| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Fuser for parallel linear projections.""" |
| 4 | |
| 5 | import ast |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING, ClassVar |
| 8 | |
| 9 | from torch import fx, nn |
| 10 | |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.layers.linear import MergedColumnParallelLinear |
| 13 | from vllm.model_executor.models.transformers.fusers.base import StackedFuser |
| 14 | from vllm.model_executor.models.transformers.fx_utils import ( |
| 15 | block_chain, |
| 16 | compile_forward, |
| 17 | is_linear, |
| 18 | recover_forward, |
| 19 | ) |
| 20 | from vllm.model_executor.models.utils import ShardId, maybe_prefix |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from vllm.config import VllmConfig |
| 24 | |
| 25 | logger = init_logger(__name__) |
| 26 | |
| 27 | |
| 28 | @dataclass |
| 29 | class MergedColumnParallelFuser(StackedFuser): |
| 30 | """Fuser for merging column-parallel linear projections.""" |
| 31 | |
| 32 | linear_names: tuple[str, ...] |
| 33 | merged_cls_name: ClassVar[str] = "MergedColumnParallelLinear" |
| 34 | |
| 35 | @property |
| 36 | def merged_name(self) -> str: |
| 37 | """Programmatic name for the merged projection, based on the original names.""" |
| 38 | # len is used to disambiguate names like `a_b` + `c` vs `a` + `b_c` |
| 39 | parts = "_".join(f"{len(name)}_{name}" for name in self.linear_names) |
| 40 | return f"merged_proj_{parts}" |
| 41 | |
| 42 | @property |
| 43 | def shards(self) -> list[tuple[str, ShardId]]: |
| 44 | return [(name, index) for index, name in enumerate(self.linear_names)] |
| 45 | |
| 46 | @staticmethod |
| 47 | def _names(group: list[fx.Node]) -> tuple[str, ...] | None: |
| 48 | names = tuple(str(node.target) for node in group) |
| 49 | if len(set(names)) != len(names) or any("." in name for name in names): |
| 50 | return None |
| 51 | return names |
| 52 | |
| 53 | @classmethod |
| 54 | def match( |
| 55 | cls, graph: fx.Graph, module: nn.Module |
| 56 | ) -> "MergedColumnParallelFuser | None": |
| 57 | """Fuse the module's sibling linears when there is only one such group.""" |
| 58 | by_input: dict[fx.Node, list[fx.Node]] = {} |
| 59 | for node in graph.nodes: |
| 60 | if ( |
| 61 | is_linear(node, module) |
| 62 | and len(node.args) == 1 |
| 63 | and not node.kwargs |
| 64 | and isinstance(node.args[0], fx.Node) |
| 65 | # Like QKVFuser/PackedQKVFuser: this fuser has no head or |
| 66 | # TP-replication awareness, so it must not absorb a QKV-shaped |
| 67 | # pattern (e.g. behind a norm) a more specific fuser declined. |
| 68 | and node.args[0].op == "placeholder" |
| 69 | ): |
| 70 | by_input.setdefault(node.args[0], []).append(node) |
| 71 | # A group whose members are not distinct direct children is dropped: |
| 72 | # the source rewrite addresses each projection as `self.<name>` once. |
| 73 | groups = [ |
| 74 | (group, names) |
| 75 | for group in by_input.values() |
| 76 | if len(group) >= 2 and (names := cls._names(group)) is not None |
| 77 | ] |
| 78 | if len(groups) > 1: |
| 79 | logger.debug( |
| 80 | "%s has %d fusable sibling-linear groups; skipping fusion " |
| 81 | "since which one to merge is ambiguous", |
| 82 | type(module).__name__, |
| 83 | len(groups), |
| 84 | ) |
| 85 | if len(groups) != 1: |
| 86 | return None |
| 87 | _, names = groups[0] |
| 88 | fuser = cls(source_cls=type(module).__name__, linear_names=names) |
| 89 | return None if hasattr(module, fuser.merged_name) else fuser |
| 90 | |
| 91 | def update_forward(self, module: nn.Module) -> None: |
| 92 | """Replace the parallel calls with one merged call and split.""" |
| 93 | funcdef, fn = recover_forward(type(module)) |
| 94 | calls = self._unguarded_calls(funcdef, self.linear_names) |
| 95 | if len(set(ast.dump(call.args[0]) for call in calls)) != 1: |
| 96 | raise ValueError("parallel linears read different inputs") |
| 97 | chains = [block_chain(funcdef.body, call) for call in calls] |
| 98 | if any(not chain for chain in chains): |
| 99 | raise ValueError("parallel linear calls not found in the function body") |
| 100 | blocks = [chain[-1] for chain in chains] |
| 101 | if len(set(id(block) for block, _ in blocks)) != 1: |
| 102 | raise ValueError("parallel linear calls are in different blocks") |
| 103 | |
| 104 | block = blocks[0][0] |
| 105 | indices = [index for _, index in blocks] |
| 106 | index = min(indices) |
| 107 | self._check_input_stable(funcdef, module, calls, block, indices) |
| 108 | |
| 109 | # l1(x), l2(x), ... -> merged(x).split(merged.output_sizes / merged.tp_size, -1) |
| 110 | self._splice_merged_split(funcdef, calls, block, index) |
| 111 | self.fused_forward = compile_forward(funcdef, fn) |
| 112 | |
| 113 | def validate(self, module: nn.Module, vllm_config: "VllmConfig") -> bool: |
| 114 | """Check that the parallel linears are compatible for merging.""" |
| 115 | linear_layers = [module.get_submodule(name) for name in self.linear_names] |
| 116 | tp_size = vllm_config.parallel_config.tensor_parallel_size |
| 117 | return ( |
| 118 | len(linear_layers) >= 2 |
| 119 | and len(set(linear.in_features for linear in linear_layers)) == 1 |
| 120 | and len(set(linear.bias is None for linear in linear_layers)) == 1 |
| 121 | and all(linear.out_features % tp_size == 0 for linear in linear_layers) |
| 122 | ) |
| 123 | |
| 124 | def update_attrs( |
| 125 | self, module: nn.Module, prefix: str, vllm_config: "VllmConfig" |
| 126 | ) -> None: |
| 127 | """Replace the module's parallel linears with one merged projection.""" |
| 128 | linear_modules = [module.get_submodule(name) for name in self.linear_names] |
| 129 | merged = MergedColumnParallelLinear( |
| 130 | input_size=linear_modules[0].in_features, |
| 131 | output_sizes=[linear.out_features for linear in linear_modules], |
| 132 | bias=linear_modules[0].bias is not None, |
| 133 | quant_config=vllm_config.quant_config, |
| 134 | prefix=maybe_prefix(prefix, self.merged_name), |
| 135 | return_bias=False, |
| 136 | ) |
| 137 | setattr(module, self.merged_name, merged) |
| 138 | for name in self.linear_names: |
| 139 | delattr(module, name) |
| 140 | logger.debug( |
| 141 | "%s -> %s: %s", |
| 142 | ", ".join(f"{n}: {m}" for n, m in zip(self.linear_names, linear_modules)), |
| 143 | self.merged_name, |
| 144 | merged, |
| 145 | ) |