[Model Loader] Keep norm layers neutral under dummy weight init (#61323)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/model_executor/test_weight_utils.py

@@ -352,5 +352,27 @@ def test_dummy_wide_integer_param_is_untouched():
     assert torch.equal(param, torch.arange(16, dtype=torch.int32))
 
 
+def test_dummy_norm_weights_are_neutral(default_vllm_config):
+    """Dummy norms stay identity-like, keeping activations at a realistic scale."""
+    from vllm.model_executor.layers.layernorm import (
+        GemmaRMSNorm,
+        LayerNorm,
+        RMSNorm,
+        RMSNormGated,
+    )
+
+    expected = {
+        RMSNorm(16): {"weight": 1.0},
+        GemmaRMSNorm(16): {"weight": 0.0},
+        LayerNorm(16): {"weight": 1.0, "bias": 0.0},
+        RMSNormGated(16): {"weight": 1.0},
+    }
+    for norm, values in expected.items():
+        for name, param in norm.named_parameters():
+            param.data.uniform_(-5, 5)
+            initialize_single_dummy_weight(param, name=name)
+            assert torch.all(param == values[name]), (type(norm).__name__, name)
+
+
 if __name__ == "__main__":
     test_download_weights_from_hf()

vllm/model_executor/layers/layernorm.py

@@ -12,6 +12,7 @@ from vllm import envs, ir
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.determinism.batch_invariant import rms_norm_batch_invariant
+from vllm.model_executor.utils import set_weight_attrs
 
 logger = init_logger(__name__)
 
@@ -63,6 +64,9 @@ class RMSNorm(CustomOp):
         self.weight = torch.ones(hidden_size, dtype=weight_dtype)
         if self.has_weight:
             self.weight = nn.Parameter(self.weight)
+            # Dummy weights keep norms neutral: random ones would shrink every
+            # activation and skew data-dependent kernels (e.g. sparse top-k).
+            set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
 
         # When has_weight=False, pass weight=None so implementations that
         # support a weightless path can skip the per-channel multiply.
@@ -154,6 +158,7 @@ class GemmaRMSNorm(CustomOp):
     ) -> None:
         super().__init__()
         self.weight = nn.Parameter(torch.zeros(hidden_size))
+        set_weight_attrs(self.weight, {"dummy_weight_value": 0.0})
         self.variance_epsilon = eps
 
     def forward_native(
@@ -239,8 +244,10 @@ class LayerNorm(CustomOp):
         self.bias: nn.Parameter | None = None
         if elementwise_affine:
             self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype))
+            set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
             if bias:
                 self.bias = nn.Parameter(torch.zeros(hidden_size, dtype=weight_dtype))
+                set_weight_attrs(self.bias, {"dummy_weight_value": 0.0})
 
     def _compute_dtype(self, x: torch.Tensor) -> torch.dtype:
         if self.weight is None:
@@ -328,6 +335,7 @@ class RMSNormGated(CustomOp):
         self.eps = eps
         self.activation = activation
         self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))
+        set_weight_attrs(self.weight, {"dummy_weight_value": 1.0})
         self.register_parameter("bias", None)
         self.group_size = group_size
         self.norm_before_gate = norm_before_gate