From 32ab06462176143d7e1e1cd9a2a89ef3dac0ef96 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Tue, 7 Jul 2026 11:43:08 +0800 Subject: [PATCH] [UX] Add `model_class_overrides` for development and debugging (#47148) Signed-off-by: Jee Jee Li Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/test_config.py | 24 ++++++++++++++++++++++++ vllm/config/model.py | 39 +++++++++++++++++++++++++++++++++++++++ vllm/engine/arg_utils.py | 7 +++++++ 3 files changed, 70 insertions(+) diff --git a/tests/test_config.py b/tests/test_config.py index 3837057658..1e93b610da 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -649,6 +649,30 @@ def test_nested_hf_overrides(): assert model_config.hf_config.vision_config.hidden_size == 512 +def test_model_class_overrides_registers_target(): + """`model_class_overrides` redirects an architecture to a custom class.""" + from vllm.model_executor.models import ModelRegistry + + arch = "_TestModelClassOverrideArch" + target = "vllm.model_executor.models.llama:LlamaForCausalLM" + assert arch not in ModelRegistry.models + + model_config = ModelConfig( + "facebook/opt-125m", + model_class_overrides={arch: target}, + ) + try: + # Accessing `.registry` is the chokepoint that applies the overrides; + # it has already run during construction. + registered = model_config.registry.models[arch] + assert registered.module_name == "vllm.model_executor.models.llama" + assert registered.class_name == "LlamaForCausalLM" + # Idempotent: a second access does not re-register or error out. + assert model_config.registry.models[arch] is registered + finally: + ModelRegistry.models.pop(arch, None) + + @pytest.mark.skipif( current_platform.is_rocm(), reason="Encoder Decoder models not supported on ROCm." ) diff --git a/vllm/config/model.py b/vllm/config/model.py index c53686eaf0..6e3fef0dcc 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -74,6 +74,12 @@ else: logger = init_logger(__name__) +# Process-local record of which (arch, target) model-class overrides have been +# registered in *this* process. Must not live on ModelConfig: that instance is +# pickled to each worker, so an instance flag would arrive already "registered" +# while the worker's own global ModelRegistry is still untouched. +_REGISTERED_MODEL_CLASS_OVERRIDES: set[tuple[str, str]] = set() + RunnerOption = Literal["auto", RunnerType] ConvertType = Literal["none", "embed", "classify"] ConvertOption = Literal["auto", ConvertType] @@ -274,6 +280,13 @@ class ModelConfig: hf_overrides: HfOverrides = field(default_factory=dict) """If a dictionary, contains arguments to be forwarded to the Hugging Face config. If a callable, it is called to update the HuggingFace config.""" + model_class_overrides: dict[str, str] = field(default_factory=dict) + """Override the model class used for one or more architectures, mapping the + architecture name to a `"module:class"` target (the same format accepted by + `ModelRegistry.register_model`). This registers the target class at runtime, + e.g. `{"GlmMoeDsaForCausalLM": + "vllm.models.deepseek_v32.nvidia.model:DeepseekV32ForCausalLM"}`. This + argument is for development and debugging purposes only.""" generation_config: str = "auto" """The folder path to the generation config. Defaults to `"auto"`, the generation config will be loaded from model path. If set to `"vllm"`, no @@ -812,8 +825,34 @@ class ModelConfig: @property def registry(self): + self._maybe_register_model_class_overrides() return me_models.ModelRegistry + def _maybe_register_model_class_overrides(self) -> None: + # Apply ``model_class_overrides`` here because this property is the + # single chokepoint through which every model-class inspect/resolve + # passes, in both the engine front-end and every worker process. The + # guard is process-local (see ``_REGISTERED_MODEL_CLASS_OVERRIDES``), so + # each worker re-registers into its own ModelRegistry exactly once + # rather than trusting a pickled-in instance flag. + if not self.model_class_overrides: + return + pending = [ + (arch, target) + for arch, target in self.model_class_overrides.items() + if (arch, target) not in _REGISTERED_MODEL_CLASS_OVERRIDES + ] + if not pending: + return + logger.warning_once( + "Applying model_class_overrides %s. This is intended for " + "development/debugging.", + str(self.model_class_overrides), + ) + for arch, target in pending: + me_models.ModelRegistry.register_model(arch, target) + _REGISTERED_MODEL_CLASS_OVERRIDES.add((arch, target)) + @property def architectures(self) -> list[str]: return self.model_arch_config.architectures diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index efdd7696fd..c7a9335bbe 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -536,6 +536,9 @@ class EngineArgs: code_revision: str | None = ModelConfig.code_revision hf_token: bool | str | None = ModelConfig.hf_token hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides") + model_class_overrides: dict[str, str] = get_field( + ModelConfig, "model_class_overrides" + ) tokenizer_revision: str | None = ModelConfig.tokenizer_revision quantization: QuantizationMethods | str | None = ModelConfig.quantization quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None @@ -851,6 +854,9 @@ class EngineArgs: model_group.add_argument("--config-format", **model_kwargs["config_format"]) model_group.add_argument("--hf-token", **model_kwargs["hf_token"]) model_group.add_argument("--hf-overrides", **model_kwargs["hf_overrides"]) + model_group.add_argument( + "--model-class-overrides", **model_kwargs["model_class_overrides"] + ) model_group.add_argument("--pooler-config", **model_kwargs["pooler_config"]) model_group.add_argument( "--generation-config", **model_kwargs["generation_config"] @@ -1622,6 +1628,7 @@ class EngineArgs: code_revision=self.code_revision, hf_token=self.hf_token, hf_overrides=self.hf_overrides, + model_class_overrides=self.model_class_overrides, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization,