| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for the Transformers modeling backend's RMSNorm fuser.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | import torch.nn as nn |
| 10 | import torch.nn.functional as F |
| 11 | |
| 12 | from vllm.model_executor.models.transformers.fuser import get_fuser |
| 13 | from vllm.model_executor.models.transformers.fusers import RMSNormFuser |
| 14 | |
| 15 | |
| 16 | class RMSNorm(nn.Module): |
| 17 | """The canonical HF RMSNorm: `weight * x * rsqrt(mean(x**2) + eps)`.""" |
| 18 | |
| 19 | def __init__(self, hidden: int = 16, eps: float = 1e-5, weight: bool = True): |
| 20 | super().__init__() |
| 21 | if weight: |
| 22 | self.weight = nn.Parameter(torch.ones(hidden)) |
| 23 | self.variance_epsilon = eps |
| 24 | |
| 25 | def _rms(self, x): |
| 26 | return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) |
| 27 | |
| 28 | def forward(self, x): |
| 29 | return self.weight * self._rms(x.to(torch.float32)).to(x.dtype) |
| 30 | |
| 31 | |
| 32 | class GemmaRMSNorm(RMSNorm): |
| 33 | """Zero-centered weight: `(1 + weight) * normalized`.""" |
| 34 | |
| 35 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 36 | super().__init__(hidden, eps) |
| 37 | self.weight = nn.Parameter(torch.zeros(hidden)) |
| 38 | |
| 39 | def forward(self, x): |
| 40 | return (1.0 + self.weight) * self._rms(x.to(torch.float32)).to(x.dtype) |
| 41 | |
| 42 | |
| 43 | class WeightlessRMSNorm(RMSNorm): |
| 44 | """No scale parameter (e.g. Gemma3n `with_scale=False`).""" |
| 45 | |
| 46 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 47 | super().__init__(hidden, eps, weight=False) |
| 48 | |
| 49 | def forward(self, x): |
| 50 | return self._rms(x.to(torch.float32)).to(x.dtype) |
| 51 | |
| 52 | |
| 53 | class Gemma4RMSNorm(RMSNorm): |
| 54 | """`pow(v, -0.5)` where the others write `rsqrt(v)` (HF `Gemma4RMSNorm._norm`). |
| 55 | |
| 56 | Gemma 4 spells the reciprocal square root as `torch.pow(..., -0.5)` to keep |
| 57 | Torch and JAX in agreement, and -- unlike Gemma 1-3 -- its weight is *not* |
| 58 | zero-centered, so this fuses to a plain `RMSNorm`. |
| 59 | """ |
| 60 | |
| 61 | def _rms(self, x): |
| 62 | mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon |
| 63 | return x * torch.pow(mean_squared, -0.5) |
| 64 | |
| 65 | def forward(self, x): |
| 66 | return (self._rms(x.float()) * self.weight.float()).type_as(x) |
| 67 | |
| 68 | |
| 69 | class PowOperatorRMSNorm(Gemma4RMSNorm): |
| 70 | """The `v ** -0.5` operator spelling of the same reciprocal square root.""" |
| 71 | |
| 72 | def _rms(self, x): |
| 73 | return x * (x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) ** -0.5 |
| 74 | |
| 75 | |
| 76 | class KwargPowRMSNorm(Gemma4RMSNorm): |
| 77 | """The keyword spelling `torch.pow(v, exponent=-0.5)` of the same rsqrt.""" |
| 78 | |
| 79 | def _rms(self, x): |
| 80 | mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon |
| 81 | return x * torch.pow(mean_squared, exponent=-0.5) |
| 82 | |
| 83 | |
| 84 | class KwargBasePowRMSNorm(Gemma4RMSNorm): |
| 85 | """The fully functional spelling `torch.pow(input=v, exponent=-0.5)`.""" |
| 86 | |
| 87 | def _rms(self, x): |
| 88 | mean_squared = x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon |
| 89 | return x * torch.pow(input=mean_squared, exponent=-0.5) |
| 90 | |
| 91 | |
| 92 | class WeightlessGemma4RMSNorm(Gemma4RMSNorm): |
| 93 | """Gemma 4 with `with_scale=False`: the `pow` spelling and no scale parameter.""" |
| 94 | |
| 95 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 96 | super().__init__(hidden, eps, weight=False) |
| 97 | |
| 98 | def forward(self, x): |
| 99 | return self._rms(x.float()).type_as(x) |
| 100 | |
| 101 | |
| 102 | class LayerNorm(RMSNorm): |
| 103 | """An RMSNorm not named `*RMSNorm`, keeping the input dtype (no upcast).""" |
| 104 | |
| 105 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 106 | super().__init__(hidden, eps) |
| 107 | |
| 108 | def forward(self, x): |
| 109 | return self.weight * self._rms(x) |
| 110 | |
| 111 | |
| 112 | class NamedEpsRMSNorm(nn.Module): |
| 113 | """Holds eps under `attr`, with no `variance_epsilon` to find instead.""" |
| 114 | |
| 115 | attr = "eps" |
| 116 | |
| 117 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 118 | super().__init__() |
| 119 | self.weight = nn.Parameter(torch.ones(hidden)) |
| 120 | setattr(self, self.attr, eps) |
| 121 | |
| 122 | def forward(self, x): |
| 123 | eps = getattr(self, self.attr) |
| 124 | return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + eps)) |
| 125 | |
| 126 | |
| 127 | class ToleranceRMSNorm(NamedEpsRMSNorm): |
| 128 | """Names eps something other than `eps`, so only its value identifies it.""" |
| 129 | |
| 130 | attr = "tolerance" |
| 131 | |
| 132 | |
| 133 | class LiteralEpsRMSNorm(RMSNorm): |
| 134 | """eps is a literal in the source, and an unrelated attribute happens to |
| 135 | hold the same value: matching on the value alone would bind to it.""" |
| 136 | |
| 137 | def __init__(self, hidden: int = 16, eps: float = 1e-5): |
| 138 | super().__init__(hidden, eps) |
| 139 | self.epsilon_lookalike = 1e-4 |
| 140 | |
| 141 | def forward(self, x): |
| 142 | return self.weight * (x * torch.rsqrt(x.pow(2).mean(-1, True) + 1e-4)) |
| 143 | |
| 144 | |
| 145 | class AmbiguousEpsRMSNorm(nn.Module): |
| 146 | """Two attributes hold the eps value, and the forward reads the second, so |
| 147 | their order cannot pick it.""" |
| 148 | |
| 149 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 150 | super().__init__() |
| 151 | self.weight = nn.Parameter(torch.ones(hidden)) |
| 152 | self.decoy_eps = eps |
| 153 | self.variance_epsilon = eps |
| 154 | |
| 155 | def forward(self, x): |
| 156 | var = x.pow(2).mean(-1, True) |
| 157 | return self.weight * (x * torch.rsqrt(var + self.variance_epsilon)) |
| 158 | |
| 159 | |
| 160 | class NotAnRMSNorm(RMSNorm): |
| 161 | """Mean-subtracting LayerNorm-like math -> not an RMSNorm.""" |
| 162 | |
| 163 | def __init__(self, hidden: int = 16, eps: float = 1e-6): |
| 164 | super().__init__(hidden, eps) |
| 165 | |
| 166 | def forward(self, x): |
| 167 | x = x - x.mean(-1, keepdim=True) |
| 168 | variance = x.var(-1, keepdim=True) |
| 169 | return self.weight * x / torch.sqrt(variance + self.variance_epsilon) |
| 170 | |
| 171 | |
| 172 | class GatedRMSNorm(RMSNorm): |
| 173 | """Second input and tail compute -> not an RMSNorm.""" |
| 174 | |
| 175 | def forward(self, x, gate=None): |
| 176 | normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) |
| 177 | return normed * F.silu(gate) |
| 178 | |
| 179 | |
| 180 | class GatedFusedRMSNorm(nn.Module): |
| 181 | """Same as GatedRMSNorm, but built on the fused `rms_norm` op -> not an RMSNorm.""" |
| 182 | |
| 183 | def __init__(self, hidden: int = 16, eps: float = 1e-5): |
| 184 | super().__init__() |
| 185 | self.weight = nn.Parameter(torch.ones(hidden)) |
| 186 | self.eps = eps |
| 187 | |
| 188 | def forward(self, x, gate=None): |
| 189 | return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) * F.silu(gate) |
| 190 | |
| 191 | |
| 192 | class UntraceableGatedRMSNorm(RMSNorm): |
| 193 | """Tracer can't see tail compute in forward, but still has a second input (gate).""" |
| 194 | |
| 195 | def forward(self, x, gate=None): |
| 196 | normed = self.weight * self._rms(x.to(torch.float32)).to(x.dtype) |
| 197 | if gate.sum() > 0: # untraceable -> partial graph, no visible tail |
| 198 | normed = normed * F.silu(gate) |
| 199 | return normed |
| 200 | |
| 201 | |
| 202 | @pytest.mark.parametrize( |
| 203 | "cls,eps,zero_centered", |
| 204 | [ |
| 205 | (RMSNorm, 1e-5, False), |
| 206 | (GemmaRMSNorm, 1e-6, True), |
| 207 | (WeightlessRMSNorm, 1e-6, False), |
| 208 | (Gemma4RMSNorm, 1e-6, False), |
| 209 | (PowOperatorRMSNorm, 1e-6, False), |
| 210 | (KwargPowRMSNorm, 1e-6, False), |
| 211 | (KwargBasePowRMSNorm, 1e-6, False), |
| 212 | (WeightlessGemma4RMSNorm, 1e-6, False), |
| 213 | (LayerNorm, 1e-6, False), |
| 214 | (torch.nn.RMSNorm, 1e-5, False), # fused `F.rms_norm` op |
| 215 | ], |
| 216 | ) |
| 217 | def test_detects_rms_norm_variants(cls, eps, zero_centered): |
| 218 | with torch.device("meta"): |
| 219 | fuser = get_fuser(cls(16, eps=eps), RMSNormFuser) |
| 220 | assert isinstance(fuser, RMSNormFuser) |
| 221 | assert fuser.zero_centered == zero_centered |
| 222 | |
| 223 | |
| 224 | @pytest.mark.parametrize("cls", [NotAnRMSNorm, nn.LayerNorm, nn.SiLU]) |
| 225 | def test_non_rms_norms_are_not_matched(cls): |
| 226 | with torch.device("meta"): |
| 227 | module = cls(16) if cls is nn.LayerNorm else cls() |
| 228 | assert get_fuser(module, RMSNormFuser) is None |
| 229 | |
| 230 | |
| 231 | @pytest.mark.parametrize( |
| 232 | "cls", [GatedRMSNorm, GatedFusedRMSNorm, UntraceableGatedRMSNorm] |
| 233 | ) |
| 234 | def test_gated_rms_norm_is_not_fused(cls): |
| 235 | with torch.device("meta"): |
| 236 | assert get_fuser(cls(), RMSNormFuser) is None |
| 237 | |
| 238 | |
| 239 | @pytest.mark.parametrize( |
| 240 | "cls,expected,zero_centered", |
| 241 | [ |
| 242 | (RMSNorm, "RMSNorm", False), |
| 243 | (GemmaRMSNorm, "GemmaRMSNorm", True), |
| 244 | (WeightlessRMSNorm, "RMSNorm", False), |
| 245 | (Gemma4RMSNorm, "RMSNorm", False), |
| 246 | (KwargPowRMSNorm, "RMSNorm", False), |
| 247 | (KwargBasePowRMSNorm, "RMSNorm", False), |
| 248 | (WeightlessGemma4RMSNorm, "RMSNorm", False), |
| 249 | ], |
| 250 | ) |
| 251 | def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_config): |
| 252 | from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm |
| 253 | from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm |
| 254 | |
| 255 | with torch.device("meta"): |
| 256 | module = cls() |
| 257 | fuser = get_fuser(module, RMSNormFuser) |
| 258 | built = fuser.fuse(module, "norm", default_vllm_config) |
| 259 | from vllm.model_executor.models.transformers.fusers.rms_norm import ( |
| 260 | TPAwareNormMixin, |
| 261 | ) |
| 262 | |
| 263 | types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm} |
| 264 | assert isinstance(built, types_by_name[expected]) |
| 265 | assert isinstance(built, TPAwareNormMixin) |
| 266 | assert built.variance_epsilon == module.variance_epsilon |
| 267 | # The weight states the hidden size (and is what the TP check reads); a norm |
| 268 | # without one has none to state. |
| 269 | weight = getattr(module, "weight", None) |
| 270 | assert isinstance(built.weight, nn.Parameter) == (weight is not None) |
| 271 | assert built.weight.shape[0] == (weight.size(0) if weight is not None else 0) |
| 272 | |
| 273 | |
| 274 | @pytest.mark.parametrize( |
| 275 | "cls", [Gemma4RMSNorm, PowOperatorRMSNorm, KwargPowRMSNorm, KwargBasePowRMSNorm] |
| 276 | ) |
| 277 | def test_pow_spelled_norm_fuses_to_equivalent_math(cls, default_vllm_config): |
| 278 | """Matching `pow(v, -0.5)` is only sound if it really is the reciprocal square |
| 279 | root: the fused norm must reproduce the unfused forward, not just replace it. |
| 280 | `pow` and `rsqrt` take different paths to the same value, so compare at float32 |
| 281 | tolerance rather than bit-for-bit.""" |
| 282 | torch.manual_seed(0) |
| 283 | module = cls(16) |
| 284 | with torch.no_grad(): |
| 285 | module.weight.copy_(torch.randn(16)) |
| 286 | built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config) |
| 287 | with torch.no_grad(): |
| 288 | built.weight.copy_(module.weight) |
| 289 | x = torch.randn(4, 16) |
| 290 | torch.testing.assert_close(built(x), module(x), rtol=1e-5, atol=1e-6) |
| 291 | |
| 292 | |
| 293 | def test_weightless_norm_has_no_hidden_size(default_vllm_config): |
| 294 | """A weightless norm states no hidden size, and the model's (LM) hidden size |
| 295 | would be wrong for one on a sub-dimension: Llama 4's `qk_norm` normalizes |
| 296 | head_dim, so sizing it at hidden_size made the TP gather reject its input. |
| 297 | It fuses with none, and so normalizes the width it is given at any TP size, |
| 298 | matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`. |
| 299 | """ |
| 300 | module = WeightlessRMSNorm(72) |
| 301 | built = get_fuser(module, RMSNormFuser).fuse(module, "norm", default_vllm_config) |
| 302 | assert built.hidden_size == 0 |
| 303 | |
| 304 | built.tp_size = 2 # emulate TP=2 without a real process group |
| 305 | x = torch.randn(4, 72) |
| 306 | torch.testing.assert_close(built(x), module(x)) |
| 307 | |
| 308 | |
| 309 | def test_fused_rms_norm_op_default_eps(default_vllm_config): |
| 310 | """`torch.nn.RMSNorm` (a single `F.rms_norm` call) matches via the fast path; |
| 311 | its default `eps=None` resolves to `finfo(dtype).eps` in `fuse`.""" |
| 312 | from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm |
| 313 | |
| 314 | with torch.device("meta"): |
| 315 | module = torch.nn.RMSNorm(16) # forward is a single `F.rms_norm` call |
| 316 | fuser = get_fuser(module, RMSNormFuser) |
| 317 | assert isinstance(fuser, RMSNormFuser) |
| 318 | assert not fuser.zero_centered |
| 319 | vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float32)) |
| 320 | built = fuser.fuse(module, "norm", vllm_config) |
| 321 | assert isinstance(built, VLLMRMSNorm) |
| 322 | assert built.variance_epsilon == torch.finfo(torch.float32).eps |
| 323 | |
| 324 | |
| 325 | def test_eps_is_derived_per_instance(default_vllm_config): |
| 326 | """Two instances of the same norm class with different eps must fuse to their |
| 327 | own eps: the type-cached fuser holds only structure, not this value.""" |
| 328 | with torch.device("meta"): |
| 329 | for eps in (1e-5, 1e-6): |
| 330 | module = RMSNorm(16, eps=eps) |
| 331 | fuser = get_fuser(module, RMSNormFuser) |
| 332 | built = fuser.fuse(module, "norm", default_vllm_config) |
| 333 | assert built.variance_epsilon == eps |
| 334 | |
| 335 | |
| 336 | @pytest.mark.parametrize("cls", [NamedEpsRMSNorm, ToleranceRMSNorm]) |
| 337 | def test_eps_attr_is_found_by_value_not_name(cls, default_vllm_config): |
| 338 | """The eps attribute is identified by holding the traced value, so a norm |
| 339 | stays per-instance correct whatever it names it.""" |
| 340 | with torch.device("meta"): |
| 341 | for eps in (1e-5, 1e-6): |
| 342 | module = cls(16, eps=eps) |
| 343 | fuser = get_fuser(module, RMSNormFuser) |
| 344 | assert fuser.eps_attr == cls.attr |
| 345 | built = fuser.fuse(module, "norm", default_vllm_config) |
| 346 | assert built.variance_epsilon == eps |
| 347 | |
| 348 | |
| 349 | def test_literal_eps_is_not_mistaken_for_an_attribute(default_vllm_config, caplog): |
| 350 | """A literal eps is recognised as coming from no attribute, even when one |
| 351 | holds the same value, and is taken from the traced source instead.""" |
| 352 | logger = "vllm.model_executor.models.transformers.fusers.rms_norm" |
| 353 | with caplog.at_level("DEBUG", logger=logger), torch.device("meta"): |
| 354 | module = LiteralEpsRMSNorm() |
| 355 | fuser = get_fuser(module, RMSNormFuser) |
| 356 | built = fuser.fuse(module, "norm", default_vllm_config) |
| 357 | assert fuser.eps_attr is None |
| 358 | assert built.variance_epsilon == 1e-4 |
| 359 | assert "does not hold its eps" in caplog.text |
| 360 | |
| 361 | |
| 362 | def test_ambiguous_eps_attrs_are_disambiguated(default_vllm_config): |
| 363 | """When several attributes hold the eps value, the one the forward actually |
| 364 | reads is identified, and they are all left as they were found.""" |
| 365 | with torch.device("meta"): |
| 366 | module = AmbiguousEpsRMSNorm(16, eps=1e-6) |
| 367 | before = dict(vars(module)) |
| 368 | fuser = get_fuser(module, RMSNormFuser) |
| 369 | assert fuser.eps_attr == "variance_epsilon" |
| 370 | assert vars(module) == before |
| 371 | |
| 372 | other = AmbiguousEpsRMSNorm(16, eps=1e-5) |
| 373 | other.decoy_eps = 1e-3 |
| 374 | built = fuser.fuse(other, "norm", default_vllm_config) |
| 375 | assert built.variance_epsilon == 1e-5 |
| 376 | |
| 377 | |
| 378 | def test_fused_norm_is_gather_capable(default_vllm_config): |
| 379 | """Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded |
| 380 | projection (OLMoE-style) self-corrects at runtime with no QKV-specific |
| 381 | plumbing. A full-width input skips the gather and equals a plain norm.""" |
| 382 | from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm |
| 383 | from vllm.model_executor.models.transformers.fusers import rms_norm |
| 384 | |
| 385 | torch.manual_seed(0) |
| 386 | x = torch.randn(4, 16) |
| 387 | for gathered_cls, plain_cls in [ |
| 388 | (rms_norm.TPAwareRMSNorm, RMSNorm), |
| 389 | (rms_norm.TPAwareGemmaRMSNorm, GemmaRMSNorm), |
| 390 | ]: |
| 391 | gathered = gathered_cls(hidden_size=16, eps=1e-6) |
| 392 | assert isinstance(gathered, rms_norm.TPAwareNormMixin) |
| 393 | plain = plain_cls(hidden_size=16, eps=1e-6) |
| 394 | with torch.no_grad(): |
| 395 | weight = torch.randn(16) |
| 396 | gathered.weight.copy_(weight) |
| 397 | plain.weight.copy_(weight) |
| 398 | torch.testing.assert_close(gathered(x), plain(x)) |
| 399 | |
| 400 | |
| 401 | def test_gathered_norm_rejects_uneven_sharding(default_vllm_config): |
| 402 | """A sharded input (narrower than the full-width weight) that does not tile |
| 403 | the weight evenly across ranks is rejected before any collective.""" |
| 404 | from vllm.model_executor.models.transformers.fusers import rms_norm |
| 405 | |
| 406 | norm = rms_norm.TPAwareRMSNorm(hidden_size=8, eps=1e-6) |
| 407 | norm.tp_size = 2 # emulate TP=2 without a real process group |
| 408 | with pytest.raises(ValueError, match="does not tile it evenly"): |
| 409 | norm(torch.randn(2, 3)) # 3 * 2 != 8 |