From e960ead0aa70362a57f205be831083b68c7a5bdb Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Thu, 17 Sep 2026 11:11:31 -0500 Subject: [PATCH] [Mypy] Fix mypy typing for Transformers models (#54320) Signed-off-by: Taneem Ibrahim Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/models/transformers/fusers/test_mla.py | 6 +-- tools/pre_commit/mypy.py | 1 - vllm/compilation/decorators.py | 13 +++-- .../models/transformers/__init__.py | 24 ++++----- .../models/transformers/base.py | 44 +++++++++------ .../models/transformers/causal.py | 4 +- .../models/transformers/fusers/attention.py | 12 +++-- .../models/transformers/fusers/base.py | 7 ++- .../models/transformers/fusers/mla.py | 36 ++++++------- .../models/transformers/fusers/rms_norm.py | 12 ++--- .../models/transformers/fx_utils.py | 25 +++++---- .../models/transformers/legacy.py | 31 ++++++----- .../model_executor/models/transformers/moe.py | 42 +++++++++++---- .../models/transformers/multimodal.py | 54 ++++++++++++------- .../models/transformers/pooling.py | 37 ++++++++----- .../models/transformers/utils.py | 21 ++++---- 16 files changed, 224 insertions(+), 145 deletions(-) diff --git a/tests/models/transformers/fusers/test_mla.py b/tests/models/transformers/fusers/test_mla.py index 9f51ab5662..625fc4dc08 100644 --- a/tests/models/transformers/fusers/test_mla.py +++ b/tests/models/transformers/fusers/test_mla.py @@ -53,7 +53,6 @@ def _match(q_lora_rank: int | None) -> MLAFuser | None: def test_discovers_modules_without_q_lora(): fuser = _match(q_lora_rank=None) assert isinstance(fuser, MLAFuser) - assert not fuser.has_q_lora assert fuser.q_proj_name == "q_proj" assert fuser.kv_a_proj_name == "kv_a_proj_with_mqa" assert fuser.kv_a_layernorm_name == "kv_a_layernorm" @@ -68,7 +67,6 @@ def test_discovers_modules_without_q_lora(): def test_discovers_modules_with_q_lora(): fuser = _match(q_lora_rank=64) assert isinstance(fuser, MLAFuser) - assert fuser.has_q_lora assert fuser.q_a_proj_name == "q_a_proj" assert fuser.q_a_layernorm_name == "q_a_layernorm" assert fuser.q_b_proj_name == "q_b_proj" @@ -107,7 +105,7 @@ def test_update_forward_rewrites_real_attention(q_lora_rank): assert isinstance(fuser, MLAFuser) fuser.update_forward(_attention(q_lora_rank)) names = set(fuser.fused_forward.__code__.co_names) - if fuser.has_q_lora: + if fuser.q_a_proj_name is not None: assert _FUSED_QKV_A_PROJ in names assert not {"q_a_proj", "kv_a_proj_with_mqa"} & names @@ -158,7 +156,7 @@ def test_discovers_modules_under_arbitrary_names(): module = RenamedMLA() fuser = MLAFuser.match(trace(module), module) assert isinstance(fuser, MLAFuser) - assert not fuser.has_q_lora + assert fuser.q_a_proj_name is None assert fuser.q_proj_name == "alpha" assert fuser.kv_a_proj_name == "beta" assert fuser.kv_a_layernorm_name == "gamma" diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py index b7a278b9bd..234acc0298 100755 --- a/tools/pre_commit/mypy.py +++ b/tools/pre_commit/mypy.py @@ -116,7 +116,6 @@ EXCLUDE = [ r"vllm/model_executor/models/[nN]", r"vllm/model_executor/models/[oO]", r"vllm/model_executor/models/[qQ]", - r"vllm/model_executor/models/transformers", r"vllm/model_executor/models/[uU]", r"vllm/model_executor/models/[vV]", r"vllm/model_executor/models/[wW]", diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index 242e7227e9..35c0e6dee5 100644 --- a/vllm/compilation/decorators.py +++ b/vllm/compilation/decorators.py @@ -7,7 +7,7 @@ import inspect import os import sys from collections.abc import Callable, Generator -from typing import TYPE_CHECKING, Any, TypeVar, overload +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload from unittest.mock import patch import torch @@ -32,6 +32,9 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer from .monitor import monitor_profiling_run, monitor_torch_compile +DynamicArgDims: TypeAlias = dict[str, int | list[int] | dict[int, str]] +"""Argument name -> the dynamic dimension(s) of that argument.""" + # shape_id parameter was added to mark_unbacked in PyTorch 2.11.0 _SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0") @@ -89,7 +92,7 @@ def support_torch_compile( @overload def support_torch_compile( *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None, + dynamic_arg_dims: DynamicArgDims | None, ) -> Callable[[type[_T]], type[_T]]: ... @@ -103,7 +106,7 @@ def support_torch_compile( @overload def support_torch_compile( *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None, + dynamic_arg_dims: DynamicArgDims | None = None, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, @@ -117,7 +120,7 @@ def support_torch_compile(cls: type[_T]) -> type[_T]: ... def support_torch_compile( cls: type[_T] | None = None, *, - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]] | None = None, + dynamic_arg_dims: DynamicArgDims | None = None, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, @@ -328,7 +331,7 @@ def _try_load_aot_compiled_fn( def _support_torch_compile( cls: type[_T], - dynamic_arg_dims: dict[str, int | list[int] | dict[int, str]], + dynamic_arg_dims: DynamicArgDims, mark_unbacked_dims: dict[str, int | list[int]] | None = None, enable_if: Callable[[VllmConfig], bool] | None = None, is_encoder: bool = False, diff --git a/vllm/model_executor/models/transformers/__init__.py b/vllm/model_executor/models/transformers/__init__.py index 2d127c430e..e19de6d4e8 100644 --- a/vllm/model_executor/models/transformers/__init__.py +++ b/vllm/model_executor/models/transformers/__init__.py @@ -21,7 +21,7 @@ from typing import TYPE_CHECKING import torch.nn.functional as F from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR, Base +from vllm.model_executor.models.transformers.base import VLLM_ATTN_ATTR from vllm.model_executor.models.transformers.causal import CausalMixin from vllm.model_executor.models.transformers.legacy import LegacyMixin from vllm.model_executor.models.transformers.moe import MoEMixin @@ -126,10 +126,10 @@ ALL_ATTENTION_FUNCTIONS.register("vllm_mla", vllm_mla_attention_forward) # Text only models -class TransformersForCausalLM(CausalMixin, Base): ... +class TransformersForCausalLM(CausalMixin): ... -class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... +class TransformersMoEForCausalLM(MoEMixin, CausalMixin): ... # Multimodal models @@ -138,7 +138,7 @@ class TransformersMoEForCausalLM(MoEMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... +class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin): ... @MULTIMODAL_REGISTRY.register_processor( @@ -146,16 +146,14 @@ class TransformersMultiModalForCausalLM(MultiModalMixin, CausalMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalMoEForCausalLM( - MoEMixin, MultiModalMixin, CausalMixin, Base -): ... +class TransformersMultiModalMoEForCausalLM(MoEMixin, MultiModalMixin, CausalMixin): ... # Embedding models -class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin, Base): ... +class TransformersEmbeddingModel(EmbeddingMixin, LegacyMixin): ... -class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... +class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin): ... @MULTIMODAL_REGISTRY.register_processor( @@ -163,17 +161,17 @@ class TransformersMoEEmbeddingModel(EmbeddingMixin, MoEMixin, Base): ... info=MultiModalProcessingInfo, dummy_inputs=MultiModalDummyInputsBuilder, ) -class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin, Base): ... +class TransformersMultiModalEmbeddingModel(EmbeddingMixin, MultiModalMixin): ... # Sequence classification models class TransformersForSequenceClassification( - SequenceClassificationMixin, LegacyMixin, Base + SequenceClassificationMixin, LegacyMixin ): ... class TransformersMoEForSequenceClassification( - SequenceClassificationMixin, MoEMixin, Base + SequenceClassificationMixin, MoEMixin ): ... @@ -183,7 +181,7 @@ class TransformersMoEForSequenceClassification( dummy_inputs=MultiModalDummyInputsBuilder, ) class TransformersMultiModalForSequenceClassification( - SequenceClassificationMixin, MultiModalMixin, Base + SequenceClassificationMixin, MultiModalMixin ): ... diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 62f6d684e8..ff8fc98dbc 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/vllm/model_executor/models/transformers/base.py @@ -35,7 +35,7 @@ from transformers.conversion_mapping import ( get_model_conversion_mapping, ) -from vllm.compilation.decorators import support_torch_compile +from vllm.compilation.decorators import DynamicArgDims, support_torch_compile from vllm.config.utils import getattr_iter from vllm.distributed import get_pp_group, get_tp_group from vllm.distributed.utils import get_pp_indices @@ -76,6 +76,7 @@ from vllm.model_executor.models.transformers.utils import ( from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, + ShardId, WeightsMapper, make_empty_intermediate_tensors_factory, maybe_prefix, @@ -110,10 +111,13 @@ class Base( SupportsEagle, SupportsEagle3, ): - embedding_modules = ["embed_tokens"] # TODO transformers will have a util to get it + hf_to_vllm_mapper: WeightsMapper + + # TODO transformers will have a util to get these + embedding_modules = {"embed_tokens": "input_embeddings"} def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): - super().__init__() + nn.Module.__init__(self) logger.info("Using Transformers modeling backend.") self.vllm_config = vllm_config @@ -248,7 +252,7 @@ class Base( def _decorate_cls_for_torch_compile( self, cls: type["PreTrainedModel"], - dynamic_arg_dims: dict[str, int] | None, + dynamic_arg_dims: DynamicArgDims | None, enable_if: Callable[["VllmConfig"], bool], is_encoder: bool, ): @@ -284,7 +288,7 @@ class Base( self._decorate_cls_for_torch_compile( cls=self._pre_trained_model_classes.decoder, # Applied to a PreTrainedModel so the batch dimension will exist - dynamic_arg_dims=dict[str, int]( + dynamic_arg_dims=DynamicArgDims( input_ids=1, # shape: [1, seq_len] inputs_embeds=1, # shape: [1, seq_len, hidden_size] position_ids=-1, # shape: [1, seq_len] or [3, 1, seq_len] for mrope @@ -303,9 +307,8 @@ class Base( - Checkpoints saved with no base model prefix - Any quantization config specific mappings """ - self.hf_to_vllm_mapper = WeightsMapper() - orig_to_new_renaming = self.hf_to_vllm_mapper.orig_to_new_renaming - orig_to_new_regex = self.hf_to_vllm_mapper.orig_to_new_regex + orig_to_new_renaming: list[WeightRenaming] = [] + orig_to_new_regex: dict[re.Pattern, str | None] = {} for mapping in get_model_conversion_mapping(self.model): # Handle weights which have been renamed in Transformers @@ -341,6 +344,11 @@ class Base( nested_lm_head_pattern = re.compile(r"^model\.(.+\.)*(lm_head.+)") orig_to_new_regex[nested_lm_head_pattern] = r"\2" + self.hf_to_vllm_mapper = WeightsMapper( + orig_to_new_renaming=orig_to_new_renaming, + orig_to_new_regex=orig_to_new_regex, + ) + # Apply mapping to quantization config if needed self._maybe_apply_model_mapping() @@ -371,8 +379,8 @@ class Base( "modeling backend will infer the split from the layers of %s in order " "of declaration and keep parameter-free modules on every rank. This " "may fail if the model's structure is non-standard. %s", - type(self.model), - type(module), + type(self.model).__name__, + type(module).__name__, tip, ) @@ -470,7 +478,7 @@ class Base( "backend will shard the model the best it can during graph fusion and " "replicate the rest. This may be suboptimal or fail if the model does " "not fuse cleanly. %s", - type(self.model), + type(self.model).__name__, tip, ) @@ -481,6 +489,8 @@ class Base( vocab_embeddings = self._vocab_embeddings() + orig_to_new_stacked: dict[str, tuple[str, ShardId]] = {} + def register_fusion(fuser: BaseFuser, prefix: str, module: nn.Module): """Register a fused layer's mappings just before it is built.""" self.fusers.setdefault(prefix, []).append(fuser) @@ -498,8 +508,7 @@ class Base( ) self.attention_fusers[index] = (prefix, fuser) - orig_to_new_stacked = fuser.orig_to_new_stacked(prefix) - self.hf_to_vllm_mapper.orig_to_new_stacked.update(orig_to_new_stacked) + orig_to_new_stacked.update(fuser.orig_to_new_stacked(prefix)) packed_modules_mapping = fuser.packed_modules_mapping self.packed_modules_mapping.update(packed_modules_mapping) @@ -564,6 +573,8 @@ class Base( _recursive_replace(self.model, prefix="model") + self.hf_to_vllm_mapper |= WeightsMapper(orig_to_new_stacked=orig_to_new_stacked) + def _create_attention_instances(self): """Create `Attention` instances to inform KV cache allocation.""" text_config = self.text_config @@ -589,8 +600,11 @@ class Base( for i in range(start, end): if i not in self.attention_fusers: - in_range = layer_types and i < len(layer_types) - layer = f"{i} ({layer_types[i]})" if in_range else str(i) + layer = ( + f"{i} ({layer_types[i]})" + if layer_types and i < len(layer_types) + else str(i) + ) raise ValueError( f"Layer {layer} does not dispatch through the Transformers " "attention interface and vLLM has no other way to handle it." diff --git a/vllm/model_executor/models/transformers/causal.py b/vllm/model_executor/models/transformers/causal.py index c5231f9e97..e7dba82b0f 100644 --- a/vllm/model_executor/models/transformers/causal.py +++ b/vllm/model_executor/models/transformers/causal.py @@ -27,13 +27,15 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( from vllm.model_executor.models.interfaces_base import VllmModelForTextGeneration from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from .base import Base + if TYPE_CHECKING: import torch from vllm.config import VllmConfig -class CausalMixin(VllmModelForTextGeneration): +class CausalMixin(VllmModelForTextGeneration, Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip VllmModelForTextGeneration.__init__ and call the next class in MRO super(VllmModelForTextGeneration, self).__init__( diff --git a/vllm/model_executor/models/transformers/fusers/attention.py b/vllm/model_executor/models/transformers/fusers/attention.py index fc16527f15..d8d324536d 100644 --- a/vllm/model_executor/models/transformers/fusers/attention.py +++ b/vllm/model_executor/models/transformers/fusers/attention.py @@ -5,6 +5,7 @@ import ast import inspect import textwrap +from collections.abc import Callable from dataclasses import dataclass from functools import cache from typing import TYPE_CHECKING, ClassVar @@ -36,16 +37,16 @@ def _is_interface_lookup(node: ast.expr | None) -> bool: @cache -def interface_call(cls: type[nn.Module]) -> ast.Call | None: - """The attention interface call in `cls.forward`, if it makes exactly one.""" +def interface_call(forward: Callable) -> ast.Call | None: + """The attention interface call in `forward`, if it makes exactly one.""" try: - source = inspect.getsource(inspect.unwrap(cls.forward)) + source = inspect.getsource(inspect.unwrap(forward)) tree = ast.parse(textwrap.dedent(source)) except (AttributeError, OSError, SyntaxError, TypeError): return None # Find the names of the local variables that read from the interface lookup - names = set() + names: set[str] = set() for node in ast.walk(tree): if not isinstance(node, (ast.Assign, ast.AnnAssign)): continue @@ -97,7 +98,7 @@ class AttentionFuser(BaseFuser): def match( cls, graph: fx.Graph | None, module: nn.Module ) -> "AttentionFuser | None": - if (call := interface_call(type(module))) is None: + if (call := interface_call(type(module).forward)) is None: return None scaling = [kw.value for kw in call.keywords if kw.arg == "scaling"] scale_expr = scaling[0] if len(scaling) == 1 else None @@ -120,6 +121,7 @@ class AttentionFuser(BaseFuser): self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" ) -> nn.Module: if (sinks := self.sinks(module)) is not None: + assert isinstance(self.s_aux_expr, ast.Attribute) size = sinks.numel() // vllm_config.parallel_config.tensor_parallel_size device = vllm_config.device_config.device data = torch.empty(size, dtype=sinks.dtype, device=device) diff --git a/vllm/model_executor/models/transformers/fusers/base.py b/vllm/model_executor/models/transformers/fusers/base.py index ad24f2542a..937b70c5ed 100644 --- a/vllm/model_executor/models/transformers/fusers/base.py +++ b/vllm/model_executor/models/transformers/fusers/base.py @@ -130,11 +130,14 @@ class StackedFuser(RewriteFuser): """A fuser that merges sibling projections into one stacked linear and rewrites the forward to call it.""" - merged_name: ClassVar[str] - """Attribute name of the merged module created by `update_attrs`.""" merged_cls_name: ClassVar[str] """Name of the vLLM class the merged projection becomes (for logging).""" + @property + @abstractmethod + def merged_name(self) -> str: + """Attribute name of the merged module created by `update_attrs`.""" + def info(self, name: str) -> str: sources = " + ".join(shard for shard, _ in self.shards) return ( diff --git a/vllm/model_executor/models/transformers/fusers/mla.py b/vllm/model_executor/models/transformers/fusers/mla.py index fa4b34750c..96a4a0b87b 100644 --- a/vllm/model_executor/models/transformers/fusers/mla.py +++ b/vllm/model_executor/models/transformers/fusers/mla.py @@ -24,6 +24,7 @@ from vllm.model_executor.models.transformers.fx_utils import ( upstream_linear, ) from vllm.model_executor.models.transformers.utils import ( + Style, log_replacement, replace_linear_class, ) @@ -128,30 +129,26 @@ class MLAFuser(StackedFuser): merged_name: ClassVar[str] = "fused_qkv_a_proj" merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" - @property - def has_q_lora(self) -> bool: - return self.q_a_proj_name is not None - def info(self, name: str) -> str: info_str = ( f"Fused: {name} ({self.source_cls}) -> MLAAttention (attention interface)" ) - if self.has_q_lora: + if self.q_a_proj_name is not None: info_str += "; " + super().info(name).removeprefix("Fused: ") return info_str @property def shards(self) -> list[tuple[str, ShardId]]: """`q_a_proj` and `kv_a_proj_with_mqa` stack into one down-projection.""" - if self.has_q_lora: - return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)] - return [] + if self.q_a_proj_name is None: + return [] + return [(self.q_a_proj_name, 0), (self.kv_a_proj_name, 1)] @property def packed_modules_mapping(self) -> dict[str, list[str]]: - if self.has_q_lora: - return super().packed_modules_mapping - return {} + if self.q_a_proj_name is None: + return {} + return super().packed_modules_mapping @classmethod def match(cls, graph: fx.Graph, module: nn.Module) -> "MLAFuser | None": @@ -234,10 +231,10 @@ class MLAFuser(StackedFuser): Bypass the KV expansion method so the compressed latent reaches the `vllm_mla` attention interface unexpanded.""" funcdef, fn = recover_forward(type(module)) - if self.has_q_lora: + if (q_a_proj_name := self.q_a_proj_name) is not None: # q_a_proj is usually inside the `else` of `if self.q_lora_rank is None`. # The fused call is inserted at the top-level statement preceding both. - names = [self.q_a_proj_name, self.kv_a_proj_name] + names = [q_a_proj_name, self.kv_a_proj_name] calls = self._unguarded_calls(funcdef, names) if ast.dump(calls[0].args[0]) != ast.dump(calls[1].args[0]): raise ValueError("down-projections read different inputs") @@ -258,7 +255,7 @@ class MLAFuser(StackedFuser): def update_attrs(self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"): quant_config = vllm_config.quant_config - def replace_linear_by_name(name: str, style: str): + def replace_linear_by_name(name: str, style: Style): linear = module.get_submodule(name) replacement = replace_linear_class( linear, style, quant_config, prefix=maybe_prefix(prefix, name) @@ -266,8 +263,10 @@ class MLAFuser(StackedFuser): setattr(module, name, replacement) log_replacement(maybe_prefix(prefix, name), linear, replacement) - if self.has_q_lora: - q_a = module.get_submodule(self.q_a_proj_name) + if (q_a_proj_name := self.q_a_proj_name) is not None: + # `match` sets the q-LoRA names together, or none of them + assert self.q_b_proj_name is not None + q_a = module.get_submodule(q_a_proj_name) kv_a = module.get_submodule(self.kv_a_proj_name) merged = MergedColumnParallelLinear( input_size=q_a.in_features, @@ -280,7 +279,7 @@ class MLAFuser(StackedFuser): ) logger.debug( "%s: %s, %s: %s -> %s: %s", - self.q_a_proj_name, + q_a_proj_name, q_a, self.kv_a_proj_name, kv_a, @@ -289,10 +288,11 @@ class MLAFuser(StackedFuser): ) setattr(module, self.merged_name, merged) # The rewritten forward calls the merged projection instead. - delattr(module, self.q_a_proj_name) + delattr(module, q_a_proj_name) delattr(module, self.kv_a_proj_name) replace_linear_by_name(self.q_b_proj_name, "colwise") else: + assert self.q_proj_name is not None replace_linear_by_name(self.kv_a_proj_name, "replicate") replace_linear_by_name(self.q_proj_name, "colwise") diff --git a/vllm/model_executor/models/transformers/fusers/rms_norm.py b/vllm/model_executor/models/transformers/fusers/rms_norm.py index dd2c5339ce..f7972f03a6 100644 --- a/vllm/model_executor/models/transformers/fusers/rms_norm.py +++ b/vllm/model_executor/models/transformers/fusers/rms_norm.py @@ -216,7 +216,7 @@ class RMSNormFuser(BaseFuser): finally: for name, value in candidates.items(): setattr(module, name, value) - if (name := markers.get(marked)) is not None: + if marked is not None and (name := markers.get(marked)) is not None: return name, None logger.debug_once( "%s does not hold its eps (%s) in an attribute. Every instance in this " @@ -251,18 +251,18 @@ class RMSNormFuser(BaseFuser): ) -> nn.Module: """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp.""" weight = getattr(module, "weight", None) - has_weight = weight is not None - hidden_size = weight.size(0) if has_weight else 0 + hidden_size = weight.size(0) if weight is not None else 0 eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps if not isinstance(eps, (int, float)): # If eps was not detected, match torch behaviour. - dtype = weight.dtype if has_weight else vllm_config.model_config.dtype + config_dtype = vllm_config.model_config.dtype + dtype = weight.dtype if weight is not None else config_dtype eps = torch.finfo(dtype).eps if self.zero_centered: return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps) return TPAwareRMSNorm( hidden_size=hidden_size, eps=eps, - has_weight=has_weight, - dtype=weight.dtype if has_weight else None, + has_weight=weight is not None, + dtype=weight.dtype if weight is not None else None, ) diff --git a/vllm/model_executor/models/transformers/fx_utils.py b/vllm/model_executor/models/transformers/fx_utils.py index 3cdc4e1201..cb630866a3 100644 --- a/vllm/model_executor/models/transformers/fx_utils.py +++ b/vllm/model_executor/models/transformers/fx_utils.py @@ -15,17 +15,19 @@ import operator import textwrap from collections.abc import Callable from itertools import chain +from typing import Any from unittest import mock import torch from torch import fx, nn from torch.nn import functional as F +from typing_extensions import TypeIs from vllm.logger import init_logger logger = init_logger(__name__) -_UNKNOWN = object() +_UNKNOWN: Any = object() """Sentinel meta value for proxies whose concrete value could not be inferred. Distinct from `None`, which is a valid concrete value (e.g. `attn_weights`).""" @@ -34,7 +36,7 @@ _MODULE_CALL = nn.Module.__call__ `call_module` nodes; meta execution must call modules for real.""" -def is_leaf_call(node: object) -> bool: +def is_leaf_call(node: object) -> TypeIs[fx.Node]: """Is node a call recorded by `_as_leaf_call` (e.g. an attention interface).""" return isinstance(node, fx.Node) and node.meta.get("leaf_call", False) @@ -60,7 +62,7 @@ class _MetaProxy(fx.Proxy): executing each op on the meta values, so PyTorch's meta kernels are the single source of shape inference — no per-op rules.""" - meta: object = _UNKNOWN + meta: Any = _UNKNOWN def __len__(self) -> int: if self.meta is not _UNKNOWN: @@ -149,6 +151,7 @@ class _AllLeafTracer(fx.Tracer): if unknown: return _UNKNOWN if kind == "call_function": + assert callable(target) return target(*meta_args, **meta_kwargs) if kind == "call_method": receiver, *rest = meta_args @@ -576,7 +579,7 @@ def _aliasing_reads(node: ast.expr) -> set[str]: would taint nearly every later name. """ names: set[str] = set() - stack = [node] + stack: list[ast.AST] = [node] while stack: current = stack.pop() if isinstance(current, ast.Attribute) and current.attr in _METADATA_ATTRS: @@ -646,7 +649,7 @@ def aliasing_names( """ names = set(seed) assigns = [ - node + (node, node.value) for node in ast.walk(funcdef) if isinstance(node, (ast.Assign, ast.AnnAssign)) and node.value is not None @@ -655,10 +658,10 @@ def aliasing_names( changed = True while changed: changed = False - for node in assigns: - if not _aliasing_reads(node.value) & names: + for node, value in assigns: + if not _aliasing_reads(value) & names: continue - targets = getattr(node, "targets", None) or [node.target] + targets = node.targets if isinstance(node, ast.Assign) else [node.target] stores = { child.id for target in targets @@ -770,7 +773,7 @@ def peel(node: object) -> object: return node -def is_fn(node: object, target: Callable) -> bool: +def is_fn(node: object, target: object) -> TypeIs[fx.Node]: """Is node `()`.""" return ( isinstance(node, fx.Node) @@ -779,14 +782,14 @@ def is_fn(node: object, target: Callable) -> bool: ) -def is_method(node: object, name: str) -> bool: +def is_method(node: object, name: str) -> TypeIs[fx.Node]: """Is node `.()`.""" return ( isinstance(node, fx.Node) and node.op == "call_method" and node.target == name ) -def is_op(node: object, name: str) -> bool: +def is_op(node: object, name: str) -> TypeIs[fx.Node]: """Is node `.()` for torch, F, operator, or Tensor.""" return any( is_fn(node, getattr(module, name, None)) for module in (torch, F, operator) diff --git a/vllm/model_executor/models/transformers/legacy.py b/vllm/model_executor/models/transformers/legacy.py index 9c8a5c45b6..7a256f29d5 100644 --- a/vllm/model_executor/models/transformers/legacy.py +++ b/vllm/model_executor/models/transformers/legacy.py @@ -20,35 +20,38 @@ from typing import TYPE_CHECKING import torch +from vllm.model_executor.models.utils import WeightsMapper from vllm.sequence import IntermediateTensors +from .base import Base + if TYPE_CHECKING: from vllm.config import VllmConfig -class LegacyMixin: +class LegacyMixin(Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): super().__init__(vllm_config=vllm_config, prefix=prefix) - # Drop unsupported/unwanted output embeddings layers. - self.hf_to_vllm_mapper.orig_to_new_prefix.update( - { + self.hf_to_vllm_mapper |= WeightsMapper( + # Drop unsupported/unwanted output embeddings layers. + orig_to_new_prefix={ "model.lm_head.": None, "model.predictions.": None, "model.qa_outputs.": None, "model.embeddings_project.": None, "model.discriminator_predictions.": None, - } + }, + orig_to_new_substr={ + # Some encoder models have the position_ids buffer in the checkpoint. + # vLLM always passes position_ids as an argument, so drop the buffer. + "position_ids": None, + # Some encoder models have the bias of the final classifier layer in + # the checkpoint. vLLM does not use this bias, so drop it. + "score.bias": None, + }, ) - # Some encoder models have the position_ids buffer in the checkpoint. vLLM will - # always pass position_ids as an argument, so we drop the buffer if it exists. - self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None - - # Some encoder models have the bias of the final classifier layer in the - # checkpoint. vLLM does not use this bias, so we drop it if it exists. - self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None - # roberta-like models an extra padding in positions. # FIXME(Isotr0py): This is quite hacky for roberta edge case, # we should find a better way to handle this. @@ -61,6 +64,7 @@ class LegacyMixin: positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, + **kwargs, ) -> torch.Tensor | IntermediateTensors: if self.is_roberta: # RoBERTa positions start at padding_idx + 1. @@ -72,4 +76,5 @@ class LegacyMixin: positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, + **kwargs, ) diff --git a/vllm/model_executor/models/transformers/moe.py b/vllm/model_executor/models/transformers/moe.py index 480bc7f66e..9f63e73336 100644 --- a/vllm/model_executor/models/transformers/moe.py +++ b/vllm/model_executor/models/transformers/moe.py @@ -41,6 +41,7 @@ from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser from vllm.model_executor.models.utils import extract_layer_index, maybe_prefix from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, direct_register_custom_op +from .base import Base from .utils import log_replacement, maybe_per_layer if TYPE_CHECKING: @@ -66,14 +67,15 @@ class TransformersMoERunner(MoERunner): self._moe_state = moe_state self._moe_state.is_sequence_parallel = self.moe_config.is_sequence_parallel - def forward( + def forward( # type: ignore[override] self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, **kwargs: Any, ) -> torch.Tensor: - """In Transformers `experts.forward` will have this signature. + """In Transformers `experts.forward` will have this signature, which is why + it deliberately does not match `MoERunner.forward`. We discard any extra kwargs because we cannot use them here.""" # Note: we need to forward through a custom op so the topk_ids @@ -126,10 +128,18 @@ direct_register_custom_op( class TransformersRoutedExperts(RoutedExperts): def get_expert_mapping( - self, include_fused: bool = False + self, + ckpt_gate_proj_name: str | None = None, + ckpt_down_proj_name: str | None = None, + ckpt_up_proj_name: str | None = None, + include_fused: bool = False, ) -> list[tuple[str, str, int, str]]: - common_names = ("gate_proj", "down_proj", "up_proj") - common_map = super().get_expert_mapping(*common_names, include_fused) + common_map = super().get_expert_mapping( + ckpt_gate_proj_name or "gate_proj", + ckpt_down_proj_name or "down_proj", + ckpt_up_proj_name or "up_proj", + include_fused, + ) mixtral_map = super().get_expert_mapping("w1", "w2", "w3", include_fused) if not include_fused: return common_map + mixtral_map @@ -138,7 +148,13 @@ class TransformersRoutedExperts(RoutedExperts): return common_fused + mixtral_fused + common_unfused + mixtral_unfused -class MoEMixin(MixtureOfExperts): +class MoEMixin(MixtureOfExperts, Base): + mlp_layers: list[nn.Module] + """MoE blocks whose experts were replaced, for the `MixtureOfExperts` methods.""" + + # Redeclared because `recursive_replace` assigns it below the first read + num_local_physical_experts: int + def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): self.check_version("5.0.0", "MoE models support") # Skip MixtureOfExperts.__init__ and call the next class in MRO @@ -197,7 +213,8 @@ class MoEMixin(MixtureOfExperts): # Dtype the router computes in, if it is not the activation dtype. config_router_dtype = getattr(text_config, "moe_router_dtype", None) - config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) + if config_router_dtype is not None: + config_router_dtype = STR_DTYPE_TO_TORCH_DTYPE.get(config_router_dtype) # Grouped topk routing kwargs num_expert_group = getattr(text_config, "n_group", None) @@ -222,7 +239,7 @@ class MoEMixin(MixtureOfExperts): # MixtureOfExperts mixin settings ep_size = get_ep_group().world_size - self.mlp_layers = [] # Used for MixtureOfExperts methods + self.mlp_layers = [] self.moe_layers = [] self.num_expert_groups = 1 if num_expert_group is None else num_expert_group self.num_logical_experts = num_experts @@ -321,9 +338,10 @@ class MoEMixin(MixtureOfExperts): if shared_experts is not None: hf_shared = shared_experts.shared_experts glu_fuser = get_fuser(hf_shared, GLUFuser) - down_name = glu_fuser and glu_fuser.down_name - if down_name is not None: - shared_down_projs.append((hf_shared, down_name)) + if glu_fuser is not None: + down_name = glu_fuser.down_name + if down_name is not None: + shared_down_projs.append((hf_shared, down_name)) # Prefer config, otherwise read it from fuser. router_dtype = config_router_dtype or fuser.router_dtype gate = fuser.gate(moe_block, prefix, router_dtype) @@ -374,7 +392,9 @@ class MoEMixin(MixtureOfExperts): # Handle all gather in expert parallel if topk_ids.size(0) != hidden_states.size(0): dp_metadata = get_forward_context().dp_metadata + assert dp_metadata is not None sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + assert sizes is not None is_sp = moe_state.is_sequence_parallel group = get_ep_group() if is_sp else get_dp_group() assert sizes[group.rank_in_group] == topk_ids.shape[0] diff --git a/vllm/model_executor/models/transformers/multimodal.py b/vllm/model_executor/models/transformers/multimodal.py index 855b267c02..164a4dbd83 100644 --- a/vllm/model_executor/models/transformers/multimodal.py +++ b/vllm/model_executor/models/transformers/multimodal.py @@ -17,17 +17,24 @@ """Transformers modeling backend mixin for multi-modal models.""" from collections import defaultdict -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from contextlib import ExitStack, contextmanager from typing import TYPE_CHECKING, Any import torch import transformers from packaging.version import Version +from transformers.utils.generic import ModelOutput from vllm.compilation.decorators import should_torch_compile_mm_encoder +from vllm.config.multimodal import AudioDummyOptions from vllm.config.utils import getattr_iter -from vllm.inputs import MultiModalDataDict, MultiModalInput, mm_input +from vllm.inputs import ( + MultiModalDataBuiltins, + MultiModalDataDict, + MultiModalInput, + mm_input, +) from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -62,6 +69,8 @@ from vllm.sequence import IntermediateTensors from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import async_tensor_h2d +from .base import Base + if TYPE_CHECKING: from transformers import BatchFeature, PreTrainedModel @@ -148,7 +157,7 @@ class MultiModalProcessingInfo(BaseProcessingInfo): def get_max_image_tokens(self) -> int: width, height = self.get_image_size_with_most_features() processor = self.get_hf_processor() - multimodal_config = self.ctx.model_config.multimodal_config + multimodal_config = self.ctx.model_config.get_multimodal_config() mm_processor_kwargs = multimodal_config.mm_processor_kwargs or {} mm_tokens = processor._get_num_multimodal_tokens( image_sizes=([height, width],), **mm_processor_kwargs @@ -188,7 +197,7 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn mm_counts: Mapping[str, int], mm_options: Mapping[str, "BaseDummyOptions"], ) -> MultiModalDataDict: - data: MultiModalDataDict = {} + data = MultiModalDataBuiltins() if self.info._is_audio_model() and (num_audios := mm_counts.get("audio", 0)): sampling_rate = self.info._get_audio_sampling_rate() sub = self.info._get_audio_processor() @@ -196,10 +205,14 @@ class MultiModalDummyInputsBuilder(BaseDummyInputsBuilder[MultiModalProcessingIn if chunk_length is None: chunk_length = 30 audio_len = int(chunk_length * sampling_rate) + audio_overrides = mm_options.get("audio") + assert audio_overrides is None or isinstance( + audio_overrides, AudioDummyOptions + ) data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, - overrides=mm_options.get("audio"), + overrides=audio_overrides, ) if self.info._is_image_model() and (num_images := mm_counts.get("image", 0)): target_width, target_height = self.info.get_image_size_with_most_features() @@ -359,7 +372,9 @@ class _MultiModalProcessorBase(BaseMultiModalProcessor[MultiModalProcessingInfo] itself is to process it by itself. """ audios = mm_data.get("audio") - if not audios or len({len(audio) for audio in audios}) == 1: + if not isinstance(audios, Iterable) or not audios: + return + if len({len(audio) for audio in audios}) == 1: return alone = [ @@ -444,7 +459,8 @@ class LegacyMultiModalProcessor(_MultiModalProcessorBase): token_id = getattr(processor, names[0], getattr_iter(config, names)) if token_id is None: token = getattr(processor, f"{modality}_token", None) - token_id = info.get_tokenizer().get_vocab().get(token) + if token is not None: + token_id = info.get_tokenizer().get_vocab().get(token) if token_id is None: raise ValueError( f"Cannot find {modality}_token_id on processor or model config" @@ -722,13 +738,15 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase): for modality, items in out_mm_kwargs.items(): # Popped so they are neither cached nor sent to the model; the updates # they produce are cached alongside the item instead - replacements = [ - PromptUpdateDetails.select_token_id( - (ids := item.pop(f"{modality}_replacement_ids").data).tolist(), - _get_embed_token_id(ids), + replacements = [] + for item in items: + ids = item.pop(f"{modality}_replacement_ids").data + assert isinstance(ids, torch.Tensor) + replacements.append( + PromptUpdateDetails.select_token_id( + ids.tolist(), _get_embed_token_id(ids) + ) ) - for item in items - ] token = getattr(hf_processor, f"{modality}_token") updates.append( PromptReplacement( @@ -777,7 +795,7 @@ class OffsetsMultiModalProcessor(_MultiModalProcessorBase): ) -> list[int] | None: """Ask the HF processor how many rows of image data each image produces.""" images = mm_data.get("images") - if not images: + if not isinstance(images, Iterable) or not images: return None try: sizes = [(image.height, image.width) for image in images] @@ -893,7 +911,7 @@ MultiModalProcessor = ( ) -class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): +class MultiModalMixin(SupportsMultiModal, SupportsMRoPE, Base): def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): # Skip SupportsMRoPE.__init__ and call the next class in MRO super(SupportsMRoPE, self).__init__(vllm_config=vllm_config, prefix=prefix) @@ -925,7 +943,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): if model_config.skip_tokenizer_init: # Determining the supported modalities needs the HF processor, which in # turn needs a tokenizer - mm_config = model_config.multimodal_config + mm_config = model_config.get_multimodal_config() if mm_config.mm_encoder_only or any( mm_config.get_limit_per_prompt(modality) == 0 for modality in encoder_classes @@ -1029,7 +1047,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): # Exclude MultiModalMixin itself bases = [b for b in bases if b is not MultiModalMixin] - class LanguageModel(*bases): + class LanguageModel(*bases): # type: ignore[misc] def __init__(self, multimodal_model): # Don't call super().__init__() to avoid re-initialization self.__dict__.update(multimodal_model.__dict__) @@ -1203,7 +1221,7 @@ class MultiModalMixin(SupportsMultiModal, SupportsMRoPE): # from `self.get_image_features` if isinstance(features, tuple): return features[0] - if isinstance(features, dict): + if isinstance(features, ModelOutput): return features.pooler_output return features diff --git a/vllm/model_executor/models/transformers/pooling.py b/vllm/model_executor/models/transformers/pooling.py index 47a9c23d2c..79594387b7 100644 --- a/vllm/model_executor/models/transformers/pooling.py +++ b/vllm/model_executor/models/transformers/pooling.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING import torch +import torch.nn as nn from transformers import AutoModelForSequenceClassification from vllm.config.utils import getattr_iter @@ -26,11 +27,25 @@ from vllm.model_executor.layers.pooler import DispatchPooler from vllm.model_executor.models.interfaces import SupportsCrossEncoding from vllm.model_executor.models.interfaces_base import VllmModelForPooling +from .base import Base + if TYPE_CHECKING: from vllm.config import VllmConfig -class EmbeddingMixin(VllmModelForPooling): +class ClassifierWithReshape(nn.Module): + """Token extraction has already been applied in `pooler.pooling`. + + Add dim to match expected input shape of `classifier.forward`. + """ + + def forward(self, *args, **kwargs): + if len(args) > 0: + args = (args[0].unsqueeze(1), *args[1:]) + return super().forward(*args, **kwargs) + + +class EmbeddingMixin(VllmModelForPooling, Base): default_seq_pooling_type = "CLS" def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): @@ -45,7 +60,7 @@ class EmbeddingMixin(VllmModelForPooling): self.pooler = DispatchPooler.for_embedding(pooler_config) -class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling): +class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling, Base): default_seq_pooling_type = "CLS" def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""): @@ -83,17 +98,13 @@ class SequenceClassificationMixin(SupportsCrossEncoding, VllmModelForPooling): ) self.init_parameters(self.classifier, dtype=self.model_config.head_dtype) - class ClassifierWithReshape(self.classifier.__class__): - """Token extraction has already been applied in `pooler.pooling`. - Add dim to match expected input shape of `classifier.forward`. - """ - - def forward(self, *args, **kwargs): - if len(args) > 0: - args = (args[0].unsqueeze(1), *args[1:]) - return super().forward(*args, **kwargs) - - self.classifier.__class__ = ClassifierWithReshape + # Order `ClassifierWithReshape` ahead of the classifier's own class so that + # its `super().forward(...)` reaches the original implementation. + self.classifier.__class__ = type( + "ClassifierWithReshape", + (ClassifierWithReshape, type(self.classifier)), + {}, + ) self.pooler = DispatchPooler.for_seq_cls( pooler_config, diff --git a/vllm/model_executor/models/transformers/utils.py b/vllm/model_executor/models/transformers/utils.py index ddb2f4fb71..0443f390c4 100644 --- a/vllm/model_executor/models/transformers/utils.py +++ b/vllm/model_executor/models/transformers/utils.py @@ -18,11 +18,10 @@ from collections.abc import Callable, Iterator from contextlib import contextmanager -from functools import lru_cache from itertools import chain from operator import attrgetter from pathlib import Path -from typing import TYPE_CHECKING, Literal, TypeVar +from typing import TYPE_CHECKING, Literal, TypeAlias, TypeVar import torch from torch import nn @@ -70,7 +69,7 @@ def init_on_device_without_buffers(device: torch.device): module._parameters[name].to(device), **kwargs ) - tensor_constructors_to_patch = {} + tensor_constructors_to_patch: dict[str, Callable] = {} def patch_tensor_constructor(fn): def wrapper(*args, **kwargs): @@ -147,8 +146,8 @@ def replace_linear_class( ) -TorchConv = nn.Conv2d | nn.Conv3d -VllmConv = Conv2dLayer | Conv3dLayer +TorchConv: TypeAlias = nn.Conv2d | nn.Conv3d +VllmConv: TypeAlias = Conv2dLayer | Conv3dLayer def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv: @@ -212,7 +211,9 @@ class _VocabParallelEmbeddingBase(VocabParallelEmbedding, _UninitializedEmbeddin `super().forward(...)` in an `nn.Embedding` subclass reaches vLLM's embedding.""" -@lru_cache +_rebased_embedding_classes: dict[type[nn.Embedding], type[VocabParallelEmbedding]] = {} + + def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbedding]: """Subclass `cls` so that `VocabParallelEmbedding` supersedes its `nn.Embedding`. @@ -224,7 +225,10 @@ def _rebase_on_vocab_parallel(cls: type[nn.Embedding]) -> type[VocabParallelEmbe The new class, to assign to `__class__` of an instance of `cls`. """ - return type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) + if cls not in _rebased_embedding_classes: + rebased_cls = type(cls.__name__, (cls, _VocabParallelEmbeddingBase), {}) + _rebased_embedding_classes[cls] = rebased_cls + return _rebased_embedding_classes[cls] def replace_embedding_class( @@ -278,10 +282,9 @@ def recursive_replace_linear( qual_name = maybe_prefix(prefix, child_name) # Replace modules as needed if isinstance(child_module, nn.Linear): - style = "replicate" new_module = replace_linear_class( child_module, - style, + "replicate", quant_config, prefix=qual_name, )