diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index a915f7f85d..244ff3cf27 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -411,36 +411,15 @@ def get_remat_policy(self): "kv_proj", "qkv_proj", ) - elif cfg.remat_policy == "qkv_proj_offloaded": + elif cfg.remat_policy in ("qkv_proj_offloaded", "minimal_offloaded", "custom"): + # minimal_offloaded offloads all except context. All three share a single + # source of truth for their save/offload name lists (see + # maxtext_utils.get_offload_remat_names) so that offloading configured via + # `custom` resolves identically to the named presets. + save_names, offload_names = maxtext_utils.get_offload_remat_names(cfg) policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=[], - names_which_can_be_offloaded=["query_proj", "value_proj", "key_proj", "kv_proj"], - offload_src="device", - offload_dst="pinned_host", - ) - elif cfg.remat_policy == "minimal_offloaded": - # offload all except context - policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=[], - names_which_can_be_offloaded=[ - "query_proj", - "value_proj", - "key_proj", - "kv_proj", - "qkv_proj", - "out_proj", - "mlpwi_0", - "mlpwi_1", - "mlpwi", - "mlpwo", - ], - offload_src="device", - offload_dst="pinned_host", - ) - elif cfg.remat_policy == "custom": - policy = jax.checkpoint_policies.save_and_offload_only_these_names( - names_which_can_be_saved=cfg.tensors_on_device, - names_which_can_be_offloaded=cfg.tensors_to_offload, + names_which_can_be_saved=save_names, + names_which_can_be_offloaded=offload_names, offload_src="device", offload_dst="pinned_host", ) @@ -1433,7 +1412,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 @@ -1456,7 +1440,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, @@ -1467,6 +1451,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, @@ -1477,6 +1462,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 diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 528523e438..47371b9759 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -37,7 +37,7 @@ ShardMode, ) from maxtext.layers import initializers, linears, mhc, normalizations, quantizations -from maxtext.layers import nnx_wrappers +from maxtext.layers import nnx_scan, nnx_wrappers from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding from maxtext.layers.normalizations import RMSNorm @@ -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 @@ -838,95 +848,20 @@ def _create_scanned_layers( rngs: nnx.Rngs, **layer_kwargs, ): - """Creates a scanned stack of layers using jax.lax.scan for memory-efficient initialization.""" - if length == 0: - return None - scan_axis = self.config.param_scan_axis - - # Fork rngs to get per-layer RNG states for scanning - try: - forked_rngs = rngs.fork(split=length) - except: # pylint: disable=bare-except - pass - - rngs_graphdef, rngs_state = nnx.split(forked_rngs) - - first_rng_state = jax.tree.map(lambda x: x[0], rngs_state) - ref_rngs = nnx.merge(rngs_graphdef, first_rng_state) - ref_layer = decoder_layer_class( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - rngs=ref_rngs, - **layer_kwargs, + return nnx_scan.create_scanned_layers( + lambda layer_rngs: decoder_layer_class( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + rngs=layer_rngs, + **layer_kwargs, + ), + length=length, + param_scan_axis=self.config.param_scan_axis, + metadata_axis_name=metadata_axis_name, + rngs=rngs, ) - layer_graphdef, _, _ = nnx.split(ref_layer, nnx.Param, ...) - del ref_layer - - def scan_body(carry, rng_state_slice): - layer_rngs = nnx.merge(rngs_graphdef, rng_state_slice) - layer = decoder_layer_class( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - rngs=layer_rngs, - **layer_kwargs, - ) - _, params, rest = nnx.split(layer, nnx.Param, ...) - return carry, (params, rest) - - _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state) - - if scan_axis != 0: - stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), stacked_params) - - def _add_scan_metadata(state, axis): - def _update_leaf(leaf): - if hasattr(leaf, "replace") and hasattr(leaf, "value"): - replace_kwargs = {} - if hasattr(leaf, "get_metadata"): - replace_kwargs.update(leaf.get_metadata()) - - replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name - replace_kwargs["param_scan_axis"] = axis - - for key in [ - "sharding", - "out_sharding", - "kernel_axes", - "sharding_names", - ]: - val = getattr(leaf, key, None) - if val is None and key in replace_kwargs: - val = replace_kwargs[key] - - if val is not None: - if isinstance(val, str): - val = (val,) - if isinstance(val, tuple): - l = list(val) - # Safely insert the scan axis into the logical axes string - if metadata_axis_name not in l: - insert_idx = min(axis, len(l)) - l.insert(insert_idx, metadata_axis_name) - replace_kwargs[key] = tuple(l) - - return leaf.replace(**replace_kwargs) - return leaf - - # We must use a custom is_leaf to catch the VariableState instances - return jax.tree.map( - _update_leaf, - state, - is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"), - ) - - stacked_params = _add_scan_metadata(stacked_params, scan_axis) - stacked_rest = _add_scan_metadata(stacked_rest, 0) - - return nnx.merge(layer_graphdef, stacked_params, stacked_rest) def _apply_layer_with_remat(self, layer: nnx.Module, y: jax.Array, policy: Any, prevent_cse: bool, **kwargs): """Helper to cleanly apply jax.checkpoint to a single unscanned layer or block.""" @@ -944,7 +879,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: @@ -955,6 +900,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: @@ -1037,7 +987,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 @@ -1069,7 +1024,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 @@ -2032,13 +1987,25 @@ def _apply_gemma4_scanned_blocks( if attention_metadata is not None: layer_kwargs["attention_metadata"] = attention_metadata - # 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 diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py new file mode 100644 index 0000000000..b675c65357 --- /dev/null +++ b/src/maxtext/layers/nnx_scan.py @@ -0,0 +1,137 @@ +# Copyright 2023–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. + +"""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 +import jax.numpy as jnp + + +def create_scanned_layers( + layer_factory: Callable[[nnx.Rngs], nnx.Module], + *, + length: int, + param_scan_axis: int, + metadata_axis_name: str, + rngs: nnx.Rngs, +) -> nnx.Module | None: + """Constructs an NNX layer whose variables are stacked for a layer scan.""" + if length == 0: + return None + + forked_rngs = rngs.fork(split=length) + rngs_graphdef, rngs_state = nnx.split(forked_rngs) + + first_rng_state = jax.tree.map(lambda x: x[0], rngs_state) + reference_layer = layer_factory(nnx.merge(rngs_graphdef, first_rng_state)) + layer_graphdef, _, _ = nnx.split(reference_layer, nnx.Param, ...) + del reference_layer + + def scan_body(carry, rng_state_slice): + layer = layer_factory(nnx.merge(rngs_graphdef, rng_state_slice)) + _, params, rest = nnx.split(layer, nnx.Param, ...) + return carry, (params, rest) + + _, (stacked_params, stacked_rest) = jax.lax.scan(scan_body, None, rngs_state) + + if param_scan_axis != 0: + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), stacked_params) + + def add_scan_metadata(state, axis): + def update_leaf(leaf): + if hasattr(leaf, "replace") and hasattr(leaf, "value"): + replace_kwargs = {} + if hasattr(leaf, "get_metadata"): + replace_kwargs.update(leaf.get_metadata()) + + replace_kwargs[nnx.PARTITION_NAME] = metadata_axis_name + replace_kwargs["param_scan_axis"] = axis + + for key in ["sharding", "out_sharding", "kernel_axes", "sharding_names"]: + value = getattr(leaf, key, None) + if value is None and key in replace_kwargs: + value = replace_kwargs[key] + + if value is not None: + if isinstance(value, str): + value = (value,) + if isinstance(value, tuple): + logical_axes = list(value) + if metadata_axis_name not in logical_axes: + logical_axes.insert(min(axis, len(logical_axes)), metadata_axis_name) + replace_kwargs[key] = tuple(logical_axes) + + return leaf.replace(**replace_kwargs) + return leaf + + return jax.tree.map( + update_leaf, + state, + is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"), + ) + + 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 diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index dbf8efc4de..670bf0af3b 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -15,18 +15,19 @@ """Specialized layers for Gemma 4.""" import jax +from jax.experimental import xla_metadata from jax.ad_checkpoint import checkpoint_name from jax.sharding import Mesh import jax.numpy as jnp from flax import linen as nn from flax import nnx -from typing import Optional +from typing import Optional, Any from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL from maxtext.layers import initializers from maxtext.layers import moe -from maxtext.layers import nnx_wrappers +from maxtext.layers import nnx_scan, nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock @@ -35,6 +36,7 @@ from maxtext.layers.normalizations import RMSNorm from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils GEMMA4_ATTENTION_PATTERN = ( @@ -409,7 +411,7 @@ def update_cache(cache, val): class Gemma4ScannableBlock(nnx.Module): - """A repeatable block of Gemma4 decoder layers.""" + """A repeatable block of Gemma4 decoder layers, scanning local layers.""" def __init__( self, @@ -418,7 +420,9 @@ def __init__( model_mode: str, rngs: nnx.Rngs, quant: None | Quant = None, - num_of_layers: int = 1, + num_of_layers: int = 6, + remat_policy_fn: Any = None, + apply_internal_remat: bool = False, ): """Initializes the instance. @@ -429,6 +433,14 @@ def __init__( rngs: The random number generators for initialization. quant: The quantization configuration. num_of_layers: The number of layers in the model. + remat_policy_fn: The resolved rematerialization policy function. + apply_internal_remat: When True, the block rematerializes its own local + (scanned) and global layers, and the caller must NOT also apply + block-level remat (that would double-rematerialize and make XLA treat the + whole block as one unit). Both the pure-NNX and linen decoders set this + and skip block-level remat, so remat happens per layer rather than over + the whole block. When False, the block does not self-remat and relies on + the caller's block-level remat instead. """ self.config = config self.mesh = mesh @@ -436,20 +448,95 @@ def __init__( self.quant = quant self.rngs = rngs self.num_of_layers = num_of_layers + self.remat_policy_fn = remat_policy_fn + self.apply_internal_remat = apply_internal_remat - for layer_id in range(self.num_of_layers): - attention_type = get_attention_type(layer_id) - layer_name = f"layers_{layer_id}" - layer = Gemma4DecoderLayer( + 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}." + ) + + # Pattern is 5 local, 1 global. + self.num_local = min(5, num_of_layers) + self.num_global = max(0, num_of_layers - 5) + + if self.num_local > 0: + self.local_layers = nnx_scan.create_scanned_layers( + lambda layer_rngs: Gemma4DecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + rngs=layer_rngs, + attention_type=AttentionType.LOCAL_SLIDING, + layer_idx=0, # layer_idx is not used in the class + ), + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + metadata_axis_name="local_layers", + rngs=self.rngs, + ) + else: + self.local_layers = None + + if self.num_global > 0: + self.global_layer = Gemma4DecoderLayer( config=self.config, mesh=self.mesh, model_mode=self.model_mode, rngs=self.rngs, quant=self.quant, - attention_type=attention_type, - layer_idx=layer_id, + attention_type=AttentionType.GLOBAL, + layer_idx=5, # layer_idx is not used in the class ) - setattr(self, layer_name, layer) + else: + self.global_layer = None + + def _apply_local_layers( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot=None, + previous_chunk=None, + bidirectional_mask=None, + attention_metadata=None, + ): + def apply_layer(layer, carry): + layer_out = layer( + carry, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot=slot, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + attention_metadata=attention_metadata, + ) + return layer_out[0] if isinstance(layer_out, tuple) else layer_out + + apply_remat = self.apply_internal_remat and self.config.remat_policy != "none" + if apply_remat: + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config) + remat_policy = self.remat_policy_fn + else: + prevent_cse = True + remat_policy = None + + return nnx_scan.apply_scanned_layers( + self.local_layers, + y, + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + apply_fn=apply_layer, + remat=apply_remat, + remat_policy=remat_policy, + prevent_cse=prevent_cse, + ) def __call__( self, @@ -471,22 +558,128 @@ def __call__( y = inputs updated_kvs = [] - for layer_id in range(self.num_of_layers): - current_kv = kv_cache[layer_id] if kv_cache is not None else None - y, new_kv = getattr(self, f"layers_{layer_id}")( + + if kv_cache is not None: + # External per-layer KV caches require statically slicing the scanned local layers. + if self.local_layers is not None: + graphdef, params, state = nnx.split(self.local_layers, nnx.Param, ...) + scan_axis = self.config.param_scan_axis + if scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params) + per_layer_states = [] + for i in range(self.num_local): + current_params = jax.tree.map(lambda x, i=i: x[i], params) + current_state = jax.tree.map(lambda x, i=i: x[i], state) + layer = nnx.merge(graphdef, current_params, current_state) + y, new_kv = layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + slot=slot, + previous_chunk=previous_chunk, + bidirectional_mask=bidirectional_mask, + kv_cache=kv_cache[i], + attention_metadata=attention_metadata, + ) + updated_kvs.append(new_kv) + per_layer_states.append(nnx.state(layer)) + + stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) + if scan_axis != 0: + stacked_params, stacked_other = stacked_state.split(nnx.Param, ...) + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params) + stacked_state = nnx.State.merge(stacked_params, stacked_other) + nnx.update(self.local_layers, stacked_state) + elif self.local_layers is not None: + y = self._apply_local_layers( y, decoder_segment_ids, decoder_positions, deterministic, model_mode, - previous_chunk=previous_chunk, slot=slot, + previous_chunk=previous_chunk, bidirectional_mask=bidirectional_mask, - kv_cache=current_kv, attention_metadata=attention_metadata, ) + + if self.global_layer is not None: if kv_cache is not None: + y, new_kv = self.global_layer( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + bidirectional_mask=bidirectional_mask, + kv_cache=kv_cache[self.num_local], + attention_metadata=attention_metadata, + ) updated_kvs.append(new_kv) + else: + graphdef_g, state_g = nnx.split(self.global_layer) + + def scan_global_layer(carry, _): + current_y, current_state = carry + layer = nnx.merge(graphdef_g, current_state) + layer_out = layer( + current_y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk=previous_chunk, + slot=slot, + bidirectional_mask=bidirectional_mask, + attention_metadata=attention_metadata, + ) + new_y = layer_out[0] if isinstance(layer_out, tuple) else layer_out + return (new_y, nnx.state(layer)), None + + # Remat the global layer, kept behind the trip-count-one while boundary + # below. The boundary bounds its full-sequence-attention activations to one + # layer live; without it (blocks are unrolled) XLA co-schedules every + # block's backward working set and OOMs. + # + # Offloaded (pinned-host) residuals cannot cross that boundary: XLA stacks + # them over the length-1 scan axis in host memory and unstacks them in the + # backward via a squeeze that lowers to a host<->device bitcast, which is + # rejected. So the global layer saves would-be-offloaded tensors on device + # instead; the local-layer scan (a real multi-iteration scan) still offloads. + global_remat_policy = self.remat_policy_fn + offload_names = maxtext_utils.get_offload_remat_names(cfg) + if offload_names is not None and (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) + ) + + if self.apply_internal_remat and self.config.remat_policy != "none": + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config) + scan_global_layer = jax.checkpoint( + scan_global_layer, + policy=global_remat_policy, + prevent_cse=prevent_cse, + ) + + # 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"} + ): + (y, global_state), _ = jax.lax.scan( + scan_global_layer, + (y, state_g), + xs=None, + length=1, + ) + + nnx.update(self.global_layer, global_state) if kv_cache is not None: return y, tuple(updated_kvs) diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 92a50dc887..b37cdaf7e1 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -195,6 +195,39 @@ def should_prevent_cse_in_remat(config): return True +def get_offload_remat_names(config): + """Returns ``(save_names, offload_names)`` for offloading remat policies, else ``None``. + + Single source of truth for which checkpointed tensors a remat policy saves on + device vs. offloads to pinned host. Shared by ``Decoder.get_remat_policy`` (which + builds the save-and-offload policy) and by models that must convert an offload + into a device-save at a scan boundary that cannot carry pinned-host residuals + (e.g. Gemma4's global-layer trip-count-one scan). Keeping both callers on this + helper ensures ``remat_policy=qkv_proj_offloaded`` and ``remat_policy=custom`` + with the same tensors marked ``offload`` resolve to identical name sets. + + Returns ``None`` for non-offloading policies. + """ + if config.remat_policy == "qkv_proj_offloaded": + return [], ["query_proj", "value_proj", "key_proj", "kv_proj"] + if config.remat_policy == "minimal_offloaded": + return [], [ + "query_proj", + "value_proj", + "key_proj", + "kv_proj", + "qkv_proj", + "out_proj", + "mlpwi_0", + "mlpwi_1", + "mlpwi", + "mlpwo", + ] + if config.remat_policy == "custom": + return list(config.tensors_on_device or []), list(config.tensors_to_offload or []) + return None + + def load_compiled(config, partial_train, state, execution_devices): """# Loading a serialized compiled train step function.""" diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py index 41cad97198..7c3d4eadb5 100644 --- a/tests/unit/model_test.py +++ b/tests/unit/model_test.py @@ -18,6 +18,7 @@ import jax import jax.numpy as jnp +from flax import nnx from jax.sharding import Mesh from maxtext.configs import pyconfig from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN @@ -42,19 +43,22 @@ def setUp(self): def init_pyconfig(self, **kwargs): """Init pyconfig.""" + default_kwargs = { + "per_device_batch_size": 1.0, + "run_name": "test", + "enable_checkpointing": False, + "base_num_decoder_layers": 2, + "attention": "dot_product", + "max_target_length": 16, + "base_emb_dim": 32, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "max_prefill_predict_length": 4, + } + default_kwargs.update(kwargs) config = pyconfig.initialize( [sys.argv[0], get_test_config_path()], - per_device_batch_size=1.0, - run_name="test", - enable_checkpointing=False, - base_num_decoder_layers=2, - attention="dot_product", - max_target_length=16, - base_emb_dim=32, - base_num_query_heads=2, - base_num_kv_heads=2, - max_prefill_predict_length=4, - **kwargs, + **default_kwargs, ) return config @@ -196,6 +200,88 @@ def test_train_vs_prefill_and_autoregress(self): equal_nan=False, ) + def test_gemma4_model(self): + """Test pure-NNX model execution with a scanned Gemma4 block.""" + new_config = self.init_pyconfig( + decoder_block="gemma4", + share_kv_projections=True, + use_post_attn_norm=True, + use_post_ffw_norm=True, + base_num_decoder_layers=12, + scan_layers=True, + enable_nnx=True, + pure_nnx=True, + pure_nnx_decoder=True, + ) + devices_array = maxtext_utils.create_device_mesh(new_config) + mesh = Mesh(devices_array, new_config.mesh_axes) + model = models.Transformer( + config=new_config, + mesh=mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(0), + ) + + shape = (new_config.global_batch_size_to_train_on, new_config.max_target_length) + ids = jax.random.randint(self.rng, shape, 0, new_config.vocab_size) + decoder_segment_ids = jnp.full(shape, DECODING_ACTIVE_SEQUENCE_INDICATOR) + decoder_positions = jnp.broadcast_to(jnp.arange(new_config.max_target_length, dtype=jnp.int32), shape) + + logits = model( + ids, + decoder_positions, + decoder_segment_ids, + 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)) + + def test_gemma4_model_linen(self): + """Test the shared Gemma4 scannable block on the linen (ToLinen) path. + + The block is shared with the pure-NNX path. On the linen path it is driven + through nn.scan + the ToLinen wrapper with apply_internal_remat=True (the + block rematerializes its own layers; no block-level remat), exercising the + nested transform stack and guarding against nnx-migration regressions. + """ + new_config = self.init_pyconfig( + decoder_block="gemma4", + share_kv_projections=True, + use_post_attn_norm=True, + use_post_ffw_norm=True, + base_num_decoder_layers=12, + scan_layers=True, + remat_policy="minimal", + ) + devices_array = maxtext_utils.create_device_mesh(new_config) + mesh = Mesh(devices_array, new_config.mesh_axes) + model = models.transformer_as_linen(config=new_config, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN) + + shape = (new_config.global_batch_size_to_train_on, new_config.max_target_length) + ids = jax.random.randint(self.rng, shape, 0, new_config.vocab_size) + decoder_segment_ids = jnp.full(shape, DECODING_ACTIVE_SEQUENCE_INDICATOR) + decoder_positions = jnp.broadcast_to(jnp.arange(new_config.max_target_length, dtype=jnp.int32), shape) + + transformer_vars = model.init( + {"params": self.rng, "aqt": self.rng, "dropout": self.rng}, + ids, + decoder_positions, + decoder_segment_ids, + enable_dropout=False, + ) + logits = model.apply( + transformer_vars, + ids, + decoder_positions, + decoder_segment_ids, + enable_dropout=False, + model_mode=MODEL_MODE_TRAIN, + 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)) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index 4585c9c3ee..1475fa1630 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -22,7 +22,9 @@ """ import sys +from types import SimpleNamespace import unittest +from unittest import mock import pytest import jax @@ -34,8 +36,10 @@ from maxtext.common.common_types import ( DECODING_ACTIVE_SEQUENCE_INDICATOR, + MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, + AttentionType, DecoderBlockType, MultimodalInput, ) @@ -45,7 +49,7 @@ from maxtext.layers.embeddings import Embed from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process from maxtext.layers.normalizations import RMSNorm -from maxtext.models import gemma4_small +from maxtext.models import gemma4, gemma4_small from maxtext.models.gpt3 import Gpt3LayerNorm from maxtext.models.llama2 import LlamaDecoderLayer from maxtext.utils import maxtext_utils @@ -714,6 +718,98 @@ def test_scan_layers(self): unittest.main() +class _StatefulGemma4DecoderLayer(nnx.Module): + """Small stand-in that exposes cache ordering and mutable-state updates.""" + + def __init__(self, *, attention_type, **unused_kwargs): + self.increment = 10 if attention_type == AttentionType.GLOBAL else 1 + self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32)) + self.received_attention_metadata = nnx.Intermediate(jnp.array(False)) + + def __call__( + self, + inputs, + *unused_args, + kv_cache=None, + attention_metadata=None, + **unused_kwargs, + ): + self.call_count.value += 1 + self.received_attention_metadata.value = attention_metadata is not None + output = inputs + self.increment + if kv_cache is None: + return output + return output, kv_cache + self.increment + + +class TestGemma4ScannableBlock(unittest.TestCase): + """Tests Gemma4's nested local/global decoder block behavior.""" + + def setUp(self): + super().setUp() + self.config = SimpleNamespace( + dtype=jnp.float32, + param_scan_axis=1, + remat_policy="none", + scan_layers=True, + ) + + def _make_block(self): + return gemma4.Gemma4ScannableBlock( + config=self.config, + mesh=None, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(0), + ) + + def test_updates_state_through_global_single_iteration_scan(self): + with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15)) + self.assertIsNone(updated_kvs) + np.testing.assert_array_equal( + block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32) + ) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + + def test_restores_local_state_and_preserves_kv_order(self): + attention_metadata = object() + + with mock.patch.object(gemma4, "Gemma4DecoderLayer", _StatefulGemma4DecoderLayer): + block = self._make_block() + output, updated_kvs = block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + kv_cache=tuple(jnp.array(i) for i in range(6)), + attention_metadata=attention_metadata, + ) + + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 15)) + np.testing.assert_array_equal(jnp.stack(updated_kvs), jnp.array([1, 2, 3, 4, 5, 15])) + np.testing.assert_array_equal( + block.local_layers.call_count.value, jnp.ones(5, dtype=jnp.int32) + ) + np.testing.assert_array_equal( + block.local_layers.received_attention_metadata.value, + jnp.ones(5, dtype=jnp.bool_), + ) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) + np.testing.assert_array_equal( + block.global_layer.received_attention_metadata.value, True + ) + + class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" diff --git a/tests/unit/nnx_scan_test.py b/tests/unit/nnx_scan_test.py new file mode 100644 index 0000000000..f42652a11b --- /dev/null +++ b/tests/unit/nnx_scan_test.py @@ -0,0 +1,88 @@ +# 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. + +"""Tests for NNX scan utilities.""" + +from unittest import mock + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.layers import nnx_scan + + +class _LinearLayer(nnx.Module): + def __init__(self, rngs: nnx.Rngs): + self.kernel = nnx.Param(jax.random.normal(rngs.params(), (2, 2))) + + def __call__(self, inputs): + return inputs @ self.kernel.value + + +def test_nonzero_param_scan_axis_round_trip(): + """The scan dimension is stored at param_scan_axis and restored for application.""" + length = 3 + layers = nnx_scan.create_scanned_layers( + _LinearLayer, + length=length, + param_scan_axis=1, + metadata_axis_name="layers", + rngs=nnx.Rngs(0), + ) + + assert layers.kernel.value.shape == (2, length, 2) + + inputs = jnp.array([1.0, -1.0]) + kernels = jnp.moveaxis(layers.kernel.value, 1, 0) + expected = inputs + for kernel in kernels: + expected = expected @ kernel + + actual = nnx_scan.apply_scanned_layers( + layers, + inputs, + length=length, + param_scan_axis=1, + apply_fn=lambda layer, carry: layer(carry), + ) + + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) + assert layers.kernel.value.shape == (2, length, 2) + + +def test_full_remat_checkpoints_scan_body_with_none_policy(): + """A None policy means full remat and must not disable jax.checkpoint.""" + layers = nnx_scan.create_scanned_layers( + _LinearLayer, + length=2, + param_scan_axis=1, + metadata_axis_name="layers", + rngs=nnx.Rngs(0), + ) + + with mock.patch.object(nnx_scan.jax, "checkpoint", wraps=jax.checkpoint) as checkpoint: + nnx_scan.apply_scanned_layers( + layers, + jnp.array([1.0, -1.0]), + length=2, + param_scan_axis=1, + apply_fn=lambda layer, carry: layer(carry), + remat=True, + remat_policy=None, + ) + + checkpoint.assert_called_once() + assert checkpoint.call_args.kwargs["policy"] is None