chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,18 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Concrete fusers for the Transformers modeling backend."""
|
||||
|
||||
from vllm.model_executor.models.transformers.fusers.base import BaseFuser, StackedFuser
|
||||
from vllm.model_executor.models.transformers.fusers.glu import GLUFuser
|
||||
from vllm.model_executor.models.transformers.fusers.moe import MoEBlockFuser
|
||||
from vllm.model_executor.models.transformers.fusers.qkv import QKVFuser
|
||||
from vllm.model_executor.models.transformers.fusers.rms_norm import RMSNormFuser
|
||||
|
||||
__all__ = [
|
||||
"BaseFuser",
|
||||
"StackedFuser",
|
||||
"GLUFuser",
|
||||
"MoEBlockFuser",
|
||||
"QKVFuser",
|
||||
"RMSNormFuser",
|
||||
]
|
||||
@@ -0,0 +1,146 @@
|
||||
# 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
|
||||
@@ -0,0 +1,218 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear."""
|
||||
|
||||
import ast
|
||||
import operator
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
from torch import fx, nn
|
||||
from transformers.activations import ACT2CLS
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.activation import (
|
||||
_ACTIVATION_AND_MUL_REGISTRY,
|
||||
get_act_and_mul_fn,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import MergedColumnParallelLinear
|
||||
from vllm.model_executor.models.transformers.fusers.base import StackedFuser
|
||||
from vllm.model_executor.models.transformers.fx_utils import (
|
||||
compile_forward,
|
||||
find_node,
|
||||
is_linear,
|
||||
peel,
|
||||
recover_forward,
|
||||
replace_expr,
|
||||
single_self_call,
|
||||
)
|
||||
from vllm.model_executor.models.transformers.utils import (
|
||||
log_replacement,
|
||||
replace_linear_class,
|
||||
)
|
||||
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
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
CLS2ACT: dict[type, list[str]] = {}
|
||||
for _act_name, _act_cls in ACT2CLS.items():
|
||||
if isinstance(_act_cls, tuple):
|
||||
_act_cls = _act_cls[0]
|
||||
CLS2ACT.setdefault(_act_cls, []).append(_act_name)
|
||||
|
||||
ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys())
|
||||
|
||||
|
||||
@dataclass
|
||||
class GLUFuser(StackedFuser):
|
||||
"""Fuser for the GLU pattern `act(gate(x)) * up(x)`."""
|
||||
|
||||
act_name: str
|
||||
gate_name: str
|
||||
up_name: str
|
||||
down_name: str | None
|
||||
merged_name: ClassVar[str] = "gate_up_proj"
|
||||
merged_cls: ClassVar[str] = "MergedColumnParallelLinear"
|
||||
|
||||
@property
|
||||
def shards(self) -> list[tuple[str, ShardId]]:
|
||||
return [(self.gate_name, 0), (self.up_name, 1)]
|
||||
|
||||
@classmethod
|
||||
def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool:
|
||||
"""Is node `act(gate(x))` where `gate` is linear and `act` is not linear."""
|
||||
return (
|
||||
node.op == "call_module"
|
||||
and not is_linear(node, module)
|
||||
and len(node.args) == 1
|
||||
and isinstance(node.args[0], fx.Node)
|
||||
and is_linear(node.args[0], module)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _get_glu_nodes(
|
||||
cls, graph: fx.Graph, module: nn.Module
|
||||
) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None:
|
||||
"""Search graph for the GLU pattern `act(gate(x)) * up(x)`."""
|
||||
for mul in graph.nodes:
|
||||
if (
|
||||
mul.op == "call_function"
|
||||
and mul.target == operator.mul
|
||||
and len(mul.args) == 2
|
||||
and all(isinstance(arg, fx.Node) for arg in mul.args)
|
||||
):
|
||||
a, b = mul.args
|
||||
if cls._is_act_of_gate(a, module) and is_linear(b, module):
|
||||
act, gate, up = a, a.args[0], b
|
||||
elif cls._is_act_of_gate(b, module) and is_linear(a, module):
|
||||
act, gate, up = b, b.args[0], a
|
||||
else:
|
||||
continue
|
||||
if (
|
||||
all(len(args) == 1 for args in (gate.args, up.args))
|
||||
and isinstance(x := gate.args[0], fx.Node)
|
||||
and x is up.args[0]
|
||||
):
|
||||
return act, gate, up, mul
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_act_and_mul_name(act: nn.Module) -> str | None:
|
||||
"""Get the name of `act` if it has an `...AndMul` equivalent."""
|
||||
for name in CLS2ACT.get(type(act), []):
|
||||
if name in ACT_AND_MUL_NAMES:
|
||||
return name
|
||||
# nn.GELU is not in ACT2CLS, but could be in model code
|
||||
if type(act) is nn.GELU:
|
||||
return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu"
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def _get_act_and_mul(cls, act: nn.Module) -> nn.Module:
|
||||
"""Get the `...AndMul` equivalent of a Transformers activation module."""
|
||||
if name := cls._get_act_and_mul_name(act):
|
||||
return get_act_and_mul_fn(name)
|
||||
raise ValueError(f"No AndMul equivalent for {type(act)}")
|
||||
|
||||
@classmethod
|
||||
def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None":
|
||||
if (glu_nodes := cls._get_glu_nodes(graph, module)) is None:
|
||||
return None
|
||||
act_node, gate_node, up_node, mul_node = glu_nodes
|
||||
|
||||
gate = module.get_submodule(gate_node.target)
|
||||
up = module.get_submodule(up_node.target)
|
||||
# Shapes must be compatible for a single merged GEMM.
|
||||
if gate.in_features == up.in_features and (gate.bias is None) == (
|
||||
up.bias is None
|
||||
):
|
||||
predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node
|
||||
down_node = find_node(graph, predicate)
|
||||
return cls(
|
||||
source_cls=type(module).__name__,
|
||||
act_name=act_node.target,
|
||||
gate_name=gate_node.target,
|
||||
up_name=up_node.target,
|
||||
down_name=down_node.target if down_node is not None else None,
|
||||
)
|
||||
return None
|
||||
|
||||
def update_forward(self, module: nn.Module) -> None:
|
||||
"""Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source."""
|
||||
funcdef, fn = recover_forward(type(module))
|
||||
act_call = single_self_call(funcdef, self.act_name)
|
||||
gate_call = single_self_call(funcdef, self.gate_name)
|
||||
up_call = single_self_call(funcdef, self.up_name)
|
||||
if act_call.args[0] is not gate_call:
|
||||
raise ValueError("activation does not directly wrap the gate")
|
||||
if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]):
|
||||
raise ValueError("gate and up inputs are written differently")
|
||||
muls = [
|
||||
node
|
||||
for node in ast.walk(funcdef)
|
||||
if isinstance(node, ast.BinOp)
|
||||
and isinstance(node.op, ast.Mult)
|
||||
and {id(node.left), id(node.right)} == {id(act_call), id(up_call)}
|
||||
]
|
||||
if len(muls) != 1:
|
||||
raise ValueError("no multiply of the activation and up projection")
|
||||
|
||||
# act(gate(x)) * up(x) -> act(gate_up(x))
|
||||
assert isinstance(gate_call.func, ast.Attribute)
|
||||
gate_call.func.attr = self.merged_name
|
||||
replace_expr(funcdef, muls[0], act_call)
|
||||
self.fused_forward = compile_forward(funcdef, fn)
|
||||
|
||||
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
|
||||
act = module.get_submodule(self.act_name)
|
||||
if self._get_act_and_mul_name(act) is None:
|
||||
logger.debug("No AndMul equivalent for %s; skipping fusion", type(act))
|
||||
return False
|
||||
return True
|
||||
|
||||
def update_attrs(
|
||||
self,
|
||||
module: nn.Module,
|
||||
prefix: str,
|
||||
model_config: "ModelConfig",
|
||||
quant_config: "QuantizationConfig",
|
||||
) -> None:
|
||||
act_fn = self._get_act_and_mul(module.get_submodule(self.act_name))
|
||||
gate = module.get_submodule(self.gate_name)
|
||||
up = module.get_submodule(self.up_name)
|
||||
merged = MergedColumnParallelLinear(
|
||||
input_size=gate.in_features,
|
||||
output_sizes=[gate.out_features, up.out_features],
|
||||
bias=gate.bias is not None,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, self.merged_name),
|
||||
return_bias=False,
|
||||
)
|
||||
logger.debug(
|
||||
"%s: %s, %s: %s -> %s: %s",
|
||||
self.gate_name,
|
||||
gate,
|
||||
self.up_name,
|
||||
up,
|
||||
self.merged_name,
|
||||
merged,
|
||||
)
|
||||
setattr(module, self.merged_name, merged)
|
||||
setattr(module, self.act_name, act_fn)
|
||||
# Drop the consumed submodules so their (meta) params are not expected.
|
||||
delattr(module, self.gate_name)
|
||||
delattr(module, self.up_name)
|
||||
# If there is a down projection, we know it must be rowwise.
|
||||
if self.down_name is not None:
|
||||
down_prefix = maybe_prefix(prefix, self.down_name)
|
||||
down = module.get_submodule(self.down_name)
|
||||
new_down = replace_linear_class(
|
||||
down, "rowwise", quant_config, prefix=down_prefix
|
||||
)
|
||||
setattr(module, self.down_name, new_down)
|
||||
log_replacement(down_prefix, down, new_down)
|
||||
@@ -0,0 +1,268 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""MoE fuser: route an HF MoE block through `FusedMoE` with vLLM's own routing."""
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import textwrap
|
||||
import types
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
|
||||
import torch
|
||||
from torch import fx, nn
|
||||
|
||||
from vllm.distributed import tensor_model_parallel_all_gather
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.model_executor.models.transformers.fx_utils import (
|
||||
find_node,
|
||||
is_op,
|
||||
peel,
|
||||
trace,
|
||||
)
|
||||
from vllm.model_executor.models.utils import maybe_prefix, sequence_parallel_chunk
|
||||
|
||||
|
||||
def named_state(module: nn.Module) -> Iterator[tuple[str, torch.Tensor]]:
|
||||
"""`module`'s own state (i.e. named parameters and buffers)."""
|
||||
return chain(module.named_parameters(), module.named_buffers())
|
||||
|
||||
|
||||
def _own_returns(node: ast.AST) -> Iterator[ast.Return]:
|
||||
"""`return` statements in `node`'s own scope, not in nested functions."""
|
||||
stack = list(ast.iter_child_nodes(node))
|
||||
while stack:
|
||||
child = stack.pop()
|
||||
if isinstance(child, ast.Return):
|
||||
yield child
|
||||
elif not isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
|
||||
stack.extend(ast.iter_child_nodes(child))
|
||||
|
||||
|
||||
def _returns_tuple(cls: type[nn.Module]) -> bool:
|
||||
"""Does `cls.forward()` return a tuple?"""
|
||||
try:
|
||||
source = textwrap.dedent(inspect.getsource(inspect.unwrap(cls.forward)))
|
||||
forward = ast.parse(source).body[0]
|
||||
except (OSError, SyntaxError, TypeError, IndexError):
|
||||
return True
|
||||
# Names bound to a tuple literal, e.g. `out = hidden, logits` then `return out`.
|
||||
tuple_names = {
|
||||
target.id
|
||||
for node in ast.walk(forward)
|
||||
if isinstance(node, ast.Assign) and isinstance(node.value, ast.Tuple)
|
||||
for target in node.targets
|
||||
if isinstance(target, ast.Name)
|
||||
}
|
||||
|
||||
def yields_tuple(value: ast.expr | None) -> bool:
|
||||
if isinstance(value, ast.Tuple):
|
||||
return True
|
||||
if isinstance(value, ast.Name):
|
||||
return value.id in tuple_names
|
||||
if isinstance(value, ast.IfExp):
|
||||
return yields_tuple(value.body) or yields_tuple(value.orelse)
|
||||
return False
|
||||
|
||||
return any(yields_tuple(ret.value) for ret in _own_returns(forward))
|
||||
|
||||
|
||||
def _is_scalar_gate(module: nn.Module) -> bool:
|
||||
"""A linear projecting to a single logit (the shared-expert sigmoid gate)."""
|
||||
weight = getattr(module, "weight", None)
|
||||
return (
|
||||
isinstance(module, nn.Linear)
|
||||
and weight is not None
|
||||
and weight.ndim == 2
|
||||
and weight.shape[0] == 1
|
||||
)
|
||||
|
||||
|
||||
def _reaches(node: fx.Node, key: str) -> set[fx.Node]:
|
||||
"""Returns the set of nodes reachable from `node` by following `key` edges."""
|
||||
seen: set[fx.Node] = set()
|
||||
stack = [node]
|
||||
while stack:
|
||||
n = stack.pop()
|
||||
if n in seen:
|
||||
continue
|
||||
seen.add(n)
|
||||
stack.extend(getattr(n, key))
|
||||
return seen
|
||||
|
||||
|
||||
class SharedExpertMLP(nn.Module):
|
||||
"""Wraps an HF shared expert, applying the output gating it is paired with."""
|
||||
|
||||
def __init__(self, shared_experts: nn.Module, gate: nn.Module | None = None):
|
||||
super().__init__()
|
||||
self.shared_experts = shared_experts
|
||||
self.gate = gate
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
out = self.shared_experts(hidden_states)
|
||||
if self.gate is not None:
|
||||
out = torch.sigmoid(self.gate(hidden_states)[0]) * out
|
||||
return out
|
||||
|
||||
|
||||
def _moe_block_forward(self: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
"""Standard MoE block forward.
|
||||
|
||||
Routing and any shared experts are handled inside `self.experts: MoERunner`."""
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_states = hidden_states.reshape(-1, orig_shape[-1])
|
||||
num_tokens = hidden_states.shape[0]
|
||||
is_sequence_parallel = self.experts.moe_config.is_sequence_parallel
|
||||
if is_sequence_parallel:
|
||||
hidden_states = sequence_parallel_chunk(hidden_states)
|
||||
out = self.experts(hidden_states, router_logits=hidden_states)
|
||||
if is_sequence_parallel:
|
||||
out = tensor_model_parallel_all_gather(out, 0)[:num_tokens]
|
||||
return out.reshape(orig_shape)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoEBlockFuser:
|
||||
"""Fuser for MoE block `experts`, `gate` and `shared_experts` (optional)."""
|
||||
|
||||
gate_name: str
|
||||
scoring_func: str
|
||||
shared_name: str | None
|
||||
shared_gate_name: str | None
|
||||
|
||||
@staticmethod
|
||||
def _match_router(gate: nn.Module) -> str | None:
|
||||
"""Matches `topk(score(linear(x)))`, `score` being `softmax`/`sigmoid`."""
|
||||
if [name for name, _ in named_state(gate)] != ["weight"]:
|
||||
return None
|
||||
graph = trace(gate)
|
||||
if graph is None:
|
||||
return None
|
||||
topk = find_node(graph, lambda n: is_op(n, "topk"))
|
||||
if topk is None:
|
||||
return None
|
||||
# Exactly one scoring op upstream of the top-k, fed (transitively) by a linear.
|
||||
scorers = [
|
||||
n
|
||||
for n in _reaches(topk, "all_input_nodes")
|
||||
if is_op(n, "softmax") or is_op(n, "sigmoid")
|
||||
]
|
||||
if len(scorers) != 1:
|
||||
return None
|
||||
scorer = scorers[0]
|
||||
if not any(is_op(n, "linear") for n in _reaches(scorer, "all_input_nodes")):
|
||||
return None
|
||||
return "softmax" if is_op(scorer, "softmax") else "sigmoid"
|
||||
|
||||
@staticmethod
|
||||
def _match_shared_experts(
|
||||
graph: fx.Graph, experts: str
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""Detects the shared expert and its optional gate by dataflow."""
|
||||
experts_predicate = lambda n: n.op == "call_module" and n.target == experts
|
||||
if (experts_node := find_node(graph, experts_predicate)) is None:
|
||||
return None, None
|
||||
from_experts = _reaches(experts_node, "users")
|
||||
for add in graph.nodes:
|
||||
if not is_op(add, "add"):
|
||||
continue
|
||||
operands = [a for a in add.args if isinstance(a, fx.Node)]
|
||||
# Exactly one side is the experts' output; the other is the shared path.
|
||||
sides = [a in from_experts for a in operands]
|
||||
if len(operands) != 2 or sides.count(True) != 1:
|
||||
continue
|
||||
cone = _reaches(operands[sides.index(False)], "all_input_nodes")
|
||||
modules = [n for n in cone if n.op == "call_module" and n.target != experts]
|
||||
# A sigmoid wrapping one of those modules marks the shared-expert gate.
|
||||
gate = next(
|
||||
(
|
||||
src
|
||||
for n in cone
|
||||
if is_op(n, "sigmoid")
|
||||
and isinstance(src := peel(n.args[0]), fx.Node)
|
||||
and src in modules
|
||||
),
|
||||
None,
|
||||
)
|
||||
shared = [n for n in modules if n is not gate]
|
||||
if len(shared) != 1:
|
||||
return None, None
|
||||
return shared[0].target, (gate.target if gate is not None else None)
|
||||
return None, None
|
||||
|
||||
@classmethod
|
||||
def match(cls, moe_block: nn.Module, experts_name: str) -> "MoEBlockFuser | None":
|
||||
# Standard MoE block returns a single tensor.
|
||||
if _returns_tuple(type(moe_block)):
|
||||
return None
|
||||
# Router: the child that scores + top-k selects.
|
||||
gate_name = scoring_func = None
|
||||
for name, child in moe_block.named_children():
|
||||
if name != experts_name and (func := cls._match_router(child)) is not None:
|
||||
gate_name, scoring_func = name, func
|
||||
break
|
||||
if gate_name is None or scoring_func is None:
|
||||
return None
|
||||
# Shared expert: a child the block adds to the experts' output.
|
||||
shared_name = shared_gate_name = None
|
||||
others = [
|
||||
n
|
||||
for n, _ in moe_block.named_children()
|
||||
if n not in {experts_name, gate_name}
|
||||
]
|
||||
if others:
|
||||
graph = trace(moe_block)
|
||||
if graph is None:
|
||||
return None
|
||||
shared_name, shared_gate_name = cls._match_shared_experts(
|
||||
graph, experts_name
|
||||
)
|
||||
if shared_gate_name is not None and not _is_scalar_gate(
|
||||
getattr(moe_block, shared_gate_name)
|
||||
):
|
||||
return None
|
||||
# Fail closed: `rewrite_forward` runs only the experts and the detected
|
||||
# shared expert, so any other stateful child would be dropped.
|
||||
accounted = {experts_name, gate_name, shared_name, shared_gate_name}
|
||||
for name, child in moe_block.named_children():
|
||||
if name not in accounted and next(named_state(child), None) is not None:
|
||||
return None
|
||||
return cls(gate_name, scoring_func, shared_name, shared_gate_name)
|
||||
|
||||
def gate(self, moe_block: nn.Module, prefix: str) -> ReplicatedLinear:
|
||||
"""Rebuild the HF gate as a `ReplicatedLinear` for vLLM's fused MoE."""
|
||||
num_experts, hidden_size = getattr(moe_block, self.gate_name).weight.shape
|
||||
gate = ReplicatedLinear(
|
||||
hidden_size,
|
||||
num_experts,
|
||||
bias=False,
|
||||
prefix=maybe_prefix(prefix, self.gate_name),
|
||||
)
|
||||
setattr(moe_block, self.gate_name, gate)
|
||||
return gate
|
||||
|
||||
def shared_experts(
|
||||
self, moe_block: nn.Module, prefix: str
|
||||
) -> SharedExpertMLP | None:
|
||||
"""Build the HF shared expert (and its optional gate)
|
||||
as a `SharedExpertMLP` for vLLM's fused MoE."""
|
||||
if self.shared_name is None:
|
||||
return None
|
||||
shared_experts = getattr(moe_block, self.shared_name)
|
||||
gate = None
|
||||
if self.shared_gate_name is not None:
|
||||
hf_gate = getattr(moe_block, self.shared_gate_name)
|
||||
gate = ReplicatedLinear(
|
||||
hf_gate.in_features,
|
||||
hf_gate.out_features,
|
||||
bias=hf_gate.bias is not None,
|
||||
prefix=maybe_prefix(prefix, self.shared_gate_name),
|
||||
)
|
||||
setattr(moe_block, self.shared_gate_name, gate)
|
||||
return SharedExpertMLP(shared_experts, gate)
|
||||
|
||||
def rewrite_forward(self, moe_block: nn.Module) -> None:
|
||||
"""Rewrite `moe_block.forward` to route through vLLM's fused MoE."""
|
||||
moe_block.forward = types.MethodType(_moe_block_forward, moe_block)
|
||||
@@ -0,0 +1,212 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""QKV projection fuser: `q(x), k(x), v(x)` -> a fused qkv linear + split."""
|
||||
|
||||
import ast
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, ClassVar
|
||||
|
||||
from torch import fx, nn
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import QKVParallelLinear
|
||||
from vllm.model_executor.models.transformers.fusers.base import StackedFuser
|
||||
from vllm.model_executor.models.transformers.fx_utils import (
|
||||
compile_forward,
|
||||
innermost_block,
|
||||
is_linear,
|
||||
recover_forward,
|
||||
replace_expr,
|
||||
single_self_call,
|
||||
)
|
||||
from vllm.model_executor.models.transformers.utils import (
|
||||
log_replacement,
|
||||
replace_linear_class,
|
||||
)
|
||||
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
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class QKVFuser(StackedFuser):
|
||||
"""Fuser for the attention QKV pattern `q(x), k(x), v(x)`."""
|
||||
|
||||
q_name: str
|
||||
k_name: str
|
||||
v_name: str
|
||||
o_name: str | None
|
||||
merged_name: ClassVar[str] = "qkv_proj"
|
||||
merged_cls: ClassVar[str] = "QKVParallelLinear"
|
||||
|
||||
@property
|
||||
def shards(self) -> list[tuple[str, ShardId]]:
|
||||
return [(self.q_name, "q"), (self.k_name, "k"), (self.v_name, "v")]
|
||||
|
||||
@classmethod
|
||||
def _get_qkv_nodes(
|
||||
cls, graph: fx.Graph, module: nn.Module
|
||||
) -> tuple[fx.Node, fx.Node, fx.Node] | None:
|
||||
"""Search `graph` for the QKV pattern `q(x), k(x), v(x)`."""
|
||||
by_input: dict[fx.Node, list[fx.Node]] = {}
|
||||
for node in graph.nodes:
|
||||
if (
|
||||
is_linear(node, module)
|
||||
and len(node.args) == 1
|
||||
and not node.kwargs
|
||||
and isinstance(node.args[0], fx.Node)
|
||||
and node.args[0].op == "placeholder"
|
||||
):
|
||||
by_input.setdefault(node.args[0], []).append(node)
|
||||
triples = [nodes for nodes in by_input.values() if len(nodes) == 3]
|
||||
if len(triples) != 1:
|
||||
return None
|
||||
|
||||
q_node, k_node, v_node = nodes = triples[0]
|
||||
outs = [module.get_submodule(node.target).out_features for node in nodes]
|
||||
if len(set(outs)) == 2:
|
||||
# q is identified as the larger projection (GQA)
|
||||
(q_node,) = (n for n, out in zip(nodes, outs) if outs.count(out) == 1)
|
||||
k_node, v_node = (n for n, out in zip(nodes, outs) if outs.count(out) == 2)
|
||||
if module.get_submodule(q_node.target).out_features != max(outs):
|
||||
return None
|
||||
elif len(set(outs)) != 1:
|
||||
return None
|
||||
return q_node, k_node, v_node
|
||||
|
||||
@classmethod
|
||||
def match(cls, graph: fx.Graph, module: nn.Module) -> "QKVFuser | None":
|
||||
if (qkv_nodes := cls._get_qkv_nodes(graph, module)) is None:
|
||||
return None
|
||||
q, k, v = qkv_nodes
|
||||
names = dict(q_name=q.target, k_name=k.target, v_name=v.target)
|
||||
attn_width = module.get_submodule(q.target).out_features
|
||||
candidates = [
|
||||
name
|
||||
for name, child in module.named_children()
|
||||
if isinstance(child, nn.Linear)
|
||||
and name not in names.values()
|
||||
and child.in_features == attn_width
|
||||
]
|
||||
names["o_name"] = candidates[0] if len(candidates) == 1 else None
|
||||
return cls(source_cls=type(module).__name__, **names)
|
||||
|
||||
def update_forward(self, module: nn.Module) -> None:
|
||||
"""Replace `q(x), k(x), v(x)` with `qkv(x).split(sizes, -1)` in source."""
|
||||
funcdef, fn = recover_forward(type(module))
|
||||
calls = [
|
||||
single_self_call(funcdef, name)
|
||||
for name in (self.q_name, self.k_name, self.v_name)
|
||||
]
|
||||
arg_dumps = {ast.dump(call.args[0]) for call in calls}
|
||||
if len(arg_dumps) != 1:
|
||||
raise ValueError("projection inputs are written differently")
|
||||
# The trace may be partial, so prove projection exclusivity in source:
|
||||
# no other linear child may consume the same input (else the matched
|
||||
# three may not be q, k and v)
|
||||
other_linears = {
|
||||
name
|
||||
for name, child in module.named_children()
|
||||
if isinstance(child, nn.Linear)
|
||||
} - {self.q_name, self.k_name, self.v_name}
|
||||
for node in ast.walk(funcdef):
|
||||
if (
|
||||
isinstance(node, ast.Call)
|
||||
and isinstance(node.func, ast.Attribute)
|
||||
and node.func.attr in other_linears
|
||||
and any(ast.dump(arg) in arg_dumps for arg in node.args)
|
||||
):
|
||||
raise ValueError("another linear consumes the same input")
|
||||
blocks = [innermost_block(funcdef.body, call) for call in calls]
|
||||
if any(found is None for found in blocks):
|
||||
raise ValueError("projection calls not found in the function body")
|
||||
if len({id(block) for block, _ in blocks}) != 1:
|
||||
raise ValueError("projection calls are in different blocks")
|
||||
|
||||
# q(x), k(x), v(x) -> q, k, v = qkv(x).split(qkv.output_sizes / qkv.tp_size, -1)
|
||||
names = {node.id for node in ast.walk(funcdef) if isinstance(node, ast.Name)}
|
||||
temps = [f"{name}_fused" for name in (self.q_name, self.k_name, self.v_name)]
|
||||
if names & set(temps):
|
||||
raise ValueError("fused temporaries would shadow existing names")
|
||||
merged = f"self.{self.merged_name}"
|
||||
sections = f"[s // {merged}.tp_size for s in {merged}.output_sizes]"
|
||||
template = f"{', '.join(temps)} = {merged}(__arg__).split({sections}, -1)"
|
||||
assign = ast.parse(template).body[0]
|
||||
arg = next(
|
||||
node
|
||||
for node in ast.walk(assign)
|
||||
if isinstance(node, ast.Name) and node.id == "__arg__"
|
||||
)
|
||||
replace_expr(assign, arg, calls[0].args[0])
|
||||
block, index = blocks[0]
|
||||
ast.copy_location(assign, block[index])
|
||||
block.insert(min(index for _, index in blocks), assign)
|
||||
for call, temp in zip(calls, temps):
|
||||
replace_expr(funcdef, call, ast.Name(id=temp, ctx=ast.Load()))
|
||||
self.fused_forward = compile_forward(funcdef, fn)
|
||||
|
||||
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
|
||||
"""Shapes must be compatible for a single merged, head-sharded GEMM."""
|
||||
q = module.get_submodule(self.q_name)
|
||||
k = module.get_submodule(self.k_name)
|
||||
v = module.get_submodule(self.v_name)
|
||||
head_size = model_config.get_head_size()
|
||||
compatible = (
|
||||
q.in_features == k.in_features == v.in_features
|
||||
and len({proj.bias is None for proj in (q, k, v)}) == 1
|
||||
and k.out_features == v.out_features
|
||||
and q.out_features % head_size == 0
|
||||
and k.out_features % head_size == 0
|
||||
)
|
||||
if not compatible:
|
||||
logger.debug("%s is not compatible with QKV fusion", type(module))
|
||||
return compatible
|
||||
|
||||
def update_attrs(
|
||||
self,
|
||||
module: nn.Module,
|
||||
prefix: str,
|
||||
model_config: "ModelConfig",
|
||||
quant_config: "QuantizationConfig",
|
||||
) -> None:
|
||||
head_size = model_config.get_head_size()
|
||||
q = module.get_submodule(self.q_name)
|
||||
k = module.get_submodule(self.k_name)
|
||||
merged = QKVParallelLinear(
|
||||
hidden_size=q.in_features,
|
||||
head_size=head_size,
|
||||
total_num_heads=q.out_features // head_size,
|
||||
total_num_kv_heads=k.out_features // head_size,
|
||||
bias=q.bias is not None,
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, self.merged_name),
|
||||
return_bias=False,
|
||||
)
|
||||
logger.debug(
|
||||
"%s: %s, %s: %s, %s: %s -> %s: %s",
|
||||
self.q_name,
|
||||
q,
|
||||
self.k_name,
|
||||
k,
|
||||
self.v_name,
|
||||
module.get_submodule(self.v_name),
|
||||
self.merged_name,
|
||||
merged,
|
||||
)
|
||||
setattr(module, self.merged_name, merged)
|
||||
# Drop the consumed submodules so their (meta) params are not expected.
|
||||
for name in (self.q_name, self.k_name, self.v_name):
|
||||
delattr(module, name)
|
||||
# If there is an output projection, we know it must be rowwise.
|
||||
if self.o_name is not None:
|
||||
o_proj_prefix = maybe_prefix(prefix, self.o_name)
|
||||
o_proj = module.get_submodule(self.o_name)
|
||||
new_o = replace_linear_class(
|
||||
o_proj, "rowwise", quant_config, prefix=o_proj_prefix
|
||||
)
|
||||
setattr(module, self.o_name, new_o)
|
||||
log_replacement(o_proj_prefix, o_proj, new_o)
|
||||
@@ -0,0 +1,217 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""RMSNorm fuser: detect the norm structurally and swap in vLLM's fused RMSNorm."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import fx, nn
|
||||
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_gather,
|
||||
)
|
||||
from vllm.distributed.parallel_state import model_parallel_is_initialized
|
||||
from vllm.distributed.utils import split_tensor_along_last_dim
|
||||
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
|
||||
from vllm.model_executor.models.transformers.fusers.base import BaseFuser
|
||||
from vllm.model_executor.models.transformers.fx_utils import (
|
||||
find_node,
|
||||
forward_input_count,
|
||||
is_op,
|
||||
peel,
|
||||
trace,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.config.model import ModelConfig
|
||||
from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
|
||||
|
||||
def _is_squared(node: object, x: fx.Node) -> bool:
|
||||
"""`x**2`, `x.square()` or `x * x`, through any dtype casts."""
|
||||
node = peel(node)
|
||||
if is_op(node, "pow"):
|
||||
base, exp = node.args
|
||||
return peel(base) is x and exp == 2
|
||||
if is_op(node, "square"):
|
||||
return peel(node.args[0]) is x
|
||||
if is_op(node, "mul"):
|
||||
a, b = node.args
|
||||
return peel(a) is x and peel(b) is x
|
||||
return False
|
||||
|
||||
|
||||
def _variance_eps(rsqrt: fx.Node, x: fx.Node) -> float | None:
|
||||
"""eps from `rsqrt(mean(x**2, -1) + eps)`, or `None` if not that shape."""
|
||||
add = peel(rsqrt.args[0])
|
||||
if not is_op(add, "add"):
|
||||
return None
|
||||
consts = [a for a in add.args if isinstance(a, (int, float))]
|
||||
nodes = [a for a in add.args if isinstance(a, fx.Node)]
|
||||
if len(consts) != 1 or len(nodes) != 1:
|
||||
return None
|
||||
mean = peel(nodes[0])
|
||||
if not is_op(mean, "mean"):
|
||||
return None
|
||||
if not _is_squared(mean.args[0], x):
|
||||
return None
|
||||
return float(consts[0])
|
||||
|
||||
|
||||
def _is_one_plus(node: object) -> bool:
|
||||
"""`1 + weight` in either operand order (marks a zero-centered weight)."""
|
||||
node = peel(node)
|
||||
if not is_op(node, "add"):
|
||||
return False
|
||||
return any(isinstance(a, (int, float)) and a == 1 for a in node.args)
|
||||
|
||||
|
||||
def _has_trailing_compute(graph: fx.Graph, node: fx.Node) -> bool:
|
||||
"""Does the forward compute anything after `node` before returning?"""
|
||||
output = find_node(graph, lambda n: n.op == "output")
|
||||
if output is None or not output.args:
|
||||
return False
|
||||
return peel(output.args[0]) is not node
|
||||
|
||||
|
||||
class TPAwareNormMixin(nn.Module):
|
||||
"""Mixin for RMSNorms that reconstructs a TP-sharded input before normalizing."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
if model_parallel_is_initialized():
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.tp_rank = get_tensor_model_parallel_rank()
|
||||
else:
|
||||
self.tp_size, self.tp_rank = 1, 0
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor, residual: torch.Tensor | None = None
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
if self.tp_size > 1 and x.shape[-1] < (full := self.weight.shape[0]):
|
||||
if x.shape[-1] * self.tp_size != full:
|
||||
raise ValueError(
|
||||
f"Cannot gather norm of width {full}: a TP-sharded input of "
|
||||
f"width {x.shape[-1]} does not tile it evenly across "
|
||||
f"{self.tp_size} ranks (replicated or uneven sharding)."
|
||||
)
|
||||
x = tensor_model_parallel_all_gather(x.contiguous())
|
||||
x = super().forward(x)
|
||||
splits = split_tensor_along_last_dim(x, num_partitions=self.tp_size)
|
||||
return splits[self.tp_rank]
|
||||
return super().forward(x, residual)
|
||||
|
||||
|
||||
class TPAwareRMSNorm(TPAwareNormMixin, RMSNorm):
|
||||
"""`RMSNorm` that reconstructs a TP-sharded input before normalizing."""
|
||||
|
||||
|
||||
class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm):
|
||||
"""`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class RMSNormFuser(BaseFuser):
|
||||
"""Fuser for RMSNorm patterns, including Gemma-style zero-centered weights."""
|
||||
|
||||
zero_centered: bool
|
||||
"""Gemma-style `(1 + weight)` scaling (weight initialised at zero)."""
|
||||
source_cls: str
|
||||
"""Class name of the norm this was matched from (for logging)."""
|
||||
|
||||
def info(self, name: str) -> str:
|
||||
norm = "GemmaRMSNorm" if self.zero_centered else "RMSNorm"
|
||||
return f"Fused: {name} ({self.source_cls}) -> {norm} (CustomOp)"
|
||||
|
||||
@classmethod
|
||||
def match(cls, graph: fx.Graph, module: nn.Module) -> "RMSNormFuser | None":
|
||||
"""Match a graph to the RMSNorm pattern, returning a fuser if found."""
|
||||
if forward_input_count(type(module)) != 1:
|
||||
return None
|
||||
x = find_node(graph, lambda n: n.op == "placeholder")
|
||||
if x is None:
|
||||
return None
|
||||
# Handle native torch `rms_norm` op.
|
||||
rms_norm = find_node(graph, lambda n: is_op(n, "rms_norm"))
|
||||
if rms_norm is not None and rms_norm.args and peel(rms_norm.args[0]) is x:
|
||||
if _has_trailing_compute(graph, rms_norm):
|
||||
return None
|
||||
return cls(zero_centered=False, source_cls=type(module).__name__)
|
||||
# Handle explicit `x * rsqrt(mean(x**2, -1) + eps)` pattern.
|
||||
# The rsqrt over the mean-square variance is the spine of the norm.
|
||||
rsqrt = None
|
||||
for node in graph.nodes:
|
||||
if is_op(node, "rsqrt") and _variance_eps(node, x) is not None:
|
||||
rsqrt = node
|
||||
break
|
||||
if rsqrt is None:
|
||||
return None
|
||||
# The `x * rsqrt(...)` normalize multiply.
|
||||
normalize = find_node(
|
||||
graph, lambda n: is_op(n, "mul") and rsqrt in map(peel, n.args)
|
||||
)
|
||||
if normalize is None:
|
||||
return None
|
||||
# An optional later `weight * normalized` (or `(1 + weight) * normalized`).
|
||||
tail, zero_centered = normalize, False
|
||||
for node in graph.nodes:
|
||||
if not is_op(node, "mul") or node is normalize:
|
||||
continue
|
||||
operands = [peel(a) for a in node.args if isinstance(a, fx.Node)]
|
||||
if len(operands) == 2 and normalize in operands:
|
||||
weight = next(o for o in operands if o is not normalize)
|
||||
tail, zero_centered = node, _is_one_plus(weight)
|
||||
break
|
||||
# The norm must be the last compute in forward, or it is not a pure norm.
|
||||
if _has_trailing_compute(graph, tail):
|
||||
return None
|
||||
return cls(zero_centered=zero_centered, source_cls=type(module).__name__)
|
||||
|
||||
@staticmethod
|
||||
def _eps_from_graph(graph: fx.Graph) -> float | None:
|
||||
"""Extract the `eps` constant from the graph, if present."""
|
||||
if (x := find_node(graph, lambda n: n.op == "placeholder")) is None:
|
||||
return None
|
||||
fused = find_node(graph, lambda n: is_op(n, "rms_norm"))
|
||||
if fused is not None and fused.args and peel(fused.args[0]) is x:
|
||||
args, kwargs = fused.args, fused.kwargs
|
||||
eps = args[3] if len(args) > 3 else kwargs.get("eps")
|
||||
return eps if isinstance(eps, (int, float)) else None
|
||||
for node in graph.nodes:
|
||||
if is_op(node, "rsqrt") and (eps := _variance_eps(node, x)) is not None:
|
||||
return eps
|
||||
return None
|
||||
|
||||
def validate(self, module: nn.Module, model_config: "ModelConfig") -> bool:
|
||||
return True
|
||||
|
||||
def fuse(
|
||||
self,
|
||||
module: nn.Module,
|
||||
prefix: str,
|
||||
model_config: "ModelConfig",
|
||||
quant_config: "QuantizationConfig",
|
||||
) -> nn.Module:
|
||||
"""Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
|
||||
weight = getattr(module, "weight", None)
|
||||
hidden_size = (
|
||||
weight.size(0) if weight is not None else model_config.get_hidden_size()
|
||||
)
|
||||
graph = trace(module)
|
||||
eps = self._eps_from_graph(graph) if graph is not None else None
|
||||
if eps is None:
|
||||
# If eps not in graph, match torch behaviour.
|
||||
dtype = weight.dtype if weight is not None else model_config.dtype
|
||||
eps = torch.finfo(dtype).eps
|
||||
if self.zero_centered:
|
||||
return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
|
||||
has_weight = weight is not None
|
||||
return TPAwareRMSNorm(
|
||||
hidden_size=hidden_size,
|
||||
eps=eps,
|
||||
has_weight=has_weight,
|
||||
dtype=weight.dtype if has_weight else None,
|
||||
)
|
||||
Reference in New Issue
Block a user