A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 22:30 UTC
[Bugfix] Fix `_CONFIG_REGISTRY` types getting wrong config class when on-disk model_type differs (#39554)
Signed-off-by: Misa <misaAle@users.noreply.github.com> Signed-off-by: Misael Casarez <misacasa@amazon.com> Co-authored-by: Misael Casarez <misacasa@amazon.com>
Misa committed
Apr 20, 2026 at 19:04 UTC
0e884fe638a3120a8772c3e95e71728b56db20e4
2 files changed
+91
tests/transformers_utils/test_hf_overrides_model_type.py
new
+81
@@ -0,0 +1,81 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Test that hf_overrides model_type returns the correct config class."""
4
+
5
+import json
6
+import tempfile
7
+
8
+from transformers import PretrainedConfig
9
+
10
+from vllm.transformers_utils.config import _CONFIG_REGISTRY, get_config
11
+
12
+
13
+class _TestCustomConfig(PretrainedConfig):
14
+ model_type = "test_custom_model"
15
+
16
+ def __init__(self, custom_attr=42, **kw):
17
+ super().__init__(**kw)
18
+ self.custom_attr = custom_attr
19
+
20
+
21
+def test_hf_overrides_model_type_returns_correct_config_class():
22
+ """When hf_overrides sets model_type to a registered custom type whose
23
+ checkpoint has a *different* model_type on disk, get_config() must return
24
+ an instance of the registered config class — not the class that matches
25
+ the on-disk model_type."""
26
+
27
+ # Register the custom config
28
+ _CONFIG_REGISTRY["test_custom_model"] = _TestCustomConfig
29
+
30
+ try:
31
+ with tempfile.TemporaryDirectory() as tmpdir:
32
+ # Checkpoint says model_type="mixtral" on disk
33
+ cfg = {
34
+ "model_type": "mixtral",
35
+ "hidden_size": 128,
36
+ "num_hidden_layers": 2,
37
+ "num_attention_heads": 4,
38
+ "num_key_value_heads": 4,
39
+ "intermediate_size": 256,
40
+ "num_local_experts": 4,
41
+ "num_experts_per_tok": 2,
42
+ }
43
+ with open(f"{tmpdir}/config.json", "w") as f:
44
+ json.dump(cfg, f)
45
+
46
+ config = get_config(
47
+ tmpdir,
48
+ trust_remote_code=False,
49
+ hf_overrides_kw={
50
+ "model_type": "test_custom_model",
51
+ },
52
+ )
53
+
54
+ from transformers import AutoConfig
55
+ from transformers.models.auto.configuration_auto import CONFIG_MAPPING
56
+
57
+ # get_config() returns the registered custom class
58
+ assert isinstance(config, _TestCustomConfig), (
59
+ f"Expected _TestCustomConfig, got {type(config).__name__}"
60
+ )
61
+
62
+ # AutoConfig has _TestCustomConfig registered under both
63
+ # the overridden model_type and the on-disk model_type
64
+ assert CONFIG_MAPPING["test_custom_model"] is _TestCustomConfig
65
+ assert CONFIG_MAPPING["mixtral"] is _TestCustomConfig
66
+
67
+ # AutoConfig.from_pretrained now returns _TestCustomConfig
68
+ # for this checkpoint (even though its on-disk model_type
69
+ # is "mixtral")
70
+ auto_config = AutoConfig.from_pretrained(tmpdir)
71
+ assert isinstance(auto_config, _TestCustomConfig), (
72
+ f"Expected _TestCustomConfig from AutoConfig, got "
73
+ f"{type(auto_config).__name__}"
74
+ )
75
+ finally:
76
+ _CONFIG_REGISTRY.pop("test_custom_model", None)
77
+ # Restore the original mixtral AutoConfig mapping to avoid
78
+ # side effects on other tests in the same process
79
+ from transformers import AutoConfig, MixtralConfig
80
+
81
+ AutoConfig.register("mixtral", MixtralConfig, exist_ok=True)
vllm/transformers_utils/config.py
+10
@@ -210,6 +210,16 @@ class HFConfigParser(ConfigParserBase):
210
config_class = _CONFIG_REGISTRY[model_type]
211
config_class.model_type = model_type
212
AutoConfig.register(model_type, config_class, exist_ok=True)
213
+ # If the on-disk model_type differs from the overridden
214
+ # one, register under both so AutoConfig.from_pretrained
215
+ # returns the correct class regardless of what the
216
+ # checkpoint says
217
+ if (
218
+ config_model_type := config_dict.get("model_type")
219
+ ) and config_model_type != model_type:
220
+ config_class.model_type = config_model_type
221
+ AutoConfig.register(config_model_type, config_class, exist_ok=True)
222
+ config_class.model_type = model_type
223
# Now that it is registered, it is not considered remote code anymore
224
trust_remote_code = False
225
try: