Fix zero-centered gamma parity for ET RMSNorm (pytorch#23193)

Summary:

Preserve zero-centered gamma semantics for transformer block and output RMSNorms in both ExecuTorch and rlformers ET paths while leaving Q/K norms unchanged. Preserve output_norm_gain_center_type so center-zero output weights are not shifted.

Authored with assistance from Codex.

Differential Revision: D121895059

Pull Request resolved: pytorch#23193

Author: chinnadhurai

examples/models/llama/llama_transformer.py

@@ -227,13 +227,17 @@ def __init__(
             self.attention_norm = RMSNorm(
                 args.dim,
                 eps=args.norm_eps,
-                add_unit_offset=args.rms_norm_add_unit_offset,
+                add_unit_offset=(
+                    args.rms_norm_add_unit_offset or args.norm_zero_centered_gamma
+                ),
             )
         if self.mlp_type != "skip":
             self.ffn_norm = RMSNorm(
                 args.dim,
                 eps=args.norm_eps,
-                add_unit_offset=args.rms_norm_add_unit_offset,
+                add_unit_offset=(
+                    args.rms_norm_add_unit_offset or args.norm_zero_centered_gamma
+                ),
             )
 
         if args.use_residual_gate:
@@ -346,7 +350,13 @@ def __init__(self, params: ModelArgs, layers: nn.ModuleList, rope: Rope):
         self.norm = RMSNorm(
             params.dim,
             eps=params.norm_eps,
-            add_unit_offset=params.rms_norm_add_unit_offset,
+            add_unit_offset=(
+                params.rms_norm_add_unit_offset
+                or (
+                    params.norm_zero_centered_gamma
+                    and params.output_norm_gain_center_type != "zero"
+                )
+            ),
         )
         self.output = (
             nn.Linear(params.dim, params.vocab_size, bias=False)

examples/models/llama/model_args.py

@@ -195,6 +195,7 @@ class ModelArgs:
     # 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_norm_gain_center_type: str = "one"
     output_soft_cap_temp: Optional[float] = None
     output_linear_intermediate_dim: int = 0
 

examples/models/llama/norm.py

@@ -55,7 +55,7 @@ def forward(self, x):
 
 
 class RMSNormCoreML(torch.nn.Module):
-    def __init__(self, dim: int, eps: float = 1e-6):
+    def __init__(self, dim: int, eps: float = 1e-6, add_unit_offset: bool = False):
         """
         CoreML-friendly RMSNorm — uses `torch.linalg.vector_norm` so the op is
         preserved in the CoreML graph for numerical stability.
@@ -78,6 +78,7 @@ def __init__(self, dim: int, eps: float = 1e-6):
         )
         self.dim = dim
         self.eps = eps
+        self.add_unit_offset = add_unit_offset
         self.weight = nn.Parameter(torch.ones(dim))
 
     def _norm(self, x):
@@ -99,7 +100,7 @@ def _norm(self, x):
 
     def forward(self, x):
         output = self._norm(x)
-        return output * self.weight
+        return output * (1.0 + self.weight if self.add_unit_offset else self.weight)
 
 
 class RMSNormWithInputScale(torch.nn.Module):

examples/models/llama/tests/test_transformer_block.py

@@ -8,11 +8,46 @@
 
 import torch
 from executorch.examples.models.llama.attention import AttentionSkip
-from executorch.examples.models.llama.llama_transformer import TransformerBlock
+from executorch.examples.models.llama.llama_transformer import (
+    Transformer,
+    TransformerBlock,
+)
 from executorch.examples.models.llama.model_args import ModelArgs
+from executorch.examples.models.llama.rope import Rope
 
 
 class TestTransformerBlock(unittest.TestCase):
+    def test_zero_centered_gamma_applies_to_main_norms(self) -> None:
+        args = ModelArgs(
+            dim=4,
+            hidden_dim=8,
+            n_heads=1,
+            n_kv_heads=1,
+            head_dim=4,
+            norm_zero_centered_gamma=True,
+        )
+        block = TransformerBlock(args, AttentionSkip()).eval()
+        model = Transformer(args, torch.nn.ModuleList([block]), Rope(args)).eval()
+
+        self.assertTrue(block.ffn_norm.add_unit_offset)
+        self.assertTrue(model.norm.add_unit_offset)
+
+        center_zero_args = ModelArgs(
+            dim=4,
+            hidden_dim=8,
+            n_heads=1,
+            n_kv_heads=1,
+            head_dim=4,
+            norm_zero_centered_gamma=True,
+            output_norm_gain_center_type="zero",
+        )
+        center_zero_model = Transformer(
+            center_zero_args,
+            torch.nn.ModuleList(),
+            Rope(center_zero_args),
+        ).eval()
+        self.assertFalse(center_zero_model.norm.add_unit_offset)
+
     def test_attention_skip_applies_gated_zero_branch_residual(self) -> None:
         args = ModelArgs(
             dim=4,