fix(grpo): allow FATHOM_USE_VLLM=0 to bypass vLLM rollout and avoid IS-ratio collapse under QLoRA

Gate use_vllm, vllm_mode, and vllm_gpu_memory_utilization on the
FATHOM_USE_VLLM env var (default "1" = enabled). When set to "0",
TRL falls back to HF generate() for rollouts — ~3x slower per step
but importance_sampling_ratio stays near 1.0 instead of 1e-7 to 1e-5.

Root cause of the 200-step flat-reward run: QLoRA + vLLM colocate
merge/unmerge produces token-prob drift between the rollout model
(vLLM, merged 4-bit) and the gradient model (HF, LoRA over 4-bit base),
causing TRL IS clipping to zero every gradient update.

The vllm_mode='colocate' assert is now inside the if use_vllm: branch.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

train/grpo.py

@@ -75,10 +75,15 @@ def run_grpo(
             sft_adapter_dir,
         )
 
-    # TRN-03 step 2: vllm_mode MUST be 'colocate' (STACK §10.4 / TRL #4543)
-    assert str(cfg.train.vllm_mode) == "colocate", (
-        "TRN-03 gate: vllm_mode must be 'colocate' for multi-turn OpenEnv (STACK §10.4)"
-    )
+    # TRN-03 step 2: FATHOM_USE_VLLM env var gates vLLM rollout.
+    # Set FATHOM_USE_VLLM=0 to use HF generate() for rollouts (slower but
+    # avoids IS-ratio collapse under QLoRA + vLLM merge drift).
+    use_vllm_env = os.environ.get("FATHOM_USE_VLLM", "1").strip()
+    use_vllm = use_vllm_env not in ("0", "false", "False", "")
+    if use_vllm:
+        assert str(cfg.train.vllm_mode) == "colocate", (
+            "TRN-03 gate: vllm_mode must be 'colocate' for multi-turn OpenEnv (STACK §10.4)"
+        )
 
     # TRN-03 step 2: Build GRPOConfig from cfg.train.
     # TRL minor versions have changed some GRPOConfig field names; select only
@@ -115,12 +120,20 @@ def run_grpo(
     _set_if_supported(["max_steps"], int(cfg.train.max_steps))
     _set_if_supported(["save_steps"], int(cfg.train.save_steps))
     _set_if_supported(["seed"], int(cfg.seed))
-    _set_if_supported(["vllm_mode"], str(cfg.train.vllm_mode))
-    _set_if_supported(
-        ["vllm_gpu_memory_utilization"],
-        float(cfg.train.vllm_gpu_memory_utilization),
-    )
-    import os
+    _set_if_supported(["use_vllm"], use_vllm)
+    if use_vllm:
+        _set_if_supported(["vllm_mode"], str(cfg.train.vllm_mode))
+        _set_if_supported(
+            ["vllm_gpu_memory_utilization"],
+            float(cfg.train.vllm_gpu_memory_utilization),
+        )
+        log.info("TRN-03 vLLM rollout enabled (vllm_mode=%s)", cfg.train.vllm_mode)
+    else:
+        log.warning(
+            "TRN-03 FATHOM_USE_VLLM=0 — using HF generate() for "
+            "rollouts. Slower but importance_sampling_ratio stays "
+            "near 1.0 (avoids QLoRA + vLLM merge drift bug)."
+        )
     wb_key = os.environ.get("WANDB_API_KEY", "").strip()
     use_wandb = len(wb_key) >= 40 and wb_key.isalnum()
     _set_if_supported(["report_to"], ["wandb"] if use_wandb else [])
@@ -131,8 +144,6 @@ def run_grpo(
             "W&B logging; metrics still printed to stdout.", len(wb_key)
         )
     _set_if_supported(["logging_steps"], 1)
-    # vllm_mode='colocate' alone is a no-op without `use_vllm=True` in TRL 1.2.
-    _set_if_supported(["use_vllm"], True)
     # Keep TRL from dropping our extra columns (gold_answer etc.) so they
     # reach the reward function as kwargs.
     _set_if_supported(["remove_unused_columns"], False)