diff --git a/docs/configuration.md b/docs/configuration.md index 377d274e..5d16bedb 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 29743bed..4794098c 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 9ca9c8d4..3741e80b 100644 --- a/src/server/worker_manager.py +++ b/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 diff --git a/tests/test_scheduler_worker_manager.py b/tests/test_scheduler_worker_manager.py index 40becf0e..30ca7678 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") @@ -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):