From eb7563679f24781bdf73b40fb2541ab88f33af49 Mon Sep 17 00:00:00 2001 From: mesakhcienet Date: Fri, 31 Jul 2026 07:24:39 +0000 Subject: [PATCH] fix: decoders and pipeline parity gaps of linen to nnx migrations --- src/maxtext/layers/decoders.py | 83 +-- src/maxtext/layers/nnx_decoders.py | 318 ++++++++- src/maxtext/layers/pipeline.py | 81 +-- tests/unit/nnx_decoders_test.py | 1061 ++++++++++++++++++++++++++-- tests/unit/nnx_pipeline_test.py | 525 ++++++++++++++ 5 files changed, 1910 insertions(+), 158 deletions(-) create mode 100644 tests/unit/nnx_pipeline_test.py diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index d722270c4c..d617aa3906 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -33,7 +33,12 @@ from maxtext.layers import mhc from maxtext.layers import normalizations from maxtext.layers import pipeline -from maxtext.layers.nnx_decoders import NNXDecoderLayer, NNXSequentialPipelineStage, NNXScannedPipelineStage +from maxtext.layers.nnx_decoders import ( + NNXDecoderLayer, + NNXSequentialPipelineStage, + NNXScannedPipelineStage, + _make_single_layer_remat_stage_cls, +) from maxtext.layers import quantizations from maxtext.layers.attentions import attention_as_linen from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen @@ -526,6 +531,7 @@ def get_scannable(normal_cls, scannable_cls): DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer], + DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), } @@ -567,52 +573,49 @@ def _build_nnx_pipeline_stage(self, decoder_blocks, rngs): cfg = self.config base_stage_cls = decoder_blocks[1] if cfg.decoder_block == DecoderBlockType.DEEPSEEK else decoder_blocks[0] + # Per-stage-layer remat (+ params-only host-offload inside the stage) when the flag is set. + # apply_per_stage_remat is the boolean decision; per_stage_remat is the policy value + # (None == full remat for remat_policy='full', matching Linen nn.remat(policy=None)). + apply_per_stage_remat = cfg.set_remat_policy_on_layers_per_stage + per_stage_remat = self.get_remat_policy() if apply_per_stage_remat else None + if cfg.num_layers_per_pipeline_stage == 1: + if apply_per_stage_remat: + # Linen nn.remat parity: keep params TOP-LEVEL (no 'layers_0' nesting). + stage_cls = _make_single_layer_remat_stage_cls(base_stage_cls) + return stage_cls( + config=cfg, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=True, + ) return base_stage_cls(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, rngs=rngs) elif cfg.scan_layers_per_stage: return NNXScannedPipelineStage( - base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs - ) - return NNXSequentialPipelineStage( - base_stage_cls, cfg.num_layers_per_pipeline_stage, cfg, self.mesh, self.quant, self.model_mode, rngs=rngs - ) - - def get_pipeline_stage_module(self, decoder_blocks): - """get pipeline stage module""" - - def get_layer_to_pipeline(blocks, cfg): - if cfg.decoder_block == DecoderBlockType.DEEPSEEK: - return blocks[1] # return the sparse block - else: - return blocks[0] - - cfg = self.config - base_stage = get_layer_to_pipeline(decoder_blocks, cfg) - if cfg.set_remat_policy_on_layers_per_stage: - policy = self.get_remat_policy() - base_stage = self.set_remat_policy([base_stage], policy)[0] - if cfg.num_layers_per_pipeline_stage == 1: - stage_module = base_stage(config=cfg, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode) - elif cfg.scan_layers_per_stage: - stage_module = self.scan_decoder_layers( - cfg, - base_stage, + base_stage_cls, cfg.num_layers_per_pipeline_stage, - "layers_per_stage", + cfg, self.mesh, - in_axes_tuple=(nn.broadcast,) * 4, - model_mode=self.model_mode, + self.quant, + self.model_mode, + rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=apply_per_stage_remat, ) - else: - stage_module = SequentialBlockDecoderLayers( - decoder_layer=base_stage, - num_decoder_layers=cfg.num_layers_per_pipeline_stage, - config=cfg, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - ) - return stage_module + return NNXSequentialPipelineStage( + base_stage_cls, + cfg.num_layers_per_pipeline_stage, + cfg, + self.mesh, + self.quant, + self.model_mode, + rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=apply_per_stage_remat, + ) def get_norm_layer(self, num_features: int): """get normalization layer (return type inherits from nn.Module)""" diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 6d01e9d22d..26c50736e4 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -252,6 +252,96 @@ def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): return hidden_states +def _run_stage_layer_with_remat( + graphdef, + params, + state, + inputs, + run_fn, + *, + host_offload, + remat_policy, + prevent_cse, +): + """Run one pipeline-stage layer under ``jax.checkpoint`` with params-only host-offload. + + Shared by ``NNXSequentialPipelineStage`` and ``NNXScannedPipelineStage``, whose per-layer + remat bodies were byte-identical. The layer must already be split into ``(graphdef, params, + state)`` so the host-offload targets params only. + Inside a single ``jax.checkpoint(policy=remat_policy, prevent_cse=prevent_cse)`` this: + + 1. optionally offloads params to host memory via ``jax.device_put(device_space())``, + 2. re-merges the layer with ``nnx.merge(graphdef, params, state)``, + 3. runs it through ``run_fn(merged_layer, inputs)`` -- ``run_fn`` captures the call site's + ``decoder_segment_ids`` / ``decoder_positions`` / ``deterministic`` / ``model_mode`` / kwargs, + 4. unwraps a tuple layer output to its first element. + + Returns ``(out, new_state)``; the caller round-trips ``new_state`` back via ``nnx.update``. + """ + + def pure_fn(params_in, state_in, x_in): + if host_offload: + params_in = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), params_in) + merged = nnx.merge(graphdef, params_in, state_in) + out_inner = run_fn(merged, x_in) + out_inner = out_inner[0] if isinstance(out_inner, tuple) else out_inner + return out_inner, nnx.state(merged) + + checkpointed_fn = jax.checkpoint(pure_fn, policy=remat_policy, prevent_cse=prevent_cse) + return checkpointed_fn(params, state, inputs) + + +@functools.lru_cache(maxsize=None) +def _make_single_layer_remat_stage_cls(base_stage_cls): + """Single-layer pipeline-stage subclass of base_stage_cls: applies per-stage jax.checkpoint + (+params-only host-offload) around the base layer's own __call__ WITHOUT a wrapper module, so + nnx.split(stage, nnx.Param, ...) yields param paths identical to a bare layer (top-level, no + 'layers_0') . Cached per base class for stable nnx type identity.""" + + class NNXSingleLayerRematPipelineStage(base_stage_cls): + """Single decoder layer + per-stage remat/host-offload, params kept top-level (Linen parity).""" + + def __init__(self, config, mesh, quant, model_mode, *, rngs, remat_policy=None, apply_remat=True): + super().__init__(config=config, mesh=mesh, quant=quant, model_mode=model_mode, rngs=rngs) + # Static (non-Variable) attrs -> graphdef, never the param tree. apply_remat decoupled from + # remat_policy (None == full rematerialization, matches Linen nn.remat(policy=None)). + self.apply_remat = apply_remat + self.remat_policy = remat_policy + self.prevent_cse = maxtext_utils.should_prevent_cse_in_remat(config) + + def __call__(self, inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs): + if not self.apply_remat: + return base_stage_cls.__call__( + self, inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs + ) + + def run_fn(merged_layer, x_in): + # BASE __call__ explicitly: merged_layer is a THIS-subclass instance; merged_layer(...) recurses. + return base_stage_cls.__call__( + merged_layer, x_in, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs + ) + + graphdef, params, rest = nnx.split(self, nnx.Param, ...) + out, new_state = _run_stage_layer_with_remat( + graphdef, + params, + rest, + inputs, + run_fn, + host_offload=self.config.parameter_memory_host_offload, + remat_policy=self.remat_policy, + prevent_cse=self.prevent_cse, + ) + nnx.update(self, new_state) + if self.config.scan_layers: # match bare-layer contract (both pipelines unwrap [0] iff tuple) + return out, None + return out + + NNXSingleLayerRematPipelineStage.__name__ = f"NNXSingleLayerRematStage_{base_stage_cls.__name__}" + NNXSingleLayerRematPipelineStage.__qualname__ = NNXSingleLayerRematPipelineStage.__name__ + return NNXSingleLayerRematPipelineStage + + class NNXSequentialPipelineStage(nnx.Module): """Sequential unscanned series of decoder layers formatted for a single pipeline stage.""" @@ -265,10 +355,19 @@ def __init__( model_mode: str, *, rngs: nnx.Rngs, + remat_policy: Any = None, + apply_remat: bool = False, ): self.config = config self.scan_layers = config.scan_layers self.num_layers = num_layers + + # apply_remat is decoupled from remat_policy's value: remat_policy=None is a VALID + # policy meaning "full rematerialization" (save nothing), matching Linen nn.remat(policy=None). + # Gating on `remat_policy is not None` would silently skip remat for remat_policy='full'. + self.remat_policy = remat_policy + self.apply_remat = apply_remat + self.prevent_cse = maxtext_utils.should_prevent_cse_in_remat(config) # Dynamically assign layers with explicit string names to ensure correct PyTree paths (layers_0) for i in range(num_layers): layer = layer_cls(config=config, mesh=mesh, quant=quant, model_mode=model_mode, rngs=rngs) @@ -283,17 +382,35 @@ def __call__( model_mode, **kwargs, ): + def run_fn(merged_layer, x_in): + return merged_layer(x_in, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs) + for i in range(self.num_layers): layer = getattr(self, f"layers_{i}") - out = layer( - inputs, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - **kwargs, - ) - inputs = out[0] if isinstance(out, tuple) else out + if self.apply_remat: + # Split params out so host-offload is params-only (mirrors Linen nn.map_variables(["params"])). + graphdef, params, rest = nnx.split(layer, nnx.Param, ...) + inputs, new_state = _run_stage_layer_with_remat( + graphdef, + params, + rest, + inputs, + run_fn, + host_offload=self.config.parameter_memory_host_offload, + remat_policy=self.remat_policy, + prevent_cse=self.prevent_cse, + ) + nnx.update(layer, new_state) + else: + out = layer( + inputs, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + **kwargs, + ) + inputs = out[0] if isinstance(out, tuple) else out if self.scan_layers: return inputs, None return inputs @@ -312,8 +429,14 @@ def __init__( model_mode: str, *, rngs: nnx.Rngs, + remat_policy: Any = None, + apply_remat: bool = False, ): self.config = config + # remat_policy=None is a valid "full remat" policy; gate on apply_remat, not on its value. + self.remat_policy = remat_policy + self.apply_remat = apply_remat + self.prevent_cse = maxtext_utils.should_prevent_cse_in_remat(config) def create_layer_fn(rng): return layer_cls(config=config, mesh=mesh, quant=quant, model_mode=model_mode, rngs=rng) @@ -344,8 +467,26 @@ def __call__( if scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params) + def run_fn(merged_layer, x_in): + return merged_layer(x_in, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs) + def layer_fn(carry, scanned_vars): current_params, current_state = scanned_vars + if self.apply_remat: + new_carry, new_state = _run_stage_layer_with_remat( + graphdef, + current_params, + current_state, + carry, + run_fn, + host_offload=self.config.parameter_memory_host_offload, + remat_policy=self.remat_policy, + prevent_cse=self.prevent_cse, + ) + # Avoid returning and stacking read-only parameters inside the scan body. + # This prevents huge unnecessary memory allocation (mirror the non-remat sibling). + _, updated_state = new_state.split(nnx.Param, ...) + return new_carry, updated_state layer = nnx.merge(graphdef, current_params, current_state) layer_out = layer( carry, @@ -432,12 +573,20 @@ def __init__( self.scanned_layers = None self.is_deepseek = self.config.decoder_block == DecoderBlockType.DEEPSEEK + self.is_deepseek4 = self.config.decoder_block == DecoderBlockType.DEEPSEEK4 self.is_gemma3 = self.config.decoder_block == DecoderBlockType.GEMMA3 self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4 self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL self._init_decoder_layers(decoder_block_classes, rngs, mesh) + # DeepSeek-V4 collapses its mhc_expansion_rate mHC streams with a LEARNED hyper-head + # (mhc.DeepSeek4HyperHead), not the unweighted mhc_reduce sum. + # NNX materializes params at construction; build AFTER the layers so decoder-layer RNG draws stay byte-identical. + # Attr name "hc_head" IS the checkpoint key (same convention as self.decoder_norm <-> Linen "decoder_norm"). + if self.is_deepseek4 and getattr(config, "mhc_expansion_rate", 1) > 1: + self.hc_head = mhc.DeepSeek4HyperHead(config=config, mesh=mesh, rngs=rngs) + def _init_decoder_layers(self, decoder_block_classes, rngs, mesh): """Routes layer construction through three main paths: pipeline, scanned non-pipeline, sequential.""" config = self.config @@ -532,6 +681,8 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): """Initializes decoder layers with scanning (non-pipeline).""" if self.is_deepseek: self._init_scanned_deepseek(decoder_block_classes, rngs) + elif self.is_deepseek4: + self._init_scanned_deepseek4(decoder_block_classes, rngs, mesh) elif self.is_gemma3: self._init_scanned_gemma3(decoder_block_classes, rngs, mesh) elif self.is_gemma4: @@ -683,6 +834,44 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): rngs=rngs, ) + def _init_scanned_deepseek4(self, decoder_block_classes, rngs, mesh): + """Initializes DeepSeek-V4 scanned layers. + + DeepSeek-V4 has ``first_num_hash_layers`` prefix layers (static hash routing and + heterogeneous attention) that cannot be scanned, followed by uniform alternating + HCA(compress_ratio=128)/CSA(compress_ratio=4) blocks. The prefix layers are built + individually (named ``layers_{i}`` to match the Linen reference) and the remaining + layers are paired into ``DeepSeek4ScannableBlock``s and scanned. Mirrors + ``decoders.Decoder._apply_deepseek4_scanned_blocks``. + """ + config = self.config + scannable_cls = decoder_block_classes[0] + + # 1. Prefix (hash-routing) layers: unrolled, heterogeneous attention per layer_idx. + self.num_prefix_layers = config.first_num_hash_layers + for layer_idx in range(self.num_prefix_layers): + self._create_and_register_layer(deepseek4.DeepSeek4DecoderLayer, rngs, "layers", layer_idx, layer_idx=layer_idx) + + # 2. Scanned alternating HCA(128)/CSA(4) blocks. + # The non-prefix layers are paired (HCA + CSA) into scannable blocks, so their count must be + # even; an odd count would make the `// 2` below silently drop the trailing layer. This is the + # single construction-time guard: _apply_deepseek4_scanned_blocks (which mirrors this `// 2`) + # runs only after this init -- both are gated on scan_layers + deepseek4 -- so one assert here + # covers the apply site too. + num_non_prefix_layers = config.num_decoder_layers - self.num_prefix_layers + assert num_non_prefix_layers % 2 == 0, ( + "DeepSeek-V4 scanned body pairs non-prefix layers into HCA/CSA blocks: " + f"(num_decoder_layers={config.num_decoder_layers} - " + f"first_num_hash_layers={self.num_prefix_layers}) = {num_non_prefix_layers} must be even, " + "otherwise the last decoder layer would be silently dropped." + ) + num_full_blocks = num_non_prefix_layers // 2 + self.scanned_blocks = ( + self._create_scanned_layers(scannable_cls, length=num_full_blocks, metadata_axis_name="layers", rngs=rngs) + if num_full_blocks > 0 + else None + ) + def _init_scanned_generic(self, decoder_block_classes, rngs): """Initializes scanned generic decoder layers.""" config = self.config @@ -713,13 +902,18 @@ def _init_sequential_layers(self, decoder_block_classes, rngs): self._init_sequential_generic(decoder_block_classes, rngs) def _init_sequential_deepseek(self, decoder_block_classes, rngs): - """Initializes sequential DeepSeek dense and MoE layers.""" + """Initializes sequential DeepSeek dense and MoE layers. + + Each layer receives its GLOBAL ``layer_idx`` (dense layers 0..first_num_dense_layers-1, + MoE layers continue from first_num_dense_layers). The index drives engram interleaving / hash routing + (deepseek.DeepSeekGenericLayer.is_engram_enabled). + """ config = self.config dense_cls, moe_cls = decoder_block_classes for i in range(config.first_num_dense_layers): - self._create_and_register_layer(dense_cls, rngs, "dense_layers", i) + self._create_and_register_layer(dense_cls, rngs, "dense_layers", i, layer_idx=i) for i in range(config.num_decoder_layers - config.first_num_dense_layers): - self._create_and_register_layer(moe_cls, rngs, "moe_layers", i) + self._create_and_register_layer(moe_cls, rngs, "moe_layers", i, layer_idx=config.first_num_dense_layers + i) def _init_sequential_generic(self, decoder_block_classes, rngs): """Initializes sequential generic decoder layers with per-architecture layer_kwargs.""" @@ -740,6 +934,7 @@ def _init_sequential_generic(self, decoder_block_classes, rngs): elif config.decoder_block in { DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, + DecoderBlockType.DEEPSEEK4, }: layer_kwargs = {"layer_idx": lyr} elif config.decoder_block == DecoderBlockType.GPT_OSS: @@ -783,7 +978,25 @@ def _get_pipeline_stage_module(self, decoder_blocks, rngs): cfg = self.config base_stage_cls = decoder_blocks[1] if self.is_deepseek else decoder_blocks[0] + # Per-stage-layer remat (+ params-only host-offload inside the stage) when the flag is set. + # apply_per_stage_remat is the boolean decision; per_stage_remat is the policy value + # (which may be None == full remat for remat_policy='full', matching Linen nn.remat). + apply_per_stage_remat = cfg.set_remat_policy_on_layers_per_stage + per_stage_remat = self.get_remat_policy() if apply_per_stage_remat else None + if cfg.num_layers_per_pipeline_stage == 1: + if apply_per_stage_remat: + # Linen nn.remat parity: keep params TOP-LEVEL (no 'layers_0' nesting). + stage_cls = _make_single_layer_remat_stage_cls(base_stage_cls) + return stage_cls( + config=cfg, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=True, + ) return self._create_single_layer(base_stage_cls, rngs) elif cfg.scan_layers_per_stage: return NNXScannedPipelineStage( @@ -794,6 +1007,8 @@ def _get_pipeline_stage_module(self, decoder_blocks, rngs): self.quant, self.model_mode, rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=apply_per_stage_remat, ) return NNXSequentialPipelineStage( base_stage_cls, @@ -803,6 +1018,8 @@ def _get_pipeline_stage_module(self, decoder_blocks, rngs): self.quant, self.model_mode, rngs=rngs, + remat_policy=per_stage_remat, + apply_remat=apply_per_stage_remat, ) def _create_and_register_layer(self, layer_cls, rngs, base_name, i, **layer_kwargs): @@ -1225,6 +1442,7 @@ def get_norm_layer(self, num_features: int, rngs: nnx.Rngs): DecoderBlockType.MISTRAL, DecoderBlockType.MIXTRAL, DecoderBlockType.DEEPSEEK, + DecoderBlockType.DEEPSEEK4, DecoderBlockType.GEMMA, DecoderBlockType.GEMMA2, DecoderBlockType.GEMMA3, @@ -1344,7 +1562,7 @@ def _apply_embedding( y = y.astype(cfg.dtype) if cfg.use_untrainable_positional_embedding: - y += self.positional_embedding(y, decoder_positions) + y += self.positional_embedding(y.shape[1], decoder_positions) if cfg.trainable_position_size > 0 and self.position_embedder: y += self.position_embedder(decoder_positions.astype("int32"), model_mode=model_mode) @@ -1526,10 +1744,10 @@ def __call__( model_mode=MODEL_MODE_TRAIN, previous_chunk=None, slot: None | int = None, + multimodal_input: None | MultimodalInput = None, kv_caches: list[jax.Array] | None = None, attention_metadata=None, deepstack_visual_embeds: None | list[jnp.ndarray] = None, - multimodal_input: None | MultimodalInput = None, ): cfg = self.config assert decoder_input_tokens.ndim == 2 # [batch, len] @@ -1574,7 +1792,7 @@ def __call__( if attention_metadata is not None: layer_kwargs["attention_metadata"] = attention_metadata - if cfg.engram_layers and decoder_input_tokens is not None: + if (cfg.engram_layers or self.is_deepseek4) and decoder_input_tokens is not None: layer_kwargs["decoder_input_tokens"] = decoder_input_tokens if getattr(cfg, "using_pipeline_parallelism", False): @@ -1644,7 +1862,7 @@ def __call__( kv_caches=kv_caches, ) else: - # Standard pipeline run (non-DeepSeek, incl. Gemma4 — matches Linen decoders.py). + # Standard pipeline run (non-DeepSeek, incl. Gemma4). # Gemma4 routes through the pipeline here; _apply_gemma4_scanned_blocks is # non-pipeline-only (its layers/layers_remainder are not built when # pipeline parallelism is enabled). @@ -1787,6 +2005,17 @@ def __call__( layer_kwargs, kv_caches=kv_caches, ) + elif self.is_deepseek4: + y = self._apply_deepseek4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + decoder_input_tokens, + ) else: scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) if kv_caches is not None: @@ -1897,8 +2126,12 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): # After the final transformer layer, `y` holds the raw, un-normalized hidden state. if getattr(cfg, "mhc_expansion_rate", 1) > 1: - # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) - hidden_state = mhc_reduce(y) + if self.is_deepseek4: + # DeepSeek-V4 learned weighted collapse, not the sum. + hidden_state = self.hc_head(y) + else: + # (batch, length, mhc_expansion_rate, emb_dim) --> (batch, length, emb_dim) + hidden_state = mhc_reduce(y) else: hidden_state = y @@ -1926,6 +2159,49 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in): return logits, hidden_state, kv_caches + def _apply_deepseek4_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + decoder_input_tokens, + ): + """Applies DeepSeek-V4 decoder layers: unrolled hash-routing prefix layers followed + by scanned alternating HCA(128)/CSA(4) blocks. Mirrors the Linen + ``decoders.Decoder._apply_deepseek4_scanned_blocks``. + """ + layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) + + # 1. Unrolled prefix (hash-routing) layers; each derives is_hash_routing/compress_ratio + # from its static layer_idx and consumes decoder_input_tokens for hash routing. + for layer_idx in range(self.num_prefix_layers): + prefix_layer = getattr(self, f"layers_{layer_idx}") + y, _ = prefix_layer( + y, + *layer_args, + previous_chunk=previous_chunk, + slot=slot, + decoder_input_tokens=decoder_input_tokens, + ) + + # 2. Scanned alternating HCA/CSA blocks (no hash routing -> no decoder_input_tokens). + if self.scanned_blocks is not None: + num_full_blocks = (self.config.num_decoder_layers - self.num_prefix_layers) // 2 + y, self.scanned_blocks, _ = self._apply_layers_sequentially( + self.scanned_blocks, + y, + *layer_args, + length=num_full_blocks, + previous_chunk=previous_chunk, + slot=slot, + ) + + return y + def _apply_gemma3_scanned_blocks( self, y, @@ -1965,6 +2241,9 @@ def _apply_gemma3_scanned_blocks( remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) def pure_gemma_fn(graphdef, state_in, y_in, kv_in): + # the remainder block must also be host-offloaded when parameter_memory_host_offload is enabled. + if cfg.parameter_memory_host_offload: + state_in = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), state_in) merged_layer = nnx.merge(graphdef, state_in) call_kwargs = dict(layer_kwargs) if kv_in is not None: @@ -2055,6 +2334,9 @@ def _apply_gemma4_scanned_blocks( else: def pure_gemma_fn(graphdef, state_in, y_in, kv_in): + # the remainder block must also be host-offloaded when parameter_memory_host_offload is enabled. + if cfg.parameter_memory_host_offload: + state_in = jax.tree.map(lambda x: jax.device_put(x, max_utils.device_space()), state_in) merged_layer = nnx.merge(graphdef, state_in) call_kwargs = dict(layer_kwargs) if kv_in is not None: diff --git a/src/maxtext/layers/pipeline.py b/src/maxtext/layers/pipeline.py index bf66fbcce8..e1edc3b0ed 100644 --- a/src/maxtext/layers/pipeline.py +++ b/src/maxtext/layers/pipeline.py @@ -790,7 +790,7 @@ def run_one_iteration( model_mode, ) - if self.config.scan_layers: + if isinstance(stages_output, tuple): stages_output = stages_output[0] if self.config.num_pipeline_repeats > 1: @@ -892,32 +892,18 @@ def unbox_val(x): # jax.tree.map to raise "Mismatch custom node data". Mirrors Linen # where all_gather_over_fsdp operates on # self.layers.variables (the params collection only). - _, layers_params, layers_metrics, layers_mutables = nnx.split( - layers_state, pipeline_utils.is_static_param, nnx.Intermediate, ... - ) - - # layers_mutables catch-all should contain ONLY RngState variables (RngKey/RngCount). - # If non_trainable state (e.g. BatchStat) appears here, - # it is being carried through scan instead of broadcast. - # NOTE: is_leaf stops jax.tree.leaves from traversing *into* Variable nodes, - # so we see actual Variable instances (not raw arrays). - assert all( - isinstance(v, nnx.RngState) - for v in jax.tree.leaves(layers_mutables, is_leaf=lambda x: isinstance(x, nnx.Variable)) - if isinstance(v, nnx.Variable) - ), ( - "Non-RngState variable found in layers_mutables catch-all partition. " - "Only RngState variables (RngKey/RngCount) should be present." + _, layers_params, layers_metrics, layers_rng, layers_non_trainable = nnx.split( + layers_state, pipeline_utils.is_static_param, nnx.Intermediate, nnx.RngState, ... ) if self.config.pipeline_fsdp_ag_once: layers_params = self.all_gather_over_fsdp(layers_params, logical_partition_spec) def scan_body(carry, _): - current_loop_state, current_layer_mutables = carry + current_loop_state, current_layer_rng, current_non_trainable = carry iteration = current_loop_state["loop_iteration"] - advanced_mutables = pipeline_utils.advance_rng_state(current_layer_mutables, iteration) - current_layer_state = nnx.State.merge(layers_params, layers_metrics, advanced_mutables) + advanced_rng = pipeline_utils.advance_rng_state(current_layer_rng, iteration) + current_layer_state = nnx.State.merge(layers_params, layers_metrics, advanced_rng, current_non_trainable) new_loop_state, new_layer_state = self.run_one_iteration( current_loop_state, @@ -930,10 +916,14 @@ def scan_body(carry, _): logical_partition_spec, ) - _, _, new_layer_metrics, new_layer_mutables = nnx.split( - new_layer_state, pipeline_utils.is_static_param, nnx.Intermediate, ... + # RngState and non_trainable are both carried: a layer may mutate non_trainable (Linen carries + # it whenever the collection is mutable -- 53dea32b7:898-900), and discarding the output copy + # here silently loses that update. Carrying costs no extra memory: a jax.lax.scan carry is + # threaded, not stacked -- only the ys (metrics) stack. + _, _, new_layer_metrics, new_layer_rng, new_non_trainable = nnx.split( + new_layer_state, pipeline_utils.is_static_param, nnx.Intermediate, nnx.RngState, ... ) - return (new_loop_state, new_layer_mutables), new_layer_metrics + return (new_loop_state, new_layer_rng, new_non_trainable), new_layer_metrics if self.config.set_remat_policy_on_pipeline_iterations: scan_body = jax.checkpoint( @@ -941,19 +931,19 @@ def scan_body(carry, _): ) if self.config.scan_pipeline_iterations: - (loop_state, final_layer_mutables), stacked_metrics = jax.lax.scan( - scan_body, (loop_state, layers_mutables), None, length=total_iterations + (loop_state, final_layer_rng, final_non_trainable), stacked_metrics = jax.lax.scan( + scan_body, (loop_state, layers_rng, layers_non_trainable), None, length=total_iterations ) else: - current_carry = (loop_state, layers_mutables) + current_carry = (loop_state, layers_rng, layers_non_trainable) metrics_history = [] for _ in range(total_iterations): current_carry, step_metrics = scan_body(current_carry, None) metrics_history.append(step_metrics) - loop_state, final_layer_mutables = current_carry + loop_state, final_layer_rng, final_non_trainable = current_carry stacked_metrics = jax.tree.map(lambda *xs: jnp.stack(xs), *metrics_history) if metrics_history else layers_metrics - final_layer_state = nnx.State.merge(layers_params, stacked_metrics, final_layer_mutables) + final_layer_state = nnx.State.merge(layers_params, stacked_metrics, final_layer_rng, final_non_trainable) nnx.update(self.layers, final_layer_state) final_output = self.permute_output_micro_per_stage_dim(loop_state["state_io"]) @@ -1284,7 +1274,7 @@ def run_one_iteration( model_mode, ) - if self.config.scan_layers: + if isinstance(stages_output, tuple): stages_output = stages_output[0] # Scatter-back: only update mutables (params handled by AD/gradient, metrics returned directly) @@ -1362,6 +1352,7 @@ def __call__( num_microbatches = self.config.num_pipeline_microbatches num_repeats = self.config.num_pipeline_repeats remat_policy = self.get_pipeline_remat_policy() + apply_iteration_remat = self.config.set_remat_policy_on_pipeline_iterations layers_graph, layers_state = nnx.split(self.layers) @@ -1377,16 +1368,6 @@ def unbox_val(x): layers_state, pipeline_utils.is_static_param, nnx.Intermediate, ... ) - # Validate: layers_mutables should contain ONLY RngState variables - assert all( - isinstance(v, nnx.RngState) - for v in jax.tree.leaves(layers_mutables, is_leaf=lambda x: isinstance(x, nnx.Variable)) - if isinstance(v, nnx.Variable) - ), ( - "Non-RngState variable found in layers_mutables catch-all partition. " - "Only RngState variables (RngKey/RngCount) should be present." - ) - # Pre-capture Python config values needed inside the stage function. scan_pipeline_iterations_enabled = self.config.scan_pipeline_iterations @@ -1483,7 +1464,7 @@ def _run(loop_state, bsw): logical_partition_spec_stripped, ) - _run_remat = jax.remat(_run, policy=remat_policy) + _run_remat = jax.remat(_run, policy=remat_policy) if apply_iteration_remat else _run (new_loop_state, new_layer_state), vjp_fn = jax.vjp(_run_remat, loop_state_arg, bsw_arg) return (new_loop_state, new_layer_state), vjp_fn @@ -1648,16 +1629,13 @@ def _execute_stage_bwd(stage_vjp_fn, output_grads): return (new_loop_state, w_next), inner_metrics # ---- Build lift.scan(lift.checkpoint(stage_fn)) over REPEATS ---- - if self.config.set_remat_policy_on_pipeline_iterations: - checkpointed_stage = flax_lift.checkpoint( - _stage_fn_for_scope, - variables=True, - rngs=True, - prevent_cse=not self.config.scan_pipeline_iterations, - policy=remat_policy, - ) - else: - checkpointed_stage = _stage_fn_for_scope + checkpointed_stage = flax_lift.checkpoint( + _stage_fn_for_scope, + variables=True, + rngs=True, + prevent_cse=not self.config.scan_pipeline_iterations, + policy=remat_policy, + ) scanned_stage = flax_lift.scan( checkpointed_stage, @@ -1666,6 +1644,7 @@ def _execute_stage_bwd(stage_vjp_fn, output_grads): variable_axes={}, split_rngs={}, length=num_repeats, + unroll=1 if self.config.scan_pipeline_repeats else num_repeats, ) # ---- Execute via flax.core.apply ---- @@ -1741,7 +1720,7 @@ def _run(loop_state, bsw): logical_partition_spec_stripped, ) - _run_remat = jax.remat(_run, policy=remat_policy) + _run_remat = jax.remat(_run, policy=remat_policy) if apply_iteration_remat else _run (new_loop_state, new_layer_state), vjp_fn = jax.vjp(_run_remat, loop_state_bubble, bsw_bubble) return (new_loop_state, new_layer_state), vjp_fn diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index fcd5acb5cc..62803834a7 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -34,6 +34,7 @@ import numpy as np from flax import linen as nn from flax import nnx +from flax.linen import partitioning as nn_partitioning from jax.sharding import Mesh from maxtext.common.common_types import ( @@ -46,10 +47,17 @@ MultimodalInput, ) from maxtext.configs import pyconfig -from maxtext.layers import linears +from maxtext.layers import linears, mhc from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed -from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process +from maxtext.layers.nnx_decoders import ( + NNXDecoder, + NNXDecoderLayer, + NNXScannedPipelineStage, + NNXSequentialPipelineStage, + _make_single_layer_remat_stage_cls, + deepstack_process, +) from maxtext.layers.normalizations import RMSNorm from maxtext.models import gemma4, gemma4_small from maxtext.models.gpt3 import Gpt3LayerNorm @@ -60,6 +68,22 @@ # --------------------------------------------------------------------------- # Shared minimal config overrides used across most tests # --------------------------------------------------------------------------- +# jax.checkpoint lowers to this primitive. Match it by IDENTITY, never by substring-searching the +# printed jaxpr: JAX renames primitives' printed names for cosmetics (pjit_p went "pjit" -> "jit" in a +# commit that also touched ad_checkpoint.py), and "remat2" is a leftover marker from the 2021 remat +# rewrite. A rename would break present-checks loudly and make absent-checks pass vacuously forever. +from jax._src.ad_checkpoint import remat_p as _REMAT_PRIMITIVE # pylint: disable=wrong-import-position +from jax._src import core as _jax_core # pylint: disable=wrong-import-position + + +def _jaxpr_contains_primitive(jaxpr, primitive): + """True if `primitive` appears anywhere in `jaxpr`, including nested sub-jaxprs (scan/cond bodies).""" + inner = jaxpr.jaxpr if hasattr(jaxpr, "jaxpr") else jaxpr + if any(eqn.primitive is primitive for eqn in inner.eqns): + return True + return any(_jaxpr_contains_primitive(sub, primitive) for sub in _jax_core.subjaxprs(inner)) + + _BASE_CONFIG = { "per_device_batch_size": 1.0, "run_name": "nnx_decoder_test", @@ -73,6 +97,7 @@ "base_mlp_dim": 512, "max_prefill_predict_length": 4, "scan_layers": False, + "activations_in_float32": True, } @@ -233,19 +258,42 @@ def test_has_mlp(self): # --- forward pass ---------------------------------------------------------- - def test_forward_output_shape_train(self): - """Forward pass output shape matches input shape in train mode.""" - layer = self._make_layer(MODEL_MODE_TRAIN) + def _assert_layer_matches_submodule_composition(self, model_mode): + """Forward pass output shape matches input shape, AND the layer composes its sub-modules exactly + as norm -> self_attention(lnx,lnx) + mlp(lnx) -> dropout -> +residual . + """ + layer = self._make_layer(model_mode) inputs, segment_ids, positions = self._make_inputs() out, _ = layer( inputs, segment_ids, positions, deterministic=True, - model_mode=MODEL_MODE_TRAIN, + model_mode=model_mode, ) self.assertEqual(out.shape, inputs.shape) + ref_layer = self._make_layer(model_mode) + lnx = ref_layer.pre_self_attention_norm(inputs) + attention_lnx, _ = ref_layer.self_attention( + lnx, + lnx, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=model_mode, + kv_cache=None, + attention_metadata=None, + ) + mlp_lnx = ref_layer.mlp(lnx, deterministic=True) + combined = ref_layer.dropout(mlp_lnx + attention_lnx, deterministic=True) + expected = combined + inputs + np.testing.assert_allclose(np.array(out), np.array(expected), rtol=1e-5, atol=1e-5) + + def test_forward_output_shape_train(self): + """Forward pass in train mode matches shape and the sub-module composition (see helper docstring).""" + self._assert_layer_matches_submodule_composition(MODEL_MODE_TRAIN) + def test_forward_output_dtype(self): """Output dtype matches config dtype.""" layer = self._make_layer() @@ -260,17 +308,10 @@ def test_forward_output_dtype(self): self.assertEqual(out.dtype, self.cfg.dtype) def test_forward_prefill_mode(self): - """Test forward pass in prefill mode.""" - layer = self._make_layer(MODEL_MODE_PREFILL) - inputs, segment_ids, positions = self._make_inputs() - out, _ = layer( - inputs, - segment_ids, - positions, - deterministic=True, - model_mode=MODEL_MODE_PREFILL, - ) - self.assertEqual(out.shape, inputs.shape) + """Forward pass in prefill mode matches shape and the sub-module composition (see helper + docstring). Prefill selects a different sharding-axis-name branch inside __call__ (values + unaffected) and a different attention code path; the composition still must hold.""" + self._assert_layer_matches_submodule_composition(MODEL_MODE_PREFILL) def test_record_metrics(self): """Test recording intermediate activation metrics.""" @@ -524,6 +565,32 @@ def _make_token_inputs(self): positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) return ids, segment_ids, positions + def test_logits_depend_on_input_tokens(self): + """The decoder's output must actually depend on the token ids it was given.""" + _, segment_ids, positions = self._make_token_inputs() + cfg = self.cfg + batch, seq_len = cfg.global_batch_size_to_train_on, cfg.max_target_length + key_a, _ = jax.random.split(jax.random.PRNGKey(1234)) + ids_a = jax.random.randint(key_a, (batch, seq_len), 0, cfg.vocab_size) + ids_b = (ids_a + 1) % cfg.vocab_size # guaranteed different at every position + + def run(ids): + logits, _, _ = self.decoder( + self.shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return np.array(logits) + + self.assertFalse( + bool(jnp.allclose(run(ids_a), run(ids_b), rtol=1e-4, atol=1e-4)), + "logits are identical for two different token sequences -> the decoder is ignoring its input " + "tokens (e.g. the embedding lookup is not using decoder_input_tokens)", + ) + def test_forward_pass_returns_three_tuple(self): """__call__ must return (logits, hidden_state, kv_caches).""" ids, segment_ids, positions = self._make_token_inputs() @@ -558,7 +625,9 @@ def test_logits_shape(self): self.assertEqual(logits.shape, expected) def test_hidden_state_shape(self): - """hidden_state shape: [batch, seq_len, emb_dim].""" + """hidden_state shape: [batch, seq_len, emb_dim], AND it must actually be the decoder LAYER + STACK's output, not the raw embeddings passed through untouched. + """ cfg = self.cfg ids, segment_ids, positions = self._make_token_inputs() _, hidden_state, _ = self.decoder( @@ -576,8 +645,16 @@ def test_hidden_state_shape(self): ) self.assertEqual(hidden_state.shape, expected) + embeds = self.decoder._apply_embedding( # pylint: disable=protected-access + self.shared_embedding, ids, positions, True, MODEL_MODE_TRAIN + ) + self.assertFalse( + bool(jnp.allclose(hidden_state, embeds, rtol=1e-3, atol=1e-3)), + msg="hidden_state equals the raw embeddings -> the decoder layer stack was not applied", + ) + def test_logits_are_finite(self): - """Logits must not contain NaN or Inf in a simple forward pass.""" + """Logits must not contain NaN or Inf, AND be non-degenerate.""" ids, segment_ids, positions = self._make_token_inputs() logits, _, _ = self.decoder( self.shared_embedding, @@ -588,16 +665,14 @@ def test_logits_are_finite(self): model_mode=MODEL_MODE_TRAIN, ) self.assertTrue(jnp.all(jnp.isfinite(logits))) + self.assertGreater(float(jnp.std(logits)), 1e-2) + self.assertFalse( + bool(jnp.allclose(logits[:, 0, :], logits[:, -1, :], rtol=1e-2, atol=1e-2)), + msg="logits are position-invariant -> forward is degenerate", + ) def test_multimodal_input_forwarded_to_apply_embedding(self): - """`multimodal_input` must reach `_apply_embedding` as the original struct. - - `NNXDecoder.__call__` takes a `MultimodalInput` struct and hands it to - `_apply_embedding`, which is the layer that actually unpacks the fields - and merges the embeddings. This test stubs `_apply_embedding` to capture - the forwarded struct without running the real embedding path (the test - config has `use_multimodal=False`). - """ + """`multimodal_input` must reach `_apply_embedding` as the original struct.""" ids, segment_ids, positions = self._make_token_inputs() # Distinct sentinels so each field can be traced independently. @@ -715,6 +790,43 @@ def test_scan_layers(self): ) self.assertEqual(logits.shape, (batch, seq_len, cfg.vocab_size)) + def test_scan_forward_uses_combined_scan_axis_helper_regression(self): + """Every scan_layers=True pure-NNX forward must run.""" + cfg = _make_config(scan_layers=True) + rngs = nnx.Rngs(params=0, dropout=1) + decoder = NNXDecoder(config=cfg, mesh=self.mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + shared_embedding = Embed( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + embedding_init=nn.initializers.normal(stddev=1.0), + config=cfg, + mesh=self.mesh, + rngs=rngs, + ) + batch = cfg.global_batch_size_to_train_on + seq_len = cfg.max_target_length + ids = jax.random.randint(self.rng, (batch, seq_len), 0, cfg.vocab_size) + segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) + + call_kwargs = { + "decoder_segment_ids": segment_ids, + "deterministic": True, + "model_mode": MODEL_MODE_TRAIN, + } + # Force the dynamic_graph_init rebuild path so nnx.Param leaves (with real, non-default + # param_scan_axis metadata) actually flow through nnx_add_and_sync_scan_axis -- see docstring. + decoder.disable_quant_stats_update = True + logits1, _, _ = decoder(shared_embedding, ids, positions, **call_kwargs) + logits2, _, _ = decoder(shared_embedding, ids, positions, **call_kwargs) + + self.assertEqual(logits1.shape, (batch, seq_len, cfg.vocab_size)) + self.assertTrue(jnp.all(jnp.isfinite(logits1))) + # The scanned params must survive the post-scan axis restoration + write-back round trip: a + # second forward re-reads the persisted params, so a mis-restored scan axis would crash or drift. + np.testing.assert_allclose(np.array(logits1), np.array(logits2), rtol=1e-5, atol=1e-5) + if __name__ == "__main__": unittest.main() @@ -998,34 +1110,36 @@ def test_gemma4_scanned_layers(self): shared_embedding = self._make_shared_embedding(cfg) ids, segment_ids, positions = self._make_token_inputs(cfg) - logits, _, _ = decoder( - shared_embedding, - ids, - positions, - decoder_segment_ids=segment_ids, - deterministic=True, - model_mode=MODEL_MODE_TRAIN, - ) + def _forward(): + out, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return out + + logits = _forward() self.assertEqual( logits.shape, (cfg.global_batch_size_to_train_on, cfg.max_target_length, cfg.vocab_size), ) + _, remainder_params, remainder_rest = nnx.split(decoder.layers_remainder, nnx.Param, ...) + self.assertGreater(len(jax.tree_util.tree_leaves(remainder_params)), 0, "layers_remainder has no params to perturb") + perturbed_params = jax.tree.map(lambda x: x + 10.0, remainder_params) + nnx.update(decoder.layers_remainder, nnx.State.merge(perturbed_params, remainder_rest)) + logits_perturbed = _forward() + + self.assertFalse( + bool(jnp.allclose(logits, logits_perturbed)), + msg="gemma4 remainder-only forward is invariant to layers_remainder params -> remainder block not applied", + ) + def test_gemma4_block_external_kv_cache_matches_scanned_path(self): - """External-kv-cache path must numerically match the scanned (kv=None) path. - - Guards the real stacked-parameter slice/re-stack in - ``Gemma4ScannableBlock._forward_with_external_kv_cache`` (``nnx.split`` of the - scanned local stack, ``param_scan_axis`` moveaxis, per-layer merge, re-stack, - ``nnx.update``) against the ``jax.lax.scan`` path. The mock-based - ``TestGemma4ScannableBlock`` tests cover call ordering / cache collection but - use trivial params, so they never exercise the real stacked-param mechanics. - - Uses ``model_mode=TRAIN`` with ``dot_product`` attention: attention then - ignores the external caches (passing them straight through), so both paths - compute the same forward and only the loop mechanism (scan vs static unroll) - differs -- any mismatch is a slice/re-stack bug. - """ + """External-kv-cache path must numerically match the scanned (kv=None) path.""" cfg = _make_config( decoder_block="gemma4", scan_layers=True, @@ -1105,6 +1219,9 @@ def test_gemma4_small_decoder(self): "weight_dtype=float32", "hidden_size_per_layer_input=128", "vocab_size_per_layer_input=256", + "per_device_batch_size=1.0", + "max_target_length=16", + "max_prefill_predict_length=4", "vocab_size=256", "max_target_length=128", "per_device_batch_size=1.0", @@ -1177,6 +1294,9 @@ def test_gemma4_small_decoder_with_mock_cache_and_ple(self): "weight_dtype=float32", "hidden_size_per_layer_input=128", "vocab_size_per_layer_input=256", + "per_device_batch_size=1.0", + "max_target_length=16", + "max_prefill_predict_length=4", "vocab_size=256", "max_target_length=128", "per_device_batch_size=1.0", @@ -1259,3 +1379,846 @@ def mock_donor_idx(lyr, layer_types, num_kv_shared): model_mode=MODEL_MODE_TRAIN, kv_caches=kv_caches, ) + + +def _assert_grad_parity( + test_case, ref_leaves, other_leaves, *, what, rtol=2e-2, per_leaf_rtol=5e-2, per_leaf_atol_frac=1e-3 +): + """Assert two gradient leaf-lists agree, by BOTH an aggregate and a per-leaf relative bound.""" + ref_leaves, other_leaves = list(ref_leaves), list(other_leaves) + test_case.assertEqual(len(ref_leaves), len(other_leaves), f"{what}: grad pytrees differ in leaf count") + test_case.assertGreater(len(ref_leaves), 0, f"{what}: no gradients") + for g_ref, g_other in zip(ref_leaves, other_leaves): + test_case.assertEqual(g_ref.shape, g_other.shape, f"{what}: gradient shape mismatch") + test_case.assertTrue( + all(bool(jnp.all(jnp.isfinite(g))) for g in ref_leaves + other_leaves), f"{what}: non-finite gradient" + ) + test_case.assertTrue(any(bool(jnp.any(g != 0)) for g in ref_leaves), f"{what}: reference backward is all-zero") + test_case.assertTrue(any(bool(jnp.any(g != 0)) for g in other_leaves), f"{what}: backward produced all-zero grads") + ref = jnp.concatenate([g.astype(jnp.float32).ravel() for g in ref_leaves]) + oth = jnp.concatenate([g.astype(jnp.float32).ravel() for g in other_leaves]) + rel_l2 = float(jnp.linalg.norm(oth - ref) / (jnp.linalg.norm(ref) + 1e-12)) + + test_case.assertLess(rel_l2, rtol, f"{what}: relative L2 gradient error {rel_l2:.4%} exceeds rtol={rtol:.2%}") + total_norm = float(jnp.linalg.norm(ref)) + for i, (g_ref, g_other) in enumerate(zip(ref_leaves, other_leaves)): + r = g_ref.astype(jnp.float32) + leaf_norm = float(jnp.linalg.norm(r)) + leaf_err = float(jnp.linalg.norm(g_other.astype(jnp.float32) - r)) + allowed = per_leaf_rtol * leaf_norm + per_leaf_atol_frac * max(total_norm, 1e-12) + test_case.assertLessEqual( + leaf_err, + allowed, + f"{what}: leaf {i} (shape {tuple(g_ref.shape)}, {leaf_norm / max(total_norm, 1e-12):.2%} of total " + f"gradient norm) has error {leaf_err:.3e}, exceeding the allowed " + f"{per_leaf_rtol:.0%}*leaf + {per_leaf_atol_frac:.0e}*total = {allowed:.3e}. " + "The aggregate check can miss this when the leaf is small.", + ) + + +class TestNNXDecoderDeepseek4(unittest.TestCase): + """Parity tests for DeepSeek-V4 (deepseek4) decoder-level handling in NNXDecoder.""" + + def _make_deepseek4_config( + self, + scan_layers=False, + num_decoder_layers=5, + first_num_hash_layers=3, + compress_ratios=(0, 0, 4, 128, 4), + remat_policy="full", + ): + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + override_model_config=True, + per_device_batch_size=1.0, + run_name="deepseek4_nnx_test", + enable_checkpointing=False, + model_name="deepseek4-284b", + attention="dot_product", + remat_policy=remat_policy, + # Dense MoE (sparse_matmul=False) so the forward runs on any backend; + # the megablox GMM path is a TPU-only Pallas kernel. + sparse_matmul=False, + megablox=False, + base_num_decoder_layers=num_decoder_layers, + base_emb_dim=256, + base_mlp_dim=512, + base_moe_mlp_dim=512, + base_num_query_heads=4, + base_num_kv_heads=1, + num_experts=8, + num_experts_per_tok=2, + shared_experts=1, + first_num_hash_layers=first_num_hash_layers, + compress_ratios=list(compress_ratios), + indexer_head_dim=64, + indexer_n_heads=4, + indexer_topk=8, + head_dim=64, + q_lora_rank=64, + o_lora_rank=64, + o_groups=2, + kv_lora_rank=64, + # seq_len must be >= the largest compress_ratio (128) so HCA layers produce >=1 compressed block. + max_target_length=256, + max_prefill_predict_length=64, + vocab_size=256, + scan_layers=scan_layers, + dtype="float32", + weight_dtype="float32", + activations_in_float32=True, + sliding_window_size=8, + ) + + def test_construct_non_scan_does_not_raise(self): + """NNXDecoder(deepseek4) must construct; get_norm_layer must support deepseek4 (RMSNorm).""" + cfg = self._make_deepseek4_config(scan_layers=False) + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertIsInstance(decoder.decoder_norm, RMSNorm) + self.assertTrue(decoder.is_deepseek4) + + def test_scan_construction_registers_prefix_via_existing_helper_regression(self): + """Test that NNXDecoder(deepseek4, scan_layers=True) registers the prefix layers via the existing helper""" + cfg = self._make_deepseek4_config(scan_layers=True) + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertEqual(decoder.num_prefix_layers, cfg.first_num_hash_layers) + for i in range(cfg.first_num_hash_layers): + self.assertTrue(hasattr(decoder, f"layers_{i}"), f"prefix layer layers_{i} was not registered") + + def test_get_decoder_layers_registers_deepseek4(self): + """Regression guard: NNXDecoder.get_decoder_layers layer_map MUST contain DEEPSEEK4.""" + from maxtext.models import deepseek4 # pylint: disable=import-outside-toplevel + + cfg = self._make_deepseek4_config(scan_layers=False) + dec = NNXDecoder(config=cfg, mesh=_make_mesh(cfg), model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertEqual(dec.get_decoder_layers(), [deepseek4.DeepSeek4DecoderLayer]) + + cfg_s = self._make_deepseek4_config(scan_layers=True) + dec_s = NNXDecoder( + config=cfg_s, mesh=_make_mesh(cfg_s), model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1) + ) + self.assertEqual(dec_s.get_decoder_layers(), [deepseek4.DeepSeek4ScannableBlock]) + + def test_linen_pipeline_dispatch_includes_deepseek4(self): + """Linen Decoder._get_nnx_decoder_block_classes (pipeline path) must include DEEPSEEK4.""" + from maxtext.layers import decoders # pylint: disable=import-outside-toplevel + from maxtext.models import deepseek4 # pylint: disable=import-outside-toplevel + + cfg = self._make_deepseek4_config(scan_layers=False) + mesh = _make_mesh(cfg) + dec = decoders.Decoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN) + self.assertEqual(dec._get_nnx_decoder_block_classes(), [deepseek4.DeepSeek4DecoderLayer]) # pylint: disable=protected-access + + cfg_s = self._make_deepseek4_config(scan_layers=True) + dec_s = decoders.Decoder(config=cfg_s, mesh=_make_mesh(cfg_s), model_mode=MODEL_MODE_TRAIN) + self.assertEqual(dec_s._get_nnx_decoder_block_classes(), [deepseek4.DeepSeek4ScannableBlock]) # pylint: disable=protected-access + + def _build_and_run(self, cfg): + """Builds an NNXDecoder + shared embedding for ``cfg`` and runs one train-mode forward pass.""" + mesh = _make_mesh(cfg) + rngs = nnx.Rngs(params=0, dropout=1) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + shared_embedding = Embed( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + embedding_init=nn.initializers.normal(stddev=1.0), + config=cfg, + mesh=mesh, + rngs=rngs, + ) + batch = cfg.global_batch_size_to_train_on + seq_len = cfg.max_target_length + ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq_len), 0, cfg.vocab_size) + segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) + + # AOT structural check of the forward graph -- traces the full prefix + HCA/CSA stack + # symbolically (no execution). Functionalize decoder+embedding via split/merge so eval_shape can + # trace the stateful NNX modules without a cross-trace RngCount mutation. + graphdef, state = nnx.split((decoder, shared_embedding)) + + def _forward_from_state(state_in, ids_in): + dec, emb = nnx.merge(graphdef, state_in) + out, _, _ = dec( + emb, ids_in, positions, decoder_segment_ids=segment_ids, deterministic=True, model_mode=MODEL_MODE_TRAIN + ) + return out + + aot_logits = jax.eval_shape(_forward_from_state, state, ids) + + logits, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return decoder, logits, (batch, seq_len, cfg.vocab_size), aot_logits + + def _assert_forward_is_real(self, logits, aot_logits, expected): + """Assertions that exercise the forward beyond finiteness (reviewer: isfinite is weak).""" + self.assertEqual(aot_logits.shape, expected) + self.assertEqual(aot_logits.dtype, jnp.float32) + self.assertEqual(logits.dtype, jnp.float32) + self.assertGreater(float(jnp.std(logits)), 1e-2) + self.assertFalse( + bool(jnp.allclose(logits[:, 0, :], logits[:, -1, :], rtol=1e-2, atol=1e-2)), + msg="deepseek4 logits are position-invariant -> forward is degenerate", + ) + + def _deepseek4_decoder_loss_and_grads(self, cfg): + """Build the DeepSeek-V4 decoder + shared embedding for cfg and return (loss, grads) for a + sum-of-squares loss differentiated (nnx.value_and_grad) wrt every decoder + embedding Param. + Fixed seed, so two builds differing only in remat_policy share the same params.""" + mesh = _make_mesh(cfg) + rngs = nnx.Rngs(params=0, dropout=1) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + shared_embedding = Embed( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + embedding_init=nn.initializers.normal(stddev=1.0), + config=cfg, + mesh=mesh, + rngs=rngs, + ) + batch = cfg.global_batch_size_to_train_on + seq_len = cfg.max_target_length + ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq_len), 0, cfg.vocab_size) + segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) + + def loss_fn(dec, emb): + out, _, _ = dec( + emb, ids, positions, decoder_segment_ids=segment_ids, deterministic=True, model_mode=MODEL_MODE_TRAIN + ) + return jnp.sum(out.astype(jnp.float32) ** 2) + + return nnx.value_and_grad(loss_fn, argnums=(0, 1))(decoder, shared_embedding) + + def _assert_decoder_grad_parity(self, scan_layers): + """Run the full DeepSeek-V4 decoder under two remat policies ('full' vs 'minimal') and assert + matching loss and gradients. Gradients are compared by relative L2 error (see _assert_grad_parity, + tolerant of TPU bf16 rounding); loss at rtol=1e-2.""" + loss_full, grads_full = self._deepseek4_decoder_loss_and_grads( + self._make_deepseek4_config(scan_layers=scan_layers, remat_policy="full") + ) + loss_min, grads_min = self._deepseek4_decoder_loss_and_grads( + self._make_deepseek4_config(scan_layers=scan_layers, remat_policy="minimal") + ) + np.testing.assert_allclose(np.array(loss_full), np.array(loss_min), rtol=1e-2, atol=1e-2) + _assert_grad_parity( + self, jax.tree.leaves(grads_full), jax.tree.leaves(grads_min), what="deepseek4 decoder full-vs-minimal remat" + ) + + def test_scan_init_builds_prefix_and_scanned_blocks(self): + """scan init builds first_num_hash_layers unrolled prefix layers + a scanned block stack.""" + cfg = self._make_deepseek4_config(scan_layers=True) + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertEqual(decoder.num_prefix_layers, cfg.first_num_hash_layers) + for i in range(cfg.first_num_hash_layers): + self.assertTrue(hasattr(decoder, f"layers_{i}")) + # num_decoder_layers=5, first_num_hash_layers=3 -> (5-3)//2 = 1 scanned HCA/CSA block + self.assertIsNotNone(decoder.scanned_blocks) + + def test_scan_init_odd_non_prefix_layers_raises(self): + """scan init with an ODD non-prefix layer count must raise AssertionError (no silent drop).""" + cfg = self._make_deepseek4_config(scan_layers=True, num_decoder_layers=6, first_num_hash_layers=3) + mesh = _make_mesh(cfg) + with self.assertRaises(AssertionError): + NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + + def test_scan_init_even_non_prefix_layers_constructs(self): + """scan init with an EVEN non-prefix layer count still constructs (companion to the odd guard).""" + cfg = self._make_deepseek4_config(scan_layers=True, num_decoder_layers=5, first_num_hash_layers=3) + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertEqual(decoder.num_prefix_layers, 3) + self.assertIsNotNone(decoder.scanned_blocks) + + def test_non_scan_init_builds_deepseek4_layers(self): + """non-scan init builds num_decoder_layers DeepSeek4DecoderLayer instances, each with the + correct GLOBAL layer_idx (0..num_decoder_layers-1) baked in at construction. + """ + from maxtext.models import deepseek4 # pylint: disable=import-outside-toplevel + + cfg = self._make_deepseek4_config(scan_layers=False) + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + for i in range(cfg.num_decoder_layers): + layer = getattr(decoder, f"layers_{i}") + self.assertIsInstance(layer, deepseek4.DeepSeek4DecoderLayer) + self.assertEqual(layer.layer_idx, i, f"layers_{i} has the wrong global layer_idx") + self.assertEqual( + layer.self_attention.compress_ratio, + cfg.compress_ratios[i], + f"layers_{i}.self_attention.compress_ratio does not match config.compress_ratios[{i}] -- " + "layer_idx was not correctly routed to this layer's construction", + ) + + def test_forward_non_scan(self): + """deepseek4 non-scan forward returns correct logits shape and finite values.""" + cfg = self._make_deepseek4_config(scan_layers=False) + decoder, logits, expected, aot_logits = self._build_and_run(cfg) + self.assertEqual(logits.shape, expected) + self._assert_forward_is_real(logits, aot_logits, expected) + self.assertTrue(jnp.all(jnp.isfinite(logits))) # secondary + self.assertTrue(hasattr(decoder, "layers_0")) + # _assert_decoder_grad_parity below compares remat_policy='full' vs 'minimal' on the SAME + # construction, so it is blind to any construction bug shared by both (e.g. every layer + # silently getting layer_idx=-1 -- see test_non_scan_init_builds_deepseek4_layers). Anchor this + # forward test to real per-layer construction: layers_0 must route to compress_ratios[0]. + self.assertEqual(decoder.layers_0.self_attention.compress_ratio, cfg.compress_ratios[0]) + self._assert_decoder_grad_parity(scan_layers=cfg.scan_layers) + + def test_forward_scan(self): + """deepseek4 scan forward (unrolled hash-routing prefix + scanned HCA/CSA blocks).""" + cfg = self._make_deepseek4_config(scan_layers=True) + _, logits, expected, aot_logits = self._build_and_run(cfg) + self.assertEqual(logits.shape, expected) + self._assert_forward_is_real(logits, aot_logits, expected) + self.assertTrue(jnp.all(jnp.isfinite(logits))) # secondary + self._assert_decoder_grad_parity(scan_layers=cfg.scan_layers) + + def test_deepseek4_builds_learned_hc_head_collapse(self): + """DeepSeek-V4 must build a LEARNED hyper-head to collapse its mhc_expansion_rate mHC streams.""" + cfg = self._make_deepseek4_config(scan_layers=False) + self.assertGreater(cfg.mhc_expansion_rate, 1) # deepseek4-284b -> mhc_expansion_rate=4 + mesh = _make_mesh(cfg) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + self.assertTrue(hasattr(decoder, "hc_head"), "deepseek4 decoder is missing the learned hc_head collapse module") + self.assertIsInstance(decoder.hc_head, mhc.DeepSeek4HyperHead) + for name in ("hc_fn", "hc_base", "hc_scale"): + self.assertIsInstance(getattr(decoder.hc_head, name), nnx.Param, f"hc_head.{name} must be a materialized nnx.Param") + + def test_deepseek4_collapse_is_wired_to_hc_head(self): + """The final mHC-stream collapse must FLOW THROUGH ``self.hc_head``, not the unweighted mhc_reduce.""" + cfg = self._make_deepseek4_config(scan_layers=False) + mesh = _make_mesh(cfg) + rngs = nnx.Rngs(params=0, dropout=1) + decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + shared_embedding = Embed( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + embedding_init=nn.initializers.normal(stddev=1.0), + config=cfg, + mesh=mesh, + rngs=rngs, + ) + batch = cfg.global_batch_size_to_train_on + seq_len = cfg.max_target_length + ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq_len), 0, cfg.vocab_size) + segment_ids = jnp.full((batch, seq_len), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq_len)[None], (batch, seq_len)) + + def _forward(): + out, _, _ = decoder( + shared_embedding, + ids, + positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + ) + return out + + hidden_before = _forward() + # Perturb ONLY the learned hyper-head; nothing else in the graph changes. + decoder.hc_head.hc_scale.value += 10.0 + decoder.hc_head.hc_fn.value *= 5.0 + hidden_after = _forward() + + self.assertFalse( + bool(jnp.allclose(hidden_before, hidden_after)), + msg="forward output is invariant to hc_head params -> collapse is NOT wired through hc_head", + ) + + +class TestNNXPipelineStages(unittest.TestCase): + """Tests for the NNX pipeline-stage modules (NNXSequentialPipelineStage / NNXScannedPipelineStage), + including per-stage remat + params-only host-offload (set_remat_policy_on_layers_per_stage / + parameter_memory_host_offload) that the nnx-based-pipeline migration dropped. + """ + + def setUp(self): + super().setUp() + self.cfg = _make_config() + self.mesh = _make_mesh(self.cfg) + + def _inputs(self, cfg): + batch = cfg.global_batch_size_to_train_on + seq = cfg.max_target_length + inputs = jax.random.normal(jax.random.PRNGKey(0), (batch, seq, cfg.emb_dim)).astype(cfg.dtype) + segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq)) + return inputs, segment_ids, positions + + def _run_stage(self, stage_cls, remat_policy, num_layers=2, config=None, use_mesh=False): + """Builds a pipeline stage of num_layers NNXDecoderLayers and runs one train-mode forward.""" + cfg = config if config is not None else self.cfg + mesh = _make_mesh(cfg) if config is not None else self.mesh + + def _build_and_run(): + stage = stage_cls( + NNXDecoderLayer, + num_layers, + cfg, + mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=remat_policy, + apply_remat=remat_policy is not None, + ) + inputs, segment_ids, positions = self._inputs(cfg) + out = stage(inputs, segment_ids, positions, True, MODEL_MODE_TRAIN) + return out[0] if isinstance(out, tuple) else out + + if use_mesh: + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + return _build_and_run() + return _build_and_run() + + def _assert_stage_forward_is_real(self, out, inputs): + """Validate that a pipeline stage forward returns finite values and is not a degenerate passthrough.""" + self.assertTrue(jnp.all(jnp.isfinite(out))) + self.assertGreater(float(jnp.std(out)), 1e-2) + self.assertFalse( + bool(jnp.allclose(out, inputs, rtol=1e-3, atol=1e-3)), + msg="stage output equals the raw input -> layers were not applied (stage is a passthrough)", + ) + + def test_sequential_stage_forward_shape(self): + """NNXSequentialPipelineStage forward returns [batch, seq, emb] and finite values.""" + inputs, _, _ = self._inputs(self.cfg) + out = self._run_stage(NNXSequentialPipelineStage, None) + self.assertEqual(out.shape, inputs.shape) + self._assert_stage_forward_is_real(out, inputs) + + def test_scanned_stage_forward_shape(self): + """NNXScannedPipelineStage forward returns [batch, seq, emb] and finite values.""" + inputs, _, _ = self._inputs(self.cfg) + out = self._run_stage(NNXScannedPipelineStage, None) + self.assertEqual(out.shape, inputs.shape) + self._assert_stage_forward_is_real(out, inputs) + + def test_sequential_stage_remat_is_output_transparent(self): + """Per-stage remat on a sequential stage must not change the forward output.""" + out_no_remat = self._run_stage(NNXSequentialPipelineStage, None) + out_remat = self._run_stage(NNXSequentialPipelineStage, jax.checkpoint_policies.nothing_saveable) + np.testing.assert_allclose(np.array(out_no_remat), np.array(out_remat), rtol=1e-5, atol=1e-5) + + def test_scanned_stage_remat_is_output_transparent(self): + """Per-stage remat on a scanned stage must not change the forward output.""" + out_no_remat = self._run_stage(NNXScannedPipelineStage, None) + out_remat = self._run_stage(NNXScannedPipelineStage, jax.checkpoint_policies.nothing_saveable) + np.testing.assert_allclose(np.array(out_no_remat), np.array(out_remat), rtol=1e-5, atol=1e-5) + + def test_scanned_stage_remat_does_not_stack_params(self): + """The scanned pipeline stage must NOT stack read-only params across its layers inside jax.lax.scan.""" + real_scan = jax.lax.scan + captured = {} + + def spy_scan(*args, **kwargs): + carry, ys = real_scan(*args, **kwargs) + # The pipeline-stage scan is the only one whose stacked output is an nnx.State. + if isinstance(ys, nnx.State): + captured["ys"] = ys + return carry, ys + + for policy in (None, jax.checkpoint_policies.nothing_saveable): + captured.clear() + with patch("jax.lax.scan", spy_scan): + self._run_stage(NNXScannedPipelineStage, policy) + self.assertIn("ys", captured, f"jax.lax.scan produced no nnx.State ys for policy={policy!r}") + stacked_params, _ = captured["ys"].split(nnx.Param, ...) + param_leaves = jax.tree_util.tree_leaves(stacked_params) + self.assertEqual( + len(param_leaves), + 0, + f"scanned stage stacked {len(param_leaves)} nnx.Param leaf/leaves across layers for " + f"policy={policy!r} (apply_remat={policy is not None}); the scan body must drop read-only " + f"params from its returned state to avoid a [num_layers, *param] transient allocation.", + ) + + def _stage_value_and_grad(self, stage_cls, apply_remat, remat_policy, num_layers=2, config=None, use_mesh=False): + """Builds a pipeline stage (fixed seed) and returns (loss, input_grad, param_grads) for + ``loss = sum(stage(x)**2)``. + + """ + cfg = config if config is not None else self.cfg + mesh = _make_mesh(cfg) if config is not None else self.mesh + + def _build_and_grad(): + stage = stage_cls( + NNXDecoderLayer, + num_layers, + cfg, + mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=remat_policy, + apply_remat=apply_remat, + ) + inputs, segment_ids, positions = self._inputs(cfg) + + def loss_fn(x, module): + out = module(x, segment_ids, positions, True, MODEL_MODE_TRAIN) + out = out[0] if isinstance(out, tuple) else out + return jnp.sum(out.astype(jnp.float32) ** 2) + + loss, (input_grad, param_grads) = nnx.value_and_grad(loss_fn, argnums=(0, 1))(inputs, stage) + return loss, input_grad, param_grads + + if use_mesh: + with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules): + return _build_and_grad() + return _build_and_grad() + + def _assert_stage_grad_parity(self, stage_cls, num_layers=2): + """remat vs no-remat: matching loss + gradients (wrt inputs AND params), plus a real (nonzero) + backward. jax.checkpoint recomputes activations in the backward pass but is mathematically + transparent; gradients are compared by relative L2 error (see _assert_grad_parity), which tolerates + the bf16 rounding of the rematerialized backward on TPU. Loss parity holds everywhere.""" + loss_ref, xgrad_ref, pgrad_ref = self._stage_value_and_grad( + stage_cls, apply_remat=False, remat_policy=None, num_layers=num_layers + ) + loss_remat, xgrad_remat, pgrad_remat = self._stage_value_and_grad( + stage_cls, apply_remat=True, remat_policy=jax.checkpoint_policies.nothing_saveable, num_layers=num_layers + ) + # Loss parity (reproducible on all platforms). + np.testing.assert_allclose(np.array(loss_remat), np.array(loss_ref), rtol=1e-2, atol=1e-2) + # Input-gradient parity. + _assert_grad_parity(self, [xgrad_ref], [xgrad_remat], what="stage remat input-grad") + # Per-Param gradient parity across the whole pytree. + _assert_grad_parity( + self, + jax.tree_util.tree_leaves(pgrad_ref), + jax.tree_util.tree_leaves(pgrad_remat), + what="stage remat param-grads", + ) + + def test_sequential_stage_remat_grad_parity(self): + """Backward parity: a sequential stage's per-stage remat must reproduce the no-remat loss and + gradients (wrt inputs AND params). jax.checkpoint recomputes activations in the backward pass but + is mathematically transparent, so value_and_grad must match.""" + self._assert_stage_grad_parity(NNXSequentialPipelineStage) + + def test_scanned_stage_remat_grad_parity(self): + """Backward parity for the scanned pipeline stage (remat inside the jax.lax.scan body): remat vs + no-remat must yield identical loss and gradients (wrt inputs AND params).""" + self._assert_stage_grad_parity(NNXScannedPipelineStage) + + def test_single_layer_stage_remat_is_output_transparent(self): + """num_layers_per_pipeline_stage==1: a 1-layer stage with remat must match the no-remat output.""" + out_no_remat = self._run_stage(NNXSequentialPipelineStage, None, num_layers=1) + out_remat = self._run_stage(NNXSequentialPipelineStage, jax.checkpoint_policies.nothing_saveable, num_layers=1) + np.testing.assert_allclose(np.array(out_no_remat), np.array(out_remat), rtol=1e-5, atol=1e-5) + + def test_single_layer_stage_remat_grad_parity(self): + """num_layers_per_pipeline_stage==1: BACKWARD parity for the single-layer stage remat path (a + distinct builder branch). remat vs no-remat must yield identical loss + gradients.""" + self._assert_stage_grad_parity(NNXSequentialPipelineStage, num_layers=1) + + @pytest.mark.tpu_only + def test_remat_with_host_offload_is_output_transparent(self): + """Per-stage remat + params-only host-offload (parameter_memory_host_offload) must not change output.""" + offload_cfg = _make_config(parameter_memory_host_offload=True) + # Both runs inside the mesh so params are sharded consistently with the inputs; the offloaded run + # additionally exercises jax.device_put(params, Space.Device) inside the per-stage remat. + plain = self._run_stage(NNXSequentialPipelineStage, None, config=offload_cfg, use_mesh=True) + offloaded = self._run_stage( + NNXSequentialPipelineStage, jax.checkpoint_policies.nothing_saveable, config=offload_cfg, use_mesh=True + ) + np.testing.assert_allclose(np.array(plain), np.array(offloaded), rtol=1e-5, atol=1e-5) + + @pytest.mark.tpu_only + def test_remat_with_host_offload_grad_is_transparent(self): + """BACKWARD parity for per-stage remat + params-only host-offload: the loss AND gradients (wrt + inputs and every Param) with parameter_memory_host_offload must match the no-offload/no-remat + path -- host-offload only moves where params live, it must not change the math. + + tpu_only for the same reason as the forward host-offload test: jax.device_put(..., device_space()) + targets TPU host memory; on CPU fake-multi-device it pins params to one device and breaks sharding. + """ + offload_cfg = _make_config(parameter_memory_host_offload=True) + loss_ref, xgrad_ref, pgrad_ref = self._stage_value_and_grad( + NNXSequentialPipelineStage, apply_remat=False, remat_policy=None, config=offload_cfg, use_mesh=True + ) + loss_off, xgrad_off, pgrad_off = self._stage_value_and_grad( + NNXSequentialPipelineStage, + apply_remat=True, + remat_policy=jax.checkpoint_policies.nothing_saveable, + config=offload_cfg, + use_mesh=True, + ) + np.testing.assert_allclose(np.array(loss_off), np.array(loss_ref), rtol=1e-2, atol=1e-2) + # Gradients are compared by relative L2 error, which absorbs the bf16 rounding of the rematerialized + # backward on TPU. The offload path is additionally pinned output-transparent at 1e-5 by + # test_remat_with_host_offload_is_output_transparent. + _assert_grad_parity(self, [xgrad_ref], [xgrad_off], what="host-offload input-grad") + _assert_grad_parity( + self, + jax.tree_util.tree_leaves(pgrad_ref), + jax.tree_util.tree_leaves(pgrad_off), + what="host-offload param-grads", + ) + + +class TestNNXPerStageRematApplied(unittest.TestCase): + """Guards the per-stage remat parity bug: remat_policy='full' resolves to get_remat_policy()==None, + which is a VALID 'full rematerialization' policy (matching Linen nn.remat(policy=None)). Gating on + `remat_policy is not None` silently dropped remat for the default 'full' policy. The builder must + apply per-stage remat whenever set_remat_policy_on_layers_per_stage=True, regardless of policy value. + """ + + def _build_stage(self, remat_policy, num_layers_per_pipeline_stage=2, flag_on=True): + """Build a pipeline stage via the NNXDecoder builder for the given remat policy + flag.""" + cfg = _make_config( + remat_policy=remat_policy, + set_remat_policy_on_layers_per_stage=flag_on, + num_layers_per_pipeline_stage=num_layers_per_pipeline_stage, + scan_layers_per_stage=False, + scan_layers=False, + ) + mesh = _make_mesh(cfg) + dec = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) + stage = dec._get_pipeline_stage_module(dec.get_decoder_layers(), nnx.Rngs(params=0, dropout=1)) # pylint: disable=protected-access + return cfg, stage + + def _run_forward(self, cfg, stage): + seq = cfg.max_target_length + x = jax.random.normal(jax.random.PRNGKey(0), (1, seq, cfg.emb_dim)).astype(cfg.dtype) + seg = jnp.full((1, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) + pos = jnp.broadcast_to(jnp.arange(seq)[None], (1, seq)) + out = stage(x, seg, pos, True, MODEL_MODE_TRAIN) + return out[0] if isinstance(out, tuple) else out + + def test_full_policy_applies_remat(self): + """remat_policy='full' (get_remat_policy()==None) must STILL apply per-stage remat (the bug).""" + cfg, stage = self._build_stage("full") + self.assertTrue(stage.apply_remat, "per-stage remat dropped for remat_policy='full'") + self.assertIsNone(stage.remat_policy) # 'full' -> None policy == full remat + out = self._run_forward(cfg, stage) + self.assertTrue(jnp.all(jnp.isfinite(out))) + + def test_minimal_policy_applies_remat(self): + """Sanity: a non-None policy also applies remat and runs.""" + cfg, stage = self._build_stage("minimal") + self.assertTrue(stage.apply_remat) + self.assertIsNotNone(stage.remat_policy) + out = self._run_forward(cfg, stage) + self.assertTrue(jnp.all(jnp.isfinite(out))) + + @staticmethod + def _param_key_paths(module): + """Sorted '/'-joined nnx.Param key paths (the leaves the pipeline stacks via nnx.split).""" + _, params, _ = nnx.split(module, nnx.Param, ...) + return sorted( + "/".join(str(getattr(k, "key", k)) for k in path) for path, _ in jax.tree_util.tree_flatten_with_path(params)[0] + ) + + def test_single_layer_stage_keeps_params_top_level_linen_parity(self): + """num_layers_per_pipeline_stage==1 + set_remat_policy_on_layers_per_stage: the remat-applying + stage must keep params TOP-LEVEL -- identical nnx.Param key paths to the flag-off bare layer, + NO 'layers_0' nesting -- matching Linen nn.remat (param-tree transparent). It must be a subclass + of the base decoder layer (IS-A layer), not a NNXSequentialPipelineStage wrapper. Inverts the + prior test that accepted the layers_0-nesting wrapper; also pins remat output-/grad-transparency + vs the bare layer (same seed => identical params).""" + cfg, stage_on = self._build_stage("full", num_layers_per_pipeline_stage=1, flag_on=True) + _, bare_off = self._build_stage("full", num_layers_per_pipeline_stage=1, flag_on=False) + + # Per-stage remat still applied ('full' policy -> None == full rematerialization). + self.assertTrue(stage_on.apply_remat) + self.assertIsNone(stage_on.remat_policy) + + # IS-A base layer, NOT the wrapper -> no layers_0 nesting. base_stage_cls resolves to the + # concrete model decoder layer (LlamaDecoderLayer for base.yml), which is exactly the class of + # the flag-off bare layer; the remat stage must subclass it (Linen nn.remat is a same-class wrap). + self.assertIsInstance(stage_on, type(bare_off)) + self.assertNotIsInstance(stage_on, (NNXSequentialPipelineStage, NNXScannedPipelineStage)) + + # Core parity: identical param key paths, none under layers_0. + on_paths = self._param_key_paths(stage_on) + off_paths = self._param_key_paths(bare_off) + self.assertEqual(on_paths, off_paths) + self.assertFalse( + any("layers_0" in p for p in on_paths), + f"flag-on single-layer stage nested params under layers_0: {on_paths}", + ) + + # Output transparency: remat must not change the forward result vs the bare layer (same params). + out_on = self._run_forward(cfg, stage_on) + out_off = self._run_forward(cfg, bare_off) + np.testing.assert_allclose(np.array(out_on), np.array(out_off), rtol=1e-5, atol=1e-5) + + # Grad transparency: gradients wrt every Param match the bare (no-remat) layer. + seq = cfg.max_target_length + x = jax.random.normal(jax.random.PRNGKey(0), (1, seq, cfg.emb_dim)).astype(cfg.dtype) + seg = jnp.full((1, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) + pos = jnp.broadcast_to(jnp.arange(seq)[None], (1, seq)) + + def _param_grads(module): + def loss_fn(m): + out = m(x, seg, pos, True, MODEL_MODE_TRAIN) + out = out[0] if isinstance(out, tuple) else out + return jnp.sum(out.astype(jnp.float32) ** 2) + + return jax.tree_util.tree_leaves(nnx.grad(loss_fn)(module)) + + _assert_grad_parity( + self, + _param_grads(bare_off), + _param_grads(stage_on), + what="single-layer top-level remat param-grads", + ) + + def test_flag_off_single_layer_returns_bare_layer(self): + """Flag off: num_layers==1 returns the bare layer (no stage wrapper) -- unchanged behavior.""" + _, stage = self._build_stage("full", num_layers_per_pipeline_stage=1, flag_on=False) + self.assertNotIsInstance(stage, (NNXSequentialPipelineStage, NNXScannedPipelineStage)) + + +class TestNNXStageRematAppliedInJaxpr(unittest.TestCase): + """Guards against a silently no-op jax.checkpoint wrap in the shared helper + ``_run_stage_layer_with_remat`` (used by NNXSequentialPipelineStage, NNXScannedPipelineStage, and + the single-layer remat stage from ``_make_single_layer_remat_stage_cls``). + + """ + + def setUp(self): + super().setUp() + self.cfg = _make_config() + self.mesh = _make_mesh(self.cfg) + + def _inputs(self): + cfg = self.cfg + batch = cfg.global_batch_size_to_train_on + seq = cfg.max_target_length + inputs = jax.random.normal(jax.random.PRNGKey(0), (batch, seq, cfg.emb_dim)).astype(cfg.dtype) + segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq)) + return inputs, segment_ids, positions + + def _stage_fwd_jaxpr_text(self, stage): + """``str(jax.make_jaxpr(...))`` of one stage forward pass, via nnx.split/merge -> pure fn (the + same split/merge shape _run_stage_layer_with_remat itself uses internally, so make_jaxpr traces a + real, representative call).""" + inputs, segment_ids, positions = self._inputs() + graphdef, state = nnx.split(stage) + + def fwd(state_in, x_in): + merged = nnx.merge(graphdef, state_in) + out = merged(x_in, segment_ids, positions, True, MODEL_MODE_TRAIN) + return out[0] if isinstance(out, tuple) else out + + return jax.make_jaxpr(fwd)(state, inputs) + + def _assert_remat_applied(self, stage, apply_remat): + """Asserts that the stage's forward jaxpr contains a remat primitive iff apply_remat is True.""" + has_remat = _jaxpr_contains_primitive(self._stage_fwd_jaxpr_text(stage), _REMAT_PRIMITIVE) + if apply_remat: + self.assertTrue( + has_remat, + f"{type(stage).__name__} forward jaxpr has no remat primitive with apply_remat=True -- the " + "jax.checkpoint(pure_fn, ...) wrap inside _run_stage_layer_with_remat appears to have been " + "bypassed (e.g. replaced by a direct pure_fn(params, state, inputs) call).", + ) + else: + self.assertFalse( + has_remat, + f"{type(stage).__name__} forward jaxpr unexpectedly contains a remat primitive with " "apply_remat=False.", + ) + + def test_sequential_stage_remat_applied_in_jaxpr(self): + """NNXSequentialPipelineStage: a `remat2` primitive is present iff apply_remat=True.""" + stage_on = NNXSequentialPipelineStage( + NNXDecoderLayer, + 2, + self.cfg, + self.mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=jax.checkpoint_policies.nothing_saveable, + apply_remat=True, + ) + self._assert_remat_applied(stage_on, apply_remat=True) + + stage_off = NNXSequentialPipelineStage( + NNXDecoderLayer, + 2, + self.cfg, + self.mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=None, + apply_remat=False, + ) + self._assert_remat_applied(stage_off, apply_remat=False) + + def test_scanned_stage_remat_applied_in_jaxpr(self): + """NNXScannedPipelineStage: a `remat2` primitive is present iff apply_remat=True. The checkpoint + wraps the per-layer scan body, but jax's jaxpr pretty-printer recurses into a scan's nested jaxpr, + so the primitive still shows up in the top-level jaxpr text (same reasoning the pipeline-level + TestNNXCircularRepeatRemat test in nnx_pipeline_test.py relies on for its own scan-nested remat).""" + stage_on = NNXScannedPipelineStage( + NNXDecoderLayer, + 2, + self.cfg, + self.mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=jax.checkpoint_policies.nothing_saveable, + apply_remat=True, + ) + self._assert_remat_applied(stage_on, apply_remat=True) + + stage_off = NNXScannedPipelineStage( + NNXDecoderLayer, + 2, + self.cfg, + self.mesh, + None, + MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=None, + apply_remat=False, + ) + self._assert_remat_applied(stage_off, apply_remat=False) + + def test_single_layer_stage_remat_applied_in_jaxpr(self): + """Single-layer remat stage (_make_single_layer_remat_stage_cls, the + num_layers_per_pipeline_stage==1 builder path): a `remat2` primitive is present iff + apply_remat=True.""" + stage_cls = _make_single_layer_remat_stage_cls(LlamaDecoderLayer) + stage_on = stage_cls( + config=self.cfg, + mesh=self.mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=jax.checkpoint_policies.nothing_saveable, + apply_remat=True, + ) + self._assert_remat_applied(stage_on, apply_remat=True) + + stage_off = stage_cls( + config=self.cfg, + mesh=self.mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(params=0, dropout=1), + remat_policy=None, + apply_remat=False, + ) + self._assert_remat_applied(stage_off, apply_remat=False) diff --git a/tests/unit/nnx_pipeline_test.py b/tests/unit/nnx_pipeline_test.py new file mode 100644 index 0000000000..985900c050 --- /dev/null +++ b/tests/unit/nnx_pipeline_test.py @@ -0,0 +1,525 @@ +# 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 the NNX pipeline (pipeline.py NNXPipeline / NNXCircularPipeline). + +The integration pipeline tests are all tpu_only, so the NNX pipeline __call__ paths had no unit +coverage. These need 4 devices (a single CPU split into 4 via XLA_FLAGS, or TPU chips in CI) and lock +the migration-parity fixes: + - non_trainable handling: BOTH schedules carry it through the iteration scan -- non-circular via a + 5-way split plus the jax.lax.scan carry tuple, circular via the "carry_state" collection. + It replaces the prior RngState-only assert, which crashed on any + non_trainable variable, and a later broadcast-and-discard that silently lost mutations; + - unconditional circular repeat-level remat (output transparency vs the iteration-remat flag). + +Run standalone so the 4-device flag takes effect before JAX initializes: + XLA_FLAGS=--xla_force_host_platform_device_count=4 python -m pytest tests/unit/nnx_pipeline_test.py +""" +import os + +os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=4") + +import sys +import unittest + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from flax.linen import partitioning as nn_partitioning +from jax.sharding import Mesh +import pytest + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.layers import pipeline +from maxtext.models import simple_layer +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +# jax.checkpoint lowers to this primitive. Match it by IDENTITY, never by substring-searching the +# printed jaxpr: JAX renames primitives' printed names for cosmetics (pjit_p went "pjit" -> "jit" in a +# commit that also touched ad_checkpoint.py), and "remat2" is a leftover marker from the 2021 remat +# rewrite. A rename would break present-checks loudly and make absent-checks pass vacuously forever. +from jax._src.ad_checkpoint import remat_p as _REMAT_PRIMITIVE # pylint: disable=wrong-import-position +from jax._src import core as _jax_core # pylint: disable=wrong-import-position + + +def _jaxpr_contains_primitive(jaxpr, primitive): + """True if `primitive` appears anywhere in `jaxpr`, including nested sub-jaxprs (scan/cond bodies).""" + inner = jaxpr.jaxpr if hasattr(jaxpr, "jaxpr") else jaxpr + if any(eqn.primitive is primitive for eqn in inner.eqns): + return True + return any(_jaxpr_contains_primitive(sub, primitive) for sub in _jax_core.subjaxprs(inner)) + + +_NEEDS_4_DEVICES = pytest.mark.skipif( + jax.device_count() < 4, + reason="needs 4 devices; run with XLA_FLAGS=--xla_force_host_platform_device_count=4", +) + + +def _make_pipeline_config(ag_per_repeat, num_layers, num_micro, **overrides): + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + enable_checkpointing=False, + enable_goodput_recording=False, + run_name="nnx_pipeline_unit", + max_target_length=64, + base_emb_dim=28, + ici_pipeline_parallelism=4, + base_num_decoder_layers=num_layers, + num_pipeline_microbatches=num_micro, + per_device_batch_size=4, + pipeline_fsdp_ag_per_repeat=ag_per_repeat, + **overrides, + ) + + +def _inputs(config): + bs = config.global_batch_size_to_train_on + seq = config.max_target_length + feat = config.emb_dim + inputs = jax.random.normal(jax.random.PRNGKey(2), [bs, seq, feat], dtype=jnp.float32) + positions = jnp.broadcast_to(jnp.arange(seq, dtype=jnp.int32), (bs, seq)) + seg = jnp.ones((bs, seq), dtype=jnp.int32) + return inputs, seg, positions + + +def _run_pipeline(config, stage_factory): + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + inputs, seg, positions = _inputs(config) + my_pipeline = pipeline.create_pipeline(config=config, layers=stage_factory, mesh=mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + params = my_pipeline.init(jax.random.PRNGKey(0), inputs, seg, positions, True, MODEL_MODE_TRAIN) + out = my_pipeline.apply(params, inputs, seg, positions, True, MODEL_MODE_TRAIN) + return out + + +def _simple_factory(config, mesh): + def factory(stage_rngs): + return simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=stage_rngs) + + return factory + + +def _pipeline_value_and_grad(config, mesh): + """Builds the pipeline (fixed init seed) and returns (loss, grads) for a sum-of-squares loss + differentiated wrt the pipeline params. Same seed across calls -> identical params, so two configs + that differ only by a numerically-transparent flag (e.g. remat) must yield matching loss + grads.""" + inputs, seg, positions = _inputs(config) + my_pipeline = pipeline.create_pipeline(config=config, layers=_simple_factory(config, mesh), mesh=mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + params = my_pipeline.init(jax.random.PRNGKey(0), inputs, seg, positions, True, MODEL_MODE_TRAIN) + + def loss_fn(p): + out = my_pipeline.apply(p, inputs, seg, positions, True, MODEL_MODE_TRAIN) + return jnp.sum(out.astype(jnp.float32) ** 2) + + loss, grads = jax.value_and_grad(loss_fn)(params) + return loss, grads + + +def _assert_stage_chaining_intact(test_case, config, mesh): + """Assert that the pipeline's inter-stage chaining is intact: stage 0 receives the fresh input, every + other stage receives the previous stage's output (not the fresh input).""" + raw_pipeline = pipeline.create_nnx_pipeline( + config=config, stage_factory=_simple_factory(config, mesh), mesh=mesh, rngs=nnx.Rngs(params=0) + ) + micro_size = config.micro_batch_size_to_train_on // config.num_pipeline_microbatches + activation_shape = (micro_size, config.max_target_length, config.emb_dim) + dummy_inputs = jnp.zeros((config.num_pipeline_microbatches,) + activation_shape, dtype=jnp.float32) + with jax.set_mesh(mesh): + loop_state = raw_pipeline.init_states(dummy_inputs) + + fresh_marker, shift_marker = 111.0, -222.0 + state_io = jnp.full_like(loop_state["state_io"], fresh_marker) + shift = jnp.full_like(loop_state["shift"], shift_marker) + circ_storage = jnp.zeros_like(loop_state["circ_storage"]) if loop_state["circ_storage"] is not None else None + + stages_in = raw_pipeline.get_iteration_inputs( + loop_iteration=0, state_io=state_io, circ_storage=circ_storage, shift=shift + ) + + test_case.assertTrue( + bool(jnp.all(stages_in[0] == fresh_marker)), "stage 0 must receive the fresh state_io-derived input" + ) + test_case.assertTrue( + bool(jnp.all(stages_in[1:] == shift_marker)), + "every stage other than stage 0 must receive `shift` (the previous stage's rotated output), " + "not the fresh input -- inter-stage chaining is broken", + ) + + +# A non_trainable variable type: not Param, not Intermediate, not RngState -> lands in the +# pipeline's catch-all partition (mirrors moe.Tid2EidVar, the DeepSeek-V4 hash-routing table). +class _NonTrainableVar(nnx.Variable): + pass + + +@_NEEDS_4_DEVICES +class TestNNXPipelineForward(unittest.TestCase): + """Smoke coverage: both schedules run and produce finite output of the right shape.""" + + def _assert_ok(self, config): + """Run the pipeline and assert the output is finite, right-shaped, and not trivially equal to the raw input""" + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + out = _run_pipeline(config, _simple_factory(config, mesh)) + expected = (config.global_batch_size_to_train_on, config.max_target_length, config.emb_dim) + self.assertEqual(out.shape, expected) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + self.assertGreater(float(jnp.std(out)), 1e-2) + + raw_inputs, _, _ = _inputs(config) + self.assertFalse( + bool(jnp.allclose(np.array(out), np.array(raw_inputs), rtol=1e-3, atol=1e-3)), + msg="pipeline output equals the raw input -> per-stage compute was not applied", + ) + # Shape + finiteness alone cannot see a broken inter-stage chain (every stage still runs, just + # on the wrong input), so directly check the stage-0-vs-rest routing too. + _assert_stage_chaining_intact(self, config, mesh) + + def test_noncircular_forward(self): + self._assert_ok(_make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4)) + + def test_circular_forward(self): + self._assert_ok(_make_pipeline_config(ag_per_repeat=True, num_layers=8, num_micro=8)) + + +@_NEEDS_4_DEVICES +class TestNNXPipelineBackward(unittest.TestCase): + """Backward coverage: value_and_grad through the non-circular NNX pipeline on 4 devices. + + The forward-only smoke tests never exercised autodiff through the schedule. This locks that a real + backward runs end-to-end: the loss is finite, gradients are finite, and at least one gradient is + nonzero (the cotangent actually reached the stage parameters, i.e. the backward was not DCE'd).""" + + def test_noncircular_pipeline_backward(self): + config = _make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4) + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + loss, grads = _pipeline_value_and_grad(config, mesh) + + self.assertTrue(bool(jnp.isfinite(loss))) + grad_leaves = jax.tree_util.tree_leaves(grads) + self.assertGreater(len(grad_leaves), 0) + self.assertTrue(all(bool(jnp.all(jnp.isfinite(g))) for g in grad_leaves), "pipeline gradient has non-finite entries") + self.assertTrue( + any(bool(jnp.any(g != 0)) for g in grad_leaves), "all pipeline gradients are zero -> backward did not run" + ) + # A gradient can be finite and nonzero even if every stage is silently fed the same raw input + # instead of the previous stage's output (each stage still has a live, differentiable path to + # the loss) -- so also check the exact inter-stage wiring the backward pass depends on. + _assert_stage_chaining_intact(self, config, mesh) + + +class TestNonTrainablePartitioning(unittest.TestCase): + """Unit guards for the two building blocks of the pipeline's non_trainable handling (pipeline.py), + tested directly without running a pipeline (no device mesh needed): + 1. the state split routes a non_trainable var to the catch-all partition (non-circular path), + which is then threaded through the scan carry; + 2. advance_rng_state leaves non-RngState leaves untouched (circular carry path). + Together these let a non_trainable collection (e.g. moe.Tid2EidVar, DeepSeek-V4 hash routing) flow + through the pipeline scan. End-to-end forward+backward is covered by TestNonTrainablePipelineBackward. + """ + + def test_advance_rng_state_preserves_non_rng_leaves(self): + """Circular carry safety: advance_rng_state must pass non-RngState leaves through unchanged.""" + from maxtext.utils import pipeline_utils as pu # pylint: disable=import-outside-toplevel + + state = nnx.State({"nt": _NonTrainableVar(jnp.asarray(3.0, dtype=jnp.float32))}) + out = pu.advance_rng_state(state, jnp.int32(7)) + np.testing.assert_array_equal(np.array(out["nt"][...]), np.array(3.0, dtype=np.float32)) + + def test_four_way_split_routes_non_trainable_to_catchall(self): + """Non-circular split: a non_trainable var routes to the catch-all partition, NOT into the + RngState partition; a Param routes to the static-param partition. + + This asserts BUCKET ROUTING only, deliberately not broadcast-vs-carry semantics — so it stays + valid however the catch-all is subsequently threaded through the loop.""" + from maxtext.utils import pipeline_utils as pu # pylint: disable=import-outside-toplevel + + state = nnx.State( + {"p": nnx.Param(jnp.ones((2,), dtype=jnp.float32)), "nt": _NonTrainableVar(jnp.asarray(3.0, dtype=jnp.float32))} + ) + _, params, _, rng, catchall = nnx.split(state, pu.is_static_param, nnx.Intermediate, nnx.RngState, ...) + + def _vars(s): + return jax.tree.leaves(s, is_leaf=lambda x: isinstance(x, nnx.Variable)) + + self.assertTrue(any(isinstance(v, nnx.Param) for v in _vars(params))) + self.assertTrue(any(isinstance(v, _NonTrainableVar) for v in _vars(catchall))) + self.assertFalse(any(isinstance(v, _NonTrainableVar) for v in _vars(rng))) + + +@_NEEDS_4_DEVICES +class TestNNXCircularRepeatRemat(unittest.TestCase): + """Circular repeat-level remat is unconditional (Linen parity). It must be numerically + transparent, and the pipeline must run regardless of the iteration-remat flag value.""" + + def _assert_repeat_level_remat_applied(self, config, mesh): + """The repeat-level `flax_lift.checkpoint(_stage_fn_for_scope, ...)` wrap is unconditional and, + per inspection of NNXCircularPipeline.__call__, is completely independent of + set_remat_policy_on_pipeline_iterations (that flag is only ever read by the NON-circular + NNXPipeline; get_pipeline_remat_policy never references it either). So comparing "flag on" vs + "flag off" outputs/gradients cannot catch the wrap being silently dropped: removing it changes + neither which flag reaches this code path (neither did to begin with) nor, ordinarily, the + computed values (remat is numerically transparent by construction -- that's the whole point of + remat). Assert directly on the compiled jaxpr instead: flax_lift.checkpoint must lower to a + real remat primitive; if the wrap is removed the primitive disappears entirely. Matched by + primitive IDENTITY, not by substring-searching the printed jaxpr -- see _jaxpr_contains_primitive.""" + inputs, seg, positions = _inputs(config) + my_pipeline = pipeline.create_pipeline(config=config, layers=_simple_factory(config, mesh), mesh=mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + params = my_pipeline.init(jax.random.PRNGKey(0), inputs, seg, positions, True, MODEL_MODE_TRAIN) + + def fwd(p): + return my_pipeline.apply(p, inputs, seg, positions, True, MODEL_MODE_TRAIN) + + jaxpr = jax.make_jaxpr(fwd)(params) + + self.assertTrue( + _jaxpr_contains_primitive(jaxpr, _REMAT_PRIMITIVE), + "circular pipeline forward jaxpr has no remat primitive -- the unconditional repeat-level " + "flax_lift.checkpoint(...) wrap around _stage_fn_for_scope appears to have been removed", + ) + + def test_repeat_remat_output_transparent(self): + cfg_on = _make_pipeline_config( + ag_per_repeat=True, num_layers=8, num_micro=8, set_remat_policy_on_pipeline_iterations=True + ) + cfg_off = _make_pipeline_config( + ag_per_repeat=True, num_layers=8, num_micro=8, set_remat_policy_on_pipeline_iterations=False + ) + devices_array = maxtext_utils.create_device_mesh(cfg_on) + mesh = Mesh(devices_array, cfg_on.mesh_axes) + out_on = _run_pipeline(cfg_on, _simple_factory(cfg_on, mesh)) + out_off = _run_pipeline(cfg_off, _simple_factory(cfg_off, mesh)) + np.testing.assert_allclose(np.array(out_on), np.array(out_off), rtol=1e-5, atol=1e-5) + # The on-vs-off comparison above cannot see the remat wrap being dropped (see docstring on + # _assert_repeat_level_remat_applied), so check its presence directly for both configs. + self._assert_repeat_level_remat_applied(cfg_on, mesh) + self._assert_repeat_level_remat_applied(cfg_off, mesh) + + def test_repeat_remat_grad_parity(self): + """Remat must be transparent in the BACKWARD pass too, not just the forward output: the + circular pipeline with repeat-level remat on vs off must produce matching loss AND gradients.""" + cfg_on = _make_pipeline_config( + ag_per_repeat=True, num_layers=8, num_micro=8, set_remat_policy_on_pipeline_iterations=True + ) + cfg_off = _make_pipeline_config( + ag_per_repeat=True, num_layers=8, num_micro=8, set_remat_policy_on_pipeline_iterations=False + ) + devices_array = maxtext_utils.create_device_mesh(cfg_on) + mesh = Mesh(devices_array, cfg_on.mesh_axes) + + loss_on, grads_on = _pipeline_value_and_grad(cfg_on, mesh) + loss_off, grads_off = _pipeline_value_and_grad(cfg_off, mesh) + + np.testing.assert_allclose(np.array(loss_on), np.array(loss_off), rtol=1e-4, atol=1e-4) + on_leaves = jax.tree_util.tree_leaves(grads_on) + off_leaves = jax.tree_util.tree_leaves(grads_off) + self.assertEqual(len(on_leaves), len(off_leaves)) + self.assertGreater(len(on_leaves), 0) + for g_on, g_off in zip(on_leaves, off_leaves): + np.testing.assert_allclose(np.array(g_on), np.array(g_off), rtol=1e-4, atol=1e-4) + # Guard against a vacuous pass (all-zero grads would trivially match). + self.assertTrue(any(bool(jnp.any(g != 0)) for g in on_leaves), "all gradients are zero -> backward did not run") + # loss/grad parity between the two configs is a false signal here (see docstring on + # _assert_repeat_level_remat_applied) -- confirm the remat wrap is actually present. + self._assert_repeat_level_remat_applied(cfg_on, mesh) + + +class _StageWithNonTrainable(nnx.Module): + """A pipeline stage carrying a non_trainable variable (like moe.Tid2EidVar, DeepSeek-V4 hash + routing), added into the forward. + + Must return the SAME structure as the wrapped layer: SimpleDecoderLayer returns an (output, kv) + tuple the pipeline loop-state relies on; collapsing it to a bare array trips the shard_map. So the + non_trainable is folded into tuple[0] and the rest is passed through.""" + + def __init__(self, config, mesh, value, *, rngs): + self.inner = simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + self.nt = _NonTrainableVar(jnp.asarray(value, dtype=jnp.float32)) + + def __call__(self, inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs): + res = self.inner(inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs) + if isinstance(res, tuple): + return (res[0] + self.nt[...],) + tuple(res[1:]) + return res + self.nt[...] + + +@_NEEDS_4_DEVICES +class TestNonTrainablePipelineBackward(unittest.TestCase): + """A pipeline stage carrying a non_trainable variable must run forward AND backward, with finite, + nonzero param gradients: the non_trainable collection is carried through the iteration scan on BOTH + schedules and must not break autodiff to the trainable params. + + """ + + def _assert_backward_ok(self, config): + """Init the pipeline (non_trainable stage), value_and_grad a sum-of-squares loss, assert loss + + param grads are finite and at least one is nonzero (backward ran with non_trainable present).""" + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + inputs, seg, positions = _inputs(config) + + def factory(stage_rngs): + return _StageWithNonTrainable(config, mesh, 0.5, rngs=stage_rngs) + + my_pipeline = pipeline.create_pipeline(config=config, layers=factory, mesh=mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + params = my_pipeline.init(jax.random.PRNGKey(0), inputs, seg, positions, True, MODEL_MODE_TRAIN) + + def loss_fn(p): + out = my_pipeline.apply(p, inputs, seg, positions, True, MODEL_MODE_TRAIN) + return jnp.sum(out.astype(jnp.float32) ** 2) + + loss, grads = jax.value_and_grad(loss_fn)(params) + + self.assertTrue(bool(jnp.isfinite(loss))) + grad_leaves = jax.tree_util.tree_leaves(grads) + self.assertGreater(len(grad_leaves), 0) + self.assertTrue( + all(bool(jnp.all(jnp.isfinite(g))) for g in grad_leaves), "non_trainable pipeline grad has non-finite entries" + ) + self.assertTrue( + any(bool(jnp.any(g != 0)) for g in grad_leaves), + "all grads zero -> backward did not run with a non_trainable variable present", + ) + + def test_noncircular_nontrainable_backward(self): + self._assert_backward_ok(_make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4)) + + def test_circular_nontrainable_backward(self): + self._assert_backward_ok(_make_pipeline_config(ag_per_repeat=True, num_layers=8, num_micro=8)) + + +class _MutatingNonTrainableStage(nnx.Module): + """_StageWithNonTrainable, but it MUTATES the non_trainable on every invocation.""" + + def __init__(self, config, mesh, value, *, rngs): + self.inner = simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + self.nt = _NonTrainableVar(jnp.asarray(value, dtype=jnp.float32)) + + def __call__(self, inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs): + res = self.inner(inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs) + self.nt[...] = self.nt[...] + 1.0 + if isinstance(res, tuple): + return (res[0] + self.nt[...],) + tuple(res[1:]) + return res + self.nt[...] + + +@_NEEDS_4_DEVICES +class TestNonTrainableMutationSurvivesLoop(unittest.TestCase): + """A non_trainable mutated inside the iteration loop must still be mutated when the loop ends. + The non_trainable is carried through the iteration scan on BOTH schedules, so the mutation must + survive the loop on both schedules.""" + + _INIT = 0.0 + + def _final_non_trainable(self, config): + """Run the pipeline with a mutating non_trainable stage, return the final value of the non_trainable.""" + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + inputs, seg, positions = _inputs(config) + + def factory(stage_rngs): + return _MutatingNonTrainableStage(config, mesh, self._INIT, rngs=stage_rngs) + + my_pipeline = pipeline.create_pipeline(config=config, layers=factory, mesh=mesh) + with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + variables = my_pipeline.init(jax.random.PRNGKey(0), inputs, seg, positions, True, MODEL_MODE_TRAIN) + mutable = [k for k in variables if k != "params"] + _, updated = my_pipeline.apply(variables, inputs, seg, positions, True, MODEL_MODE_TRAIN, mutable=mutable or True) + leaves = jax.tree_util.tree_leaves(updated.get(_NonTrainableVar.__name__, {})) + return [float(v) for leaf in leaves for v in jnp.ravel(leaf)] + + def _assert_mutation_kept(self, config): + values = self._final_non_trainable(config) + self.assertTrue(values, "no non_trainable leaf surfaced; test cannot conclude") + self.assertTrue( + all(v > self._INIT for v in values), + f"non_trainable mutation was discarded by the iteration loop: {values} (init={self._INIT}). " + "The stage incremented it on every invocation, so every entry must exceed the initial value.", + ) + + def test_noncircular_preserves_non_trainable_mutation(self): + self._assert_mutation_kept(_make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4)) + + def test_circular_preserves_non_trainable_mutation(self): + self._assert_mutation_kept(_make_pipeline_config(ag_per_repeat=True, num_layers=8, num_micro=8)) + + +class _FlagIndependentReturnStage(nnx.Module): + """ + A stage whose return SHAPE deliberately does not track ``config.scan_layers``. + """ + + def __init__(self, config, mesh, *, rngs, always_tuple): + self.inner = simple_layer.SimpleDecoderLayer(config=config, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=rngs) + self.always_tuple = always_tuple + + def __call__(self, inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs): + res = self.inner(inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, **kwargs) + out = res[0] if isinstance(res, tuple) else res + return (out, None) if self.always_tuple else out + + +@_NEEDS_4_DEVICES +class TestStageOutputUnwrapIsFlagIndependent(unittest.TestCase): + """The stage-output unwrap must key off the actual return shape, not off config.scan_layers. + + Both schedules must tolerate a stage whose return shape stops tracking the flag.""" + + def _assert_pipeline_ok(self, config, always_tuple): + """Run the pipeline with a stage whose return shape does not track config.scan_layers, assert the + output is finite and right-shaped.""" + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + def factory(stage_rngs): + return _FlagIndependentReturnStage(config, mesh, rngs=stage_rngs, always_tuple=always_tuple) + + out = _run_pipeline(config, factory) + expected = (config.global_batch_size_to_train_on, config.max_target_length, config.emb_dim) + self.assertEqual( + out.shape, + expected, + f"unwrap did not track the stage's real return shape (always_tuple={always_tuple}, " + f"scan_layers={config.scan_layers})", + ) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_noncircular_tuple_return_when_scan_layers_off(self): + config = _make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4, scan_layers=False) + self._assert_pipeline_ok(config, always_tuple=True) + + def test_noncircular_bare_return_when_scan_layers_on(self): + config = _make_pipeline_config(ag_per_repeat=False, num_layers=4, num_micro=4, scan_layers=True) + self._assert_pipeline_ok(config, always_tuple=False) + + def test_circular_tuple_return_when_scan_layers_off(self): + config = _make_pipeline_config(ag_per_repeat=True, num_layers=8, num_micro=8, scan_layers=False) + self._assert_pipeline_ok(config, always_tuple=True) + + def test_circular_bare_return_when_scan_layers_on(self): + config = _make_pipeline_config(ag_per_repeat=True, num_layers=8, num_micro=8, scan_layers=True) + self._assert_pipeline_ok(config, always_tuple=False) + + +if __name__ == "__main__": + unittest.main()