Rewrite Gemma4 scannable block

[aireenmei]

Description

Currently theimplementation of Gemma4ScannableBlock makes XLA treat the whole 6-layer block as a "big layer". The all-gathers for all 6 layers are run at the beginning of each block. Similarly, in backward pass, remat runs the whole block. This causes elevated HBM usage and more exposed collectives.

After this rewrite:

• 5 local layers via jax.lax.scan, each per-layer jax.checkpoint
• length-1 scan for global layer, also wrapped in jax.checkpoint
• remove block-level remat and unroll outer block scan

Tests

• added unit test
• Before rewrite: xprof for pdbs=4. xprof for pdbs=8
• After rewrite: xprof for pdbs=4. xprof for pdbs=8 + vocab_tiling=8 (OOM without vocab_tiling).
• Before rewrite, the HBM peak is not on vocab peak(screenshot). After rewrite, vocab peak becomes the highest (screenshot) due to other peaks become smaller, also by having remat boundary at layer level adds HBM usage and overlap with the vocab peak. Therefore, without vocab tiling, the peak HBM usage doesn't change much compared with original. But when combined with vocab_tiling, we can effectively lower the HBM peak a lot (70 GB -> 40 GB for pdbs=4).

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.

[github-actions]

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

• Job/Matrix: code_quality_check / Static code-quality checkers
• Failing Test: pre-commit checks on just the files that have changed
• Error: Pylint and pyink code style violations

🪵 Error Details & Stack Trace

# Missing function docstring in Gemma4ScannableBlock helper method
src/maxtext/models/gemma4.py:496:2: C0116: Missing function or method docstring (missing-function-docstring)

# Line too long violations in model_test.py
tests/unit/model_test.py:238:0: C0301: Line too long (131/125) (line-too-long)
tests/unit/model_test.py:283:0: C0301: Line too long (131/125) (line-too-long)

# Formatting non-compliance reported by pyink
would reformat tests/unit/model_test.py
would reformat src/maxtext/models/gemma4.py
would reformat tests/unit/nnx_scan_test.py
would reformat tests/unit/nnx_decoders_test.py

💡 Root Cause Analysis & Context

Confidence: high (confirmed cause)

The failure is caused by code style and quality violations introduced in this PR. Specifically:

1. In src/maxtext/models/gemma4.py, the new helper method _apply_local_layers was added without a docstring, triggering the pylint missing-function-docstring check.
2. In tests/unit/model_test.py, the lines asserting logits shapes exceed the maximum length of 125 characters configured for pylint (each was 131 characters).
3. The files tests/unit/model_test.py, src/maxtext/models/gemma4.py, tests/unit/nnx_scan_test.py, and tests/unit/nnx_decoders_test.py had spacing and line wraps that did not comply with the pyink formatting rules (indentation 2, line length 122).

These issues are purely style/quality regressions and do not affect behavioral correctness. They have been fully corrected and verified locally using pre-commit run.

🛠️ Recommended Fix

Apply the formatting changes and add the missing docstring to src/maxtext/models/gemma4.py. The required diffs are:

diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py
index b974ec5..f626b8e 100644
--- a/src/maxtext/models/gemma4.py
+++ b/src/maxtext/models/gemma4.py
@@ -453,9 +453,7 @@ class Gemma4ScannableBlock(nnx.Module):
 
     pattern_length = len(GEMMA4_ATTENTION_PATTERN)
     if not 0 <= num_of_layers <= pattern_length:
-      raise ValueError(
-          f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}."
-      )
+      raise ValueError(f"Gemma4ScannableBlock must contain between 0 and {pattern_length} layers; got {num_of_layers}.")
 
     # Pattern is 5 local, 1 global.
     self.num_local = min(5, num_of_layers)
@@ -505,6 +503,8 @@ class Gemma4ScannableBlock(nnx.Module):
       bidirectional_mask=None,
       attention_metadata=None,
   ):
+    """Applies local attention layers sequentially using scan."""
+
     def apply_layer(layer, carry):
       layer_out = layer(
           carry,
@@ -654,9 +652,7 @@ class Gemma4ScannableBlock(nnx.Module):
         offload_names = maxtext_utils.get_save_and_offload_names(cfg)
         if offload_names[0] or offload_names[1]:
           save_names, offload_to_device = offload_names
-          global_remat_policy = jax.checkpoint_policies.save_only_these_names(
-              *(save_names + offload_to_device)
-          )
+          global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device))
 
         if self.apply_internal_remat and self.config.remat_policy != "none":
           prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config)
@@ -669,9 +667,7 @@ class Gemma4ScannableBlock(nnx.Module):
         # Carry state through the loop instead of returning a stacked [1, ...]
         # scan result: slicing that result previously introduced a bitcast
         # between device and pinned-host memory under offload remat.
-        with xla_metadata.set_xla_metadata(
-            **{"skip-simplify-while-loops_trip-count-one": "true"}
-        ):
+        with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}):
           (y, global_state), _ = jax.lax.scan(
               scan_global_layer,
               (y, state_g),
diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py
index 7c3d4ea..dd679ff 100644
--- a/tests/unit/model_test.py
+++ b/tests/unit/model_test.py
@@ -235,7 +235,9 @@ class TestModel(unittest.TestCase):
         enable_dropout=False,
         model_mode=MODEL_MODE_TRAIN,
     )
-    self.assertEqual(logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size))
+    self.assertEqual(
+        logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size)
+    )
 
   def test_gemma4_model_linen(self):
     """Test the shared Gemma4 scannable block on the linen (ToLinen) path.
@@ -280,7 +282,9 @@ class TestModel(unittest.TestCase):
         rngs={"aqt": self.rng},
     )
     logits = logits[0] if isinstance(logits, tuple) else logits
-    self.assertEqual(logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size))
+    self.assertEqual(
+        logits.shape, (new_config.global_batch_size_to_train_on, new_config.max_target_length, new_config.vocab_size)
+    )

[codecov]

Codecov Report

❌ Patch coverage is 79.79798% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/models/gemma4.py	80.00%	6 Missing and 8 partials ⚠️
src/maxtext/layers/nnx_scan.py	80.00%	1 Missing and 3 partials ⚠️
src/maxtext/layers/nnx_decoders.py	77.77%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[aireenmei]

I add pull ready to generate internal cl for testing. Won't submit until PR approved

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @RissyRan, but I was unable to process your request. Please see the logs for more details.

[RissyRan]

Thanks for the change! Overall it looks good. I am a little bit lost in gemma4.py if/else branches, wondering if we could simplify a little bit?

[RissyRan]

It seems the pattern has been defined. Will something like this work?

    active_pattern = GEMMA4_ATTENTION_PATTERN[:num_of_layers]
    self.num_local = sum(1 for t in active_pattern if t == AttentionType.LOCAL_SLIDING)
    self.num_global = sum(1 for t in active_pattern if t == AttentionType.GLOBAL)

[RissyRan]

Will this be triggered for Gemma4?

[RissyRan]

I guess you meet issue when calling kv_cache[layer_id]?

Wondering if you tested for decoding with the change end-to-end and see reasonable outputs?

[RissyRan]

Is this the same as line 564?