vllm.model_executor.models.transformers.fusers.rms_norm ¶

@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)."""
    eps_attr: str | None = None
    """Attribute holding `eps`, read per instance in `fuse`."""
    eps: float | None = None
    """`eps` itself, when it is not held in an attribute (see `_eps_source`)."""

    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
            eps_attr, eps = cls._eps_source(graph, module)
            return cls(
                zero_centered=False,
                source_cls=type(module).__name__,
                eps_attr=eps_attr,
                eps=eps,
            )
        # 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
        eps_attr, eps = cls._eps_source(graph, module)
        return cls(
            zero_centered=zero_centered,
            source_cls=type(module).__name__,
            eps_attr=eps_attr,
            eps=eps,
        )

    @classmethod
    def _eps_source(
        cls, graph: fx.Graph, module: nn.Module
    ) -> tuple[str | None, float | None]:
        """Where `fuse` should read `eps` from, resolved once per class."""
        eps = cls._eps_from_graph(graph)
        if eps is None:
            return None, None
        # Whatever supplied the constant must still equal it.
        candidates = {
            name: value
            for name, value in vars(module).items()
            if isinstance(value, float) and value == eps
        }
        # Use unique markers and retrace to verify exactly which attribute is eps.
        markers = {float(-index - 1): name for index, name in enumerate(candidates)}
        marked = None
        if markers:
            try:
                for marker, name in markers.items():
                    setattr(module, name, marker)
                if (remarked := trace(module)) is not None:
                    marked = cls._eps_from_graph(remarked)
            finally:
                for name, value in candidates.items():
                    setattr(module, name, value)
        if (name := markers.get(marked)) is not None:
            return name, None
        logger.debug_once(
            "%s does not hold its eps (%s) in an attribute. Every instance in this "
            "model will use the value traced from this instance. If this is not "
            "desired, consider storing and reading eps using attribute of %s.",
            type(module).__name__,
            eps,
            type(module).__name__,
        )
        return None, eps

    @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 float(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, vllm_config: "VllmConfig") -> bool:
        return True

    def fuse(
        self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
    ) -> nn.Module:
        """Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
        weight = getattr(module, "weight", None)
        has_weight = weight is not None
        hidden_size = weight.size(0) if has_weight else 0
        eps = getattr(module, self.eps_attr, None) if self.eps_attr else self.eps
        if not isinstance(eps, (int, float)):
            # If eps was not detected, match torch behaviour.
            dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
            eps = torch.finfo(dtype).eps
        if self.zero_centered:
            return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
        return TPAwareRMSNorm(
            hidden_size=hidden_size,
            eps=eps,
            has_weight=has_weight,
            dtype=weight.dtype if has_weight else None,
        )

class TPAwareGemmaRMSNorm(TPAwareNormMixin, GemmaRMSNorm):
    """`GemmaRMSNorm` that reconstructs a TP-sharded input before normalizing."""

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."""