iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Mypy] Fix mypy typing for Transformers models (#54320)

Browse Source 
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Taneem Ibrahim
2026-09-17 12:11:31 -04:00
committed by GitHub
co-authored by Harry Mellor
parent b78fbab4be
commit e960ead0aa
16 changed files with 224 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/models/transformers/fusers/test_mla.py
+2 -4
View File
@@ -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"
tools/pre_commit/mypy.py
-1
View File
@@ -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]",
vllm/compilation/decorators.py
+8 -5
View File
@@ -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,
vllm/model_executor/models/transformers/__init__.py
+11 -13
View File
@@ -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
): ...
vllm/model_executor/models/transformers/base.py
+29 -15
View File
@@ -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."
vllm/model_executor/models/transformers/causal.py
+3 -1
View File
@@ -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__(
vllm/model_executor/models/transformers/fusers/attention.py
+7 -5
View File
@@ -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)
vllm/model_executor/models/transformers/fusers/base.py
+5 -2
View File
@@ -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 (
vllm/model_executor/models/transformers/fusers/mla.py
+18 -18
View File
@@ -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")
vllm/model_executor/models/transformers/fusers/rms_norm.py
+6 -6
View File
@@ -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,
)
vllm/model_executor/models/transformers/fx_utils.py
+14 -11
View File
@@ -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 `<target>()`."""
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 `.<name>()`."""
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 `<mod>.<name>()` for torch, F, operator, or Tensor."""
return any(
is_fn(node, getattr(module, name, None)) for module in (torch, F, operator)
vllm/model_executor/models/transformers/legacy.py
+18 -13
View File
@@ -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,
)
vllm/model_executor/models/transformers/moe.py
+31 -11
View File
@@ -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]
vllm/model_executor/models/transformers/multimodal.py
+36 -18
View File
@@ -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
vllm/model_executor/models/transformers/pooling.py
+24 -13
View File
@@ -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,
vllm/model_executor/models/transformers/utils.py
+12 -9
View File
@@ -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,
)