From 1301fdbe973a9ce2f7f51977fe3044b5d8f65931 Mon Sep 17 00:00:00 2001 From: Tarek Elgamal Date: Fri, 4 Sep 2026 16:29:43 -0700 Subject: [PATCH] Map norm_zero_centered_gamma onto the existing rms_norm_add_unit_offset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: ALTERNATIVE to D117233461, implementing minguo's review suggestion. Same bug, one fewer concept. rlformers stores gamma offset by -1 when `norm_zero_centered_gamma` is set, so the effective scale is `weight + 1`. ET dropped the flag on the params -> ModelArgs conversion, so a checkpoint carrying it loaded without error and produced systematically wrong activations. D117233461 fixes that by adding a **new** `ModelArgs.norm_zero_centered_gamma` and wiring it into `RMSNormWithInputScale`. But ET already has a field meaning exactly this: `ModelArgs.rms_norm_add_unit_offset` (`model_args.py:56`), and `RMSNorm.forward` already implements it as `output * (1 + weight)` (`norm.py:37-41`). It is already wired into `attention_norm` and `ffn_norm` (`llama_transformer.py:226,232`), `q_norm`/`k_norm` (`attention.py:468,474`) and the final norm (`llama_transformer.py:345`) — nothing ever populated it. On the rlformers side the single `norm_zero_centered_gamma` drives all of those too (`transformer.py:3618,3634,3710,3715,5468`). It is one checkpoint property, not two. So the two-field version leaves a checkpoint with `norm_zero_centered_gamma: true` and affine norms computing `rms_norm((weight + 1) * x)` in `post_ffn_norm` and `rms_norm(weight * x)` in the other five — the same silent wrong-activation bug, relocated rather than fixed. `f1129404401` reaches `inf` under either version only because its other norms happen to be unaffected; that is a property of that checkpoint, not of the fix. This version instead populates the existing field from the existing checkpoint key and passes it into `RMSNormWithInputScale`. One name, one source of truth, six norms covered instead of one, and no new `ModelArgs` surface: 10 files instead of 12. It also removes the reason `patch_rms_unit_offset` exists in `utils/omni_patches.py:76`, whose own docstring says ET "silently crushes activations" because "the omni params->ModelArgs conversion never sets it". That monkeypatch is left in place here — deleting it belongs in its own change. Two things a reviewer should decide, both raised by minguo and NOT resolved here: - rlformers gates the **final** norm on `output_norm_gain_center_type` ("one" -> +1, "zero" -> no shift), independent of the per-layer flag. ET has one flag covering the final norm too, so mapping them 1:1 is a deliberate simplification that needs sign-off. - Widening the mapping changes numerics for any already-validated checkpoint whose params.json carries `norm_zero_centered_gamma: true`. The conversion scripts bake the +1 in and then set the flag false, so most are unaffected — but this wants an A/B on the existing `backbone_cuda_test` checkpoints rather than an assumption. Differential Revision: D118485339 --- examples/models/llama/llama_transformer.py | 10 +++++++- examples/models/llama/norm.py | 29 ++++++++++++++++++---- 2 files changed, 33 insertions(+), 6 deletions(-) diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py index 73117826708..c0837d32920 100644 --- a/examples/models/llama/llama_transformer.py +++ b/examples/models/llama/llama_transformer.py @@ -244,7 +244,15 @@ 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, + # Same checkpoint property that drives attention_norm, ffn_norm, + # q_norm/k_norm and the final norm. rlformers has one flag + # (`norm_zero_centered_gamma`) covering all of them, so ET must not + # grow a second name for it. + zero_centered_gamma=args.rms_norm_add_unit_offset, + ) @classmethod def from_type(cls, layer_id, args, rope) -> "TransformerBlock": diff --git a/examples/models/llama/norm.py b/examples/models/llama/norm.py index 6777a29bbae..bb242d8dfd9 100644 --- a/examples/models/llama/norm.py +++ b/examples/models/llama/norm.py @@ -103,19 +103,30 @@ 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)``. + + ``zero_centered_gamma``: the checkpoint stores gamma offset by -1, so the + effective scale is ``weight + 1``. + """ 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. @@ -131,6 +142,8 @@ def __init__(self, dim: int, eps: float = 1e-5): dim (int): The dimension of the input tensor. eps (float, optional): Floor on the L2-norm denominator (`clamp_min(‖x‖₂, √(dim·eps))`), matching RMSNormCoreML. Must be > 0. + zero_centered_gamma (bool, optional): checkpoint stores gamma offset + by -1, so the effective scale is `weight + 1`. """ super().__init__() assert eps > 0, ( @@ -139,6 +152,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 +167,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 +212,11 @@ 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) + 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]`.