| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GLU projection fuser: `act(gate(x)) * up(x)` -> a fused gate/up linear.""" |
| 4 | |
| 5 | import ast |
| 6 | import operator |
| 7 | from dataclasses import dataclass |
| 8 | from typing import TYPE_CHECKING, ClassVar |
| 9 | |
| 10 | from torch import fx, nn |
| 11 | from transformers.activations import ACT2CLS |
| 12 | |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.layers.activation import _ACTIVATION_AND_MUL_REGISTRY |
| 15 | from vllm.model_executor.models.transformers.fusers.merged_column import ( |
| 16 | MergedColumnParallelFuser, |
| 17 | ) |
| 18 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 19 | compile_forward, |
| 20 | find_node, |
| 21 | is_linear, |
| 22 | peel, |
| 23 | recover_forward, |
| 24 | replace_expr, |
| 25 | single_self_call, |
| 26 | ) |
| 27 | from vllm.model_executor.models.transformers.layers import get_act_and_mul_fn |
| 28 | from vllm.model_executor.models.transformers.utils import ( |
| 29 | log_replacement, |
| 30 | replace_linear_class, |
| 31 | ) |
| 32 | from vllm.model_executor.models.utils import maybe_prefix |
| 33 | |
| 34 | if TYPE_CHECKING: |
| 35 | from vllm.config import VllmConfig |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | |
| 40 | CLS2ACT: dict[type, list[str]] = {} |
| 41 | for _act_name, _act_cls in ACT2CLS.items(): |
| 42 | if isinstance(_act_cls, tuple): |
| 43 | _act_cls = _act_cls[0] |
| 44 | CLS2ACT.setdefault(_act_cls, []).append(_act_name) |
| 45 | |
| 46 | ACT_AND_MUL_NAMES = frozenset(_ACTIVATION_AND_MUL_REGISTRY.keys()) |
| 47 | |
| 48 | |
| 49 | @dataclass |
| 50 | class GLUFuser(MergedColumnParallelFuser): |
| 51 | """Fuser for the GLU pattern `act(gate(x)) * up(x)`.""" |
| 52 | |
| 53 | act_name: str |
| 54 | down_name: str | None |
| 55 | merged_name: ClassVar[str] = "gate_up_proj" |
| 56 | |
| 57 | @property |
| 58 | def gate_name(self) -> str: |
| 59 | return self.linear_names[0] |
| 60 | |
| 61 | @property |
| 62 | def up_name(self) -> str: |
| 63 | return self.linear_names[1] |
| 64 | |
| 65 | @classmethod |
| 66 | def _is_act_of_gate(cls, node: fx.Node, module: nn.Module) -> bool: |
| 67 | """Is node `act(gate(x))` where `gate` is linear and `act` is not linear.""" |
| 68 | return ( |
| 69 | node.op == "call_module" |
| 70 | and not is_linear(node, module) |
| 71 | and len(node.args) == 1 |
| 72 | and isinstance(node.args[0], fx.Node) |
| 73 | and is_linear(node.args[0], module) |
| 74 | ) |
| 75 | |
| 76 | @classmethod |
| 77 | def _get_glu_nodes( |
| 78 | cls, graph: fx.Graph, module: nn.Module |
| 79 | ) -> tuple[fx.Node, fx.Node, fx.Node, fx.Node] | None: |
| 80 | """Search graph for the GLU pattern `act(gate(x)) * up(x)`.""" |
| 81 | for mul in graph.nodes: |
| 82 | if ( |
| 83 | mul.op == "call_function" |
| 84 | and mul.target == operator.mul |
| 85 | and len(mul.args) == 2 |
| 86 | and all(isinstance(arg, fx.Node) for arg in mul.args) |
| 87 | ): |
| 88 | a, b = mul.args |
| 89 | if cls._is_act_of_gate(a, module) and is_linear(b, module): |
| 90 | act, gate, up = a, a.args[0], b |
| 91 | elif cls._is_act_of_gate(b, module) and is_linear(a, module): |
| 92 | act, gate, up = b, b.args[0], a |
| 93 | else: |
| 94 | continue |
| 95 | if ( |
| 96 | all(len(args) == 1 for args in (gate.args, up.args)) |
| 97 | and isinstance(x := gate.args[0], fx.Node) |
| 98 | and x is up.args[0] |
| 99 | ): |
| 100 | return act, gate, up, mul |
| 101 | return None |
| 102 | |
| 103 | @staticmethod |
| 104 | def _get_act_and_mul_name(act: nn.Module) -> str | None: |
| 105 | """Get the name of `act` if it has an `...AndMul` equivalent.""" |
| 106 | for name in CLS2ACT.get(type(act), []): |
| 107 | if name in ACT_AND_MUL_NAMES: |
| 108 | return name |
| 109 | # nn.GELU is not in ACT2CLS, but could be in model code |
| 110 | if type(act) is nn.GELU: |
| 111 | return "gelu_pytorch_tanh" if act.approximate == "tanh" else "gelu" |
| 112 | return None |
| 113 | |
| 114 | @classmethod |
| 115 | def _get_act_and_mul(cls, act: nn.Module) -> nn.Module: |
| 116 | """Get the `...AndMul` equivalent of a Transformers activation module.""" |
| 117 | if name := cls._get_act_and_mul_name(act): |
| 118 | return get_act_and_mul_fn(name) |
| 119 | raise ValueError(f"No AndMul equivalent for {type(act)}") |
| 120 | |
| 121 | @classmethod |
| 122 | def match(cls, graph: fx.Graph, module: nn.Module) -> "GLUFuser | None": |
| 123 | if (glu_nodes := cls._get_glu_nodes(graph, module)) is None: |
| 124 | return None |
| 125 | act_node, gate_node, up_node, mul_node = glu_nodes |
| 126 | |
| 127 | predicate = lambda n: is_linear(n, module) and peel(n.args[0]) is mul_node |
| 128 | down_node = find_node(graph, predicate) |
| 129 | return cls( |
| 130 | source_cls=type(module).__name__, |
| 131 | act_name=act_node.target, |
| 132 | linear_names=(gate_node.target, up_node.target), |
| 133 | down_name=down_node.target if down_node is not None else None, |
| 134 | ) |
| 135 | |
| 136 | def update_forward(self, module: nn.Module) -> None: |
| 137 | """Replace `act(gate(x)) * up(x)` with `act(gate_up(x))` in source.""" |
| 138 | funcdef, fn = recover_forward(type(module)) |
| 139 | act_call = single_self_call(funcdef, self.act_name) |
| 140 | calls = self._unguarded_calls(funcdef, (self.gate_name, self.up_name)) |
| 141 | gate_call, up_call = calls |
| 142 | if act_call.args[0] is not gate_call: |
| 143 | raise ValueError("activation does not directly wrap the gate") |
| 144 | if ast.dump(gate_call.args[0]) != ast.dump(up_call.args[0]): |
| 145 | raise ValueError("gate and up inputs are written differently") |
| 146 | muls = [ |
| 147 | node |
| 148 | for node in ast.walk(funcdef) |
| 149 | if isinstance(node, ast.BinOp) |
| 150 | and isinstance(node.op, ast.Mult) |
| 151 | and {id(node.left), id(node.right)} == {id(act_call), id(up_call)} |
| 152 | ] |
| 153 | if len(muls) != 1: |
| 154 | raise ValueError("no multiply of the activation and up projection") |
| 155 | |
| 156 | # act(gate(x)) * up(x) -> act(gate_up(x)) |
| 157 | assert isinstance(gate_call.func, ast.Attribute) |
| 158 | gate_call.func.attr = self.merged_name |
| 159 | replace_expr(funcdef, muls[0], act_call) |
| 160 | self.fused_forward = compile_forward(funcdef, fn) |
| 161 | |
| 162 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 163 | if not super().validate(module, vllm_config): |
| 164 | return False |
| 165 | # An AndMul kernel splits its input in half, so gate and up must match. |
| 166 | gate = module.get_submodule(self.gate_name) |
| 167 | up = module.get_submodule(self.up_name) |
| 168 | if gate.out_features != up.out_features: |
| 169 | logger.debug("%s and %s differ in size; skipping fusion", gate, up) |
| 170 | return False |
| 171 | act = module.get_submodule(self.act_name) |
| 172 | if self._get_act_and_mul_name(act) is None: |
| 173 | logger.debug("%s has no AndMul equivalent; skipping fusion", act) |
| 174 | return False |
| 175 | return True |
| 176 | |
| 177 | def update_attrs( |
| 178 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 179 | ) -> None: |
| 180 | super().update_attrs(module, prefix, vllm_config) |
| 181 | act_fn = self._get_act_and_mul(module.get_submodule(self.act_name)) |
| 182 | setattr(module, self.act_name, act_fn) |
| 183 | # If there is a down projection, we know it must be rowwise. |
| 184 | if self.down_name is not None: |
| 185 | down_prefix = maybe_prefix(prefix, self.down_name) |
| 186 | down = module.get_submodule(self.down_name) |
| 187 | new_down = replace_linear_class( |
| 188 | down, "rowwise", vllm_config.quant_config, prefix=down_prefix |
| 189 | ) |
| 190 | setattr(module, self.down_name, new_down) |
| 191 | log_replacement(down_prefix, down, new_down) |