fix: fuse weightless RMSNorms at their declared width (#50867)
Signed-off-by: Anuj Bolewar <anujbolewar@gmail.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Anuj Bolewar <anujbolewar@gmail.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Anuj Bolewar committed
Aug 4, 2026 at 21:33 UTC
166f4e2dc3a6ed1dfc3641611c9a07d8f0e87f9c
2 files changed
+29
-23
tests/models/transformers/fusers/test_rms_norm.py
+26
-17
@@ -146,26 +146,39 @@ def test_rms_norm_builds_vllm_class(cls, expected, zero_centered, default_vllm_c
146
from vllm.model_executor.layers.layernorm import GemmaRMSNorm as VLLMGemmaRMSNorm
147
from vllm.model_executor.layers.layernorm import RMSNorm as VLLMRMSNorm
148
149
- # `default_vllm_config` supplies the config context the CustomOp needs; the
150
- # weightless path reads hidden size from the model config, so stub it.
151
- vllm_config = SimpleNamespace(
152
- model_config=SimpleNamespace(get_hidden_size=lambda: 16)
153
- )
149
with torch.device("meta"):
150
module = cls()
151
fuser = get_fuser(module)
157
- built = fuser.fuse(module, "norm", vllm_config)
152
+ built = fuser.fuse(module, "norm", default_vllm_config)
153
from vllm.model_executor.models.transformers.fusers.rms_norm import (
154
TPAwareNormMixin,
155
)
156
157
types_by_name = {"RMSNorm": VLLMRMSNorm, "GemmaRMSNorm": VLLMGemmaRMSNorm}
158
assert isinstance(built, types_by_name[expected])
164
- assert isinstance(built, TPAwareNormMixin) # fused norms self-correct under TP
159
+ assert isinstance(built, TPAwareNormMixin)
160
assert built.variance_epsilon == module.variance_epsilon
166
- assert isinstance(built.weight, nn.Parameter) == (
167
- getattr(module, "weight", None) is not None
168
- )
161
+ # The weight states the hidden size (and is what the TP check reads); a norm
162
+ # without one has none to state.
163
+ weight = getattr(module, "weight", None)
164
+ assert isinstance(built.weight, nn.Parameter) == (weight is not None)
165
+ assert built.weight.shape[0] == (weight.size(0) if weight is not None else 0)
166
+
167
+
168
+def test_weightless_norm_has_no_hidden_size(default_vllm_config):
169
+ """A weightless norm states no hidden size, and the model's (LM) hidden size
170
+ would be wrong for one on a sub-dimension: Llama 4's `qk_norm` normalizes
171
+ head_dim, so sizing it at hidden_size made the TP gather reject its input.
172
+ It fuses with none, and so normalizes the width it is given at any TP size,
173
+ matching the unfused norm and vLLM's native `RMSNorm(hidden_size=head_dim)`.
174
+ """
175
+ module = WeightlessRMSNorm(72)
176
+ built = get_fuser(module).fuse(module, "norm", default_vllm_config)
177
+ assert built.hidden_size == 0
178
+
179
+ built.tp_size = 2 # emulate TP=2 without a real process group
180
+ x = torch.randn(4, 72)
181
+ torch.testing.assert_close(built(x), module(x))
182
183
184
def test_fused_rms_norm_op_default_eps(default_vllm_config):
@@ -178,8 +191,7 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):
191
fuser = get_fuser(module)
192
assert isinstance(fuser, RMSNormFuser)
193
assert not fuser.zero_centered
181
- mc = SimpleNamespace(get_hidden_size=lambda: 16, dtype=torch.float32)
182
- vllm_config = SimpleNamespace(model_config=mc)
194
+ vllm_config = SimpleNamespace(model_config=SimpleNamespace(dtype=torch.float32))
195
built = fuser.fuse(module, "norm", vllm_config)
196
assert isinstance(built, VLLMRMSNorm)
197
assert built.variance_epsilon == torch.finfo(torch.float32).eps
@@ -188,18 +200,15 @@ def test_fused_rms_norm_op_default_eps(default_vllm_config):
200
def test_eps_is_derived_per_instance(default_vllm_config):
201
"""Two instances of the same norm class with different eps must fuse to their
202
own eps: the type-cached fuser holds only structure, not this value."""
191
- vllm_config = SimpleNamespace(
192
- model_config=SimpleNamespace(get_hidden_size=lambda: 16)
193
- )
203
with torch.device("meta"):
204
for eps in (1e-5, 1e-6):
205
module = RMSNorm(16, eps=eps)
197
- built = get_fuser(module).fuse(module, "norm", vllm_config)
206
+ built = get_fuser(module).fuse(module, "norm", default_vllm_config)
207
assert built.variance_epsilon == eps
208
209
210
def test_fused_norm_is_gather_capable(default_vllm_config):
202
- """Every fused norm is emitted gather-capable, so a norm on a head-sharded
211
+ """Every weighted fused norm is emitted gather-capable, so a norm on a head-sharded
212
projection (OLMoE-style) self-corrects at runtime with no QKV-specific
213
plumbing. A full-width input skips the gather and equals a plain norm."""
214
from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
vllm/model_executor/models/transformers/fusers/rms_norm.py
+3
-6
@@ -190,20 +190,17 @@ class RMSNormFuser(BaseFuser):
190
self, module: nn.Module, prefix: str, vllm_config: "VllmConfig"
191
) -> nn.Module:
192
"""Fuse the matched RMSNorm pattern into a vLLM fused RMSNorm CustomOp."""
193
- model_config = vllm_config.model_config
193
weight = getattr(module, "weight", None)
195
- hidden_size = (
196
- weight.size(0) if weight is not None else model_config.get_hidden_size()
197
- )
194
+ has_weight = weight is not None
195
+ hidden_size = weight.size(0) if has_weight else 0
196
graph = trace(module)
197
eps = self._eps_from_graph(graph) if graph is not None else None
198
if eps is None:
199
# If eps not in graph, match torch behaviour.
202
- dtype = weight.dtype if weight is not None else model_config.dtype
200
+ dtype = weight.dtype if has_weight else vllm_config.model_config.dtype
201
eps = torch.finfo(dtype).eps
202
if self.zero_centered:
203
return TPAwareGemmaRMSNorm(hidden_size=hidden_size, eps=eps)
206
- has_weight = weight is not None
204
return TPAwareRMSNorm(
205
hidden_size=hidden_size,
206
eps=eps,