[Gemma4] Fix the bug of use_qk_norm accidently set to False in Gemma 4 configs

[Dr-Left]

Description

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.

Tests

Smoke test on V5p-256 x 2 slices (Data parallel training of Gemma4-26B-A4B)
Before fix: http://tb/281018541
After fix: http://tb/281083608 log_link

Added unit test:

• tests/unit/gemma4_canonical_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request enables QK normalization (use_qk_norm: true) across several Gemma 4 model configurations (26b, 31b, e2b, and e4b), disables layer scanning for the smaller variants, updates the attention layer's weight initialization logic, and adds a new unit test suite. The review feedback highlights a potential silent correctness bug where a mismatch can occur between weight initialization and the forward pass if self.use_qk_norm and self.config.use_qk_norm are inconsistent. Additionally, the feedback points out that Flax NNX variables in the new unit tests should access their underlying JAX arrays using .value instead of slicing with [...].

[gemini-code-assist]

There is a potential silent correctness bug/inconsistency here.

If a model configuration has use_qk_norm: true (so self.config.use_qk_norm is True), but the attention layer is instantiated without explicitly passing use_qk_norm=True (meaning self.use_qk_norm defaults to False in Attention.__init__), then:

1. self.query_norm and self.key_norm will be None (QK normalization is NOT applied during the forward pass).
2. However, init_query_w will see getattr(self.config, "use_qk_norm", False) as True and set depth_scaling = 1.0 (disabling depth scaling during weight initialization).

This results in a silent mismatch where weights are initialized as if QK normalization is active, but the forward pass does not actually apply QK normalization.

To prevent this, we should ensure that self.use_qk_norm in Attention.__init__ defaults to config.use_qk_norm when not explicitly provided, or at least raise an error/warning if they are inconsistent. Since Attention.__init__ is outside of this diff, please consider updating Attention.__init__ in a follow-up or expanding this PR's scope to make self.use_qk_norm the single source of truth (e.g., self.use_qk_norm = use_qk_norm or getattr(config, "use_qk_norm", False)).

[Dr-Left]

Fixed in new push. default to None for all attentions and override by config.use_qk_norm

[gemini-code-assist]

In Flax NNX, nnx.Variable (and its subclass nnx.Param) does not support direct indexing/slicing via [...]. To access the underlying JAX array, you should use the .value attribute.

Please update this to use kernel.value instead.

Suggested change

	q_kernel = layer.self_attention.query.kernel[...]
	q_kernel = layer.self_attention.query.kernel.value

[Dr-Left]

Replacing kernel[...] with kernel.value would re-introduce DeprecationWarnings across the codebase.

[gemini-code-assist]

In Flax NNX, nnx.Variable (and its subclass nnx.Param) does not support direct indexing/slicing via [...]. To access the underlying JAX array, you should use the .value attribute.

Please update this to use kernel.value instead.

Suggested change

	q_kernel = layer.self_attention.query.kernel[...]
	q_kernel = layer.self_attention.query.kernel.value

[Dr-Left]

Replacing kernel[...] with kernel.value would re-introduce DeprecationWarnings across the codebase.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!