Rewrite Gemma4 scannable block

Author: aireenmei

src/maxtext/layers/decoders.py

@@ -1426,7 +1426,12 @@ def _apply_gemma4_scanned_blocks(
     if num_full_blocks > 0:
       ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
       policy = self.get_remat_policy()
-      RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0]
+      # Gemma4ScannableBlock rematerializes its own local (scanned) and global
+      # layers when apply_internal_remat=True, so we do NOT wrap it in
+      # block-level remat here (that would double-rematerialize and make XLA
+      # treat the whole block as one unit). Unrolling the block scan lets XLA
+      # free each block's activations across iterations instead of keeping the
+      # block live as a unit.
 
       kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
           kv_caches, num_full_blocks, block_pattern_len, stack=True
@@ -1449,7 +1454,7 @@ def _apply_gemma4_scanned_blocks(
 
       # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks
       y, returned_kv_cache = nn.scan(
-          RemattedGemma4Block,
+          ScannableBlockToLinen,
           variable_axes={
               "params": cfg.param_scan_axis,
               "cache": 0,
@@ -1460,6 +1465,7 @@ def _apply_gemma4_scanned_blocks(
           split_rngs={"params": True, "dropout": cfg.enable_dropout},
           in_axes=in_axes_tuple,
           length=num_full_blocks,
+          unroll=num_full_blocks,
           metadata_params={
               nn.PARTITION_NAME: "layers",
               "abstract_init": False,
@@ -1470,6 +1476,8 @@ def _apply_gemma4_scanned_blocks(
           quant=self.quant,
           model_mode=model_mode,
           num_of_layers=block_pattern_len,
+          remat_policy_fn=policy,
+          apply_internal_remat=True,
           name="scanned_blocks",
       )(
           y, *broadcast_args

src/maxtext/layers/nnx_decoders.py

@@ -646,9 +646,19 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh):
     attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN)
     scan_length = config.num_decoder_layers // attention_pattern_length
     num_remaining_layers = config.num_decoder_layers % attention_pattern_length
-    layer_kwargs = {"num_of_layers": attention_pattern_length}
-
-    rem_layer_kwargs = {"num_of_layers": num_remaining_layers}
+    policy = self.get_remat_policy()
+    # The pure-NNX decoder skips block-level remat (skip_block_remat=True below),
+    # so the block rematerializes its own local/global layers instead.
+    layer_kwargs = {
+        "num_of_layers": attention_pattern_length,
+        "remat_policy_fn": policy,
+        "apply_internal_remat": True,
+    }
+    rem_layer_kwargs = {
+        "num_of_layers": num_remaining_layers,
+        "remat_policy_fn": policy,
+        "apply_internal_remat": True,
+    }
 
     RemattedGemma4Block = gemma4.Gemma4ScannableBlock
 
@@ -870,7 +880,17 @@ def pure_layer_fn(state_in, y_in):
 
     return out
 
-  def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches_stacked=None, **kwargs):
+  def _apply_layers_sequentially(
+      self,
+      layers,
+      x_in,
+      *args,
+      length: int,
+      kv_caches_stacked=None,
+      skip_block_remat: bool = False,
+      unroll: int = 1,
+      **kwargs,
+  ):
     """Runs the layer stack using nnx.scan.
 
     Args:
@@ -881,6 +901,11 @@ def _apply_layers_sequentially(self, layers, x_in, *args, length: int, kv_caches
       kv_caches_stacked: Optional pytree whose leaves have shape [num_layers, ...].
         When provided, the i-th slice is passed as `kv_cache=` to layer i and the
         updated caches are returned as a third element of the tuple.
+      skip_block_remat: When True, do not wrap the scanned body in jax.checkpoint.
+        Used when the scanned module already applies its own (finer-grained,
+        e.g. per-layer) remat internally, to avoid double rematerialization.
+      unroll: Number of scan iterations to unroll into straight-line code
+        (forwarded to jax.lax.scan). unroll >= length fully unrolls the loop.
       **kwargs: Keyword args forwarded to the layer (filtered by the layer signature).
 
     Returns:
@@ -963,7 +988,12 @@ def layer_fn(carry, scanned_vars):
         return new_carry, (new_current_state, updated_kv)
       return new_carry, new_current_state
 
-    layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
+    if skip_block_remat:
+      # The scanned module applies its own remat internally; wrapping the whole
+      # body again would double-remat and recompute the entire block.
+      layer_fn_wrapped = layer_fn
+    else:
+      layer_fn_wrapped = jax.checkpoint(layer_fn, policy=policy, prevent_cse=prevent_cse)
 
     if use_kv:
       # If kv_caches is provided (e.g., from vLLM), we CANNOT use jax.lax.scan
@@ -995,7 +1025,7 @@ def layer_fn(carry, scanned_vars):
       params = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(params, length)
       state = maxtext_utils_nnx.nnx_ensure_scan_leading_axis(state, length)
 
-      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state))
+      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, (params, state), unroll=unroll)
       returned_kv_stacked = None
 
       # Ensure metadata rank matches the stacked values
@@ -1938,13 +1968,25 @@ def _apply_gemma4_scanned_blocks(
     attention_pattern_length = len(gemma4.GEMMA4_ATTENTION_PATTERN)
     scan_length = cfg.num_decoder_layers // attention_pattern_length
 
-    # Apply the main scan over the full blocks
+    # Apply the main scan over the full blocks. Gemma4ScannableBlock applies
+    # per-layer remat internally (local scan + global layer), so skip the
+    # block-level remat here to avoid double rematerialization. Unrolling the
+    # block loop (one iteration per repeated block) lets XLA pipeline/free block
+    # activations across iterations (memory + overlap knob).
+    block_unroll = max(1, scan_length)
     if scan_length > 0:
       grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(
           kv_caches, scan_length, attention_pattern_length, stack=False
       )
       y, self.scanned_blocks, _ = self._apply_layers_sequentially(
-          self.scanned_blocks, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs
+          self.scanned_blocks,
+          y,
+          *layer_args,
+          length=scan_length,
+          kv_caches_stacked=grouped_kv_caches,
+          skip_block_remat=True,
+          unroll=block_unroll,
+          **layer_kwargs,
       )
       maxtext_utils.update_kv_caches_after_scan(
           kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False

src/maxtext/layers/nnx_scan.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Utilities for constructing stacks of scanned NNX layers."""
+"""Utilities for constructing and applying stacks of scanned NNX layers."""
 
 from collections.abc import Callable
+from typing import Any
 
 from flax import nnx
 import jax
@@ -87,3 +88,50 @@ def update_leaf(leaf):
   stacked_params = add_scan_metadata(stacked_params, param_scan_axis)
   stacked_rest = add_scan_metadata(stacked_rest, 0)
   return nnx.merge(layer_graphdef, stacked_params, stacked_rest)
+
+
+def apply_scanned_layers(
+    layers: nnx.Module,
+    carry: Any,
+    *,
+    length: int,
+    param_scan_axis: int,
+    apply_fn: Callable[[nnx.Module, Any], Any],
+    remat: bool = False,
+    remat_policy: Callable[..., Any] | None = None,
+    prevent_cse: bool = True,
+    unroll: int = 1,
+) -> Any:
+  """Applies stacked NNX layers using ``jax.lax.scan``.
+
+  This helper owns the generic NNX state and scan-axis mechanics. ``apply_fn``
+  defines the model-specific module invocation and must return the next carry.
+  ``remat`` is separate from ``remat_policy`` because ``None`` is JAX's full
+  rematerialization policy, not an indication that rematerialization is off.
+
+  Externally managed per-layer state, such as KV caches, is not supported by
+  this scan path.
+  """
+  if length <= 0:
+    return carry
+
+  layer_graphdef, params, state = nnx.split(layers, nnx.Param, ...)
+  if param_scan_axis != 0:
+    params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params)
+
+  def scan_body(current_carry, scanned_state):
+    current_params, current_state = scanned_state
+    current_layer = nnx.merge(layer_graphdef, current_params, current_state)
+    next_carry = apply_fn(current_layer, current_carry)
+    return next_carry, nnx.state(current_layer)
+
+  scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body
+  final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), unroll=unroll)
+
+  if param_scan_axis != 0:
+    scanned_params, scanned_other = scanned_state.split(nnx.Param, ...)
+    scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_params)
+    scanned_state = nnx.State.merge(scanned_params, scanned_other)
+
+  nnx.update(layers, scanned_state)
+  return final_carry