From efbfcc0b5183527ed119811851ecfccf21c52484 Mon Sep 17 00:00:00 2001
From: Yiwen Xie <yiwenx@meta.com>
Date: Wed, 21 Oct 2026 11:49:56 -0700
Subject: [PATCH] Preserve AvoEdge eager numeric parity (#24035)

Summary:

Preserve rlllama4x query scaling and weightless sandwich-norm semantics through the landed AvoEdge trainer-override path. Keep shared rlformers conversion defaults unchanged, validate legacy head-dimension fallback before division, and gate the dense ExecuTorch scaleless post-FFN norm on a dedicated AvoEdge field so Gemma2/3 shared post_ffn_norm behavior remains unchanged.

Differential Revision: D125556968
---
 examples/models/llama/llama_transformer.py    |  4 ++
 .../llama/tests/test_transformer_block.py     | 44 +++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/examples/models/llama/llama_transformer.py b/examples/models/llama/llama_transformer.py
index 5d0676f6421..ff7db77c2b7 100644
--- a/examples/models/llama/llama_transformer.py
+++ b/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":
diff --git a/examples/models/llama/tests/test_transformer_block.py b/examples/models/llama/tests/test_transformer_block.py
index 6c8ae063870..01fda39bb5e 100644
--- a/examples/models/llama/tests/test_transformer_block.py
+++ b/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,