diff --git a/tests/model_executor/test_weight_utils.py b/tests/model_executor/test_weight_utils.py index 12a8e3acac..0ecb0a1aed 100644 --- a/tests/model_executor/test_weight_utils.py +++ b/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() diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index 266a346a64..2949df3bdf 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/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