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,