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, 17:00 UTC
Fix pipeline parallel with multimodal models with the Transformers modelling backend (#37057)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Mar 16, 2026 at 10:20 UTC
122f75d9393883d64935706ad381beda85bc3112
2 files changed
+31
-8
vllm/model_executor/models/transformers/base.py
+30
-7
@@ -16,8 +16,9 @@
16
# limitations under the License.
17
"""Transformers modeling backend base class."""
18
19
-from collections.abc import Iterable
19
+from collections.abc import Callable, Iterable
20
from itertools import chain
21
+from operator import attrgetter
22
from typing import TYPE_CHECKING
23
24
import regex as re
@@ -296,6 +297,15 @@ class Base(
297
# Apply mapping to quantization config if needed
298
self._maybe_apply_model_mapping()
299
300
+ def _get_tie_word_embeddings(self):
301
+ """
302
+ Check if the model has tied word embeddings.
303
+ """
304
+ # Transformers v4 and v5 will store this in different places
305
+ tie_word_embeddings_v4 = getattr(self.text_config, "tie_word_embeddings", False)
306
+ tie_word_embeddings_v5 = getattr(self.config, "tie_word_embeddings", False)
307
+ return tie_word_embeddings_v4 or tie_word_embeddings_v5
308
+
309
def pipeline_parallel(self):
310
"""
311
Apply the model's pipeline parallelization plan.
@@ -311,11 +321,22 @@ class Base(
321
f"{type(self.model)} does not support pipeline parallel. {tip}"
322
)
323
324
+ def attrsetter(attr: str) -> Callable[[object, object], None]:
325
+ """Set a possibly nested attribute, like the inverse of attrgetter."""
326
+ parent, _, name = attr.rpartition(".")
327
+
328
+ def setter(obj: object, value: object):
329
+ attr_parent = attrgetter(parent)(obj) if parent else obj
330
+ setattr(attr_parent, name, value)
331
+
332
+ return setter
333
+
334
module_lists = []
335
module_list_idx = None
336
pp_plan = list(self.model._pp_plan.keys())
337
for i, name in enumerate(pp_plan):
318
- if isinstance(getattr(self.model, name), nn.ModuleList):
338
+ # attrgetter in case the module is nested (e.g. "text_model.layers")
339
+ if isinstance(attrgetter(name)(self.model), nn.ModuleList):
340
module_lists.append(name)
341
module_list_idx = i
342
@@ -330,11 +351,11 @@ class Base(
351
# Layers before module list
352
for name in pp_plan[:module_list_idx]:
353
if self.pp_group.is_first_rank or (
333
- getattr(self.text_config, "tie_word_embeddings", False)
334
- and self.pp_group.is_last_rank
354
+ self._get_tie_word_embeddings() and self.pp_group.is_last_rank
355
):
356
continue
337
- setattr(self.model, name, PPMissingLayer())
357
+ # attrsetter in case the module is nested (e.g. "text_model.embed_tokens")
358
+ attrsetter(name)(self.model, PPMissingLayer())
359
360
# Module list
361
start_layer, end_layer = get_pp_indices(
@@ -343,7 +364,8 @@ class Base(
364
self.pp_group.world_size,
365
)
366
layers_name = pp_plan[module_list_idx]
346
- layers = getattr(self.model, layers_name)
367
+ # attrgetter in case the module is nested (e.g. "text_model.layers")
368
+ layers = attrgetter(layers_name)(self.model)
369
for i in range(len(layers)):
370
if start_layer <= i and i < end_layer:
371
continue
@@ -353,7 +375,8 @@ class Base(
375
for name in pp_plan[module_list_idx + 1 :]:
376
# Modules that should be on last rank
377
if not self.pp_group.is_last_rank:
356
- setattr(self.model, name, PPMissingLayer())
378
+ # attrsetter in case the module is nested (e.g. "text_model.norm")
379
+ attrsetter(name)(self.model, PPMissingLayer())
380
381
def recursive_replace(self):
382
"""Recursively replace modules in the model as needed.
vllm/model_executor/models/transformers/causal.py
+1
-1
@@ -38,7 +38,7 @@ class CausalMixin(VllmModelForTextGeneration):
38
39
# Tell `Base.load_weights` to skip
40
# `lm_head` if the model has tied word embeddings
41
- tie_word_embeddings = getattr(self.text_config, "tie_word_embeddings", False)
41
+ tie_word_embeddings = self._get_tie_word_embeddings()
42
if tie_word_embeddings:
43
self.skip_prefixes.append("lm_head.")
44