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 6, 19:15 UTC
Use Transformers v5 `WeightRenaming` for Transformers modeling backend (#31545)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 13, 2026 at 20:49 UTC
0005d2a3c9ed8cf8bab4018b7064ceb4fd9548d1
7 files changed
+163
-90
tests/models/multimodal/generation/test_common.py
+1
-3
@@ -206,9 +206,7 @@ VLM_TEST_SETTINGS = {
206
"model_impl": "transformers",
207
"default_torch_num_threads": 1,
208
},
209
- # FIXME: Investigate why the test hangs
210
- # when processing the 3rd prompt in vLLM
211
- marks=[pytest.mark.core_model, pytest.mark.skip(reason="Test hangs")],
209
+ marks=[pytest.mark.core_model],
210
),
211
# Gemma3 has bidirectional mask on images
212
"gemma3-transformers": VLMTestInfo(
tests/models/multimodal/test_mapping.py
+31
-4
@@ -5,9 +5,10 @@ from collections.abc import Iterable
5
import pytest
6
import torch
7
import transformers
8
-from transformers import AutoConfig, PreTrainedModel
8
+from transformers import AutoConfig, AutoModel, PreTrainedModel
9
10
from vllm.config import ModelConfig
11
+from vllm.model_executor.models.transformers.base import Base as TransformersBase
12
from vllm.model_executor.models.utils import WeightsMapper
13
from vllm.multimodal import MULTIMODAL_REGISTRY
14
from vllm.transformers_utils.config import try_get_safetensors_metadata
@@ -23,6 +24,16 @@ def create_repo_dummy_weights(repo: str) -> Iterable[tuple[str, torch.Tensor]]:
24
return ((name, torch.empty(0)) for name in weight_names)
25
26
27
+def create_dummy_base_model(repo: str, model_arch: str) -> PreTrainedModel:
28
+ """
29
+ Create weights from a dummy meta deserialized hf base model with name conversion
30
+ """
31
+ config = AutoConfig.from_pretrained(repo)
32
+ with torch.device("meta"):
33
+ model = AutoModel.from_config(config)
34
+ return model
35
+
36
+
37
def create_dummy_model(repo: str, model_arch: str) -> PreTrainedModel:
38
"""
39
Create weights from a dummy meta deserialized hf model with name conversion
@@ -79,6 +90,19 @@ def test_hf_model_weights_mapper(model_arch: str):
90
dtype=model_info.dtype,
91
)
92
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
93
+ if issubclass(model_cls, TransformersBase):
94
+ # Transformers backend models create their mapper during __init__
95
+ # by inspecting the HF model instance. We simulate this by calling
96
+ # _create_hf_to_vllm_mapper with a minimal proxy object.
97
+ model_cls = type(
98
+ "ProxyModelCls",
99
+ (),
100
+ {
101
+ "model": create_dummy_base_model(model_id, model_arch),
102
+ "_maybe_apply_model_mapping": lambda self: None,
103
+ },
104
+ )()
105
+ TransformersBase._create_hf_to_vllm_mapper(model_cls)
106
107
original_weights = create_repo_dummy_weights(model_id)
108
hf_dummy_model = create_dummy_model(model_id, model_arch)
@@ -102,9 +126,12 @@ def test_hf_model_weights_mapper(model_arch: str):
126
# after they are tied in the model, so the mapper will not be able to map them.
127
# We exclude them from the reference weight names for this test.
128
if isinstance(tied := getattr(hf_dummy_model, "_tied_weights_keys", None), dict):
105
- mapped_tied_weights = mapper.apply((k, None) for k in tied)
106
- tied_weight_names = set(map(lambda x: x[0], mapped_tied_weights))
107
- ref_weight_names -= tied_weight_names
129
+ config = hf_dummy_model.config
130
+ key = "tie_word_embeddings"
131
+ if getattr(config.get_text_config(), key, False) or getattr(config, key, False):
132
+ mapped_tied_weights = mapper.apply((k, None) for k in tied)
133
+ tied_weight_names = set(map(lambda x: x[0], mapped_tied_weights))
134
+ ref_weight_names -= tied_weight_names
135
136
weights_missing = ref_weight_names - weight_names
137
weights_unmapped = weight_names - ref_weight_names
vllm/model_executor/models/interfaces.py
+13
-13
@@ -995,19 +995,10 @@ class SupportsQuant:
995
def __new__(cls, *args, **kwargs) -> Self:
996
instance = super().__new__(cls)
997
998
- # find config passed in arguments
999
- quant_config = cls._find_quant_config(*args, **kwargs)
1000
- if quant_config is not None:
1001
- # attach config to model for general use
1002
- instance.quant_config = quant_config
1003
-
1004
- # apply model mappings to config for proper config-model matching
1005
- if (hf_to_vllm_mapper := instance.hf_to_vllm_mapper) is not None:
1006
- instance.quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
1007
- if instance.packed_modules_mapping is not None:
1008
- instance.quant_config.packed_modules_mapping.update(
1009
- instance.packed_modules_mapping
1010
- )
998
+ # find config passed in arguments and attach it to model for general use
999
+ instance.quant_config = cls._find_quant_config(*args, **kwargs)
1000
+
1001
+ cls._maybe_apply_model_mapping(instance)
1002
1003
return instance
1004
@@ -1026,6 +1017,15 @@ class SupportsQuant:
1017
1018
return None
1019
1020
+ def _maybe_apply_model_mapping(self):
1021
+ """Apply model mappings to config for proper config-model matching"""
1022
+ if self.quant_config is None:
1023
+ return
1024
+ if (hf_to_vllm_mapper := self.hf_to_vllm_mapper) is not None:
1025
+ self.quant_config.apply_vllm_mapper(hf_to_vllm_mapper)
1026
+ if self.packed_modules_mapping is not None:
1027
+ self.quant_config.packed_modules_mapping.update(self.packed_modules_mapping)
1028
+
1029
1030
@runtime_checkable
1031
class SupportsRealtime(Protocol):
vllm/model_executor/models/transformers/base.py
+103
-23
@@ -17,6 +17,7 @@
17
"""Transformers modeling backend base class."""
18
19
from collections.abc import Iterable
20
+from itertools import chain
21
from typing import TYPE_CHECKING
22
23
import regex as re
@@ -107,27 +108,6 @@ class Base(
108
SupportsEagle3,
109
):
110
embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it
110
- hf_to_vllm_mapper = WeightsMapper(
111
- orig_to_new_prefix={
112
- # Add `model.` prefix for base model checkpoints,
113
- # handling the case where it is already present
114
- "": "model.",
115
- "model.model.": "model.",
116
- # Heads will be adjacent to `model` (pooling included because of adapters)
117
- "model.lm_head.": "lm_head.",
118
- "model.score.": "classifier.",
119
- "model.classifier.": "classifier.",
120
- }
121
- )
122
-
123
- def __init_subclass__(cls, *args, **kwargs):
124
- """Merge hf_to_vllm_mapper in MRO from most specific to least specific."""
125
- super().__init_subclass__(*args, **kwargs)
126
- hf_to_vllm_mapper = WeightsMapper()
127
- for base in cls.__mro__:
128
- if base_hf_to_vllm_mapper := getattr(base, "hf_to_vllm_mapper", None):
129
- hf_to_vllm_mapper |= base_hf_to_vllm_mapper
130
- cls.hf_to_vllm_mapper = hf_to_vllm_mapper
111
112
def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
113
super().__init__()
@@ -174,8 +154,8 @@ class Base(
154
if "gptq" in quant_method_name:
155
self.ignore_unexpected_suffixes.append(".bias")
156
177
- # Set correct attn and init on "meta" to delay allocating GPU tensors
178
- self.text_config._attn_implementation = "vllm"
157
+ # Patch config and init on "meta" to delay allocating GPU tensors
158
+ self._patch_config()
159
with init_on_device_without_buffers("meta"):
160
self.model: PreTrainedModel = AutoModel.from_config(
161
self.config,
@@ -183,6 +163,8 @@ class Base(
163
trust_remote_code=self.model_config.trust_remote_code,
164
)
165
166
+ # Create weight name to module qualname mapper
167
+ self._create_hf_to_vllm_mapper()
168
# Remove layers not on this pipeline parallel rank
169
self.pipeline_parallel()
170
# Substitute remaining layers with vLLM's layers as needed
@@ -216,6 +198,104 @@ class Base(
198
["hidden_states"], self.text_config.hidden_size
199
)
200
201
+ def _patch_config(self):
202
+ """
203
+ Patch the config to ensure that the model is created correctly:
204
+
205
+ - Sets the attention implementation to "vllm" so the attention instances from
206
+ `create_attention_instances` are used
207
+ - Sets the dtype to the default torch dtype set by vLLM because Transformers
208
+ uses the config dtype when creating the model
209
+ - Propagates this dtype to any sub-configs because Transformers model
210
+ implementations do not support/use different dtypes in sub-models
211
+ """
212
+ self.text_config._attn_implementation = "vllm"
213
+ self.config.dtype = torch.get_default_dtype()
214
+ # TODO(hmellor): Remove this when Transformers v4 support is dropped
215
+ for sub_config_name in getattr(self.config, "sub_configs", {}):
216
+ sub_config = getattr(self.config, sub_config_name)
217
+ if sub_config.dtype != (dtype := self.config.dtype):
218
+ sub_config.dtype = dtype
219
+
220
+ def _create_hf_to_vllm_mapper(self):
221
+ """
222
+ Create a WeightsMapper to map checkpoint weight names to module qualnames.
223
+
224
+ This handles:
225
+
226
+ - Transformers weight renaming:
227
+ - from `WeightRenaming` in Transformers v5
228
+ - from `_checkpoint_conversion_mapping` in Transformers v4
229
+ - Checkpoints saved with a base model prefix that is not `model`
230
+ - Checkpoints saved with no base model prefix
231
+ - Any quantization config specific mappings
232
+ """
233
+ self.hf_to_vllm_mapper = WeightsMapper()
234
+ orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex
235
+
236
+ if Version(transformers.__version__) >= Version("5.0.0"):
237
+ from transformers.conversion_mapping import (
238
+ WeightRenaming,
239
+ get_model_conversion_mapping,
240
+ )
241
+
242
+ for mapping in get_model_conversion_mapping(self.model):
243
+ # Handle weights which have been renamed in Transformers
244
+ if isinstance(mapping, WeightRenaming):
245
+ # Recompile using regex (Transformers used re)
246
+ compiled_sources = re.compile(
247
+ mapping.compiled_sources.pattern, mapping.compiled_sources.flags
248
+ )
249
+ target_pattern = mapping.target_patterns[0]
250
+ orig_to_new_regex[compiled_sources] = target_pattern
251
+ # TODO: Handle WeightConverter to enable layer merging
252
+ else:
253
+ # Replace legacy suffixes used for norms
254
+ # TODO(hmellor): Remove this when Transformers v4 support is dropped
255
+ orig_to_new_regex.update(
256
+ {
257
+ re.compile(r"\.gamma$"): ".weight",
258
+ re.compile(r"\.beta$"): ".bias",
259
+ }
260
+ )
261
+
262
+ # Handle weights which have been renamed in Transformers
263
+ # TODO(hmellor): Remove this when Transformers v4 support is dropped
264
+ ccm = getattr(self.model, "_checkpoint_conversion_mapping", {})
265
+ for source, target in ccm.items():
266
+ orig_to_new_regex[re.compile(source)] = target
267
+
268
+ # Handle unexpected weights which should be ignored
269
+ if self.model._keys_to_ignore_on_load_unexpected is not None:
270
+ for key in self.model._keys_to_ignore_on_load_unexpected:
271
+ orig_to_new_regex[re.compile(key)] = None
272
+
273
+ # Standardise base model prefix
274
+ bmp = self.model.base_model_prefix
275
+ expected_bmp = r"model.\1"
276
+ # Handle checkpoints saved with different base model prefix
277
+ if bmp and bmp != "model":
278
+ different_bmp_pattern = re.compile(rf"^{bmp}\.(.+)")
279
+ orig_to_new_regex[different_bmp_pattern] = expected_bmp
280
+ # Handle direct children of self.model which were saved without the model prefix
281
+ direct_children = chain(
282
+ self.model.named_children(),
283
+ self.model.named_parameters(recurse=False),
284
+ self.model.named_buffers(recurse=False),
285
+ )
286
+ model_children = "|".join(name for name, _ in direct_children)
287
+ missing_bmp_pattern = re.compile(rf"^(?!model\.)(({model_children}).*)")
288
+ orig_to_new_regex[missing_bmp_pattern] = expected_bmp
289
+ # Handle weights saved as direct children of self.model which no longer are
290
+ unexpected_bmp_pattern = re.compile(rf"^(model\.)((?!{model_children}).+)")
291
+ orig_to_new_regex[unexpected_bmp_pattern] = r"\2"
292
+ # Handle lm_head which was saved inside the base model
293
+ nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)")
294
+ orig_to_new_regex[nested_lm_head_pattern] = r"\2"
295
+
296
+ # Apply mapping to quantization config if needed
297
+ self._maybe_apply_model_mapping()
298
+
299
def pipeline_parallel(self):
300
"""
301
Apply the model's pipeline parallelization plan.
vllm/model_executor/models/transformers/legacy.py
-15
@@ -20,7 +20,6 @@ from typing import TYPE_CHECKING
20
21
import torch
22
23
-from vllm.model_executor.models.utils import WeightsMapper
23
from vllm.sequence import IntermediateTensors
24
25
if TYPE_CHECKING:
@@ -28,20 +27,6 @@ if TYPE_CHECKING:
27
28
29
class LegacyMixin:
31
- hf_to_vllm_mapper = WeightsMapper(
32
- # These are applied in order, so the order matters!
33
- orig_to_new_prefix={
34
- # Handle BERT-like models
35
- "roberta": "model",
36
- "bert": "model",
37
- },
38
- orig_to_new_suffix={
39
- # Replace legacy suffixes used for norms
40
- ".gamma": ".weight",
41
- ".beta": ".bias",
42
- },
43
- )
44
-
30
def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
31
super().__init__(vllm_config=vllm_config, prefix=prefix)
32
vllm/model_executor/models/transformers/multimodal.py
-25
@@ -24,7 +24,6 @@ import torch
24
from vllm.config.utils import getattr_iter
25
from vllm.logger import init_logger
26
from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal
27
-from vllm.model_executor.models.utils import WeightsMapper
27
from vllm.multimodal import MultiModalKwargsItems
28
from vllm.multimodal.inputs import (
29
MultiModalDataDict,
@@ -273,30 +272,6 @@ class MultiModalProcessor(BaseMultiModalProcessor[MultiModalProcessingInfo]):
272
class MultiModalMixin(SupportsMultiModal, SupportsMRoPE):
273
supports_multimodal_raw_input_only = True
274
276
- # Backwards compatibility for prev released models. State dicts back then
277
- # had different formats and cannot be loaded with `AutoModel` mapping as is
278
- hf_to_vllm_mapper = WeightsMapper(
279
- orig_to_new_prefix={
280
- "language_model.model": "model.language_model",
281
- "text_model.model": "model.text_model",
282
- "vision_tower": "model.vision_tower",
283
- "vqmodel": "model.vqmodel",
284
- "visual": "model.visual",
285
- "vision_model": "model.vision_model",
286
- "vision_embed_tokens": "model.vision_embed_tokens",
287
- "image_newline": "model.image_newline",
288
- "multi_modal_projector": "model.multi_modal_projector",
289
- "text_model.lm_head": "lm_head",
290
- "language_model.lm_head": "lm_head",
291
- # Qwen models used "model" as the name for the language model.
292
- # Therefore, we must map each of submodule explicitly to avoid
293
- # conflicts with newer models that use "model.language_model".
294
- "model.embed_tokens": "model.language_model.embed_tokens",
295
- "model.layers": "model.language_model.layers",
296
- "model.norm": "model.language_model.norm",
297
- }
298
- )
299
-
275
def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
276
# Skip SupportsMRoPE.__init__ and call the next class in MRO
277
super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix)
vllm/model_executor/models/utils.py
+15
-7
@@ -7,6 +7,7 @@ from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
from typing import Any, Literal, Protocol, overload
9
10
+import regex as re
11
import torch
12
import torch.nn as nn
13
from torch.nn.modules.module import register_module_module_registration_hook
@@ -38,17 +39,17 @@ from vllm.utils.torch_utils import (
39
40
logger = init_logger(__name__)
41
41
-WeightsMapping = Mapping[str, str | None]
42
-"""If a key maps to a value of `None`, the corresponding weight is ignored."""
43
-
42
43
@dataclass
44
class WeightsMapper:
47
- """Maps the name of each weight if they match the following patterns."""
45
+ """Maps the name of each weight if they match the following patterns.
46
+
47
+ If a key maps to a value of `None`, the corresponding weight is ignored."""
48
49
- orig_to_new_substr: WeightsMapping = field(default_factory=dict)
50
- orig_to_new_prefix: WeightsMapping = field(default_factory=dict)
51
- orig_to_new_suffix: WeightsMapping = field(default_factory=dict)
49
+ orig_to_new_regex: Mapping[re.Pattern, str | None] = field(default_factory=dict)
50
+ orig_to_new_substr: Mapping[str, str | None] = field(default_factory=dict)
51
+ orig_to_new_prefix: Mapping[str, str | None] = field(default_factory=dict)
52
+ orig_to_new_suffix: Mapping[str, str | None] = field(default_factory=dict)
53
54
def __or__(self, other: "WeightsMapper") -> "WeightsMapper":
55
"""Combine two `WeightsMapper`s by merging their mappings."""
@@ -59,6 +60,13 @@ class WeightsMapper:
60
)
61
62
def _map_name(self, key: str) -> str | None:
63
+ for pattern, new_key in self.orig_to_new_regex.items():
64
+ if pattern.search(key):
65
+ if new_key is None:
66
+ return None
67
+
68
+ key = pattern.sub(new_key, key)
69
+
70
for substr, new_key in self.orig_to_new_substr.items():
71
if substr in key:
72
if new_key is None: