From f85570dc192d6803e2de2dc2837f3d2c9efc2253 Mon Sep 17 00:00:00 2001 From: Chris Zuo Date: Tue, 18 Aug 2026 23:18:27 +0000 Subject: [PATCH] [Gemma4] Enable use_qk_norm in Gemma 4 configs and fix attention weight init depth_scaling This change resolves an issue with Gemma 4 attention weight initialization and QK normalization: 1. In `src/maxtext/layers/attentions.py`, `init_query_w` now checks `self.use_qk_norm` in addition to `self.config.use_qk_norm`. When `use_qk_norm` is enabled (or `query_pre_attn_scalar != 1.0`), `depth_scaling` is set to 1.0 instead of `sqrt(head_dim)`. This prevents dividing the initial query projection weights by `sqrt(head_dim)`, ensuring unscaled canonical weight initialization for Gemma 4. 2. In Gemma 4 model configs (`gemma4-26b.yml`, `gemma4-31b.yml`, `gemma4-e2b.yml`, `gemma4-e4b.yml`), explicitly enabled `use_qk_norm: true`. 3. In `gemma4-e2b.yml` and `gemma4-e4b.yml`, set `scan_layers: false` since per-layer embeddings and variable layer configs are not compatible with layer scanning. 4. Added `tests/unit/gemma4_canonical_test.py` to verify that canonical attention scaling (`query_pre_attn_scalar = 1.0`), `use_qk_norm`, and unscaled weight initialization hold for both dense/MoE Gemma 4 and small Gemma 4 models. Added unit test: - `tests/unit/gemma4_canonical_test.py` TAG=agy CONV=f731b620-dff8-4a9f-b791-657890c4604f --- src/maxtext/configs/models/gemma4-26b.yml | 1 + src/maxtext/configs/models/gemma4-31b.yml | 1 + src/maxtext/configs/models/gemma4-e2b.yml | 2 + src/maxtext/configs/models/gemma4-e4b.yml | 2 + src/maxtext/layers/attention_compressed.py | 4 +- src/maxtext/layers/attention_mla.py | 4 +- src/maxtext/layers/attentions.py | 10 +-- tests/unit/gemma4_canonical_test.py | 91 ++++++++++++++++++++++ 8 files changed, 105 insertions(+), 10 deletions(-) create mode 100644 tests/unit/gemma4_canonical_test.py diff --git a/src/maxtext/configs/models/gemma4-26b.yml b/src/maxtext/configs/models/gemma4-26b.yml index b20397dee7..8f764e5caf 100644 --- a/src/maxtext/configs/models/gemma4-26b.yml +++ b/src/maxtext/configs/models/gemma4-26b.yml @@ -39,6 +39,7 @@ global_rope_proportion: 0.25 local_rope_proportion: 1.0 v_norm_with_scale: false final_logits_soft_cap: 30.0 +use_qk_norm: true # MoE configuration num_experts: 128 diff --git a/src/maxtext/configs/models/gemma4-31b.yml b/src/maxtext/configs/models/gemma4-31b.yml index 7e97f9dba5..61e7324662 100644 --- a/src/maxtext/configs/models/gemma4-31b.yml +++ b/src/maxtext/configs/models/gemma4-31b.yml @@ -42,6 +42,7 @@ rope_max_timescale: 1000000 global_rope_proportion: 0.25 local_rope_proportion: 1.0 final_logits_soft_cap: 30.0 +use_qk_norm: true # Multimodal flags (need to set use_multimodal=true) vision_encoder_block: "gemma4" diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml index 81c8d4ea66..ca3e1c66e6 100644 --- a/src/maxtext/configs/models/gemma4-e2b.yml +++ b/src/maxtext/configs/models/gemma4-e2b.yml @@ -43,6 +43,8 @@ rope_max_timescale: 1000000 global_rope_proportion: 0.25 local_rope_proportion: 1.0 final_logits_soft_cap: 30.0 +use_qk_norm: true +scan_layers: false # Vision encoder flags — multimodal not yet supported for E2B / E4B. vision_encoder_block: "gemma4" diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml index b8da0d1d46..3ebca9b1ba 100644 --- a/src/maxtext/configs/models/gemma4-e4b.yml +++ b/src/maxtext/configs/models/gemma4-e4b.yml @@ -44,6 +44,8 @@ rope_max_timescale: 1000000 global_rope_proportion: 0.25 local_rope_proportion: 1.0 final_logits_soft_cap: 30.0 +use_qk_norm: true +scan_layers: false # Vision encoder flags — multimodal not yet supported for E2B / E4B. vision_encoder_block: "gemma4" diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index 23a1e3b539..30c0d5eba9 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -1110,7 +1110,7 @@ def __init__( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, # Compression Specific Parameters: @@ -1606,7 +1606,7 @@ def compressed_attention( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, q_lora_rank: int = 1536, diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index 3fc1a3c69d..6ed5b23049 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -477,7 +477,7 @@ def mla_as_linen( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections # Temperature tuning parameters used for Llama4 @@ -615,7 +615,7 @@ def __init__( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections # Temperature tuning parameters used for Llama4 diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index a2f4b48afd..fccd2ddec1 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -126,7 +126,7 @@ def attention_as_linen( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections share_kv_projections: bool = False, # If true, Key and Value use the same projection @@ -283,7 +283,7 @@ def __init__( sliding_window_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, - use_qk_norm: bool = False, + use_qk_norm: bool | None = None, query_pre_attn_scalar: float | None = None, use_bias_in_projections: bool = False, # Set to True will enable bias in q, k, v, o projections share_kv_projections: bool = False, # If true, Key and Value use the same projection @@ -394,7 +394,7 @@ def __init__( self.sliding_window_size = sliding_window_size self.use_ragged_attention = use_ragged_attention self.ragged_block_size = ragged_block_size - self.use_qk_norm = use_qk_norm + self.use_qk_norm = getattr(self.config, "use_qk_norm", False) if use_qk_norm is None else use_qk_norm self.query_pre_attn_scalar = query_pre_attn_scalar self.use_bias_in_projections = use_bias_in_projections self.share_kv_projections = share_kv_projections @@ -665,9 +665,7 @@ def init_query_w(self, inputs_q_shape: Tuple) -> nnx.Module: # linear transformations, which is equivalent under Adafactor. # We disable depth_scaling when using qk_norm or a query_pre_attn_scalar # to avoid applying scaling twice. - if getattr(self.config, "use_qk_norm", False) or ( - self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar != 1.0 - ): + if self.use_qk_norm or (self.query_pre_attn_scalar is not None and self.query_pre_attn_scalar != 1.0): depth_scaling = 1.0 else: depth_scaling = jnp.sqrt(self.head_dim).astype(self.dtype) diff --git a/tests/unit/gemma4_canonical_test.py b/tests/unit/gemma4_canonical_test.py new file mode 100644 index 0000000000..615053da41 --- /dev/null +++ b/tests/unit/gemma4_canonical_test.py @@ -0,0 +1,91 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for canonical Gemma 4 attention scaling and weight initialization.""" + +import os +import unittest +import jax +import jax.numpy as jnp +from flax import nnx + +from maxtext.configs import pyconfig +from maxtext.common import common_types +from maxtext.models import gemma4, gemma4_small +from maxtext.utils.globals import MAXTEXT_REPO_ROOT + + +class Gemma4CanonicalAttentionTest(unittest.TestCase): + """Tests that Gemma 4 models follow canonical attention scaling and unscaled weight init.""" + + def setUp(self): + super().setUp() + self.base_config_path = os.path.join(MAXTEXT_REPO_ROOT, "src", "maxtext", "configs", "base.yml") + + def test_gemma4_26b_attention_config(self): + config = pyconfig.initialize( + ["", self.base_config_path], + model_name="gemma4-26b", + enable_dropout=False, + ) + self.assertTrue(config.use_qk_norm, "gemma4-26b should enable use_qk_norm in config") + + mesh = jax.sharding.Mesh(jax.devices()[:1], ("data",)) + rngs = nnx.Rngs(0) + layer = gemma4.Gemma4DecoderLayer( + config=config, + mesh=mesh, + model_mode=common_types.MODEL_MODE_PREFILL, + rngs=rngs, + attention_type=gemma4.AttentionType.LOCAL_SLIDING, + layer_idx=0, + ) + # Canonical Gemma 4 attention uses query_pre_attn_scalar = 1.0 (unscaled logits) + self.assertEqual(layer.self_attention.query_pre_attn_scalar, 1.0) + self.assertTrue(layer.self_attention.use_qk_norm) + + # Initial query weights should NOT be divided by sqrt(head_dim) = 16.0 + # Expected standard deviation for fan_in=2816 is 1/sqrt(2816) ~= 0.0188 + q_kernel = layer.self_attention.query.kernel[...] + std_q = float(jnp.std(q_kernel)) + self.assertGreater(std_q, 0.01, f"Query kernel std ({std_q}) should not be divided by depth_scaling") + + def test_gemma4_small_attention_config(self): + for model_name in ["gemma4-e2b", "gemma4-e4b"]: + config = pyconfig.initialize( + ["", self.base_config_path], + model_name=model_name, + enable_dropout=False, + ) + self.assertTrue(config.use_qk_norm, f"{model_name} should enable use_qk_norm in config") + + mesh = jax.sharding.Mesh(jax.devices()[:1], ("data",)) + rngs = nnx.Rngs(0) + layer = gemma4_small.Gemma4SmallDecoderLayer( + config=config, + mesh=mesh, + model_mode=common_types.MODEL_MODE_PREFILL, + layer_idx=0, + rngs=rngs, + ) + self.assertEqual(layer.self_attention.query_pre_attn_scalar, 1.0) + self.assertTrue(layer.self_attention.use_qk_norm) + + q_kernel = layer.self_attention.query.kernel[...] + std_q = float(jnp.std(q_kernel)) + self.assertGreater(std_q, 0.01, f"Query kernel std ({std_q}) should not be divided by depth_scaling") + + +if __name__ == "__main__": + unittest.main()