Preserve eager numeric parity (#24035)

Differential Revision: D125556968

Pull Request resolved: #24035

Author: YIWENX14

examples/models/llama/llama_transformer.py

@@ -257,6 +257,10 @@ def __init__(
                 eps=args.norm_eps,
                 zero_centered_gamma=args.norm_zero_centered_gamma,
             )
+        elif (
+            getattr(args, "weightless_post_ffn_norm", False) and self.mlp_type != "skip"
+        ):
+            self.post_ffn_norm = ScalelessRMSNorm(args.dim, eps=args.norm_eps)
 
     @classmethod
     def from_type(cls, layer_id, args, rope) -> "TransformerBlock":

examples/models/llama/tests/test_transformer_block.py

@@ -13,10 +13,54 @@
     TransformerBlock,
 )
 from executorch.examples.models.llama.model_args import ModelArgs
+from executorch.examples.models.llama.norm import (
+    RMSNormWithInputScale,
+    ScalelessRMSNorm,
+)
 from executorch.examples.models.llama.rope import Rope
 
 
 class TestTransformerBlock(unittest.TestCase):
+    @staticmethod
+    def _model_args(**overrides) -> ModelArgs:
+        args = {
+            "dim": 4,
+            "hidden_dim": 8,
+            "n_heads": 1,
+            "n_kv_heads": 1,
+            "head_dim": 4,
+            **overrides,
+        }
+        return ModelArgs(**args)
+
+    def test_dedicated_weightless_post_ffn_norm(self) -> None:
+        args = self._model_args()
+        args.weightless_post_ffn_norm = True  # pyrefly: ignore[missing-attribute]
+
+        block = TransformerBlock(args, AttentionSkip()).eval()
+
+        self.assertIsInstance(block.post_ffn_norm, ScalelessRMSNorm)
+
+    def test_learnable_post_ffn_norm_takes_precedence(self) -> None:
+        args = self._model_args(use_ffn_learnable_scales=True)
+        args.weightless_post_ffn_norm = True  # pyrefly: ignore[missing-attribute]
+
+        block = TransformerBlock(args, AttentionSkip()).eval()
+
+        self.assertIsInstance(block.post_ffn_norm, RMSNormWithInputScale)
+
+    def test_shared_post_ffn_norm_does_not_select_weightless_norm(self) -> None:
+        args = self._model_args(post_ffn_norm=True)
+
+        block = TransformerBlock(args, AttentionSkip()).eval()
+
+        self.assertFalse(hasattr(block, "post_ffn_norm"))
+
+    def test_default_has_no_post_ffn_norm(self) -> None:
+        block = TransformerBlock(self._model_args(), AttentionSkip()).eval()
+
+        self.assertFalse(hasattr(block, "post_ffn_norm"))
+
     def test_zero_centered_gamma_applies_to_main_norms(self) -> None:
         args = ModelArgs(
             dim=4,