[BugFix] Forward CLI --tokenizer to per-stage engine configs (#3120)

Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>

tests/engine/test_async_omni_engine_stage_init.py

@@ -563,6 +563,78 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
     assert captured_timeout == 302
 
 
+def test_async_omni_engine_reads_tokenizer_from_engine_args(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    class DummyThread:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def start(self):
+            pass
+
+        def is_alive(self):
+            return False
+
+    fake_engine_args = types.SimpleNamespace(
+        _explicit_fields=frozenset({"tokenizer"}),
+        explicit_kwargs=lambda: {"tokenizer": "/tokenizer/from-engine-args"},
+    )
+
+    monkeypatch.setattr(engine_mod.threading, "Thread", DummyThread)
+    monkeypatch.setattr(
+        AsyncOmniEngine,
+        "_resolve_stage_configs",
+        lambda self, model, kwargs: ("dummy-config", [types.SimpleNamespace(engine_args={})]),
+    )
+    monkeypatch.setattr(AsyncOmniEngine, "_wait_for_orchestrator_init", lambda *_, **__: None)
+
+    engine = AsyncOmniEngine("dummy-model", engine_args=fake_engine_args)
+
+    assert engine.tokenizer == "/tokenizer/from-engine-args"
+    engine.shutdown()
+
+
+def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "llm"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(
+        stage_cfg,
+        str(tmp_path),
+        cli_tokenizer="/external/tokenizer",
+    )
+
+    assert engine_args["model"] == os.path.join(str(tmp_path), "llm")
+    assert engine_args["tokenizer"] == "/external/tokenizer"
+
+
+def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path):
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict
+
+    stage_cfg = types.SimpleNamespace(
+        stage_id=0,
+        stage_type="llm",
+        engine_args={"model_subdir": "llm", "tokenizer_subdir": "tokenizer"},
+        default_sampling_params={},
+    )
+
+    engine_args = build_engine_args_dict(
+        stage_cfg,
+        str(tmp_path),
+        cli_tokenizer="/external/tokenizer",
+    )
+
+    assert engine_args["model"] == os.path.join(str(tmp_path), "llm")
+    assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer")
+
+
 def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch):
     import vllm_omni.engine.stage_init_utils as init_mod
 

vllm_omni/engine/arg_utils.py

@@ -438,10 +438,9 @@ class OrchestratorArgs:
     parallel_config: Any = None
 
     # === Multi-stage guards ===
-    # --tokenizer is captured here so it does not propagate to every stage
-    # uniformly (different stages often need different tokenizers, e.g.
-    # qwen3_omni thinker vs talker). Users wanting a per-stage tokenizer
-    # should set it in the deploy YAML.
+    # --tokenizer is captured by the orchestrator and forwarded to stages
+    # only when the stage does not define tokenizer/tokenizer_subdir itself.
+    # Users wanting a per-stage tokenizer should set it in the deploy YAML.
     tokenizer: str | None = None
 
 

vllm_omni/engine/async_omni_engine.py

@@ -261,6 +261,8 @@ class AsyncOmniEngine:
             ea_dict.pop("model", None)
             kwargs = {**ea_dict, **kwargs}
 
+        self.tokenizer: str | None = kwargs.get("tokenizer")
+
         # ------------------------------------------------------------------ #
         # Single-stage mode detection                                        #
         # ------------------------------------------------------------------ #
@@ -485,6 +487,7 @@ class AsyncOmniEngine:
                     stage_cfg,
                     self.model,
                     stage_connector_spec=stage_connector_spec,
+                    cli_tokenizer=getattr(self, "tokenizer", None),
                 )
                 omni_conn_cfg, omni_from, omni_to = omni_kv_connector
                 if omni_conn_cfg:
@@ -654,6 +657,7 @@ class AsyncOmniEngine:
                                 stage_cfg,
                                 self.model,
                                 stage_connector_spec=plan.stage_connector_spec,
+                                cli_tokenizer=getattr(self, "tokenizer", None),
                             )
                             lock_fds = acquire_device_locks(
                                 plan.metadata.stage_id,

vllm_omni/engine/stage_init_utils.py

@@ -104,6 +104,18 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s
     return model
 
 
+def apply_cli_tokenizer(
+    engine_args: dict[str, Any],
+    *,
+    cli_tokenizer: str | None,
+    stage_defines_tokenizer: bool,
+) -> None:
+    """Forward CLI tokenizer unless the stage config defines its own."""
+    if cli_tokenizer is None or stage_defines_tokenizer:
+        return
+    engine_args["tokenizer"] = cli_tokenizer
+
+
 def terminate_alive_proc(proc, timeout=5):
     if proc.is_alive():
         proc.terminate()
@@ -531,6 +543,7 @@ def build_engine_args_dict(
     stage_config: Any,
     model: str,
     stage_connector_spec: dict[str, Any] | None = None,
+    cli_tokenizer: str | None = None,
 ) -> dict[str, Any]:
     """Build the normalized engine args dict for one stage."""
     engine_args = stage_config.engine_args
@@ -545,7 +558,15 @@ def build_engine_args_dict(
     stage_id = stage_config.stage_id
 
     engine_args_dict = _to_dict(engine_args)
+    stage_defines_tokenizer = (
+        engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None
+    )
     model = _resolve_model_tokenizer_paths(model, engine_args_dict)
+    apply_cli_tokenizer(
+        engine_args_dict,
+        cli_tokenizer=cli_tokenizer,
+        stage_defines_tokenizer=stage_defines_tokenizer,
+    )
     engine_args_dict["model"] = model
     # Stage id must come from stage config instead of inherited CLI kwargs
     # (e.g. `--stage-id` defaulting to None).

vllm_omni/entrypoints/cli/serve.py

@@ -628,6 +628,7 @@ def run_headless(args: argparse.Namespace) -> None:
         stage_cfg,
         model,
         stage_connector_spec=stage_connector_spec,
+        cli_tokenizer=getattr(args, "tokenizer", None),
     )
 
     # Inject omni KV connector config so the engine runner can initialize the