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 9, 19:30 UTC
[Model] Extend `collect_children` and `no_init_weights` contexts (#32757)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Jan 22, 2026 at 16:20 UTC
2b8a38b6d6be4f6e09cc20381c7027e7c35cb6c9
20 files changed
+442
-255
tests/utils_/test_collection_utils.py
+18
-2
@@ -2,11 +2,27 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
5
-from vllm.utils.collection_utils import swap_dict_values
5
+from vllm.utils.collection_utils import common_prefix, swap_dict_values
6
7
8
@pytest.mark.parametrize(
9
- "obj,key1,key2",
9
+ ("inputs", "expected_output"),
10
+ [
11
+ ([""], ""),
12
+ (["a"], "a"),
13
+ (["a", "b"], ""),
14
+ (["a", "ab"], "a"),
15
+ (["a", "ab", "b"], ""),
16
+ (["abc", "a", "ab"], "a"),
17
+ (["aba", "abc", "ab"], "ab"),
18
+ ],
19
+)
20
+def test_common_prefix(inputs, expected_output):
21
+ assert common_prefix(inputs) == expected_output
22
+
23
+
24
+@pytest.mark.parametrize(
25
+ ("obj", "key1", "key2"),
26
[
27
# Tests for both keys exist
28
({1: "a", 2: "b"}, 1, 2),
vllm/model_executor/model_loader/utils.py
+2
-15
@@ -18,7 +18,7 @@ from vllm.model_executor.layers.quantization.base_config import (
18
QuantizationConfig,
19
QuantizeMethodBase,
20
)
21
-from vllm.model_executor.models.interfaces import SupportsQuant, supports_multimodal
21
+from vllm.model_executor.models.interfaces import SupportsQuant
22
from vllm.utils.platform_utils import is_pin_memory_available
23
24
logger = init_logger(__name__)
@@ -165,11 +165,7 @@ _MODEL_ARCH_BY_HASH = dict[int, tuple[type[nn.Module], str]]()
165
166
167
def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], str]:
168
- from vllm.model_executor.models.adapters import (
169
- as_embedding_model,
170
- as_seq_cls_model,
171
- try_create_mm_pooling_model_cls,
172
- )
168
+ from vllm.model_executor.models.adapters import as_embedding_model, as_seq_cls_model
169
170
architectures = getattr(model_config.hf_config, "architectures", [])
171
@@ -189,15 +185,6 @@ def _get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module],
185
)
186
187
convert_type = model_config.convert_type
192
- if convert_type != "none" and supports_multimodal(model_cls):
193
- logger.debug_once("Detected conversion of Multi Modal model.")
194
- converted = try_create_mm_pooling_model_cls(model_cls)
195
- if converted is not None:
196
- logger.debug_once("Creating wrapper class to forward pooler.")
197
- return converted, arch
198
- else:
199
- logger.debug_once("Attempting direct conversion.")
200
-
188
if convert_type == "none":
189
pass
190
elif convert_type == "embed":
vllm/model_executor/models/adapters.py
+81
-89
@@ -1,8 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import ast
5
-import inspect
4
from collections.abc import Iterable
5
from typing import TYPE_CHECKING, Any, TypeVar, cast
6
@@ -18,10 +16,12 @@ from vllm.transformers_utils.config import (
16
)
17
from vllm.transformers_utils.repo_utils import get_hf_file_bytes
18
19
+from .interfaces import supports_multimodal
20
from .interfaces_base import VllmModelForPooling, is_pooling_model
21
22
if TYPE_CHECKING:
23
from vllm.config import ModelConfig, VllmConfig
24
+ from vllm.model_executor.layers.pooler import Pooler
25
26
_T = TypeVar("_T", bound=type[nn.Module])
27
@@ -124,20 +124,12 @@ def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
124
return model_name + pooling_suffix
125
126
127
-def try_create_mm_pooling_model_cls(orig_cls: _T) -> _T:
128
- class CallVisitor(ast.NodeVisitor):
129
- def __init__(self):
130
- self.calls = []
131
-
132
- def visit_Call(self, node):
133
- if isinstance(node.func, ast.Name):
134
- self.calls.append(node.func.id)
135
- self.generic_visit(node)
127
+def _create_pooling_model_cls(orig_cls: _T) -> _T:
128
+ # Lazy import
129
+ from vllm.model_executor.layers.logits_processor import LogitsProcessor
130
+ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
131
137
- visitor = CallVisitor()
138
- visitor.visit(ast.parse(inspect.getsource(orig_cls)))
139
- if "init_vllm_registered_model" not in visitor.calls:
140
- return None
132
+ from .utils import AutoWeightsLoader, StageMissingLayer, no_init_weights
133
134
class ModelForPooling(orig_cls, VllmModelForPooling):
135
is_pooling_model = True
@@ -149,90 +141,84 @@ def try_create_mm_pooling_model_cls(orig_cls: _T) -> _T:
141
prefix: str = "",
142
**kwargs: Any,
143
) -> None:
152
- super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
153
-
154
- self.pooler = self.get_language_model().pooler
155
-
156
- return ModelForPooling # type: ignore
144
+ with no_init_weights(
145
+ self,
146
+ lambda mod: StageMissingLayer("output", mod),
147
+ targets=(LogitsProcessor, ParallelLMHead),
148
+ ):
149
+ super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
150
+
151
+ # Used by SEQ_CLS_LOAD_METHODS
152
+ self.vllm_config = vllm_config
153
154
+ # If the model already defines a pooler instance, don't overwrite it
155
+ pooler = getattr(self, "pooler", None)
156
+ if not pooler and supports_multimodal(self):
157
+ # Try to get the pooler from the LM backbone
158
+ language_model = self.get_language_model()
159
+ if hasattr(language_model, "pooler"):
160
+ pooler = language_model.pooler
161
159
-def _create_pooling_model_cls(orig_cls: _T) -> _T:
160
- # Lazy import
161
- from .utils import AutoWeightsLoader, WeightsMapper
162
+ if not pooler:
163
+ pooler = self._init_pooler(vllm_config, prefix=prefix)
164
163
- class ModelForPooling(orig_cls, VllmModelForPooling):
164
- is_pooling_model = True
165
+ self.pooler = pooler
166
166
- def __init__(
167
+ def _init_pooler(
168
self,
168
- *,
169
vllm_config: "VllmConfig",
170
prefix: str = "",
171
- **kwargs: Any,
172
- ) -> None:
173
- super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
174
-
175
- self.vllm_config = vllm_config
176
-
177
- # These are not used in pooling models
178
- objects_to_clean = [self]
179
- if language_model := getattr(self, "language_model", None):
180
- objects_to_clean.append(language_model)
181
-
182
- for obj in objects_to_clean:
183
- for attr in ("lm_head", "logits_processor"):
184
- if hasattr(obj, attr):
185
- delattr(obj, attr)
171
+ ) -> "Pooler":
172
+ raise NotImplementedError
173
187
- # If the model already defines a pooler instance, don't overwrite it
188
- if not getattr(self, "pooler", None):
189
- self._init_pooler(vllm_config, prefix=prefix)
174
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
175
+ params_dict = dict(self.named_parameters())
176
191
- def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
192
- raise NotImplementedError
177
+ # We support loading from both `*ForCausalLM` and `*Model`
178
+ candidate_prefixes = ["", "model."]
179
+ target_prefix = ""
180
194
- def load_weights(
195
- self,
196
- weights: Iterable[tuple[str, torch.Tensor]],
197
- load_lm_head: bool = False,
198
- ):
199
- # TODO: Support uninitialized params tracking
200
-
201
- # For most pooling models: We have deleted this attribute, so don't load it.
202
- # For converting an LLM into a seq cls model, we need the lm_head.
203
- if not load_lm_head:
204
- weights = (
205
- (name, data)
206
- for name, data in weights
207
- if not name.startswith("lm_head.")
208
- )
181
+ seen_weights = list[tuple[str, torch.Tensor]]()
182
+ for name, loaded_weight in weights:
183
+ seen_weights.append((name, loaded_weight))
184
210
- # If `*ForCausalLM` defines `load_weights` on the inner model
211
- # and there are no other inner modules with parameters,
212
- # we support loading from both `*Model` and `*ForCausalLM`
213
- if hasattr(self, "model") and hasattr(self.model, "load_weights"):
214
- # Whether only `self.model` contains parameters
215
- model_is_only_param = all(
216
- name == "model" or next(child.parameters(), None) is None
217
- for name, child in self.named_children()
185
+ try:
186
+ target_prefix = next(
187
+ prefix
188
+ for prefix in candidate_prefixes
189
+ if prefix + name in params_dict
190
+ )
191
+ break
192
+ except StopIteration:
193
+ # The weight might not exist on the model
194
+ # (to be handled by AutoWeightsLoader)
195
+ pass
196
+
197
+ if target_prefix:
198
+ target_model = self
199
+ for attr in target_prefix.split("."):
200
+ if attr:
201
+ target_model = getattr(self, attr)
202
+
203
+ logger.info(
204
+ "Mapping weights to %s as they are "
205
+ "relative to this model instead of %s.",
206
+ target_model._get_name(),
207
+ self._get_name(),
208
)
209
220
- if model_is_only_param:
221
- mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
222
- weights = mapper.apply(weights)
223
-
224
- loaded_params = self.model.load_weights(weights)
225
- loaded_params = {f"model.{name}" for name in loaded_params}
226
- return loaded_params
210
+ mapped_weights = (
211
+ (target_prefix + name, weight)
212
+ for name, weight in (*seen_weights, *weights)
213
+ )
214
228
- # For most other models
229
- if hasattr(orig_cls, "load_weights"):
230
- return orig_cls.load_weights(self, weights) # type: ignore
231
- # Fallback
232
- else:
215
+ def default_load_weights(weights):
216
loader = AutoWeightsLoader(self)
217
return loader.load_weights(weights)
218
219
+ load_weights = getattr(super(), "load_weights", default_load_weights)
220
+ return load_weights(mapped_weights)
221
+
222
return ModelForPooling # type: ignore
223
224
@@ -255,11 +241,15 @@ def as_embedding_model(cls: _T) -> _T:
241
from vllm.model_executor.layers.pooler import DispatchPooler
242
243
class ModelForEmbedding(_create_pooling_model_cls(cls)):
258
- def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
244
+ def _init_pooler(
245
+ self,
246
+ vllm_config: "VllmConfig",
247
+ prefix: str = "",
248
+ ) -> "Pooler":
249
pooler_config = vllm_config.model_config.pooler_config
250
assert pooler_config is not None
251
262
- self.pooler = DispatchPooler.for_embedding(pooler_config)
252
+ return DispatchPooler.for_embedding(pooler_config)
253
254
ModelForEmbedding.__name__ = _get_pooling_model_name(cls.__name__, "ForEmbedding")
255
@@ -292,7 +282,11 @@ def as_seq_cls_model(cls: _T) -> _T:
282
class ModelForSequenceClassification(
283
_create_pooling_model_cls(cls), SupportsCrossEncoding
284
):
295
- def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
285
+ def _init_pooler(
286
+ self,
287
+ vllm_config: "VllmConfig",
288
+ prefix: str = "",
289
+ ) -> "Pooler":
290
text_config = vllm_config.model_config.hf_config.get_text_config()
291
model_config = vllm_config.model_config
292
quant_config = vllm_config.quant_config
@@ -310,9 +304,7 @@ def as_seq_cls_model(cls: _T) -> _T:
304
pooler_config = vllm_config.model_config.pooler_config
305
assert pooler_config is not None
306
313
- self.pooler = DispatchPooler.for_seq_cls(
314
- pooler_config, classifier=self.score
315
- )
307
+ return DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
308
309
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
310
hf_config = self.config
@@ -424,7 +416,7 @@ def load_weights_using_from_2_way_softmax(
416
pooling_model_cls = next(
417
x for x in type(model).__mro__ if x.__name__ == "ModelForPooling"
418
)
427
- loaded_weights = pooling_model_cls.load_weights(model, weights, load_lm_head=True)
419
+ loaded_weights = pooling_model_cls.load_weights(model, weights)
420
421
from vllm.tokenizers import get_tokenizer
422
vllm/model_executor/models/bagel.py
+4
-4
@@ -44,11 +44,11 @@ from .interfaces import (
44
SupportsLoRA,
45
SupportsMultiModal,
46
SupportsPP,
47
- TowerMissingLayer,
47
)
48
from .siglip import SiglipVisionModel
49
from .utils import (
50
AutoWeightsLoader,
51
+ StageMissingLayer,
52
WeightsMapper,
53
init_vllm_registered_model,
54
maybe_prefix,
@@ -426,9 +426,9 @@ class BagelForConditionalGeneration(
426
hidden_size=llm_hidden_size,
427
)
428
else:
429
- self.vit_model = TowerMissingLayer("image")
430
- self.connector = TowerMissingLayer("image")
431
- self.vit_pos_embed = TowerMissingLayer("image")
429
+ self.vit_model = StageMissingLayer("image_tower")
430
+ self.connector = StageMissingLayer("image_tower")
431
+ self.vit_pos_embed = StageMissingLayer("image_tower")
432
433
self.make_empty_intermediate_tensors = (
434
self.language_model.make_empty_intermediate_tensors
vllm/model_executor/models/chameleon.py
+14
-6
@@ -935,9 +935,20 @@ class ChameleonForConditionalGeneration(
935
multimodal_config = vllm_config.model_config.multimodal_config
936
self.config = config
937
self.multimodal_config = multimodal_config
938
- self.model = ChameleonModel(
939
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
940
- )
938
+
939
+ with self._mark_composite_model(
940
+ vllm_config,
941
+ language_targets=(
942
+ ChameleonDecoderLayer
943
+ if not self.config.swin_norm
944
+ else ChameleonSwinDecoderLayer
945
+ ),
946
+ tower_targets={"image": ChameleonVQVAE},
947
+ ):
948
+ self.model = ChameleonModel(
949
+ vllm_config=vllm_config,
950
+ prefix=maybe_prefix(prefix, "model"),
951
+ )
952
953
self.lm_head = ParallelLMHead(
954
config.vocab_size,
@@ -970,9 +981,6 @@ class ChameleonForConditionalGeneration(
981
resolve_bindings={"h": expected_h, "w": expected_w},
982
)
983
973
- def get_language_model(self) -> torch.nn.Module:
974
- return self.model
975
-
984
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
985
image_input = self._parse_and_validate_image_input(**kwargs)
986
if image_input is None:
vllm/model_executor/models/gemma3_mm.py
+1
-4
@@ -539,10 +539,7 @@ class Gemma3ForConditionalGeneration(
539
)
540
541
logit_scale = getattr(config, "logit_scale", 1.0)
542
- if hasattr(self.language_model, "logits_processor"):
543
- # The logits processor can be unset if we're using
544
- # automatic conversion to pooling model.
545
- self.language_model.logits_processor.scale *= logit_scale
542
+ self.language_model.logits_processor.scale *= logit_scale
543
544
self.make_empty_intermediate_tensors = (
545
self.language_model.make_empty_intermediate_tensors
vllm/model_executor/models/glm4v.py
+11
-9
@@ -53,7 +53,7 @@ from vllm.sequence import IntermediateTensors
53
from vllm.transformers_utils.configs import ChatGLMConfig
54
from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
56
-from .chatglm import ChatGLMBaseModel, ChatGLMModel
56
+from .chatglm import ChatGLMBaseModel, ChatGLMModel, GLMTransformer
57
from .interfaces import (
58
MultiModalEmbeddings,
59
SupportsLoRA,
@@ -591,11 +591,16 @@ class GLM4VForCausalLM(
591
prefix: str = "",
592
transformer_type: type[GLM4VModel] = GLM4VModel,
593
) -> None:
594
- super().__init__(
595
- vllm_config=vllm_config,
596
- prefix=prefix,
597
- transformer_type=transformer_type,
598
- )
594
+ with self._mark_composite_model(
595
+ vllm_config,
596
+ language_targets=GLMTransformer,
597
+ tower_targets={"image": EVA2CLIPModel},
598
+ ):
599
+ super().__init__(
600
+ vllm_config=vllm_config,
601
+ prefix=prefix,
602
+ transformer_type=transformer_type,
603
+ )
604
605
self.transformer: GLM4VModel
606
@@ -752,9 +757,6 @@ class GLM4VForCausalLM(
757
mrope_position_delta = (llm_positions.max() + 1 - len(input_tokens)).item()
758
return llm_positions, mrope_position_delta
759
755
- def get_language_model(self) -> torch.nn.Module:
756
- return self.transformer
757
-
760
embed_input_ids = SupportsMultiModal.embed_input_ids
761
762
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
vllm/model_executor/models/idefics3.py
+15
-7
@@ -57,7 +57,11 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape
57
from .idefics2_vision_model import (
58
Idefics2VisionTransformer as Idefics3VisionTransformer,
59
)
60
-from .interfaces import MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal
60
+from .interfaces import (
61
+ MultiModalEmbeddings,
62
+ SupportsLoRA,
63
+ SupportsMultiModal,
64
+)
65
from .llama import LlamaModel
66
from .utils import AutoWeightsLoader, maybe_prefix
67
@@ -604,9 +608,16 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo
608
self.config = config
609
self.multimodal_config = multimodal_config
610
607
- self.model = Idefics3Model(
608
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
609
- )
611
+ with self._mark_composite_model(
612
+ vllm_config,
613
+ language_targets=LlamaModel,
614
+ tower_targets={"image": (Idefics3VisionTransformer, Idefics3Connector)},
615
+ ):
616
+ self.model = Idefics3Model(
617
+ vllm_config=vllm_config,
618
+ prefix=maybe_prefix(prefix, "model"),
619
+ )
620
+
621
self.image_token_id = self.config.image_token_id
622
623
self.lm_head = ParallelLMHead(
@@ -669,9 +680,6 @@ class Idefics3ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLo
680
num_patches = image_input["num_patches"]
681
return [e.flatten(0, 1) for e in image_features.split(num_patches.tolist())]
682
672
- def get_language_model(self) -> torch.nn.Module:
673
- return self.model
674
-
683
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
684
image_input = self._parse_and_validate_image_input(**kwargs)
685
if image_input is None:
vllm/model_executor/models/interfaces.py
+106
-65
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable, Iterable, Mapping, MutableSequence
5
-from contextlib import contextmanager, nullcontext
5
+from contextlib import ExitStack, contextmanager, nullcontext
6
from typing import (
7
TYPE_CHECKING,
8
ClassVar,
@@ -25,6 +25,7 @@ from vllm.inputs import TokensPrompt
25
from vllm.inputs.data import PromptType
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.quantization import QuantizationConfig
28
+from vllm.utils.collection_utils import common_prefix
29
from vllm.utils.func_utils import supports_kw
30
31
from .interfaces_base import VllmModel, is_pooling_model
@@ -70,46 +71,8 @@ def _require_is_multimodal(is_multimodal: Tensor | None) -> Tensor:
71
return is_multimodal
72
73
73
-class LMMissingLayer(nn.Module):
74
- def make_empty_intermediate_tensors(self, *args, **kwargs):
75
- raise RuntimeError("This module should not be called in MM encoder-only mode")
76
-
77
- def __call__(self, *args, **kwargs):
78
- raise RuntimeError("This module should not be called in MM encoder-only mode")
79
-
80
-
81
-class TowerMissingLayer(nn.Module):
82
- def __init__(self, modalities: set[str] | str) -> None:
83
- if isinstance(modalities, str):
84
- modalities = {modalities}
85
-
86
- super().__init__()
87
-
88
- self.modalities = modalities
89
-
90
- def __call__(self, *args, **kwargs):
91
- raise RuntimeError(
92
- f"This module should not be called when the following "
93
- f"modalities are disabled: {self.modalities}"
94
- )
95
-
96
-
97
-@contextmanager
98
-def _no_init_weights(module: nn.Module, placeholder: Callable[[], nn.Module]):
99
- """
100
- Within this context, prevent weight initialization from using device memory and
101
- replace direct child assignments to `module` with the result of `placeholder()`.
102
- """
103
-
104
- def callback(module_, name, submodule):
105
- if module_ is module:
106
- return placeholder()
107
-
108
- return submodule
109
-
110
- with torch.nn.modules.module.register_module_module_registration_hook(callback): # noqa: E501,SIM117
111
- with torch.device("meta"):
112
- yield
74
+# Cache results of `SupportsMultiModal.get_language_model`
75
+_language_model_by_module = dict[nn.Module, VllmModel]()
76
77
78
@runtime_checkable
@@ -187,31 +150,61 @@ class SupportsMultiModal(Protocol):
150
Returns:
151
torch.nn.Module: The core language model component.
152
"""
153
+ # Cached
154
+ if self in _language_model_by_module:
155
+ return _language_model_by_module[self]
156
+
157
if self._language_model_names:
191
- return getattr(self, self._language_model_names[0])
158
+ mod = self
159
+ for attr in common_prefix(
160
+ [name.split(".") for name in self._language_model_names]
161
+ ):
162
+ if attr:
163
+ mod = getattr(mod, attr)
164
+
165
+ if mod is not self and hasattr(mod, "embed_input_ids"):
166
+ _language_model_by_module[self] = mod
167
+ return mod
168
+
169
+ # Fallback
170
+ for mod in self.children():
171
+ if hasattr(mod, "embed_input_ids"):
172
+ _language_model_by_module[self] = mod
173
+ return mod
174
175
raise NotImplementedError(
176
f"No language model found in {type(self).__name__}! "
195
- "You should initialize it inside `_mark_language_model`."
177
+ "You should initialize it via `_mark_language_model`."
178
)
179
180
@contextmanager
199
- def _mark_language_model(self, vllm_config: VllmConfig):
200
- """
201
- Mark each child module that was assigned to this model
202
- during this context as a language model component.
181
+ def _mark_language_model(
182
+ self,
183
+ vllm_config: VllmConfig,
184
+ *,
185
+ targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
186
+ ):
187
"""
204
- mm_config = vllm_config.model_config.multimodal_config
188
+ Mark each child module that was assigned to this model during this context
189
+ as a language model component.
190
+
191
+ Language model components are automatically skipped in `--mm-encoder-only`
192
+ mode.
193
206
- children_names = list[str]()
194
+ If `targets` is set, instead include descendants that are an instance
195
+ of `targets`, even if they aren't direct children.
196
+ """
197
+ from .utils import StageMissingLayer, collect_children, no_init_weights
198
208
- def callback(module_, name, submodule):
209
- if module_ is self:
210
- children_names.append(name)
199
+ mm_config = vllm_config.model_config.multimodal_config
200
212
- with torch.nn.modules.module.register_module_module_registration_hook(callback): # noqa: E501,SIM117
201
+ with collect_children(self, targets=targets) as children_names: # noqa: SIM117
202
with (
214
- _no_init_weights(self, LMMissingLayer)
203
+ no_init_weights(
204
+ self,
205
+ lambda mod: StageMissingLayer("language_model", mod),
206
+ targets=targets,
207
+ )
208
if mm_config.mm_encoder_only
209
else nullcontext()
210
):
@@ -220,25 +213,42 @@ class SupportsMultiModal(Protocol):
213
self._language_model_names = children_names
214
215
@contextmanager
223
- def _mark_tower_model(self, vllm_config: VllmConfig, modalities: set[str] | str):
216
+ def _mark_tower_model(
217
+ self,
218
+ vllm_config: VllmConfig,
219
+ modalities: set[str] | str,
220
+ *,
221
+ targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
222
+ ):
223
"""
225
- Mark each child module that was assigned to this model
226
- during this context as a tower model component.
224
+ Mark each child module that was assigned to this model during this context
225
+ as a tower model component.
226
+
227
+ Tower model components are automatically skipped when `--limit-mm-per-prompt`
228
+ is set to zero for all of their modalities.
229
+
230
+ If `targets` is set, instead include descendants that are an instance
231
+ of `targets`, even if they aren't direct children.
232
"""
233
+ from .utils import StageMissingLayer, collect_children, no_init_weights
234
+
235
if isinstance(modalities, str):
236
modalities = {modalities}
237
231
- mm_config = vllm_config.model_config.multimodal_config
232
-
233
- children_names = list[str]()
238
+ if modalities == {"image", "video"}:
239
+ stage_name = "vision_tower"
240
+ else:
241
+ stage_name = "_".join([*modalities, "tower"])
242
235
- def callback(module_, name, submodule):
236
- if module_ is self:
237
- children_names.append(name)
243
+ mm_config = vllm_config.model_config.multimodal_config
244
239
- with torch.nn.modules.module.register_module_module_registration_hook(callback): # noqa: E501,SIM117
245
+ with collect_children(self, targets=targets) as children_names: # noqa: SIM117
246
with (
241
- _no_init_weights(self, lambda: TowerMissingLayer(modalities))
247
+ no_init_weights(
248
+ self,
249
+ lambda mod: StageMissingLayer(stage_name, mod),
250
+ targets=targets,
251
+ )
252
if all(mm_config.get_limit_per_prompt(m) == 0 for m in modalities)
253
else nullcontext()
254
):
@@ -246,6 +256,37 @@ class SupportsMultiModal(Protocol):
256
257
self._tower_model_names = children_names
258
259
+ @contextmanager
260
+ def _mark_composite_model(
261
+ self,
262
+ vllm_config: VllmConfig,
263
+ *,
264
+ language_targets: type[nn.Module] | tuple[type[nn.Module], ...],
265
+ tower_targets: dict[str, type[nn.Module] | tuple[type[nn.Module], ...]],
266
+ ):
267
+ """
268
+ Composite wrapper over `_mark_language_model` and
269
+ `_mark_tower_model` by modality.
270
+ """
271
+ with ExitStack() as stack:
272
+ stack.enter_context(
273
+ self._mark_language_model(
274
+ vllm_config,
275
+ targets=language_targets,
276
+ )
277
+ )
278
+
279
+ for modality, modality_targets in tower_targets.items():
280
+ stack.enter_context(
281
+ self._mark_tower_model(
282
+ vllm_config,
283
+ modality,
284
+ targets=modality_targets,
285
+ )
286
+ )
287
+
288
+ yield
289
+
290
def get_num_mm_encoder_tokens(self, num_image_tokens: int) -> int:
291
"""
292
Implement this function to enable LoRA support
vllm/model_executor/models/internlm2.py
+12
-4
@@ -41,10 +41,12 @@ from vllm.sequence import IntermediateTensors
41
from .interfaces import SupportsLoRA, SupportsPP
42
from .interfaces_base import default_pooling_type
43
from .utils import (
44
+ StageMissingLayer,
45
is_pp_missing_parameter,
46
make_empty_intermediate_tensors_factory,
47
make_layers,
48
maybe_prefix,
49
+ no_init_weights,
50
)
51
52
@@ -413,10 +415,16 @@ class InternLM2ForRewardModel(InternLM2ForCausalLM):
415
prefix: str = "",
416
model_type: type[InternLM2Model] = InternLM2Model,
417
):
416
- super().__init__(vllm_config=vllm_config, prefix=prefix, model_type=model_type)
417
-
418
- for attr in ("output", "logits_processor"):
419
- delattr(self, attr)
418
+ with no_init_weights(
419
+ self,
420
+ lambda mod: StageMissingLayer("output", mod),
421
+ targets=(LogitsProcessor, ParallelLMHead),
422
+ ):
423
+ super().__init__(
424
+ vllm_config=vllm_config,
425
+ prefix=prefix,
426
+ model_type=model_type,
427
+ )
428
429
config = vllm_config.model_config.hf_config
430
self.head_dtype = vllm_config.model_config.head_dtype
vllm/model_executor/models/minicpmv.py
+4
-2
@@ -1035,11 +1035,13 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
1035
)
1036
1037
with self._mark_tower_model(vllm_config, {"image", "video"}):
1038
- self.vpm = vpm = self.init_vision_module(
1038
+ self.vpm = self.init_vision_module(
1039
config, quant_config, prefix=maybe_prefix(prefix, "vpm")
1040
)
1041
self.vision_dim = (
1042
- vpm.embed_dim if self.version == (2, 0) else vpm.embeddings.embed_dim
1042
+ self.vpm.embed_dim
1043
+ if self.version == (2, 0)
1044
+ else self.vpm.embeddings.embed_dim
1045
)
1046
self.embed_dim = self.config.hidden_size
1047
vllm/model_executor/models/mllama4.py
+2
-7
@@ -70,20 +70,15 @@ from vllm.sequence import IntermediateTensors
70
from vllm.utils.tensor_schema import TensorSchema, TensorShape
71
72
from .interfaces import (
73
- LMMissingLayer,
73
MixtureOfExperts,
74
MultiModalEmbeddings,
75
SupportsEagle3,
76
SupportsLoRA,
77
SupportsMultiModal,
78
SupportsPP,
80
- TowerMissingLayer,
79
)
80
from .llama4 import Llama4ForCausalLM
83
-from .utils import (
84
- AutoWeightsLoader,
85
- maybe_prefix,
86
-)
81
+from .utils import AutoWeightsLoader, StageMissingLayer, maybe_prefix
82
from .vision import run_dp_sharded_vision_model
83
84
@@ -1024,7 +1019,7 @@ class Llama4ForConditionalGeneration(
1019
renamed = self._rename_weight_for_modelopt_checkpoint(name)
1020
1021
attr = renamed.split(".", 1)[0]
1027
- if isinstance(getattr(self, attr), (LMMissingLayer, TowerMissingLayer)):
1022
+ if isinstance(getattr(self, attr), StageMissingLayer):
1023
continue
1024
1025
if renamed.startswith("language_model."):
vllm/model_executor/models/nano_nemotron_vl.py
+2
-2
@@ -1513,7 +1513,7 @@ class NemotronH_Nano_VL_V2(
1513
self.video_pruning_rate = multimodal_config.video_pruning_rate
1514
1515
with self._mark_language_model(vllm_config):
1516
- self.language_model = language_model = init_vllm_registered_model(
1516
+ self.language_model = init_vllm_registered_model(
1517
vllm_config=vllm_config,
1518
hf_config=config.text_config,
1519
prefix=maybe_prefix(prefix, "language_model"),
@@ -1542,7 +1542,7 @@ class NemotronH_Nano_VL_V2(
1542
ReLUSquaredActivation(),
1543
nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False),
1544
)
1545
- self.mlp1 = mlp1.to(language_model.config.dtype)
1545
+ self.mlp1 = mlp1.to(self.language_model.config.dtype)
1546
1547
self.config = config
1548
self.model_config = vllm_config.model_config
vllm/model_executor/models/paddleocr_vl.py
+2
-2
@@ -1025,12 +1025,12 @@ class PaddleOCRVLForConditionalGeneration(nn.Module, SupportsMultiModal, Support
1025
self.mlp_AR = Projector(config, config.vision_config)
1026
1027
with self._mark_language_model(vllm_config):
1028
- self.language_model = language_model = Ernie4_5ForCausalLM(
1028
+ self.language_model = Ernie4_5ForCausalLM(
1029
vllm_config=vllm_config,
1030
prefix=maybe_prefix(prefix, "language_model"),
1031
)
1032
1033
- for layer in language_model.model.layers:
1033
+ for layer in self.language_model.model.layers:
1034
if not isinstance(layer, PPMissingLayer):
1035
layer.self_attn.rotary_emb.is_neox_style = True
1036
vllm/model_executor/models/paligemma.py
+3
-2
@@ -314,13 +314,14 @@ class PaliGemmaForConditionalGeneration(
314
config.text_config.architectures = ["Gemma2ForCausalLM"]
315
316
with self._mark_language_model(vllm_config):
317
- self.language_model = language_model = init_vllm_registered_model(
317
+ self.language_model = init_vllm_registered_model(
318
vllm_config=vllm_config,
319
hf_config=config.text_config,
320
prefix=maybe_prefix(prefix, "language_model"),
321
)
322
+
323
logit_scale = getattr(config, "logit_scale", 1.0)
323
- language_model.logits_processor.scale *= logit_scale
324
+ self.language_model.logits_processor.scale *= logit_scale
325
326
self.make_empty_intermediate_tensors = (
327
self.language_model.make_empty_intermediate_tensors
vllm/model_executor/models/qwen3_vl_moe.py
+8
-7
@@ -461,15 +461,16 @@ class Qwen3VLMoeForConditionalGeneration(
461
]
462
463
with self._mark_language_model(vllm_config):
464
- self.language_model = language_model = Qwen3MoeLLMForCausalLM(
465
- vllm_config=vllm_config, prefix=maybe_prefix(prefix, "language_model")
464
+ self.language_model = Qwen3MoeLLMForCausalLM(
465
+ vllm_config=vllm_config,
466
+ prefix=maybe_prefix(prefix, "language_model"),
467
)
468
468
- # Whether to include the gate_up_proj mapping is determined by
469
- # the language model.
470
- self.packed_modules_mapping = (
471
- self.packed_modules_mapping | language_model.packed_modules_mapping
472
- )
469
+ # Whether to include the gate_up_proj mapping is determined by
470
+ # the language model.
471
+ self.packed_modules_mapping = (
472
+ self.packed_modules_mapping | self.language_model.packed_modules_mapping
473
+ )
474
475
self.make_empty_intermediate_tensors = (
476
self.language_model.make_empty_intermediate_tensors
vllm/model_executor/models/qwen_vl.py
+11
-9
@@ -58,7 +58,7 @@ from .interfaces import (
58
SupportsMultiModal,
59
SupportsPP,
60
)
61
-from .qwen import QWenBaseModel, QWenModel
61
+from .qwen import QWenBaseModel, QWenBlock, QWenModel
62
63
64
class QwenImagePixelInputs(TensorSchema):
@@ -757,11 +757,16 @@ class QwenVLForConditionalGeneration(
757
prefix: str = "",
758
transformer_type: type[QwenVLModel] = QwenVLModel,
759
) -> None:
760
- super().__init__(
761
- vllm_config=vllm_config,
762
- prefix=prefix,
763
- transformer_type=transformer_type,
764
- )
760
+ with self._mark_composite_model(
761
+ vllm_config,
762
+ language_targets=QWenBlock,
763
+ tower_targets={"image": VisionTransformer},
764
+ ):
765
+ super().__init__(
766
+ vllm_config=vllm_config,
767
+ prefix=prefix,
768
+ transformer_type=transformer_type,
769
+ )
770
771
self.transformer: QwenVLModel
772
@@ -795,9 +800,6 @@ class QwenVLForConditionalGeneration(
800
801
return self.transformer.visual(image_input["data"])
802
798
- def get_language_model(self) -> torch.nn.Module:
799
- return self.transformer
800
-
803
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
804
image_input = self._parse_and_validate_image_input(**kwargs)
805
if image_input is None:
vllm/model_executor/models/utils.py
+110
-13
@@ -2,13 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import itertools
5
-from collections.abc import Iterable, Mapping
5
+from collections.abc import Callable, Iterable, Mapping
6
+from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
from typing import Any, Literal, Protocol, overload
9
10
import torch
11
import torch.nn as nn
12
from torch.func import functional_call
13
+from torch.nn.modules.module import register_module_module_registration_hook
14
from transformers import PretrainedConfig
15
16
from vllm.config import VllmConfig
@@ -24,11 +26,7 @@ from vllm.model_executor.model_loader.online_quantization import (
26
support_quantized_model_reload_from_hp_weights,
27
)
28
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
27
-from vllm.model_executor.models.interfaces import (
28
- LMMissingLayer,
29
- TowerMissingLayer,
30
- supports_any_eagle,
31
-)
29
+from vllm.model_executor.models.interfaces import supports_any_eagle
30
from vllm.multimodal import NestedTensors
31
from vllm.sequence import IntermediateTensors
32
from vllm.utils.math_utils import cdiv
@@ -214,8 +212,8 @@ class AutoWeightsLoader:
212
continue
213
214
raise ValueError(
217
- f"Attempted to load nested weight '{weight_qualname}' "
218
- f"into a single parameter '{base_prefix}'"
215
+ f"Attempted to load nested weight {weight_qualname!r} "
216
+ f"into a single parameter {base_prefix!r}"
217
)
218
219
weight_loader = getattr(param, "weight_loader", default_weight_loader)
@@ -254,7 +252,7 @@ class AutoWeightsLoader:
252
module: nn.Module,
253
weights: Iterable[tuple[str, torch.Tensor]],
254
) -> Iterable[str]:
257
- if isinstance(module, (LMMissingLayer, TowerMissingLayer, PPMissingLayer)):
255
+ if isinstance(module, (StageMissingLayer, PPMissingLayer)):
256
return
257
258
# Avoid infinite recursion since this function is typically
@@ -316,9 +314,14 @@ class AutoWeightsLoader:
314
315
continue
316
317
+ desc_param_keys = {
318
+ base_prefix + k for k, _ in module.named_parameters(recurse=True)
319
+ }
320
msg = (
320
- f"There is no module or parameter named '{prefix}' "
321
- f"in {type(self.module).__name__}"
321
+ f"There is no module or parameter named {prefix!r} "
322
+ f"in {self.module._get_name()}. "
323
+ f"The available parameters belonging to {base_prefix} "
324
+ f"({module._get_name()}) are: {desc_param_keys}"
325
)
326
raise ValueError(msg)
327
@@ -496,6 +499,100 @@ def isin_list(
499
return torch.isin(elements, test_elements)
500
501
502
+class StageMissingLayer(nn.Module):
503
+ def __init__(self, stage_name: str, module: nn.Module | None = None) -> None:
504
+ super().__init__()
505
+
506
+ self.stage_name = stage_name
507
+
508
+ # Don't register this as a child module in order to
509
+ # avoid missing keys when loading weights
510
+ self.__dict__["module"] = module
511
+
512
+ def __getattr__(self, name: str):
513
+ return getattr(self.__dict__["module"], name)
514
+
515
+ def __call__(self, *args, **kwargs):
516
+ raise RuntimeError(f"{self} should not be called")
517
+
518
+ def extra_repr(self) -> str:
519
+ return f"stage_name={self.stage_name!r}"
520
+
521
+
522
+@contextmanager
523
+def collect_children(
524
+ module: nn.Module,
525
+ *,
526
+ targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
527
+):
528
+ """
529
+ Within this context, collect all direct child assignments to `module`,
530
+ returning a list of children names that is internally updated until the
531
+ context is exited.
532
+
533
+ If `targets` is set, instead collect descendents of `module`
534
+ that are an instance of `targets`, even if they aren't direct children.
535
+ """
536
+ children_names = list[str]()
537
+
538
+ if targets is None:
539
+
540
+ def hook(module_: nn.Module, name: str, submodule: nn.Module):
541
+ if module_ is module:
542
+ children_names.append(name)
543
+
544
+ with register_module_module_registration_hook(hook):
545
+ yield children_names
546
+ else:
547
+ yield children_names
548
+
549
+ for name, module_ in module.named_modules():
550
+ if isinstance(module_, targets):
551
+ children_names.append(name)
552
+
553
+
554
+@contextmanager
555
+def no_init_weights(
556
+ module: nn.Module,
557
+ placeholder: Callable[[nn.Module], nn.Module],
558
+ *,
559
+ targets: type[nn.Module] | tuple[type[nn.Module], ...] | None = None,
560
+):
561
+ """
562
+ Within this context, prevent weight initialization from using device memory and
563
+ replace direct child assignments to `module` with the result of `placeholder()`.
564
+
565
+ If `targets` is set, instead prevent weight initialization and
566
+ replace assignments where the child is an instance of `targets`,
567
+ even if they aren't direct children of `module`.
568
+ """
569
+ if targets is None:
570
+
571
+ def hook(module_: nn.Module, name: str, submodule: nn.Module):
572
+ if module_ is module:
573
+ return placeholder(submodule)
574
+
575
+ return submodule
576
+
577
+ with register_module_module_registration_hook(hook), torch.device("meta"):
578
+ yield
579
+ else:
580
+
581
+ def hook(module_: nn.Module, name: str, submodule: nn.Module):
582
+ if isinstance(module_, targets):
583
+ submodule.to("meta") # Free memory
584
+ if isinstance(submodule, targets):
585
+ submodule.to("meta") # Free memory
586
+ return placeholder(submodule)
587
+
588
+ return submodule
589
+
590
+ # Not all descendents are targeted, so we can't use a blanket
591
+ # `torch.device("meta")` context
592
+ with register_module_module_registration_hook(hook):
593
+ yield
594
+
595
+
596
class LayerFn(Protocol):
597
def __call__(self, prefix: str) -> torch.nn.Module: ...
598
@@ -627,7 +724,7 @@ def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:
724
725
missing_layer_names = []
726
for name, module in model.named_modules():
630
- if isinstance(module, PPMissingLayer):
727
+ if isinstance(module, (StageMissingLayer, PPMissingLayer)):
728
# NOTE: the trailing dot is used to match the prefix of the layer.
729
# without the dot, we could match a layer that is not missing,
730
# e.g., 'encoder.layer.1' would match 'encoder.layer.11'
@@ -639,7 +736,7 @@ def get_pp_missing_layer_names(model: torch.nn.Module) -> list[str]:
736
737
def is_pp_missing_parameter(name: str, model: torch.nn.Module) -> bool:
738
"""Check if a parameter is missing in a pipeline parallel model."""
642
- if isinstance(model, PPMissingLayer):
739
+ if isinstance(model, (StageMissingLayer, PPMissingLayer)):
740
return True
741
742
return any(
vllm/model_executor/models/whisper.py
+6
-4
@@ -909,7 +909,12 @@ class WhisperForConditionalGeneration(
909
self.config = config
910
self.dtype = vllm_config.model_config.dtype
911
912
- self.model = WhisperModel(vllm_config=vllm_config, prefix=prefix)
912
+ with self._mark_composite_model(
913
+ vllm_config,
914
+ language_targets=WhisperDecoder,
915
+ tower_targets={"audio": WhisperEncoder},
916
+ ):
917
+ self.model = WhisperModel(vllm_config=vllm_config, prefix=prefix)
918
919
self.proj_out = ParallelLMHead(
920
config.vocab_size,
@@ -937,9 +942,6 @@ class WhisperForConditionalGeneration(
942
)
943
return decoder_outputs
944
940
- def get_language_model(self) -> torch.nn.Module:
941
- return self.model.decoder
942
-
945
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
946
# Required as part of SupportsMultiModal interface.
947
audio_input = self._parse_and_validate_audio_input(**kwargs)
vllm/utils/collection_utils.py
+30
-2
@@ -7,10 +7,10 @@ This is similar in concept to the `collections` module.
7
"""
8
9
from collections import defaultdict
10
-from collections.abc import Callable, Generator, Hashable, Iterable, Mapping
10
+from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence
11
from typing import Generic, Literal, TypeVar
12
13
-from typing_extensions import TypeIs, assert_never
13
+from typing_extensions import TypeIs, assert_never, overload
14
15
T = TypeVar("T")
16
@@ -74,6 +74,34 @@ def is_list_of(
74
assert_never(check)
75
76
77
+@overload
78
+def common_prefix(items: Sequence[str]) -> str: ...
79
+
80
+
81
+@overload
82
+def common_prefix(items: Sequence[Sequence[T]]) -> Sequence[T]: ...
83
+
84
+
85
+def common_prefix(items: Sequence[Sequence[T] | str]) -> Sequence[T] | str:
86
+ """Find the longest prefix common to all items."""
87
+ if len(items) == 0:
88
+ return []
89
+ if len(items) == 1:
90
+ return items[0]
91
+
92
+ shortest = min(items, key=len)
93
+ if not shortest:
94
+ return shortest[:0]
95
+
96
+ for match_len in range(1, len(shortest) + 1):
97
+ match = shortest[:match_len]
98
+ for item in items:
99
+ if item[:match_len] != match:
100
+ return shortest[: match_len - 1]
101
+
102
+ return shortest
103
+
104
+
105
def chunk_list(lst: list[T], chunk_size: int) -> Generator[list[T]]:
106
"""Yield successive chunk_size chunks from lst."""
107
for i in range(0, len(lst), chunk_size):