perf(sampler): load the gemma-4 LoRA sampler as the text-only model

docs/configuration.md

@@ -81,7 +81,7 @@ backend by default for physical checkpoint/restore.
 | Env var | Default | What it does |
 | --- | --- | --- |
 | `VLLM_ARCHITECTURE_OVERRIDE` | unset | Optional architecture override passed to the in-repo vLLM worker. Gemma 4 examples use `Gemma4ForCausalLM`. |
-| `VLLM_ENABLE_MULTIMODAL` | `0` | By default the samplers pass `limit_mm_per_prompt={"image": 0, "video": 0}`. Text checkpoints published as `*ForConditionalGeneration` otherwise make vLLM reserve a multi-GiB encoder cache during startup that no OpenRL code path can use, which can OOM engine init. Set to `1` to restore stock vLLM behaviour. |
+| `VLLM_ENABLE_MULTIMODAL` | `0` | By default the samplers pass `limit_mm_per_prompt={"image": 0, "video": 0}`. Text checkpoints published as `*ForConditionalGeneration` otherwise make vLLM reserve a multi-GiB encoder cache during startup that no OpenRL code path can use, which can OOM engine init. Gemma 4 samplers also load the text-only `Gemma4ForCausalLM`. Set to `1` to restore stock vLLM behaviour and keep the full model. |
 
 ## Client variables
 

src/server/scheduler_worker_manager.py

@@ -120,7 +120,7 @@ def pod_env(worker: Worker) -> list[dict[str, Any]]:
   }
   # MAX_JOBS caps FlashInfer's JIT build, which otherwise runs one ~3GB
   # compiler per core and blows through the pod's host memory limit.
-  for name in ("VLLM_GPU_MEMORY_UTILIZATION", "VLLM_MAX_MODEL_LEN", "OPEN_RL_TRAIN_TOKEN_BUDGET", "MAX_JOBS"):
+  for name in ("VLLM_GPU_MEMORY_UTILIZATION", "VLLM_MAX_MODEL_LEN", "VLLM_ENABLE_MULTIMODAL", "OPEN_RL_TRAIN_TOKEN_BUDGET", "MAX_JOBS"):
     if os.getenv(name):
       values[name] = os.environ[name]
   # No other worker shares an exclusive worker's GPUs, so it never parks.

src/server/worker_manager.py

@@ -111,7 +111,12 @@ def worker_env(meta: TrainingModelMetadata, base_model: str, runtime: str, is_lo
     env["OPEN_RL_MODEL_ID"] = runtime
     env["VLLM_SERVER_DEV_MODE"] = "1"
     env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
-    if not is_lora and ("gemma-4" in base_model.lower() or "gemma4" in base_model.lower()):
+    # The text-only class skips the multimodal front end vLLM otherwise builds
+    # at start-up (about 45 s); OpenRL only sends token ids. LoRA adapters
+    # trained on the full model apply to it unchanged. VLLM_ENABLE_MULTIMODAL=1
+    # keeps the full model.
+    gemma_4 = "gemma-4" in base_model.lower() or "gemma4" in base_model.lower()
+    if gemma_4 and os.getenv("VLLM_ENABLE_MULTIMODAL", "0") != "1":
       env["VLLM_ARCHITECTURE_OVERRIDE"] = "Gemma4ForCausalLM"
   return env
 

tests/test_scheduler_worker_manager.py

@@ -132,7 +132,7 @@ def test_mutable_worker_images_use_the_requested_pull_policy(self) -> None:
 
   def test_workers_get_the_deployment_settings(self) -> None:
     s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"})
-    settings = {"VLLM_MAX_MODEL_LEN": "131072", "OPEN_RL_TRAIN_TOKEN_BUDGET": "131072", "MAX_JOBS": "4"}
+    settings = {"VLLM_MAX_MODEL_LEN": "131072", "VLLM_ENABLE_MULTIMODAL": "1", "OPEN_RL_TRAIN_TOKEN_BUDGET": "131072", "MAX_JOBS": "4"}
     with patch.dict(os.environ, settings), patch("server.worker_manager.get_state_store", return_value=s):
       self.manager.ensure("job-lora-1", "sampler")
       self.manager.ensure("job-lora-1", "trainer")
@@ -149,6 +149,27 @@ def test_launch_is_idempotent(self) -> None:
       self.manager.ensure("job-lora-1", "trainer")
     self.assertEqual(len(self.api.created), 1)
 
+  def test_gemma_4_samplers_load_the_text_only_model(self) -> None:
+    for model_id, tuning in (("job-lora-1", "lora"), ("job-fft-1", "full")):
+      s = self.store_with(model_id, {"base_model": "google/gemma-4-E2B", "fine_tuning_type": tuning})
+      with patch.dict(os.environ, {"OPEN_RL_ENABLE_FFT": "true"}), patch("server.worker_manager.get_state_store", return_value=s):
+        self.manager.ensure(model_id, "sampler")
+
+    self.assertEqual(len(self.api.created), 2)
+    for pod in self.api.created:
+      env = {e["name"]: e.get("value") for e in pod["spec"]["template"]["spec"]["containers"][0]["env"]}
+      self.assertEqual(env.get("VLLM_ARCHITECTURE_OVERRIDE"), "Gemma4ForCausalLM")
+
+  def test_multimodal_gemma_4_samplers_keep_the_full_model(self) -> None:
+    s = self.store_with("job-lora-1", {"base_model": "google/gemma-4-E2B", "fine_tuning_type": "lora"})
+    with patch.dict(os.environ, {"VLLM_ENABLE_MULTIMODAL": "1"}), patch("server.worker_manager.get_state_store", return_value=s):
+      self.manager.ensure("job-lora-1", "sampler")
+
+    (pod,) = self.api.created
+    env = {e["name"]: e.get("value") for e in pod["spec"]["template"]["spec"]["containers"][0]["env"]}
+    self.assertNotIn("VLLM_ARCHITECTURE_OVERRIDE", env)
+    self.assertEqual(env.get("VLLM_ENABLE_MULTIMODAL"), "1")
+
   def test_placement_knowledge_stays_out_of_the_template(self) -> None:
     s = self.store_with("job-lora-1", {"base_model": "Qwen/Qwen2.5-0.5B", "fine_tuning_type": "lora"})
     with patch("server.worker_manager.get_state_store", return_value=s):