wehub-resource-sync
219 lines
8.2 KiB
Python
219 lines
8.2 KiB
Python
# 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)
|