[UX] Add `model_class_overrides` for development and debugging (#47148)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jee Jee Li committed
Jul 7, 2026 at 11:43 UTC
32ab06462176143d7e1e1cd9a2a89ef3dac0ef96
3 files changed
+70
tests/test_config.py
+24
@@ -649,6 +649,30 @@ def test_nested_hf_overrides():
649
assert model_config.hf_config.vision_config.hidden_size == 512
650
651
652
+def test_model_class_overrides_registers_target():
653
+ """`model_class_overrides` redirects an architecture to a custom class."""
654
+ from vllm.model_executor.models import ModelRegistry
655
+
656
+ arch = "_TestModelClassOverrideArch"
657
+ target = "vllm.model_executor.models.llama:LlamaForCausalLM"
658
+ assert arch not in ModelRegistry.models
659
+
660
+ model_config = ModelConfig(
661
+ "facebook/opt-125m",
662
+ model_class_overrides={arch: target},
663
+ )
664
+ try:
665
+ # Accessing `.registry` is the chokepoint that applies the overrides;
666
+ # it has already run during construction.
667
+ registered = model_config.registry.models[arch]
668
+ assert registered.module_name == "vllm.model_executor.models.llama"
669
+ assert registered.class_name == "LlamaForCausalLM"
670
+ # Idempotent: a second access does not re-register or error out.
671
+ assert model_config.registry.models[arch] is registered
672
+ finally:
673
+ ModelRegistry.models.pop(arch, None)
674
+
675
+
676
@pytest.mark.skipif(
677
current_platform.is_rocm(), reason="Encoder Decoder models not supported on ROCm."
678
)
vllm/config/model.py
+39
@@ -74,6 +74,12 @@ else:
74
75
logger = init_logger(__name__)
76
77
+# Process-local record of which (arch, target) model-class overrides have been
78
+# registered in *this* process. Must not live on ModelConfig: that instance is
79
+# pickled to each worker, so an instance flag would arrive already "registered"
80
+# while the worker's own global ModelRegistry is still untouched.
81
+_REGISTERED_MODEL_CLASS_OVERRIDES: set[tuple[str, str]] = set()
82
+
83
RunnerOption = Literal["auto", RunnerType]
84
ConvertType = Literal["none", "embed", "classify"]
85
ConvertOption = Literal["auto", ConvertType]
@@ -274,6 +280,13 @@ class ModelConfig:
280
hf_overrides: HfOverrides = field(default_factory=dict)
281
"""If a dictionary, contains arguments to be forwarded to the Hugging Face
282
config. If a callable, it is called to update the HuggingFace config."""
283
+ model_class_overrides: dict[str, str] = field(default_factory=dict)
284
+ """Override the model class used for one or more architectures, mapping the
285
+ architecture name to a `"module:class"` target (the same format accepted by
286
+ `ModelRegistry.register_model`). This registers the target class at runtime,
287
+ e.g. `{"GlmMoeDsaForCausalLM":
288
+ "vllm.models.deepseek_v32.nvidia.model:DeepseekV32ForCausalLM"}`. This
289
+ argument is for development and debugging purposes only."""
290
generation_config: str = "auto"
291
"""The folder path to the generation config. Defaults to `"auto"`, the
292
generation config will be loaded from model path. If set to `"vllm"`, no
@@ -812,8 +825,34 @@ class ModelConfig:
825
826
@property
827
def registry(self):
828
+ self._maybe_register_model_class_overrides()
829
return me_models.ModelRegistry
830
831
+ def _maybe_register_model_class_overrides(self) -> None:
832
+ # Apply ``model_class_overrides`` here because this property is the
833
+ # single chokepoint through which every model-class inspect/resolve
834
+ # passes, in both the engine front-end and every worker process. The
835
+ # guard is process-local (see ``_REGISTERED_MODEL_CLASS_OVERRIDES``), so
836
+ # each worker re-registers into its own ModelRegistry exactly once
837
+ # rather than trusting a pickled-in instance flag.
838
+ if not self.model_class_overrides:
839
+ return
840
+ pending = [
841
+ (arch, target)
842
+ for arch, target in self.model_class_overrides.items()
843
+ if (arch, target) not in _REGISTERED_MODEL_CLASS_OVERRIDES
844
+ ]
845
+ if not pending:
846
+ return
847
+ logger.warning_once(
848
+ "Applying model_class_overrides %s. This is intended for "
849
+ "development/debugging.",
850
+ str(self.model_class_overrides),
851
+ )
852
+ for arch, target in pending:
853
+ me_models.ModelRegistry.register_model(arch, target)
854
+ _REGISTERED_MODEL_CLASS_OVERRIDES.add((arch, target))
855
+
856
@property
857
def architectures(self) -> list[str]:
858
return self.model_arch_config.architectures
vllm/engine/arg_utils.py
+7
@@ -536,6 +536,9 @@ class EngineArgs:
536
code_revision: str | None = ModelConfig.code_revision
537
hf_token: bool | str | None = ModelConfig.hf_token
538
hf_overrides: HfOverrides = get_field(ModelConfig, "hf_overrides")
539
+ model_class_overrides: dict[str, str] = get_field(
540
+ ModelConfig, "model_class_overrides"
541
+ )
542
tokenizer_revision: str | None = ModelConfig.tokenizer_revision
543
quantization: QuantizationMethods | str | None = ModelConfig.quantization
544
quantization_config: "dict[str, Any] | QuantizationConfigArgs | None" = None
@@ -851,6 +854,9 @@ class EngineArgs:
854
model_group.add_argument("--config-format", **model_kwargs["config_format"])
855
model_group.add_argument("--hf-token", **model_kwargs["hf_token"])
856
model_group.add_argument("--hf-overrides", **model_kwargs["hf_overrides"])
857
+ model_group.add_argument(
858
+ "--model-class-overrides", **model_kwargs["model_class_overrides"]
859
+ )
860
model_group.add_argument("--pooler-config", **model_kwargs["pooler_config"])
861
model_group.add_argument(
862
"--generation-config", **model_kwargs["generation_config"]
@@ -1622,6 +1628,7 @@ class EngineArgs:
1628
code_revision=self.code_revision,
1629
hf_token=self.hf_token,
1630
hf_overrides=self.hf_overrides,
1631
+ model_class_overrides=self.model_class_overrides,
1632
tokenizer_revision=self.tokenizer_revision,
1633
max_model_len=self.max_model_len,
1634
quantization=self.quantization,