Merge pull request #2148 from AI-Hypercomputer:chengnuojin-add-blocks

Google-ML-Automation · Google-ML-Automation · commit 252d1fc492ef · 2025-08-11T10:34:56.000-07:00
PiperOrigin-RevId: 793707146
diff --git a/MaxText/layers/linears.py b/MaxText/layers/linears.py
@@ -406,6 +406,9 @@ def get_norm_layer(self, num_features: int):
         DecoderBlockType.MISTRAL,
         DecoderBlockType.MIXTRAL,
         DecoderBlockType.GEMMA,
+        DecoderBlockType.GEMMA2,
+        DecoderBlockType.GEMMA3,
+        DecoderBlockType.QWEN3,
         DecoderBlockType.DEEPSEEK,
         DecoderBlockType.LLAMA4,
     ):
