wehub-resource-sync
147 lines
5.2 KiB
Python
147 lines
5.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Base classes for the Transformers backend fusers."""
|
|
|
|
import types
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, ClassVar
|
|
|
|
from torch import fx, nn
|
|
|
|
from vllm.model_executor.models.utils import ShardId, maybe_prefix
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config.model import ModelConfig
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
|
|
|
|
@dataclass
|
|
class BaseFuser(ABC):
|
|
"""A detected fusion and how to apply it.
|
|
|
|
`match` analyses the module *class* once (cached, see `get_fuser`); `fuse`
|
|
then applies the fusion to an instance in `recursive_replace`, returning the
|
|
module to install in its place.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def info(self, name: str) -> str:
|
|
"""A human-readable description of the fusion at `name`, for logging."""
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def match(cls, graph: fx.Graph, module: nn.Module) -> "BaseFuser | None":
|
|
"""Match the pattern in `graph`, returning a fuser if found."""
|
|
|
|
@abstractmethod
|
|
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
|
|
"""Whether this fuser can be applied to this `module` instance."""
|
|
|
|
@abstractmethod
|
|
def fuse(
|
|
self,
|
|
module: nn.Module,
|
|
prefix: str,
|
|
model_config: "ModelConfig",
|
|
quant_config: "QuantizationConfig",
|
|
) -> nn.Module:
|
|
"""Apply the fusion to an already-validated `module`, returning the
|
|
module to install in its place (mutated in place, or freshly built)."""
|
|
|
|
def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
|
|
"""`WeightsMapper.orig_to_new_stacked` entries this fuser contributes
|
|
(none unless it stacks weights)."""
|
|
return {}
|
|
|
|
@property
|
|
def packed_modules_mapping(self) -> dict[str, list[str]]:
|
|
"""`packed_modules_mapping` entries this fuser contributes (none unless
|
|
it stacks weights)."""
|
|
return {}
|
|
|
|
|
|
@dataclass
|
|
class StackedFuser(BaseFuser):
|
|
"""A fuser that merges sibling projections into one stacked linear and
|
|
rewrites the forward to call it.
|
|
|
|
`match` and `update_forward` analyse the class once; `fuse` builds the merged
|
|
submodule and binds the compiled forward on an instance in place, so it keeps
|
|
its class and any attribute the fusion does not consume.
|
|
"""
|
|
|
|
merged_name: ClassVar[str]
|
|
"""Attribute name of the merged module created by `update_attrs`."""
|
|
merged_cls: ClassVar[str]
|
|
"""Name of the vLLM class the merged projection becomes (for logging)."""
|
|
|
|
source_cls: str
|
|
"""Class of the HF module the fused projections belonged to (for logging)."""
|
|
|
|
fused_forward: Callable = field(init=False, repr=False)
|
|
"""The compiled rewritten forward, set by `update_forward`."""
|
|
|
|
def info(self, name: str) -> str:
|
|
sources = " + ".join(shard for shard, _ in self.shards)
|
|
return (
|
|
f"Fused: {sources} ({name}: {self.source_cls}) -> "
|
|
f"{self.merged_name} ({self.merged_cls})"
|
|
)
|
|
|
|
@property
|
|
@abstractmethod
|
|
def shards(self) -> list[tuple[str, ShardId]]:
|
|
"""Each projection's original name and its shard id in the merged module.
|
|
|
|
Source for both `orig_to_new_stacked` and `packed_modules_mapping`."""
|
|
|
|
def orig_to_new_stacked(self, prefix: str) -> dict[str, tuple[str, ShardId]]:
|
|
"""`WeightsMapper.orig_to_new_stacked` entries for one fused instance.
|
|
|
|
Maps each checkpoint name to `(merged_name, shard_id)`, keyed by qualname
|
|
so only this exact layer is remapped, never a same-named projection
|
|
elsewhere (e.g. an unfused MoE expert's `gate_proj`)."""
|
|
merged = maybe_prefix(prefix, self.merged_name)
|
|
return {
|
|
maybe_prefix(prefix, name): (merged, shard) for name, shard in self.shards
|
|
}
|
|
|
|
@property
|
|
def packed_modules_mapping(self) -> dict[str, list[str]]:
|
|
"""`{merged_name: [projection names]}` so quantization can unpack the
|
|
fused layer into its per-shard configs."""
|
|
return {self.merged_name: [name for name, _ in self.shards]}
|
|
|
|
@abstractmethod
|
|
def update_forward(self, module: nn.Module) -> None:
|
|
"""Rewrite and compile `type(module)`'s forward source.
|
|
|
|
Raises if the source does not admit the rewrite (fusion is then skipped).
|
|
"""
|
|
|
|
@abstractmethod
|
|
def update_attrs(
|
|
self,
|
|
module: nn.Module,
|
|
prefix: str,
|
|
model_config: "ModelConfig",
|
|
quant_config: "QuantizationConfig",
|
|
) -> None:
|
|
"""Replace `module`'s submodules with the merged module."""
|
|
|
|
def fuse(
|
|
self,
|
|
module: nn.Module,
|
|
prefix: str,
|
|
model_config: "ModelConfig",
|
|
quant_config: "QuantizationConfig",
|
|
) -> nn.Module:
|
|
"""Fuse an already-validated `module` in place (see `Fusers.__getitem__`).
|
|
|
|
Builds the merged submodule and binds the compiled forward."""
|
|
self.update_attrs(module, prefix, model_config, quant_config)
|
|
module.forward = types.MethodType(self.fused_forward, module)
|
|
return module
|