From a9e46927e6647ecdf50d7a381496488f71b181f2 Mon Sep 17 00:00:00 2001 From: Hoda Mohammadi Date: Wed, 21 Oct 2026 22:51:21 -0700 Subject: [PATCH 1/2] perf(sampler): load the gemma-4 LoRA sampler as the text-only model --- src/server/worker_manager.py | 5 ++++- tests/test_scheduler_worker_manager.py | 11 +++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index 9ca9c8d4f..432317e08 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -111,7 +111,10 @@ 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. + if "gemma-4" in base_model.lower() or "gemma4" in base_model.lower(): env["VLLM_ARCHITECTURE_OVERRIDE"] = "Gemma4ForCausalLM" return env diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 40becf0e9..0621bc1f2 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/tests/test_scheduler_worker_manager.py @@ -149,6 +149,17 @@ 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_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): From cba2422b1a4199f7ad469dc545de30f88191ddd9 Mon Sep 17 00:00:00 2001 From: Hoda Mohammadi Date: Thu, 22 Oct 2026 13:41:00 -0700 Subject: [PATCH 2/2] perf(sampler): keep the full gemma-4 model when VLLM_ENABLE_MULTIMODAL=1 VLLM_ENABLE_MULTIMODAL=1 now also skips the Gemma4ForCausalLM override, and the API server forwards it to Kubernetes worker pods so the sampler sees the same setting. --- docs/configuration.md | 2 +- src/server/scheduler_worker_manager.py | 2 +- src/server/worker_manager.py | 6 ++++-- tests/test_scheduler_worker_manager.py | 12 +++++++++++- 4 files changed, 17 insertions(+), 5 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 377d274e6..5d16bedbe 100644 --- a/docs/configuration.md +++ b/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 diff --git a/src/server/scheduler_worker_manager.py b/src/server/scheduler_worker_manager.py index 29743bed8..4794098c2 100644 --- a/src/server/scheduler_worker_manager.py +++ b/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. diff --git a/src/server/worker_manager.py b/src/server/worker_manager.py index 432317e08..3741e80b5 100644 --- a/src/server/worker_manager.py +++ b/src/server/worker_manager.py @@ -113,8 +113,10 @@ def worker_env(meta: TrainingModelMetadata, base_model: str, runtime: str, is_lo env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1" # 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. - if "gemma-4" in base_model.lower() or "gemma4" in base_model.lower(): + # 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 diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 0621bc1f2..30ca7678a 100644 --- a/tests/test_scheduler_worker_manager.py +++ b/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") @@ -160,6 +160,16 @@ def test_gemma_4_samplers_load_the_text_only_model(self) -> None: 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):