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]`.