Honor norm_zero_centered_gamma in RMSNormWithInputScale (pytorch#22444)

Summary:

rlformers stores gamma offset by -1 when `norm_zero_centered_gamma` is set, so its `RMSNormWithInputScale` computes `rms_norm((weight + 1.0) * x)` (`core/model/transformer.py:208-211`). ET computed `rms_norm(weight * x)` unconditionally, and `norm_zero_centered_gamma` appeared in zero lines of ET — not on `ModelArgs`, not in `model_args_conversion.py`, not in `norm.py`. The flag was silently dropped, so a checkpoint carrying it loaded without error and produced systematically wrong activations from the first block onward.

Add `zero_centered_gamma` to `RMSNormWithInputScale` and `RMSNormWithInputScaleCoreML`, add `norm_zero_centered_gamma` to `ModelArgs`, plumb it through `get_et_model_args_from_genai_model_args`, and pass it where `post_ffn_norm` is built. It defaults to False, so no checkpoint without the flag changes behavior.

Reviewed By: Zhiwen-Owen-Jiang

Differential Revision: D117233461

Pull Request resolved: pytorch#22444

Author: telgamal-1

examples/models/llama/llama_transformer.py

@@ -244,7 +244,11 @@ def __init__(
             self.post_attn_norm = ScalelessRMSNorm(args.dim, eps=args.norm_eps)
 
         if args.use_ffn_learnable_scales and self.mlp_type != "skip":
-            self.post_ffn_norm = RMSNormWithInputScale(args.dim, eps=args.norm_eps)
+            self.post_ffn_norm = RMSNormWithInputScale(
+                args.dim,
+                eps=args.norm_eps,
+                zero_centered_gamma=args.norm_zero_centered_gamma,
+            )
 
     @classmethod
     def from_type(cls, layer_id, args, rope) -> "TransformerBlock":

examples/models/llama/model_args.py

@@ -187,6 +187,10 @@ class ModelArgs:
     use_attn_o_norm: bool = False
     use_residual_gate: bool = False
     use_ffn_learnable_scales: bool = False
+    # Zero-centered gamma: the checkpoint stores gamma offset by -1, so the
+    # effective scale is ``weight + 1``. rlformers applies this in
+    # RMSNormWithInputScale; ignoring it silently rescales every post-FFN norm.
+    norm_zero_centered_gamma: bool = False
     output_soft_cap_temp: Optional[float] = None
     output_linear_intermediate_dim: int = 0
 

examples/models/llama/norm.py

@@ -103,19 +103,28 @@ def forward(self, x):
 
 
 class RMSNormWithInputScale(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-5):
+    def __init__(self, dim: int, eps: float = 1e-5, zero_centered_gamma: bool = False):
+        """RMSNorm with gamma applied to the input: ``rms_norm(gamma * x)``.
+
+        Args:
+            zero_centered_gamma: when True the checkpoint stores gamma offset by
+                -1, so the effective scale is ``weight + 1``. Mirrors
+                ``RMSNormWithInputScale`` in rlformers' ``core/model/transformer.py``.
+        """
         super().__init__()
         self.eps = eps
         self.dim = dim
+        self.zero_centered_gamma = zero_centered_gamma
         self.weight = torch.nn.Parameter(torch.ones(dim))
 
     def forward(self, x):
-        scaled = self.weight * x
+        w = self.weight + 1.0 if self.zero_centered_gamma else self.weight
+        scaled = w * x
         return F.rms_norm(scaled, (self.dim,), None, self.eps)
 
 
 class RMSNormWithInputScaleCoreML(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-5):
+    def __init__(self, dim: int, eps: float = 1e-5, zero_centered_gamma: bool = False):
         """
         CoreML-friendly RMSNormWithInputScale.
 
@@ -139,6 +148,7 @@ def __init__(self, dim: int, eps: float = 1e-5):
         )
         self.eps = eps
         self.dim = dim
+        self.zero_centered_gamma = zero_centered_gamma
         self.weight = torch.nn.Parameter(torch.ones(dim))
 
     def _norm(self, x):
@@ -153,7 +163,8 @@ def _norm(self, x):
         )
 
     def forward(self, x):
-        scaled = self.weight * x
+        w = self.weight + 1.0 if self.zero_centered_gamma else self.weight
+        scaled = w * x
         return self._norm(scaled)
 
 
@@ -197,7 +208,16 @@ def replace_rms_norm_for_coreml_(model: torch.nn.Module) -> torch.nn.Module:
             # applies its scale post-norm, which would change the math here).
             dim = getattr(mod, "dim", None) or mod.normalized_shape[-1]
             eps = getattr(mod, "eps", 1e-6) or 1e-6
-            new = RMSNormWithInputScaleCoreML(dim, eps=eps)
+            # Carry `zero_centered_gamma` across. Dropping it here would silently
+            # downgrade `rms_norm((weight + 1) * x)` to `rms_norm(weight * x)` on the
+            # CoreML path only -- the same wrong-activation bug this flag exists to
+            # prevent, and one that would show up as a CoreML-vs-XNNPACK divergence
+            # rather than as an error.
+            new = RMSNormWithInputScaleCoreML(
+                dim,
+                eps=eps,
+                zero_centered_gamma=getattr(mod, "zero_centered_gamma", False),
+            )
             new.weight = mod.weight
         elif isinstance(mod, (RMSNorm, ScalelessRMSNorm, torch.nn.RMSNorm)):
             # All three carry the normalized dim either as `dim` or in `normalized_shape[-1]`.