From e463229c46a62b6ac5bff03c9a7089f1583375c6 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 21 Aug 2026 20:43:36 +0000 Subject: [PATCH 1/8] perf(qwen3-next): nested scan over hybrid attention with per-layer remat Qwen3-Next's scannable block instantiated its four heterogeneous sub-layers flat and ran them in a Python loop, so the whole block was rematerialized as a unit and all four sub-layers' activations were live together. Restructure it along the same lines as Gemma4: stack the three linear-attention (GatedDeltaNet) layers and run them through nnx_scan.apply_scanned_layers, and run the single full-attention layer inside a trip-count-one jax.lax.scan that acts as an XLA scheduling barrier. Each sub-layer is now rematerialized on its own, so the outer apply skips block-level remat. Also stop apply_scanned_layers from returning parameters out of its scan body: lax.scan stacks every carry output, so returning full state made XLA materialize a second copy of the stacked layer weights on every call. Add full_attention_layer_offset to place the full-attention layer within each cycle; it defaults to -1 (last in the cycle), reproducing the previous (layer_idx + 1) % cycle == 0 schedule. --- src/maxtext/configs/base.yml | 3 + src/maxtext/configs/types.py | 7 + src/maxtext/layers/decoders.py | 127 ++++++++++++++ src/maxtext/layers/nnx_decoders.py | 47 ++++++ src/maxtext/layers/nnx_scan.py | 21 +-- src/maxtext/models/qwen3.py | 263 ++++++++++++++++++++++++----- 6 files changed, 415 insertions(+), 53 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 31c4f0d165..b46cb9d58f 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -305,6 +305,9 @@ batch_split_factor: 1 # the factor by which to split the batch. Only used if use # inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope] # which can only be scanned together in one large block of inhomogeneous_layer_cycle_interval=4 layers. inhomogeneous_layer_cycle_interval: 1 +# Position within each inhomogeneous cycle that holds the full-attention layer, for +# hybrid stacks such as Qwen3-Next. -1 (the default) puts it last in the cycle. +full_attention_layer_offset: -1 # pipeline parallelism # The number of decoder layers is equal to the product of num_stages, num_layers_per_pipeline_stage and num_pipeline_repeats. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index d7d56469d1..31e9e4bbb8 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1116,6 +1116,13 @@ class HardwareAndMesh(BaseModel): ) shard_mode: ShardMode = Field("auto", description="can be either auto or explicit") inhomogeneous_layer_cycle_interval: int = Field(1, description="The interval of repeated inhomogeneous layer patterns.") + full_attention_layer_offset: int = Field( + -1, + description=( + "Position within each inhomogeneous cycle that holds the full-attention layer, for hybrid stacks such as " + "Qwen3-Next. -1 places it last in the cycle." + ), + ) scan_layers: bool = Field( True, description=( diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 648392738e..65224adb9b 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1071,6 +1071,18 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, ) + elif cfg.decoder_block == DecoderBlockType.QWEN3_NEXT: + y = self._apply_qwen3_next_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: y = self._apply_deepseek4_scanned_blocks( y, @@ -1423,6 +1435,121 @@ def _apply_gemma3_scanned_blocks( return y + def _apply_qwen3_next_scanned_blocks( + self, + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_caches=None, + attention_metadata=None, + ): + """Applies Qwen3-Next scanned decoder blocks, handling main scan and remainders.""" + + cfg = self.config + mesh = self.mesh + + # Define the repeating pattern length and calculate how many full blocks to scan + block_pattern_len = cfg.inhomogeneous_layer_cycle_interval + num_full_blocks = cfg.num_decoder_layers // block_pattern_len + remainder_layers = cfg.num_decoder_layers % block_pattern_len + + if num_full_blocks > 0: + ScannableBlockToLinen = qwen3.Qwen3NextScannableBlockToLinen + policy = self.get_remat_policy() + + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, num_full_blocks, block_pattern_len, stack=True + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (None, nn.broadcast), # bidirectional_mask + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + # For a fully scanned block, apply it inside an nn.scan over the calculated number of full blocks + y, returned_kv_cache = nn.scan( + ScannableBlockToLinen, + variable_axes={ + "params": cfg.param_scan_axis, + "cache": 0, + "intermediates": 0, + "aqt": 0, + "_overwrite_with_gradient": 0, + }, + split_rngs={"params": True, "dropout": cfg.enable_dropout}, + in_axes=in_axes_tuple, + length=num_full_blocks, + unroll=1, + metadata_params={ + nn.PARTITION_NAME: "layers", + "abstract_init": False, + }, + )( + config=cfg, + mesh=mesh, + quant=self.quant, + model_mode=model_mode, + num_of_layers=block_pattern_len, + remat_policy_fn=policy, + apply_internal_remat=True, + name="scanned_blocks", + )( + y, *broadcast_args + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True + ) + + # Process any remaining layers that don't fit into a full scanned block + for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): + layer = qwen3.Qwen3NextDecoderLayerToLinen( + config=cfg, + mesh=mesh, + model_mode=model_mode, + quant=self.quant, + layer_idx=layer_id, + ) + kv_cache = kv_caches[layer_id] if kv_caches is not None else None + + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + slot, + kv_cache, + attention_metadata, + ) + + y_and_kv = layer(y, *remainder_args) + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + kv_caches[layer_id] = new_kv + + return y + def _apply_gemma4_scanned_blocks( self, y, diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1a9fdd48b0..9c69e96793 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -437,6 +437,7 @@ def __init__( 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.is_qwen3_next = self.config.decoder_block == DecoderBlockType.QWEN3_NEXT if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: self.hc_head = mhc.DeepSeek4HyperHead( @@ -547,6 +548,8 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): self._init_scanned_gemma3(decoder_block_classes, rngs, mesh) elif self.is_gemma4: self._init_scanned_gemma4(decoder_block_classes, rngs, mesh) + elif self.is_qwen3_next: + self._init_scanned_qwen3_next(rngs) else: self._init_scanned_generic(decoder_block_classes, rngs) @@ -717,6 +720,27 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): rngs=rngs, ) + def _init_scanned_qwen3_next(self, rngs): + """Initializes scanned Qwen3-Next blocks with per-layer (rather than per-block) remat. + + Mirrors _init_scanned_gemma4: each block covers one period of the hybrid + attention pattern and rematerializes its own sub-layers, so the outer apply + skips block-level remat. + """ + config = self.config + block_length = config.inhomogeneous_layer_cycle_interval + scan_length = config.num_decoder_layers // block_length + if scan_length > 0: + self.layers = self._create_scanned_layers( + qwen3.Qwen3NextScannableBlock, + length=scan_length, + metadata_axis_name="layers", + rngs=rngs, + num_of_layers=block_length, + remat_policy_fn=self.get_remat_policy(), + apply_internal_remat=True, + ) + def _init_scanned_generic(self, decoder_block_classes, rngs): """Initializes scanned generic decoder layers.""" config = self.config @@ -1858,6 +1882,8 @@ def __call__( layer_kwargs, kv_caches=kv_caches, ) + elif self.is_qwen3_next and kv_caches is None: + y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs) else: scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) if kv_caches is not None: @@ -2140,6 +2166,27 @@ def pure_gemma_fn(graphdef, state_in, y_in, kv_in): return y + def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs): + """Applies the Qwen3-Next scanned blocks. + + Qwen3NextScannableBlock rematerializes its own sub-layers (a scan over the + linear-attention layers plus a trip-count-one scan over the full-attention + layer), so block-level remat is skipped here to avoid rematerializing twice. + """ + cfg = self.config + scan_length = cfg.num_decoder_layers // cfg.inhomogeneous_layer_cycle_interval + if scan_length == 0: + return y + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, + y, + *layer_args, + length=scan_length, + skip_block_remat=True, + **layer_kwargs, + ) + return y + def _apply_gemma4_scanned_blocks( self, y, diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index 1198ef172e..e5c2155e9a 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -122,7 +122,7 @@ def apply_scanned_layers( if length <= 0: return carry - layer_graphdef, params, state = nnx.split(layers, nnx.Param, ...) + layer_graphdef, params, rest = nnx.split(layers, nnx.Param, ...) if param_scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params) @@ -131,7 +131,7 @@ def _ensure_stacked(x): return jnp.broadcast_to(x, (length,)) return x - state = jax.tree.map(_ensure_stacked, state) + rest = jax.tree.map(_ensure_stacked, rest) def _strip_scan_metadata(leaf): if hasattr(leaf, "replace") and hasattr(leaf, "value"): # pylint: disable=too-many-nested-blocks @@ -164,23 +164,20 @@ def _strip_scan_metadata(leaf): return leaf def scan_body(current_carry, scanned_state): - current_params, current_state = scanned_state + current_params, current_rest = scanned_state current_params = jax.tree.map( _strip_scan_metadata, current_params, is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"), ) - current_layer = nnx.merge(layer_graphdef, current_params, current_state) + current_layer = nnx.merge(layer_graphdef, current_params, current_rest) next_carry = apply_fn(current_layer, current_carry) - return next_carry, nnx.state(current_layer) + # Avoid returning and stacking read-only parameters inside the scan body. + _, _, updated_rest = nnx.split(current_layer, nnx.Param, ...) + return next_carry, updated_rest scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body - final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), length=length, unroll=unroll) + final_carry, scanned_rest = jax.lax.scan(scan_fn, carry, (params, rest), length=length, unroll=unroll) - if param_scan_axis != 0: - scanned_params, scanned_other = scanned_state.split(nnx.Param, ...) - scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_params) - scanned_state = nnx.State.merge(scanned_params, scanned_other) - - nnx.update(layers, scanned_state) + nnx.update(layers, scanned_rest) return final_carry diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e46a41dda9..b83308a2a3 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -43,6 +43,8 @@ from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations +from maxtext.layers import nnx_scan +from jax.experimental import xla_metadata from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated from maxtext.layers.quantizations import AqtQuantization as Quant @@ -51,6 +53,7 @@ from maxtext.layers.moe import RoutedMoE from maxtext.layers.initializers import nd_dense_init, variable_to_logically_partitioned from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils from maxtext.inference import kvcache @@ -1194,12 +1197,18 @@ def __call__(self, hidden_states: Array, deterministic: bool) -> tuple[Array, Ar class Qwen3NextScannableBlock(nnx.Module): - """A scannable block of Qwen3-Next decoder layers. + """A scannable block of Qwen3-Next decoder layers with hierarchical nested scans. - This module contains a fixed number of heterogeneous decoder layers that form - a repeating pattern, as defined by `config.inhomogeneous_layer_cycle_interval`. It is - intended to be the body of an `nn.scan` transformation to construct the full - decoder stack efficiently. + One block covers a single period of the attention pattern defined by + `config.inhomogeneous_layer_cycle_interval`: several linear-attention + (GatedDeltaNet) layers plus one full-attention layer. The linear-attention + layers are homogeneous, so they are stacked and run through + `nnx_scan.apply_scanned_layers`; the lone full-attention layer runs inside a + trip-count-one `jax.lax.scan` that acts as an XLA scheduling barrier. + + Nesting the scans this way lets each sub-layer be rematerialized on its own + (`apply_internal_remat`) instead of rematerializing the whole block, so only + one sub-layer's activations are live at a time. Attributes: config: The model configuration object. @@ -1208,35 +1217,190 @@ class Qwen3NextScannableBlock(nnx.Module): quant: Optional quantization configuration. """ - def __init__(self, config: Config, mesh: Mesh, model_mode: str, quant: None | Quant = None, *, rngs: nnx.Rngs): + def __init__( + self, + config: Config, + mesh: Mesh, + model_mode: str, + quant: None | Quant = None, + *, + num_of_layers: int | None = None, + layer_idx_offset: int = 0, + remat_policy_fn: Any | None = None, + apply_internal_remat: bool = False, + rngs: nnx.Rngs, + ): self.config = config self.mesh = mesh self.model_mode = model_mode self.quant = quant self.rngs = rngs + self.remat_policy_fn = remat_policy_fn + self.apply_internal_remat = apply_internal_remat cfg = self.config + if num_of_layers is None: + num_of_layers = cfg.inhomogeneous_layer_cycle_interval + self.num_of_layers = num_of_layers + self.layer_idx_offset = layer_idx_offset + + cycle_interval = cfg.inhomogeneous_layer_cycle_interval + full_attention_offset = cfg.full_attention_layer_offset % cycle_interval + + positions = [(layer_idx_offset + i) % cycle_interval for i in range(num_of_layers)] + self.num_local = sum(1 for p in positions if p != full_attention_offset) + self.num_global = sum(1 for p in positions if p == full_attention_offset) + if self.num_global > 1: + raise ValueError( + f"A Qwen3-Next scannable block spans {num_of_layers} layers starting at offset {layer_idx_offset}, which " + f"covers {self.num_global} full-attention layers; the block supports at most one." + ) + # The local scan runs before the global layer, so the block only reproduces the + # model's layer order when the full-attention layer is last in the period. + if self.num_global == 1 and positions[-1] != full_attention_offset: + raise ValueError( + f"Qwen3-Next scannable block expects the full-attention layer last in the block, but with " + f"full_attention_layer_offset={cfg.full_attention_layer_offset} and layer_idx_offset={layer_idx_offset} it " + f"lands at block position {positions.index(full_attention_offset)} of {num_of_layers}." + ) + + if self.num_local > 0: + self.local_layers = nnx_scan.create_scanned_layers( + lambda layer_rngs: Qwen3NextDecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + layer_idx=0, + is_full_attention_layer=False, + rngs=layer_rngs, + ), + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + metadata_axis_name="local_layers", + rngs=self.rngs, + ) + else: + self.local_layers = None - # Instantiate each layer within the block in __init__ - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer_rngs = self.rngs.fork() # Fork RNGs for each layer - layer_name = f"layer_{i}" - layer = Qwen3NextDecoderLayer( + if self.num_global > 0: + self.global_layer = Qwen3NextDecoderLayer( config=self.config, mesh=self.mesh, quant=self.quant, model_mode=self.model_mode, - layer_idx=i, - rngs=layer_rngs, + layer_idx=full_attention_offset, + is_full_attention_layer=True, + rngs=self.rngs, ) - setattr(self, layer_name, layer) + else: + self.global_layer = None + + def _run_layer(self, layer, y, layer_kwargs, kv_cache=None): + """Invokes one Qwen3NextDecoderLayer, returning (output, updated_kv_cache).""" + out = layer(y, **layer_kwargs, kv_cache=kv_cache) + return out if isinstance(out, tuple) else (out, None) + + @property + def _remat_enabled(self): + """Whether the block rematerializes its own layers.""" + return self.apply_internal_remat and self.config.remat_policy != "none" + + def _scan_local_layers(self, y, layer_kwargs): + """Runs the local (linear attention / GatedDeltaNet) layers via a per-layer rematerialized jax.lax.scan.""" + remat = self._remat_enabled + return nnx_scan.apply_scanned_layers( + self.local_layers, + y, + length=self.num_local, + param_scan_axis=self.config.param_scan_axis, + apply_fn=lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0], + remat=remat, + remat_policy=self.remat_policy_fn if remat else None, + prevent_cse=maxtext_utils.should_prevent_cse_in_remat(self.config) if remat else True, + ) + + def _scan_global_layer(self, y, layer_kwargs): + """Runs the single global-attention layer inside a length-1 jax.lax.scan.""" + cfg = self.config + graphdef_g, intermediate_g, other_g = nnx.split(self.global_layer, nnx.Intermediate, ...) + intermediate_xs = jax.tree.map(lambda x: x[None], intermediate_g) + + def run_global_layer(carry, intermediate_slice): + hidden_states, other = carry + layer = nnx.merge(graphdef_g, intermediate_slice, other) + new_hidden_states = self._run_layer(layer, hidden_states, layer_kwargs)[0] + _, new_intermediate, new_other = nnx.split(layer, nnx.Intermediate, ...) + return (new_hidden_states, new_other), new_intermediate + + global_remat_policy = self.remat_policy_fn + offload_names = maxtext_utils.get_save_and_offload_names(cfg) + if offload_names[0] or offload_names[1]: + save_names, offload_to_device = offload_names + global_remat_policy = jax.checkpoint_policies.save_only_these_names(*(save_names + offload_to_device)) + + if self._remat_enabled: + prevent_cse = maxtext_utils.should_prevent_cse_in_remat(self.config) + run_global_layer = jax.checkpoint( + run_global_layer, + policy=global_remat_policy, + prevent_cse=prevent_cse, + ) + + with xla_metadata.set_xla_metadata(**{"skip-simplify-while-loops_trip-count-one": "true"}): + (y, final_other), stacked_intermediate = jax.lax.scan( + run_global_layer, + (y, other_g), + intermediate_xs, + length=1, + ) + + intermediate_state = jax.tree.map(lambda x: x[0], stacked_intermediate) + nnx.update(self.global_layer, final_other, intermediate_state) + return y + + def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): + """Runs the block with externally-supplied per-layer kv caches. + + Inference KV caches are a Python list of per-layer entries, so this path + unrolls the local layers statically rather than scanning them. + """ + updated_kvs = [] + if self.local_layers is not None: + graphdef, params, state = nnx.split(self.local_layers, nnx.Param, ...) + scan_axis = self.config.param_scan_axis + if scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0), params) + per_layer_states = [] + for i in range(self.num_local): + current_params = jax.tree.map(lambda x, i=i: x[i], params) + current_state = jax.tree.map(lambda x, i=i: x[i], state) + layer = nnx.merge(graphdef, current_params, current_state) + current_kv = kv_cache[i] if (kv_cache is not None and i < len(kv_cache)) else None + y, new_kv = self._run_layer(layer, y, layer_kwargs, current_kv) + updated_kvs.append(new_kv) + per_layer_states.append(nnx.state(layer)) + + stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) + if scan_axis != 0: + stacked_params, stacked_other = stacked_state.split(nnx.Param, ...) + stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params) + stacked_state = nnx.State.merge(stacked_params, stacked_other) + nnx.update(self.local_layers, stacked_state) + + if self.global_layer is not None: + global_kv = kv_cache[self.num_local] if (kv_cache is not None and self.num_local < len(kv_cache)) else None + y, new_kv = self._run_layer(self.global_layer, y, layer_kwargs, global_kv) + updated_kvs.append(new_kv) + + return y, tuple(updated_kvs) def __call__( self, carry: jnp.ndarray, - decoder_segment_ids: None | jnp.ndarray, - decoder_positions: None | jnp.ndarray, - deterministic: bool, - model_mode: str, + decoder_segment_ids: None | jnp.ndarray = None, + decoder_positions: None | jnp.ndarray = None, + deterministic: bool = False, + model_mode: str = "train", previous_chunk=None, slot: None | int = None, kv_cache=None, @@ -1253,27 +1417,31 @@ def __call__( value for the scan's `y` collection. """ cfg = self.config - x = carry - - # Loop over the number of sub-layers that make up one repeating pattern. - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer = getattr(self, f"layer_{i}") - # The second return value is kv_cache, which we ignore here because - # it is not passed as a carry in scannable layers. - x, _ = layer( - x, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - ) + inputs = carry + inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) + + layer_kwargs = { + "decoder_segment_ids": decoder_segment_ids, + "decoder_positions": decoder_positions, + "deterministic": deterministic, + "model_mode": model_mode, + "slot": slot, + "previous_chunk": previous_chunk, + "attention_metadata": attention_metadata, + } - # The output of the block is the carry for the next scan iteration. - return x, None + if kv_cache is not None: + return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs) + + y = inputs + if self.local_layers is not None: + y = self._scan_local_layers(y, layer_kwargs) + if self.global_layer is not None: + y = self._scan_global_layer(y, layer_kwargs) + + if cfg.scan_layers: + return y, None + return y class Qwen3NextDecoderLayer(nnx.Module): @@ -1295,7 +1463,15 @@ class Qwen3NextDecoderLayer(nnx.Module): """ def __init__( - self, config: Config, mesh: Mesh, model_mode: str, layer_idx: int, quant: None | Quant = None, *, rngs: nnx.Rngs + self, + config: Config, + mesh: Mesh, + model_mode: str, + layer_idx: int, + quant: None | Quant = None, + *, + is_full_attention_layer: bool | None = None, + rngs: nnx.Rngs, ): self.config = config self.mesh = mesh @@ -1314,8 +1490,13 @@ def __init__( rngs=rngs, ) - # Determine the type of attention mechanism for the current layer. - is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 + # Determine the type of attention mechanism for the current layer. A scanned block + # knows each sub-layer's role up front and passes it explicitly, because inside a + # scan the layer's position is not recoverable from layer_idx. + if is_full_attention_layer is None: + offset = cfg.full_attention_layer_offset % cfg.inhomogeneous_layer_cycle_interval + is_full_attention_layer = self.layer_idx % cfg.inhomogeneous_layer_cycle_interval == offset + self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: From 8a4d2520c51e2a3800d944fe1a8a0770e73b6b24 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 21 Aug 2026 21:45:19 +0000 Subject: [PATCH 2/8] test(qwen3-next): cover the scannable block's nested scans Asserts the block splits a cycle into a stacked local scan plus one global layer, that the local params really are stacked along param_scan_axis, that the nested scans reproduce a sequential unroll of the same weights, and that a block whose full-attention layer is not last is rejected. --- tests/unit/qwen3_next_scannable_block_test.py | 131 ++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 tests/unit/qwen3_next_scannable_block_test.py diff --git a/tests/unit/qwen3_next_scannable_block_test.py b/tests/unit/qwen3_next_scannable_block_test.py new file mode 100644 index 0000000000..20892cfd06 --- /dev/null +++ b/tests/unit/qwen3_next_scannable_block_test.py @@ -0,0 +1,131 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the nested scans inside Qwen3NextScannableBlock.""" + +import sys +import unittest + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.common.common_types import MODEL_MODE_TRAIN +from maxtext.configs import pyconfig +from maxtext.models import qwen3 +from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path + +_CONFIG_OVERRIDES = { + "run_name": "qwen3_next_scannable_block_test", + "model_name": "qwen3-next-80b-a3b", + "enable_checkpointing": False, + "per_device_batch_size": 1.0, + "max_target_length": 8, + "max_prefill_predict_length": 4, + "attention": "dot_product", + "base_emb_dim": 64, + "base_num_decoder_layers": 4, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "head_dim": 32, + "base_mlp_dim": 128, + "base_moe_mlp_dim": 32, + "num_experts": 4, + "num_experts_per_tok": 2, + "vocab_size": 32, + "gdn_num_key_heads": 2, + "gdn_num_value_heads": 4, + "gdn_key_head_dim": 16, + "gdn_value_head_dim": 16, + "gdn_chunk_size": 4, + "sparse_matmul": True, + "megablox": False, + "dtype": "float32", + "weight_dtype": "float32", + "scan_layers": False, +} + + +def _make_config(**overrides): + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], override_model_config=True, **{**_CONFIG_OVERRIDES, **overrides} + ) + + +class Qwen3NextScannableBlockTest(unittest.TestCase): + """The block's nested scans must reproduce a plain sequential unroll.""" + + def _build(self, **overrides): + cfg = _make_config(**overrides) + mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) + block = qwen3.Qwen3NextScannableBlock( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(0), + ) + return cfg, mesh, block + + def test_block_splits_cycle_into_local_stack_plus_one_global(self): + """A block covers one attention period: cycle-1 linear layers and one full-attention layer.""" + cfg, _, block = self._build() + self.assertEqual(block.num_local, cfg.inhomogeneous_layer_cycle_interval - 1) + self.assertEqual(block.num_global, 1) + self.assertIsNotNone(block.local_layers) + self.assertIsNotNone(block.global_layer) + + def test_local_params_are_stacked(self): + """The linear-attention layers are stacked along param_scan_axis, not stored per layer.""" + cfg, _, block = self._build() + _, params, _ = nnx.split(block.local_layers, nnx.Param, ...) + leaves = [v.value for _, v in params.flat_state()] + self.assertTrue(leaves) + for leaf in leaves: + self.assertEqual(leaf.shape[cfg.param_scan_axis], block.num_local) + + def test_nested_scan_matches_sequential_unroll(self): + """Scanning the local layers then the global layer equals applying them one by one.""" + cfg, _, block = self._build() + inputs = jax.random.normal(jax.random.PRNGKey(1), (1, cfg.max_target_length, cfg.emb_dim), dtype=jnp.float32) + positions = jnp.arange(cfg.max_target_length)[None, :] + segment_ids = jnp.ones((1, cfg.max_target_length), dtype=jnp.int32) + + scanned = block(inputs, segment_ids, positions, True, MODEL_MODE_TRAIN) + + # Reference: pull each stacked local layer out by index and run it, then the global layer. + graphdef, params, rest = nnx.split(block.local_layers, nnx.Param, ...) + if cfg.param_scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, cfg.param_scan_axis, 0), params) + y = inputs + for i in range(block.num_local): + layer = nnx.merge( + graphdef, + jax.tree.map(lambda x, i=i: x[i], params), + jax.tree.map(lambda x, i=i: x[i], rest), + ) + y = layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] + expected = block.global_layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] + + np.testing.assert_allclose(np.asarray(scanned), np.asarray(expected), rtol=1e-5, atol=1e-5) + + def test_rejects_block_whose_global_layer_is_not_last(self): + """The local scan runs before the global layer, so any other ordering must be refused.""" + with self.assertRaisesRegex(ValueError, "full-attention layer last"): + self._build(full_attention_layer_offset=0) + + +if __name__ == "__main__": + unittest.main() From 0857e11f2489715ae58676dc9fa3044c77cf4779 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 21 Aug 2026 22:20:01 +0000 Subject: [PATCH 3/8] fix(qwen3-next): keep per-layer remat on the external KV cache path The scanned-block apply was guarded on `kv_caches is None`, so an inference call fell through to the generic branch and rematerialized the whole block on top of the block's own per-layer remat. Route Qwen3-Next through `_apply_qwen3_next_scanned_blocks` unconditionally and regroup the flat per-layer cache list into per-block tuples before the scan, writing it back afterwards -- the `prepare_kv_caches_for_scan` / `update_kv_caches_after_scan` pair Gemma4 already uses. The scan runs over blocks, so the flat list would otherwise hand block i only `kv_caches[i]`. Training (kv_caches=None) is unaffected: the grouping helpers pass None through, so that path is byte-identical to before. Move the scannable-block tests out of their own file into tests/unit/nnx_decoders_test.py alongside TestGemma4ScannableBlock, and add coverage for both halves of the KV path. --- src/maxtext/layers/nnx_decoders.py | 24 ++- tests/unit/nnx_decoders_test.py | 167 +++++++++++++++++- tests/unit/qwen3_next_scannable_block_test.py | 131 -------------- 3 files changed, 186 insertions(+), 136 deletions(-) delete mode 100644 tests/unit/qwen3_next_scannable_block_test.py diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 9c69e96793..bd835e9e56 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1882,8 +1882,13 @@ def __call__( layer_kwargs, kv_caches=kv_caches, ) - elif self.is_qwen3_next and kv_caches is None: - y = self._apply_qwen3_next_scanned_blocks(y, layer_args, layer_kwargs) + elif self.is_qwen3_next: + y = self._apply_qwen3_next_scanned_blocks( + y, + layer_args, + layer_kwargs, + kv_caches=kv_caches, + ) else: scan_length = int(cfg.num_decoder_layers / cfg.inhomogeneous_layer_cycle_interval) if kv_caches is not None: @@ -2166,25 +2171,36 @@ def pure_gemma_fn(graphdef, state_in, y_in, kv_in): return y - def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs): + def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs, kv_caches=None): """Applies the Qwen3-Next scanned blocks. Qwen3NextScannableBlock rematerializes its own sub-layers (a scan over the linear-attention layers plus a trip-count-one scan over the full-attention layer), so block-level remat is skipped here to avoid rematerializing twice. + This holds for the external-KV-cache path too, hence the regrouping below + rather than falling back to the generic (block-rematerialized) branch. + + External (vLLM) caches arrive as a flat list with one entry per decoder + layer, but the scan runs over blocks, so they are grouped into per-block + tuples before the scan and written back to the flat list afterwards, as in + _apply_gemma4_scanned_blocks. """ cfg = self.config - scan_length = cfg.num_decoder_layers // cfg.inhomogeneous_layer_cycle_interval + block_length = cfg.inhomogeneous_layer_cycle_interval + scan_length = cfg.num_decoder_layers // block_length if scan_length == 0: return y + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, block_length, stack=False) y, self.layers, _ = self._apply_layers_sequentially( self.layers, y, *layer_args, length=scan_length, + kv_caches_stacked=grouped_kv_caches, skip_block_remat=True, **layer_kwargs, ) + maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped_kv_caches, scan_length, block_length, stacked=False) return y def _apply_gemma4_scanned_blocks( diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index c2832f4ce9..ac4f411a7b 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -51,7 +51,7 @@ from maxtext.layers.embeddings import Embed from maxtext.layers.nnx_decoders import NNXDecoder, NNXDecoderLayer, deepstack_process from maxtext.layers.normalizations import RMSNorm -from maxtext.models import gemma4, gemma4_small +from maxtext.models import gemma4, gemma4_small, qwen3 from maxtext.models.gpt3 import Gpt3LayerNorm from maxtext.models.llama2 import LlamaDecoderLayer from maxtext.utils import maxtext_utils, maxtext_utils_nnx @@ -843,6 +843,171 @@ def test_restores_local_state_and_preserves_kv_order(self): np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True) +# Qwen3-Next blocks read enough of the config (GatedDeltaNet head dims, MoE sizing) +# that a SimpleNamespace stand-in is not workable, so these use a real tiny config. +_QWEN3_NEXT_CONFIG = { + "run_name": "qwen3_next_scannable_block_test", + "model_name": "qwen3-next-80b-a3b", + "max_target_length": 8, + "base_emb_dim": 64, + "base_num_decoder_layers": 4, + "base_num_query_heads": 2, + "base_num_kv_heads": 2, + "head_dim": 32, + "base_mlp_dim": 128, + "base_moe_mlp_dim": 32, + "num_experts": 4, + "num_experts_per_tok": 2, + "vocab_size": 32, + "gdn_num_key_heads": 2, + "gdn_num_value_heads": 4, + "gdn_key_head_dim": 16, + "gdn_value_head_dim": 16, + "gdn_chunk_size": 4, + "sparse_matmul": True, + "megablox": False, + "dtype": "float32", + "weight_dtype": "float32", +} + + +class TestQwen3NextScannableBlock(unittest.TestCase): + """Tests Qwen3-Next's nested local(scan)/global(length-1 scan) decoder block.""" + + def _build(self, **overrides): + cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, **overrides}) + mesh = _make_mesh(cfg) + block = qwen3.Qwen3NextScannableBlock( + config=cfg, + mesh=mesh, + model_mode=MODEL_MODE_TRAIN, + rngs=nnx.Rngs(0), + ) + return cfg, mesh, block + + def _inputs(self, cfg): + inputs = jax.random.normal(jax.random.PRNGKey(1), (1, cfg.max_target_length, cfg.emb_dim), dtype=jnp.float32) + positions = jnp.arange(cfg.max_target_length)[None, :] + segment_ids = jnp.ones((1, cfg.max_target_length), dtype=jnp.int32) + return inputs, segment_ids, positions + + def test_block_splits_cycle_into_local_stack_plus_one_global(self): + """A block covers one attention period: cycle-1 linear layers and one full-attention layer.""" + cfg, _, block = self._build() + self.assertEqual(block.num_local, cfg.inhomogeneous_layer_cycle_interval - 1) + self.assertEqual(block.num_global, 1) + self.assertIsNotNone(block.local_layers) + self.assertIsNotNone(block.global_layer) + + def test_local_params_are_stacked(self): + """The linear-attention layers are stacked along param_scan_axis, not stored per layer.""" + cfg, _, block = self._build() + _, params, _ = nnx.split(block.local_layers, nnx.Param, ...) + leaves = [v.value for _, v in params.flat_state()] + self.assertTrue(leaves) + for leaf in leaves: + self.assertEqual(leaf.shape[cfg.param_scan_axis], block.num_local) + + def test_nested_scan_matches_sequential_unroll(self): + """Scanning the local layers then the global layer equals applying them one by one.""" + cfg, _, block = self._build() + inputs, segment_ids, positions = self._inputs(cfg) + + scanned = block(inputs, segment_ids, positions, True, MODEL_MODE_TRAIN) + + # Reference: pull each stacked local layer out by index and run it, then the global layer. + graphdef, params, rest = nnx.split(block.local_layers, nnx.Param, ...) + if cfg.param_scan_axis != 0: + params = jax.tree.map(lambda x: jnp.moveaxis(x, cfg.param_scan_axis, 0), params) + y = inputs + for i in range(block.num_local): + layer = nnx.merge( + graphdef, + jax.tree.map(lambda x, i=i: x[i], params), + jax.tree.map(lambda x, i=i: x[i], rest), + ) + y = layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] + expected = block.global_layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] + + np.testing.assert_allclose(np.asarray(scanned), np.asarray(expected), rtol=1e-5, atol=1e-5) + + def test_external_kv_cache_matches_scanned_path(self): + """The external-kv-cache path must match the scanned path and return one cache per sub-layer. + + In TRAIN mode with dot_product attention the caches pass straight through, so + the two paths differ only in the loop mechanism (nested scans vs static + unroll of the stacked local params) -- any mismatch is a slice/re-stack bug. + """ + cfg, _, block = self._build(matmul_precision="highest") + inputs, segment_ids, positions = self._inputs(cfg) + call_args = (segment_ids, positions, True, MODEL_MODE_TRAIN) + + y_scanned = block(inputs, *call_args) + + num_layers = block.num_local + block.num_global + external_kv = tuple(jnp.full((1, cfg.max_target_length), float(i)) for i in range(num_layers)) + y_external, updated_kvs = block(inputs, *call_args, kv_cache=external_kv) + + self.assertEqual(len(updated_kvs), num_layers) + # Order must stay local[0..n-1] then global, matching the flat per-layer cache list. + for i, updated in enumerate(updated_kvs): + np.testing.assert_array_equal(np.asarray(updated), np.asarray(external_kv[i])) + np.testing.assert_allclose(np.asarray(y_external), np.asarray(y_scanned), rtol=1e-5, atol=1e-5) + + def test_rejects_block_whose_global_layer_is_not_last(self): + """The local scan runs before the global layer, so any other ordering must be refused.""" + with self.assertRaisesRegex(ValueError, "full-attention layer last"): + self._build(full_attention_layer_offset=0) + + +class TestNNXDecoderQwen3Next(unittest.TestCase): + """Tests the NNXDecoder-level wiring of the Qwen3-Next scanned blocks.""" + + def test_decoder_regroups_flat_kv_caches_per_block(self): + """A flat per-layer kv cache list must be regrouped per block and written back in order. + + The scan runs over blocks, not layers, so passing the flat list straight + through would hand block i only ``kv_caches[i]``. Guards + ``_apply_qwen3_next_scanned_blocks``, which must also keep + ``skip_block_remat=True`` on this path rather than falling back to the + generic (block-rematerialized) branch. + """ + cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, "base_num_decoder_layers": 8, "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)) + shared_embedding = Embed( + num_embeddings=cfg.vocab_size, + num_features=cfg.emb_dim, + dtype=cfg.dtype, + config=cfg, + mesh=mesh, + rngs=nnx.Rngs(params=0), + ) + + batch = cfg.global_batch_size_to_train_on + seq = cfg.max_target_length + ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq), 0, cfg.vocab_size) + segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) + positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq)) + + # Distinct sentinel per layer: regrouping errors show up as caches landing on + # the wrong layer, which the pass-through in TRAIN mode makes visible. + kv_caches = [jnp.full((batch, seq), float(i)) for i in range(cfg.num_decoder_layers)] + decoder( + shared_embedding, + ids, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + deterministic=True, + model_mode=MODEL_MODE_TRAIN, + kv_caches=kv_caches, + ) + + self.assertEqual(len(kv_caches), cfg.num_decoder_layers) + for i, cache in enumerate(kv_caches): + np.testing.assert_array_equal(np.asarray(cache), np.full((batch, seq), float(i))) + + class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" diff --git a/tests/unit/qwen3_next_scannable_block_test.py b/tests/unit/qwen3_next_scannable_block_test.py deleted file mode 100644 index 20892cfd06..0000000000 --- a/tests/unit/qwen3_next_scannable_block_test.py +++ /dev/null @@ -1,131 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Tests for the nested scans inside Qwen3NextScannableBlock.""" - -import sys -import unittest - -from flax import nnx -import jax -import jax.numpy as jnp -import numpy as np - -from maxtext.common.common_types import MODEL_MODE_TRAIN -from maxtext.configs import pyconfig -from maxtext.models import qwen3 -from maxtext.utils import maxtext_utils -from tests.utils.test_helpers import get_test_config_path - -_CONFIG_OVERRIDES = { - "run_name": "qwen3_next_scannable_block_test", - "model_name": "qwen3-next-80b-a3b", - "enable_checkpointing": False, - "per_device_batch_size": 1.0, - "max_target_length": 8, - "max_prefill_predict_length": 4, - "attention": "dot_product", - "base_emb_dim": 64, - "base_num_decoder_layers": 4, - "base_num_query_heads": 2, - "base_num_kv_heads": 2, - "head_dim": 32, - "base_mlp_dim": 128, - "base_moe_mlp_dim": 32, - "num_experts": 4, - "num_experts_per_tok": 2, - "vocab_size": 32, - "gdn_num_key_heads": 2, - "gdn_num_value_heads": 4, - "gdn_key_head_dim": 16, - "gdn_value_head_dim": 16, - "gdn_chunk_size": 4, - "sparse_matmul": True, - "megablox": False, - "dtype": "float32", - "weight_dtype": "float32", - "scan_layers": False, -} - - -def _make_config(**overrides): - return pyconfig.initialize( - [sys.argv[0], get_test_config_path()], override_model_config=True, **{**_CONFIG_OVERRIDES, **overrides} - ) - - -class Qwen3NextScannableBlockTest(unittest.TestCase): - """The block's nested scans must reproduce a plain sequential unroll.""" - - def _build(self, **overrides): - cfg = _make_config(**overrides) - mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes) - block = qwen3.Qwen3NextScannableBlock( - config=cfg, - mesh=mesh, - model_mode=MODEL_MODE_TRAIN, - rngs=nnx.Rngs(0), - ) - return cfg, mesh, block - - def test_block_splits_cycle_into_local_stack_plus_one_global(self): - """A block covers one attention period: cycle-1 linear layers and one full-attention layer.""" - cfg, _, block = self._build() - self.assertEqual(block.num_local, cfg.inhomogeneous_layer_cycle_interval - 1) - self.assertEqual(block.num_global, 1) - self.assertIsNotNone(block.local_layers) - self.assertIsNotNone(block.global_layer) - - def test_local_params_are_stacked(self): - """The linear-attention layers are stacked along param_scan_axis, not stored per layer.""" - cfg, _, block = self._build() - _, params, _ = nnx.split(block.local_layers, nnx.Param, ...) - leaves = [v.value for _, v in params.flat_state()] - self.assertTrue(leaves) - for leaf in leaves: - self.assertEqual(leaf.shape[cfg.param_scan_axis], block.num_local) - - def test_nested_scan_matches_sequential_unroll(self): - """Scanning the local layers then the global layer equals applying them one by one.""" - cfg, _, block = self._build() - inputs = jax.random.normal(jax.random.PRNGKey(1), (1, cfg.max_target_length, cfg.emb_dim), dtype=jnp.float32) - positions = jnp.arange(cfg.max_target_length)[None, :] - segment_ids = jnp.ones((1, cfg.max_target_length), dtype=jnp.int32) - - scanned = block(inputs, segment_ids, positions, True, MODEL_MODE_TRAIN) - - # Reference: pull each stacked local layer out by index and run it, then the global layer. - graphdef, params, rest = nnx.split(block.local_layers, nnx.Param, ...) - if cfg.param_scan_axis != 0: - params = jax.tree.map(lambda x: jnp.moveaxis(x, cfg.param_scan_axis, 0), params) - y = inputs - for i in range(block.num_local): - layer = nnx.merge( - graphdef, - jax.tree.map(lambda x, i=i: x[i], params), - jax.tree.map(lambda x, i=i: x[i], rest), - ) - y = layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] - expected = block.global_layer(y, segment_ids, positions, True, MODEL_MODE_TRAIN)[0] - - np.testing.assert_allclose(np.asarray(scanned), np.asarray(expected), rtol=1e-5, atol=1e-5) - - def test_rejects_block_whose_global_layer_is_not_last(self): - """The local scan runs before the global layer, so any other ordering must be refused.""" - with self.assertRaisesRegex(ValueError, "full-attention layer last"): - self._build(full_attention_layer_offset=0) - - -if __name__ == "__main__": - unittest.main() From 99cbb62f65586b591c1d707351650bcfcff4369a Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 21 Aug 2026 22:38:48 +0000 Subject: [PATCH 4/8] refactor(qwen3-next): drop param stacking from the external KV cache path Parameters are read-only in the forward pass, so stacking them back into self.local_layers allocated a second copy of every layer weight. Collect only the non-Param state, which also removes the param_scan_axis round-trip: unlike Params, non-Param state is stacked on axis 0 by nnx_scan.create_scanned_layers. Same reasoning as the scan-body change in nnx_scan.apply_scanned_layers, applied to the static-unroll path. Also move the xla_metadata import into the jax group (PEP 8). --- src/maxtext/models/qwen3.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index b83308a2a3..548bba2681 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -25,6 +25,7 @@ import jax.nn from jax import lax from jax.ad_checkpoint import checkpoint_name +from jax.experimental import xla_metadata from jax.sharding import Mesh import jax.numpy as jnp @@ -44,7 +45,6 @@ from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers import nnx_scan -from jax.experimental import xla_metadata from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated from maxtext.layers.quantizations import AqtQuantization as Quant @@ -1378,14 +1378,13 @@ def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): current_kv = kv_cache[i] if (kv_cache is not None and i < len(kv_cache)) else None y, new_kv = self._run_layer(layer, y, layer_kwargs, current_kv) updated_kvs.append(new_kv) - per_layer_states.append(nnx.state(layer)) + # Collect only non-Param state: parameters are read-only here, so stacking + # them back would allocate a second copy of every layer weight. Non-Param + # state is stacked on axis 0, matching nnx_scan.create_scanned_layers. + _, _, updated_state = nnx.split(layer, nnx.Param, ...) + per_layer_states.append(updated_state) - stacked_state = jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states) - if scan_axis != 0: - stacked_params, stacked_other = stacked_state.split(nnx.Param, ...) - stacked_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), stacked_params) - stacked_state = nnx.State.merge(stacked_params, stacked_other) - nnx.update(self.local_layers, stacked_state) + nnx.update(self.local_layers, jax.tree.map(lambda *xs: jnp.stack(xs), *per_layer_states)) if self.global_layer is not None: global_kv = kv_cache[self.num_local] if (kv_cache is not None and self.num_local < len(kv_cache)) else None From 449a81ed6d13e709e9cfadfcc72d44186826e5a2 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Sat, 22 Aug 2026 05:38:33 +0000 Subject: [PATCH 5/8] fix(nnx): keep params created inside the layer scan `apply_scanned_layers` dropped every `nnx.Param` from the scan output to avoid materializing a second copy of the stacked layer weights. That also discarded parameters created *while tracing the body*, most notably the `nnx.LoRAParam` adapters Qwix materializes, which are an `nnx.Param` subclass -- LoRA setup then failed with "LoRA module path matched target modules, but nnx.LoRAParam is still missing". Only drop the params that were fed in as scan inputs; anything else the body produces still leaves the scan. --- src/maxtext/layers/nnx_scan.py | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index e5c2155e9a..2b736f59c4 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -126,6 +126,12 @@ def apply_scanned_layers( if param_scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params) + # Parameters fed in as scan inputs come back out unchanged, so they must not be + # re-emitted as scan outputs (see scan_body). Anything else the body produces -- + # including parameters materialized while tracing, such as Qwix LoRA adapters, + # which are ``nnx.Param`` subclasses -- still has to leave the scan. + carried_param_paths = {path for path, _ in nnx.to_flat_state(params)} + def _ensure_stacked(x): if hasattr(x, "ndim") and x.ndim == 0: return jnp.broadcast_to(x, (length,)) @@ -172,12 +178,22 @@ def scan_body(current_carry, scanned_state): ) current_layer = nnx.merge(layer_graphdef, current_params, current_rest) next_carry = apply_fn(current_layer, current_carry) - # Avoid returning and stacking read-only parameters inside the scan body. - _, _, updated_rest = nnx.split(current_layer, nnx.Param, ...) - return next_carry, updated_rest + # Drop the parameters that were carried in: ``jax.lax.scan`` stacks every + # output, so returning them would materialize a second copy of the stacked + # layer weights. Parameters created inside the body are still returned. + _, updated_params, updated_rest = nnx.split(current_layer, nnx.Param, ...) + new_params = nnx.from_flat_state( + [(path, value) for path, value in nnx.to_flat_state(updated_params) if path not in carried_param_paths] + ) + return next_carry, (new_params, updated_rest) scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body - final_carry, scanned_rest = jax.lax.scan(scan_fn, carry, (params, rest), length=length, unroll=unroll) + final_carry, (scanned_new_params, scanned_rest) = jax.lax.scan( + scan_fn, carry, (params, rest), length=length, unroll=unroll + ) + + if param_scan_axis != 0: + scanned_new_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, param_scan_axis), scanned_new_params) - nnx.update(layers, scanned_rest) + nnx.update(layers, scanned_new_params, scanned_rest) return final_carry From ef4f304266dda1b0bf5795c3773b540df832f126 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Sat, 22 Aug 2026 05:38:55 +0000 Subject: [PATCH 6/8] feat(qwen3-next): support checkpoint conversion for the scanned block layout The scannable-block rewrite replaced the per-cycle-position `layer_{i}` params with a nested layout -- `layers-local_layers-*` (an inner scan over the linear-attention layers, nested in the block scan) and `layers-global_layer-*` (the block scan only). The HF param mapping still described the old layout, so conversion produced keys the model does not have. Rewrite the scanned branch of the qwen3-next mapping (and its hooks) for the new layout. Routed-expert weights are expert-stacked *inside* the nested scan, giving `[expert][block][local]` -- a third stacked axis, which the conversion helpers did not support. Generalize `_build_multi_axis_stacked_tensor` and the inverse in `process_maxtext_param` from two axes to N, with the axis placement factored into a shared `stacked_axes` helper, and broaden the nested-scan detection from `scanned_blocks-local_layers` (gemma4's module name) to `-local_layers` so qwen3-next's `layers-local_layers` matches too. Also fix the Linen `_apply_qwen3_next_scanned_blocks`: its broadcast-arg spec had been copied from gemma4 and no longer matched `Qwen3NextScannableBlock.__call__`, and it named the scanned module `scanned_blocks` where the pure-NNX decoder uses `layers`. Both decoder paths now emit byte-identical parameter names and shapes, so one mapping serves both. --- .../checkpoint_conversion/to_maxtext.py | 73 ++++--- .../utils/load_dynamic.py | 1 + .../utils/param_mapping.py | 187 ++++++++---------- .../utils/tensor_handling.py | 124 ++++++++---- .../checkpoint_conversion/utils/utils.py | 58 +++--- src/maxtext/layers/decoders.py | 9 +- tests/unit/param_mapping_test.py | 141 ++++++++++++- 7 files changed, 375 insertions(+), 218 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index ce2eecb678..17ced7f095 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -67,7 +67,7 @@ from maxtext.common.common_types import MODEL_MODE_TRAIN from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING -from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns +from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns, nesting_depth, slice_shape, stacked_axes from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys from maxtext.inference.inference_utils import str2bool from maxtext.layers import quantizations @@ -344,66 +344,59 @@ def get_maxtext_model_info(config): def _build_multi_axis_stacked_tensor( - hf_source_keys: List[List[str]], + hf_source_keys: List[Any], tensor_getter_fn: Callable[[str], np.ndarray], hook_fns: Any, target_shape: tuple, config, mt_key: str = "", ) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along two axes. + """Builds a MaxText tensor by stacking HF weights along several axes. - Two cases share this helper: + Three layouts share this helper, distinguished by how deeply ``hf_source_keys`` + is nested and by whether ``mt_key`` names a nested block scan: * MoE expert stacking: outer=experts, inner=layers, placed at the LEADING two axes -> shape (num_experts, num_layers, ...). - * Gemma4 nested block scan (``scanned_blocks-local_layers``): the block's local - layers are an inner scan nested inside the outer block scan, so the two - stacked axes live at ``(param_scan_axis, param_scan_axis + 1)`` rather than - the leading axes, e.g. (emb, blocks, local, ...). + * Nested block scan (``...-local_layers-...``, gemma4 and qwen3-next): the + block's local layers are an inner scan nested inside the outer block scan, + so the two stacked axes live at ``(param_scan_axis, param_scan_axis + 1)`` + rather than the leading axes, e.g. (emb, blocks, local, ...). + * Both at once (qwen3-next's routed experts, which are expert-stacked inside + a nested block scan): three axes at ``(0, param_scan_axis, + param_scan_axis + 1)``, e.g. (experts, blocks, local, ...). Args: - hf_source_keys: A nested (2D) list of Hugging Face parameter names. - The outer list is the outer stack axis, the inner list the - inner stack axis. + hf_source_keys: A nested list of Hugging Face parameter names, one level of + nesting per stacked axis, outermost axis first. A ``tuple`` of names is a + composite HF source that the hook fuses into a single leaf, not an axis. tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array). hook_fns: The hook function(s) to apply to each individual weight. target_shape: The final shape of the target MaxText tensor. config: The MaxText pyconfig object. - mt_key: The MaxText parameter key, used to detect the gemma4 nested-scan case. + mt_key: The MaxText parameter key, used to detect the nested-scan case. Returns: The final, assembled NumPy array for the MaxText parameter. """ - # Gemma4's local layers are an inner scan nested inside the block scan, so the - # two stacked axes go at (param_scan_axis, param_scan_axis + 1); everything else - # (MoE experts/layers) stacks at the leading axes. - if isinstance(mt_key, str) and "scanned_blocks-local_layers" in mt_key: - outer_axis, inner_axis = config.param_scan_axis, config.param_scan_axis + 1 - else: - outer_axis, inner_axis = 0, 1 - + depth = nesting_depth(hf_source_keys) + axes = stacked_axes(mt_key, config, depth) # The hook function needs the shape of an individual slice: target_shape with - # the two stacked axes removed. - stacked_axes = {outer_axis, inner_axis} - mt_slice_shape = tuple(d for i, d in enumerate(target_shape) if i not in stacked_axes) - - all_outer_tensors = [] - # Outer loop iterates the outer stack axis (experts, or blocks for gemma4 local) - for inner_keys in hf_source_keys: - inner_tensors = [] - # Inner loop iterates the inner stack axis (layers, or local layers for gemma4) - for hf_key_single in inner_keys: - if isinstance(hf_key_single, (list, tuple)): - hf_tensor_numpy = tuple(tensor_getter_fn(k) for k in hf_key_single) + # the stacked axes removed. + mt_slice_shape = slice_shape(target_shape, axes) + + def gather(keys, level): + if level == depth: + if isinstance(keys, (list, tuple)): + raw = tuple(tensor_getter_fn(k) for k in keys) else: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - inner_tensors.append(apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns)) - all_outer_tensors.append(np.stack(inner_tensors, axis=0)) - stacked = np.stack(all_outer_tensors, axis=0) # (outer, inner, *slice) - - # Move the (outer, inner) axes from leading positions to their targets. - if (outer_axis, inner_axis) != (0, 1): - stacked = np.moveaxis(stacked, (0, 1), (outer_axis, inner_axis)) + raw = tensor_getter_fn(keys) + return apply_hook_fns(raw, mt_slice_shape, hook_fns) + # Stack with the axes leading, outermost first; they are moved into place below. + return np.stack([gather(sub, level + 1) for sub in keys], axis=0) + + stacked = gather(hf_source_keys, 0) + if axes != tuple(range(depth)): + stacked = np.moveaxis(stacked, tuple(range(depth)), axes) return stacked diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py index 079c35c00d..1de07b76d6 100644 --- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py +++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py @@ -231,6 +231,7 @@ def tensor_getter(key): hook_fn, target_leaf, maxtext_config, + mt_key, ) # Execute transformation and assign to flat_restored diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 94e96173f1..d7b80149b7 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -1402,100 +1402,81 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F } if scan_layers: - # 2. Scan over block cycles - for block_idx in range(layer_cycle_interval): - hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval)) - prefix = f"params-decoder-layers-layer_{block_idx}" + # 2. Scanned blocks. One block covers a single period of the hybrid attention + # pattern: `layer_cycle_interval - 1` linear-attention (GatedDeltaNet) layers + # run as an inner scan, then one full-attention layer. The resulting params are: + # layers-local_layers-* -> nested [block][local] (doubly scanned) + # layers-global_layer-* -> flat [block] (block scan only; the + # length-1 _scan_global_layer scan is a runtime memory boundary, not a + # param stack) + # Routed-expert weights carry an additional leading expert axis, so they are + # nested one level deeper: [expert][block][local] and [expert][block]. + # Qwen3NextScannableBlock requires the full-attention layer to be last in the + # period, so the local positions are 0..cycle-2 and the global position is + # cycle-1. + num_blocks = num_main_layers // layer_cycle_interval + local_positions = list(range(layer_cycle_interval - 1)) + global_position = layer_cycle_interval - 1 + + def hf_layer(block_idx, position, suffix): + return f"model.layers.{block_idx * layer_cycle_interval + position}.{suffix}" + + # (maxtext subkey, hf suffix) pairs shared by both the local and global layers. + shared_specs = [ + ("input_layernorm-scale", "input_layernorm.weight"), + ("post_attention_layernorm-scale", "post_attention_layernorm.weight"), + ("mlp-routed_experts-gate-kernel", "mlp.gate.weight"), + ("mlp-shared_expert-wi_0-kernel", "mlp.shared_expert.gate_proj.weight"), + ("mlp-shared_expert-wi_1-kernel", "mlp.shared_expert.up_proj.weight"), + ("mlp-shared_expert-wo-kernel", "mlp.shared_expert.down_proj.weight"), + ("mlp-shared_expert_gate-kernel", "mlp.shared_expert_gate.weight"), + ] + # Linear (GatedDeltaNet) attention: only ever on the local layers. + local_specs = shared_specs + [ + ("attention-in_proj_qkvz-kernel", "linear_attn.in_proj_qkvz.weight"), + ("attention-in_proj_ba-kernel", "linear_attn.in_proj_ba.weight"), + ("attention-conv1d-kernel", "linear_attn.conv1d.weight"), + ("attention-A_log", "linear_attn.A_log"), + ("attention-dt_bias", "linear_attn.dt_bias"), + ("attention-norm-rms_norm-scale", "linear_attn.norm.weight"), + ("attention-out_proj-kernel", "linear_attn.out_proj.weight"), + ] + # Full attention: only ever on the global layer. + global_specs = shared_specs + [ + ("attention-attention-query-kernel", "self_attn.q_proj.weight"), + ("attention-attention-key-kernel", "self_attn.k_proj.weight"), + ("attention-attention-value-kernel", "self_attn.v_proj.weight"), + ("attention-attention-out-kernel", "self_attn.o_proj.weight"), + ("attention-attention-query_norm-scale", "self_attn.q_norm.weight"), + ("attention-attention-key_norm-scale", "self_attn.k_norm.weight"), + ] + expert_specs = [ + ("mlp-routed_experts-wi_0", "gate_proj.weight"), + ("mlp-routed_experts-wi_1", "up_proj.weight"), + ("mlp-routed_experts-wo", "down_proj.weight"), + ] - # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.input_layernorm.weight" for i in hf_indices - ] # pyrefly: ignore[bad-assignment] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices + local_prefix = "params-decoder-layers-local_layers" + for subkey, suffix in local_specs: + mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + [hf_layer(b, p, suffix) for p in local_positions] for b in range(num_blocks) + ] + for subkey, suffix in expert_specs: + mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + [[hf_layer(b, p, f"mlp.experts.{e}.{suffix}") for p in local_positions] for b in range(num_blocks)] + for e in range(num_experts) ] - # Handle Interleaved Attention (Linear vs Full) - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 - - if is_full_attention_layer: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-attention-attention-query-kernel": [ - f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key-kernel": [ - f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-value-kernel": [ - f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-out-kernel": [ - f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-query_norm-scale": [ - f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key_norm-scale": [ - f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices - ], - } - ) - else: - # Linear/Hybrid Attention Block - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-attention-in_proj_qkvz-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices - ], - f"{prefix}-attention-in_proj_ba-kernel": [ - f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices - ], - f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices], - f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices], - f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices], - f"{prefix}-attention-norm-rms_norm-scale": [ - f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices - ], - f"{prefix}-attention-out_proj-kernel": [ - f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices - ], - } - ) - - # 3. Handle MLP: Gates and Shared Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices], - f"{prefix}-mlp-shared_expert-wi_0-kernel": [ - f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel": [ - f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel": [ - f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert_gate-kernel": [ - f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices - ], - } - ) - - # 4. Handle MoE Routed Experts - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-wi_0": [ - [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1": [ - [f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo": [ - [f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts) - ], - } - ) + global_prefix = "params-decoder-layers-global_layer" + for subkey, suffix in global_specs: + mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + hf_layer(b, global_position, suffix) for b in range(num_blocks) + ] + for subkey, suffix in expert_specs: + mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + [hf_layer(b, global_position, f"mlp.experts.{e}.{suffix}") for b in range(num_blocks)] + for e in range(num_experts) + ] else: # Unscanned layer mapping for i in range(num_main_layers): @@ -1589,17 +1570,21 @@ def permute_conv(input_tensor, target_shape=None): layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval num_main_layers = config["num_hidden_layers"] - loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers) - - for i in loop_indices: - if scan_layers: - prefix = f"params-decoder-layers-layer_{i}" - block_idx = i - else: - prefix = f"params-decoder-layers_{i}" - block_idx = i % layer_cycle_interval - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 + # Scanned blocks expose two prefixes -- the stacked local (linear-attention) layers + # and the single global (full-attention) layer -- rather than one prefix per position + # in the cycle. Unscanned models keep one prefix per decoder layer. + if scan_layers: + layer_prefixes = [ + ("params-decoder-layers-local_layers", False), + ("params-decoder-layers-global_layer", True), + ] + else: + layer_prefixes = [ + (f"params-decoder-layers_{i}", (i % layer_cycle_interval + 1) % layer_cycle_interval == 0) + for i in range(num_main_layers) + ] + for prefix, is_full_attention_layer in layer_prefixes: if is_full_attention_layer: for key in ["query", "key", "value", "out"]: hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] diff --git a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py index 508697624b..039948b836 100644 --- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py +++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py @@ -33,6 +33,52 @@ def apply_hook_fns(weight, target_shape, hook_fns): return weight +def nesting_depth(hf_source_keys: Any) -> int: + """Counts how many axes a (possibly nested) list of HF keys stacks over. + + Only ``list`` nesting counts: a ``tuple`` of HF keys is the composite-key + convention (several HF tensors combined into one MaxText leaf by a hook), not + a stacking axis. + """ + depth = 0 + while isinstance(hf_source_keys, list): + depth += 1 + hf_source_keys = hf_source_keys[0] + return depth + + +def stacked_axes(mt_key: str, config, depth: int) -> tuple: + """Returns where each of the ``depth`` stacked axes lands in the MaxText tensor. + + The outer-to-inner ordering of the returned axes matches the outer-to-inner + nesting of the HF key list. Three layouts occur: + + * MoE expert stacking -- ``(experts, layers, ...)``: every stacked axis is + leading, so the axes are ``0, 1, ...``. + * A nested block scan (``...-local_layers-...``, used by gemma4 and + qwen3-next): the block's local layers are an inner scan nested inside the + outer block scan, so ``(blocks, local)`` sit at + ``(param_scan_axis, param_scan_axis + 1)`` rather than at the leading axes. + * A nested block scan whose weights are *also* expert-stacked (qwen3-next's + routed experts): the expert axis still leads, giving + ``(0, param_scan_axis, param_scan_axis + 1)``. + """ + if isinstance(mt_key, str) and "-local_layers" in mt_key: + param_scan_axis = config.param_scan_axis + nested_axes = (param_scan_axis, param_scan_axis + 1) + return nested_axes if depth == 2 else (0,) + nested_axes + return tuple(range(depth)) + + +def slice_shape(target_shape: tuple, axes: tuple) -> tuple: + """Returns ``target_shape`` with the stacked ``axes`` removed. + + Hook functions operate on a single un-stacked slice, so they need this shape + rather than the shape of the fully assembled tensor. + """ + return tuple(dim for i, dim in enumerate(target_shape) if i not in set(axes)) + + def _binary_chunked_stack(tensors: List[np.ndarray], axis: int) -> np.ndarray: """Stacks JAX arrays along axis by binary division to limit memory usage from JAX compiler.""" if not tensors: @@ -49,13 +95,20 @@ def _binary_chunked_stack(tensors: List[np.ndarray], axis: int) -> np.ndarray: def _build_multi_axis_stacked_tensor( - hf_source_keys: List[List[str]], + hf_source_keys: List[Any], tensor_getter_fn: Callable[[str], np.ndarray], hook_fns: Any, target_leaf: Any, config, + mt_key: str = "", ) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along two axes (experts and layers) directly in place on device.""" + """Builds a MaxText tensor by stacking HF weights along several axes, in place on device. + + ``hf_source_keys`` is nested one level per stacked axis, outermost first (see + ``nesting_depth``), and ``stacked_axes`` decides where those axes land in the + target: leading for MoE expert stacking, or at ``param_scan_axis`` and beyond + for a nested block scan. + """ if hasattr(target_leaf, "sharding"): target_shape = target_leaf.shape target_sharding = target_leaf.sharding @@ -65,39 +118,37 @@ def _build_multi_axis_stacked_tensor( target_sharding = None target_dtype = target_leaf.dtype if hasattr(target_leaf, "dtype") else np.float32 - mt_slice_shape = target_shape[2:] - - # Pre-derive the compatible sharding specs to avoid rank mismatches - if target_sharding is not None and hasattr(target_sharding, "spec"): - # Target shape is (experts, layers, ...) -> slice is from index 2 onwards - spec_list = list(target_sharding.spec)[2:] - slice_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*spec_list)) - # Stacking layer shards - layer_spec_list = list(target_sharding.spec)[1:] - layer_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*layer_spec_list)) - else: - slice_sharding = target_sharding - layer_sharding = target_sharding - - all_expert_tensors = [] - # Outer loop iterates through experts - for layer_keys_for_expert in hf_source_keys: - layer_tensors_for_expert = [] - # Inner loop iterates through layers for the current expert - for hf_key_single in layer_keys_for_expert: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) - - if target_sharding is not None: - processed_hf_tensor = jax.device_put(processed_hf_tensor, slice_sharding) - layer_tensors_for_expert.append(processed_hf_tensor) - - expert_tensor = _binary_chunked_stack(layer_tensors_for_expert, axis=0) - if target_sharding is not None: - expert_tensor = jax.device_put(expert_tensor, layer_sharding) - all_expert_tensors.append(expert_tensor) - - stacked_array = _binary_chunked_stack(all_expert_tensors, axis=0).astype(target_dtype) + depth = nesting_depth(hf_source_keys) + axes = stacked_axes(mt_key, config, depth) + mt_slice_shape = slice_shape(target_shape, axes) + + # Pre-derive the compatible sharding specs to avoid rank mismatches. The tensor + # is assembled with its stacked axes leading and only moved into place at the + # end, so level ``l`` carries the specs of the axes not yet stacked, followed by + # the specs of the slice itself. + target_spec = list(target_sharding.spec) if target_sharding is not None and hasattr(target_sharding, "spec") else None + + def sharding_at(level): + if target_spec is None: + return target_sharding + stacked_spec = [target_spec[axis] for axis in axes[level:]] + slice_spec = [spec for i, spec in enumerate(target_spec) if i not in set(axes)] + return jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*stacked_spec, *slice_spec)) + + def gather(keys, level): + if level == depth: + # A tuple of keys is a composite HF source that the hook fuses into one leaf. + raw = tuple(tensor_getter_fn(k) for k in keys) if isinstance(keys, tuple) else tensor_getter_fn(keys) + tensor = apply_hook_fns(raw, mt_slice_shape, hook_fns) + else: + tensor = _binary_chunked_stack([gather(sub, level + 1) for sub in keys], axis=0) + if target_sharding is not None and level > 0: + tensor = jax.device_put(tensor, sharding_at(level)) + return tensor + + stacked_array = gather(hf_source_keys, 0).astype(target_dtype) + if axes != tuple(range(depth)): + stacked_array = np.moveaxis(stacked_array, tuple(range(depth)), axes) if target_sharding is not None: stacked_array = jax.device_put(stacked_array, target_sharding) return stacked_array @@ -155,7 +206,7 @@ def _build_single_axis_stacked_tensor( return stacked_array -def get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_leaf, config): +def get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_leaf, config, mt_key=""): """Determine the loading function for HF keys.""" if not isinstance(hf_source_keys_or_key, list): # Case 1: Single hf key (str) @@ -194,4 +245,5 @@ def _loader(getter, key, leaf, hook): hook_fn, mt_target_leaf, config, + mt_key, ) diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 26677a0cbb..a7aeb39d42 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -51,6 +51,7 @@ from flax.training import train_state from maxtext.common import checkpointing from maxtext.common.gcloud_stub import gcs_storage +from maxtext.checkpoint_conversion.utils.tensor_handling import nesting_depth, stacked_axes from maxtext.utils import max_logging import orbax.checkpoint as ocp @@ -312,47 +313,38 @@ def process_maxtext_param( return output_weights - # Case 4: Multi-axis stacked. Two sub-cases (the inverse of _build_multi_axis_stacked_tensor): - # - Scanned MoE: the tensor is stacked on (experts, layers) at the LEADING two axes, so we - # slice axis 0 (experts) then axis 0 again (layers, after the expert axis is removed). - # - Gemma4 nested block scan (scanned_blocks-local_layers): the block's local layers are an - # inner scan nested in the block scan, so the two axes are at (param_scan_axis, - # param_scan_axis + 1) -- outer = blocks, inner = local. We slice param_scan_axis (blocks), - # then param_scan_axis again (local shifts down into that slot once blocks is removed). + # Case 4: Multi-axis stacked -- the exact inverse of _build_multi_axis_stacked_tensor. + # `stacked_axes` says where the stacked axes sit in the MaxText tensor: leading for + # scanned MoE (experts, layers), at (param_scan_axis, param_scan_axis + 1) for a + # nested block scan's (blocks, local), or all three for weights that are both. + # Slicing removes one axis at a time, so once the first `level` axes are gone the + # next one has shifted down to `axes[level] - level`. key_str = maxtext_param_key[0] if isinstance(maxtext_param_key, tuple) else maxtext_param_key - if isinstance(key_str, str) and "scanned_blocks-local_layers" in key_str: - max_logging.log("\tscan gemma4 local") - outer_axis_to_slice = maxtext_config.param_scan_axis - inner_axis_to_slice = maxtext_config.param_scan_axis - else: - max_logging.log("\tscan moe") - outer_axis_to_slice = 0 - inner_axis_to_slice = 0 - - # Outer loop (experts for MoE, blocks for gemma4 local) - for outer_idx, inner_paths in enumerate(hf_target_paths): - if isinstance(maxtext_param_weight, list): - outer_slice = [ - jax.lax.index_in_dim(x, outer_idx, axis=outer_axis_to_slice, keepdims=False) for x in maxtext_param_weight - ] - else: - outer_slice = jax.lax.index_in_dim(maxtext_param_weight, outer_idx, axis=outer_axis_to_slice, keepdims=False) - - # Inner loop (layers for MoE, local layers for gemma4) - for inner_idx, hf_path in enumerate(inner_paths): - if isinstance(outer_slice, list): - inner_slice = [jax.lax.index_in_dim(x, inner_idx, axis=inner_axis_to_slice, keepdims=False) for x in outer_slice] - else: - inner_slice = jax.lax.index_in_dim(outer_slice, inner_idx, axis=inner_axis_to_slice, keepdims=False) + depth = nesting_depth(hf_target_paths) + axes = stacked_axes(key_str, maxtext_config, depth) + max_logging.log("\tscan nested local" if axes != tuple(range(depth)) else "\tscan moe") + def _emit_slices(weight, hf_paths, level): + if level == depth: _process( - hf_path, - inner_slice, + hf_paths, + weight, output_weights, current_hook_fns, hf_shape_map, save_dtype=maxtext_config.weight_dtype, ) + return + axis_to_slice = axes[level] - level + for idx, sub_paths in enumerate(hf_paths): + if isinstance(weight, list): + # Handles `composite_mt_key` mappings where weight is a list of tensors. + weight_slice = [jax.lax.index_in_dim(x, idx, axis=axis_to_slice, keepdims=False) for x in weight] + else: + weight_slice = jax.lax.index_in_dim(weight, idx, axis=axis_to_slice, keepdims=False) + _emit_slices(weight_slice, sub_paths, level + 1) + + _emit_slices(maxtext_param_weight, hf_target_paths, 0) return output_weights diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 65224adb9b..a4008686d9 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1465,15 +1465,14 @@ def _apply_qwen3_next_scanned_blocks( kv_caches, num_full_blocks, block_pattern_len, stack=True ) + # Positional order must match Qwen3NextScannableBlock.__call__. broadcast_args_spec = [ (decoder_segment_ids, nn.broadcast), (decoder_positions, nn.broadcast), (deterministic, nn.broadcast), (model_mode, nn.broadcast), - (slot, nn.broadcast), - (None, nn.broadcast), # page_state (previous_chunk, nn.broadcast), - (None, nn.broadcast), # bidirectional_mask + (slot, nn.broadcast), (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), (attention_metadata, nn.broadcast), ] @@ -1506,7 +1505,9 @@ def _apply_qwen3_next_scanned_blocks( num_of_layers=block_pattern_len, remat_policy_fn=policy, apply_internal_remat=True, - name="scanned_blocks", + # Keep the Linen parameter path identical to the pure-NNX decoder's + # `self.layers`, so a single checkpoint mapping serves both. + name="layers", )( y, *broadcast_args ) diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index cea6485817..8631f89855 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -102,14 +102,147 @@ def test_qwen3_next_mapping(self): self.assertIn("params-token_embedder-embedding", mapping) def test_qwen3_next_mapping_scanned(self): + num_layers, cycle, num_experts = 8, 4, 2 config = { - "num_hidden_layers": 4, - "num_experts": 2, + "num_hidden_layers": num_layers, + "num_experts": num_experts, } maxtext_config = mock.Mock() - maxtext_config.inhomogeneous_layer_cycle_interval = 2 + maxtext_config.inhomogeneous_layer_cycle_interval = cycle mapping = param_mapping.QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) - self.assertIn("params-decoder-layers-layer_0-input_layernorm-scale", mapping) + + # A block covers one period of the hybrid pattern: `cycle - 1` linear-attention + # layers as an inner scan, then one full-attention layer. + num_blocks, num_local = num_layers // cycle, cycle - 1 + local_prefix = "params-decoder-layers-local_layers" + global_prefix = "params-decoder-layers-global_layer" + self.assertIn(f"{local_prefix}-attention-in_proj_qkvz-kernel", mapping) + self.assertIn(f"{global_prefix}-attention-attention-query-kernel", mapping) + # Linear attention only exists on the local layers, full attention only on the global one. + self.assertNotIn(f"{global_prefix}-attention-in_proj_qkvz-kernel", mapping) + self.assertNotIn(f"{local_prefix}-attention-attention-query-kernel", mapping) + + # local_layers values are nested [block][local]; global_layer is flat over blocks. + local_val = mapping[f"{local_prefix}-attention-in_proj_qkvz-kernel"] + self.assertEqual(len(local_val), num_blocks) + self.assertEqual(len(local_val[0]), num_local) + self.assertEqual(local_val[0][0], "model.layers.0.linear_attn.in_proj_qkvz.weight") + global_val = mapping[f"{global_prefix}-attention-attention-query-kernel"] + self.assertEqual(len(global_val), num_blocks) + # The full-attention layer is last in the period. + self.assertEqual(global_val[0], f"model.layers.{cycle - 1}.self_attn.q_proj.weight") + + # Routed experts add a leading expert axis: [expert][block][local] and [expert][block]. + local_experts = mapping[f"{local_prefix}-mlp-routed_experts-wi_0"] + self.assertEqual(len(local_experts), num_experts) + self.assertEqual(len(local_experts[0]), num_blocks) + self.assertEqual(len(local_experts[0][0]), num_local) + self.assertEqual(local_experts[1][0][0], "model.layers.0.mlp.experts.1.gate_proj.weight") + global_experts = mapping[f"{global_prefix}-mlp-routed_experts-wi_0"] + self.assertEqual(len(global_experts), num_experts) + self.assertEqual(len(global_experts[0]), num_blocks) + + @staticmethod + def _indices_from_name(name): + """Parses a synthetic HF key such as "e1_b0_l2" back into its stacked-axis indices.""" + return tuple(int(part[1:]) for part in name.split("_")) + + def _qwen3_next_conversion_config(self): + cfg = mock.Mock() + cfg.param_scan_axis = 1 + cfg.scan_layers = True + cfg.weight_dtype = "float32" + cfg.rope_type = "" + cfg.model_name = "qwen3-next-80b-a3b" + return cfg + + def _assert_stack_unstack_roundtrip(self, mt_key, hf_names, target_shape, slice_shape, value_of, cfg): + """Stacking HF weights into `target_shape` (to_maxtext) then un-stacking them back + (to_huggingface) must be the identity, with every stacked axis in the right place.""" + + def getter(name): + return value_of(*self._indices_from_name(name)) + + stacked = _build_multi_axis_stacked_tensor(hf_names, getter, None, target_shape, cfg, mt_key) + self.assertEqual(stacked.shape, target_shape) + + param_map = {mt_key: hf_names} + flat_names = [] + + def flatten(keys): + if isinstance(keys, list): + for sub in keys: + flatten(sub) + else: + flat_names.append(keys) + + flatten(hf_names) + hf_shape_map = {name: slice_shape for name in flat_names} + out = dict(process_maxtext_param(mt_key, stacked, param_map, {}, hf_shape_map, cfg)) + self.assertEqual(len(out), len(flat_names)) + for name in flat_names: + np.testing.assert_array_equal(out[name], value_of(*self._indices_from_name(name))) + return stacked + + def test_qwen3_next_local_layers_stack_unstack_roundtrip(self): + """The local layers of a qwen3-next block are an inner scan, so their two stacked axes + land at (param_scan_axis, param_scan_axis + 1) -- e.g. (emb, blocks, local).""" + num_blocks, num_local = 2, 3 + slice_shape = (4, 3) # per-(block, local) HF weight shape + cfg = self._qwen3_next_conversion_config() + mt_key = "params-decoder-layers-local_layers-attention-in_proj_qkvz-kernel" + + def value_of(b, l): + return np.full(slice_shape, b * 100 + l, dtype=np.float32) + + hf_names = [[f"b{b}_l{l}" for l in range(num_local)] for b in range(num_blocks)] + # blocks at axis 1, local at axis 2. + target_shape = (slice_shape[0], num_blocks, num_local, slice_shape[1]) + + stacked = self._assert_stack_unstack_roundtrip(mt_key, hf_names, target_shape, slice_shape, value_of, cfg) + for b in range(num_blocks): + for l in range(num_local): + np.testing.assert_array_equal(stacked[:, b, l, :], value_of(b, l)) + + def test_qwen3_next_routed_experts_stack_unstack_roundtrip(self): + """qwen3-next's routed experts are expert-stacked *inside* the nested block scan, so they + need three stacked axes: the expert axis still leads, then (blocks, local).""" + num_experts, num_blocks, num_local = 2, 2, 3 + slice_shape = (4, 3) # per-(expert, block, local) HF weight shape + cfg = self._qwen3_next_conversion_config() + mt_key = "params-decoder-layers-local_layers-mlp-routed_experts-wi_0" + + def value_of(e, b, l): + return np.full(slice_shape, e * 10000 + b * 100 + l, dtype=np.float32) + + hf_names = [[[f"e{e}_b{b}_l{l}" for l in range(num_local)] for b in range(num_blocks)] for e in range(num_experts)] + # experts at axis 0, blocks at axis 1, local at axis 2. + target_shape = (num_experts, num_blocks, num_local, *slice_shape) + + stacked = self._assert_stack_unstack_roundtrip(mt_key, hf_names, target_shape, slice_shape, value_of, cfg) + for e in range(num_experts): + for b in range(num_blocks): + for l in range(num_local): + np.testing.assert_array_equal(stacked[e, b, l], value_of(e, b, l)) + + def test_qwen3_next_global_layer_experts_stack_unstack_roundtrip(self): + """The global (full-attention) layer is not inside the inner scan, so its routed experts + keep the plain scanned-MoE layout with both stacked axes leading: (experts, blocks).""" + num_experts, num_blocks = 2, 3 + slice_shape = (4, 3) + cfg = self._qwen3_next_conversion_config() + mt_key = "params-decoder-layers-global_layer-mlp-routed_experts-wi_0" + + def value_of(e, b): + return np.full(slice_shape, e * 100 + b, dtype=np.float32) + + hf_names = [[f"e{e}_b{b}" for b in range(num_blocks)] for e in range(num_experts)] + target_shape = (num_experts, num_blocks, *slice_shape) + + stacked = self._assert_stack_unstack_roundtrip(mt_key, hf_names, target_shape, slice_shape, value_of, cfg) + for e in range(num_experts): + for b in range(num_blocks): + np.testing.assert_array_equal(stacked[e, b], value_of(e, b)) def test_deepseek_mapping(self): config = { From 2968036dd8b3e0d0b4b1e34605fd6574fe90c644 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Mon, 24 Aug 2026 19:51:15 +0000 Subject: [PATCH 7/8] Keep trailing Qwen3-Next layers, unify remainder handling, add decode test Address the follow-ups from the scanned Qwen3-Next block PR that fit in it: - The pure-NNX decoder built `num_decoder_layers // cycle` blocks and stopped there, so a layer count that is not a multiple of the cycle silently dropped its trailing layers. They now go into a `layers_remainder` block that is built and applied like Gemma4's. - The Linen decoder spelled its remainder out as individual `Qwen3NextDecoderLayerToLinen` layers, which named parameters differently from the NNX side. It now uses a `Qwen3NextScannableBlockToLinen` named `layers_remainder` too, so one checkpoint mapping still serves both decoders. `TestQwen3NextDecoderParity` checks the two parameter trees match, with and without a remainder. - Gemma4 and Qwen3-Next now share `_apply_remainder_block` instead of keeping two copies of the same split/checkpoint/merge dance. - Deleted the unreachable third copy of `_build_multi_axis_stacked_tensor` and friends in checkpoint_conversion/utils/utils.py; nothing calls that module's `_get_hf_loading_function`. Decode coverage: `deepseek_decode_consistency_test.py` gains three Qwen3-Next cases (unscanned, scanned, scanned with a remainder). Both attention kinds in the hybrid stack carry state across a decode step, so a block that mis-threads its sub-layers' state diverges from the teacher-forced rollout. --- .../checkpoint_conversion/utils/utils.py | 135 +------------ src/maxtext/layers/decoders.py | 30 +-- src/maxtext/layers/nnx_decoders.py | 185 ++++++++++++------ .../unit/deepseek_decode_consistency_test.py | 29 +++ tests/unit/nnx_decoders_test.py | 103 ++++++++-- 5 files changed, 264 insertions(+), 218 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index a7aeb39d42..298d78d407 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -23,8 +23,7 @@ import time import json from concurrent.futures import ThreadPoolExecutor -from functools import partial -from typing import Any, Callable, List +from typing import Any from tqdm import tqdm import resource import numpy as np @@ -1297,135 +1296,3 @@ def save_weights_to_checkpoint( checkpointing.wait_until_finished(checkpoint_manager) max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min") - - -def _build_multi_axis_stacked_tensor( - hf_source_keys: List[List[str]], - tensor_getter_fn: Callable[[str], np.ndarray], - hook_fns: Any, - target_shape: tuple, - config, -) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along two axes (experts and layers). - - This function handles the complex case for scanned MoE layers, producing a tensor - with the shape (num_experts, num_layers, ...). - - Args: - hf_source_keys: A nested (2D) list of Hugging Face parameter names. - Outer list iterates experts, inner list iterates layers. - tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array). - hook_fns: The hook function(s) to apply to each individual weight. - target_shape: The final shape of the target MaxText tensor. - config: The MaxText pyconfig object. - - Returns: - The final, assembled NumPy array for the MaxText parameter. - """ - all_expert_tensors = [] - # The hook function needs the shape of an individual slice, not the full stacked tensor. - # For multi-axis stacking (experts, layers, ...), the slice shape is target_shape[2:] - mt_slice_shape = target_shape[2:] - - # Outer loop iterates through experts - for layer_keys_for_expert in hf_source_keys: - layer_tensors_for_expert = [] - # Inner loop iterates through layers for the current expert - for hf_key_single in layer_keys_for_expert: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) - layer_tensors_for_expert.append(processed_hf_tensor) - all_expert_tensors.append(np.stack(layer_tensors_for_expert, axis=0)) - return np.stack(all_expert_tensors, axis=0) - - -def _build_single_axis_stacked_tensor( - hf_source_keys: List[str], - tensor_getter_fn: Callable[[str], np.ndarray], - hook_fns: Any, - target_shape: tuple, - config, -) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along a single axis. - - This function handles both standard scanned layers (e.g., attention) and - unscanned MoE layers (which are stacked along the expert axis). - - Args: - hf_source_keys: A 1D list of Hugging Face parameter names. - tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array). - hook_fns: The hook function(s) to apply to each individual weight. - target_shape: The final shape of the target MaxText tensor. - config: The MaxText pyconfig object. - - Returns: - The final, assembled NumPy array for the MaxText parameter. - """ - tensors_to_stack = [] - - if config.scan_layers: - # If it's a standard scanned layer, we use the configured param_scan_axis. - axis_to_stack = config.param_scan_axis - else: - # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). - axis_to_stack = 0 - - # The hook function needs the shape of an individual slice, not the full stacked tensor. - # We calculate it by removing the stacking dimension from the final target shape. - mt_slice_shape_list = list(target_shape) - del mt_slice_shape_list[axis_to_stack] - mt_slice_shape = tuple(mt_slice_shape_list) - - for hf_key_single in hf_source_keys: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) - tensors_to_stack.append(processed_hf_tensor) - - # Stack all processed tensors along the determined axis. - return np.stack(tensors_to_stack, axis=axis_to_stack) - - -def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_shape_or_shapes, config): - """Determine the loading function for HF keys. - HF keys can take four forms: - Case 1: Unscanned (single string) - Case 2: Scanned (list of strings) - Case 3: Unscanned with expert stacking (list of strings) - Case 4: Scanned with expert stacking (nested list of strings) - """ - load_fn = None - if not isinstance(hf_source_keys_or_key, list): - # Case 1: Single hf key (str) - def _loader(getter, key, shape, hook): - return apply_hook_fns(getter(key), shape, hook) - - load_fn = partial( - _loader, - tensor_getter, - hf_source_keys_or_key, - mt_target_shape_or_shapes, - hook_fn, - ) - # Stacked mapping - elif not isinstance(hf_source_keys_or_key[0], list): - # Case 2 or 3: Single-Axis Stacked hf keys (un-nested list) - load_fn = partial( - _build_single_axis_stacked_tensor, - hf_source_keys_or_key, - tensor_getter, - hook_fn, - mt_target_shape_or_shapes, - config, - ) - else: - # isinstance(hf_source_keys_or_key[0], list) - # Case 4: Multi-Axis Stacked hf keys (nested list) - load_fn = partial( - _build_multi_axis_stacked_tensor, - hf_source_keys_or_key, - tensor_getter, - hook_fn, - mt_target_shape_or_shapes, - config, - ) - return load_fn diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index a4008686d9..ec8309b2ec 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1516,29 +1516,34 @@ def _apply_qwen3_next_scanned_blocks( kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True ) - # Process any remaining layers that don't fit into a full scanned block - for layer_id in range(cfg.num_decoder_layers - remainder_layers, cfg.num_decoder_layers): - layer = qwen3.Qwen3NextDecoderLayerToLinen( + # Layers past the last whole block go into a short block of their own. It is a + # Qwen3NextScannableBlock like the scanned ones, named `layers_remainder`, so the + # parameter paths still match the pure-NNX decoder's and one mapping serves both. + if remainder_layers > 0: + start_idx = cfg.num_decoder_layers - remainder_layers + remainder_kv = tuple(kv_caches[start_idx:]) if kv_caches is not None else None + y_and_kv = qwen3.Qwen3NextScannableBlockToLinen( config=cfg, mesh=mesh, - model_mode=model_mode, quant=self.quant, - layer_idx=layer_id, - ) - kv_cache = kv_caches[layer_id] if kv_caches is not None else None - - remainder_args = ( + model_mode=model_mode, + num_of_layers=remainder_layers, + layer_idx_offset=start_idx, + remat_policy_fn=self.get_remat_policy(), + apply_internal_remat=True, + name="layers_remainder", + )( + y, decoder_segment_ids, decoder_positions, deterministic, model_mode, previous_chunk, slot, - kv_cache, + remainder_kv, attention_metadata, ) - y_and_kv = layer(y, *remainder_args) if isinstance(y_and_kv, tuple): y = y_and_kv[0] new_kv = y_and_kv[1] @@ -1547,7 +1552,8 @@ def _apply_qwen3_next_scanned_blocks( new_kv = None if kv_caches is not None and new_kv is not None: - kv_caches[layer_id] = new_kv + for offset, updated_item in enumerate(new_kv): + kv_caches[start_idx + offset] = updated_item return y diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index bd835e9e56..e72b6b6ef8 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -549,7 +549,7 @@ def _init_scanned_layers(self, decoder_block_classes, rngs, mesh): elif self.is_gemma4: self._init_scanned_gemma4(decoder_block_classes, rngs, mesh) elif self.is_qwen3_next: - self._init_scanned_qwen3_next(rngs) + self._init_scanned_qwen3_next(rngs, mesh) else: self._init_scanned_generic(decoder_block_classes, rngs) @@ -720,16 +720,20 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): rngs=rngs, ) - def _init_scanned_qwen3_next(self, rngs): + def _init_scanned_qwen3_next(self, rngs, mesh): """Initializes scanned Qwen3-Next blocks with per-layer (rather than per-block) remat. Mirrors _init_scanned_gemma4: each block covers one period of the hybrid attention pattern and rematerializes its own sub-layers, so the outer apply - skips block-level remat. + skips block-level remat. Layers past the last whole period go into a short + remainder block, named the same as the Linen `Decoder`'s so both decoders + keep producing the same parameter tree. """ config = self.config block_length = config.inhomogeneous_layer_cycle_interval scan_length = config.num_decoder_layers // block_length + num_remaining_layers = config.num_decoder_layers % block_length + policy = self.get_remat_policy() if scan_length > 0: self.layers = self._create_scanned_layers( qwen3.Qwen3NextScannableBlock, @@ -737,9 +741,26 @@ def _init_scanned_qwen3_next(self, rngs): metadata_axis_name="layers", rngs=rngs, num_of_layers=block_length, - remat_policy_fn=self.get_remat_policy(), + remat_policy_fn=policy, apply_internal_remat=True, ) + if num_remaining_layers > 0: + # The remainder starts on a period boundary, so it holds only the leading + # linear-attention layers of a period -- the full-attention layer is last. + self.layers_remainder = qwen3.Qwen3NextScannableBlock( + config=config, + mesh=mesh, + quant=self.quant, + model_mode=self.model_mode, + num_of_layers=num_remaining_layers, + layer_idx_offset=scan_length * block_length, + remat_policy_fn=policy, + # _apply_remainder_block already wraps the call in jax.checkpoint, and + # the remainder is shorter than one period and runs once, so there is + # nothing to gain from a second, finer remat boundary inside it. + apply_internal_remat=False, + rngs=rngs, + ) def _init_scanned_generic(self, decoder_block_classes, rngs): """Initializes scanned generic decoder layers.""" @@ -2184,23 +2205,38 @@ def _apply_qwen3_next_scanned_blocks(self, y, layer_args, layer_kwargs, kv_cache layer, but the scan runs over blocks, so they are grouped into per-block tuples before the scan and written back to the flat list afterwards, as in _apply_gemma4_scanned_blocks. + + Layers past the last whole block are applied afterwards, in a short + remainder block, so a layer count that is not a multiple of the cycle keeps + all its layers. """ cfg = self.config block_length = cfg.inhomogeneous_layer_cycle_interval scan_length = cfg.num_decoder_layers // block_length - if scan_length == 0: - return y - grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, block_length, stack=False) - y, self.layers, _ = self._apply_layers_sequentially( - self.layers, - y, - *layer_args, - length=scan_length, - kv_caches_stacked=grouped_kv_caches, - skip_block_remat=True, - **layer_kwargs, - ) - maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped_kv_caches, scan_length, block_length, stacked=False) + if scan_length > 0: + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, block_length, stack=False) + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, + y, + *layer_args, + length=scan_length, + kv_caches_stacked=grouped_kv_caches, + skip_block_remat=True, + **layer_kwargs, + ) + maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped_kv_caches, scan_length, block_length, stacked=False) + + num_remaining_layers = cfg.num_decoder_layers % block_length + if num_remaining_layers > 0: + y = self._apply_remainder_block( + self.layers_remainder, + y, + layer_args, + layer_kwargs, + kv_caches=kv_caches, + start_idx=scan_length * block_length, + num_remaining_layers=num_remaining_layers, + ) return y def _apply_gemma4_scanned_blocks( @@ -2245,52 +2281,87 @@ def _apply_gemma4_scanned_blocks( # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length if num_remaining_layers > 0: - policy = self.get_remat_policy() - prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) + y = self._apply_remainder_block( + self.layers_remainder, + y, + layer_args, + layer_kwargs, + kv_caches=kv_caches, + start_idx=scan_length * attention_pattern_length, + num_remaining_layers=num_remaining_layers, + ) - remainder_kv = None - if kv_caches is not None: - start_idx = scan_length * attention_pattern_length - remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + return y - if cfg.use_qwix_quantization or cfg.lora.lora_weight_qtype: - call_kwargs = dict(layer_kwargs) - if remainder_kv is not None: - call_kwargs["kv_cache"] = remainder_kv - out_res = self.layers_remainder(y, *layer_args, **call_kwargs) - if isinstance(out_res, tuple): - y = out_res[0] - updated_remainder_kv = out_res[1] if len(out_res) > 1 else None - else: - y = out_res - updated_remainder_kv = None - else: + def _apply_remainder_block( + self, + block, + y, + layer_args, + layer_kwargs, + kv_caches, + start_idx, + num_remaining_layers, + ): + """Applies the trailing scannable block holding the layers left over by the main scan. - def pure_gemma_fn(graphdef, state_in, y_in, kv_in): - merged_layer = nnx.merge(graphdef, state_in) - call_kwargs = dict(layer_kwargs) - if kv_in is not None: - call_kwargs["kv_cache"] = kv_in - out_res = merged_layer(y_in, *layer_args, **call_kwargs) - if isinstance(out_res, tuple): - out_y = out_res[0] - out_kv = out_res[1] if len(out_res) > 1 else None - else: - out_y = out_res - out_kv = None - nnx.pop(merged_layer, (nnx.RngState, nnx.Intermediate)) - return out_y, out_kv, nnx.state(merged_layer) + The outer scan runs over whole blocks, so `num_decoder_layers % + block_length` layers are left over. They live in a single short block, + applied here behind one block-level remat boundary. The main scan skips + block-level remat, because there its blocks rematerialize their own + sub-layers; this one runs once, so a coarse boundary costs nothing. - checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) + Args: + block: The remainder block module; updated in place with its new state. + y: Input activations. + layer_args: Positional arguments forwarded to the block. + layer_kwargs: Keyword arguments forwarded to the block. + kv_caches: Flat per-layer external cache list, mutated in place, or None. + start_idx: Index of the first remainder layer in `kv_caches`. + num_remaining_layers: How many layers the remainder block covers. - graphdef, state = nnx.split(self.layers_remainder) - y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) - nnx.update(self.layers_remainder, new_state) + Returns: + The output activations. + """ + cfg = self.config + remainder_kv = None + if kv_caches is not None: + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + def split_output(out_res): + if isinstance(out_res, tuple): + return out_res[0], (out_res[1] if len(out_res) > 1 else None) + return out_res, None + + if cfg.use_qwix_quantization or cfg.lora.lora_weight_qtype: + call_kwargs = dict(layer_kwargs) + if remainder_kv is not None: + call_kwargs["kv_cache"] = remainder_kv + y, updated_remainder_kv = split_output(block(y, *layer_args, **call_kwargs)) + else: - if kv_caches is not None and updated_remainder_kv is not None: - start_idx = scan_length * attention_pattern_length - for offset, updated_item in enumerate(updated_remainder_kv): - kv_caches[start_idx + offset] = updated_item + def pure_block_fn(graphdef, state_in, y_in, kv_in): + merged_layer = nnx.merge(graphdef, state_in) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + out_y, out_kv = split_output(merged_layer(y_in, *layer_args, **call_kwargs)) + nnx.pop(merged_layer, (nnx.RngState, nnx.Intermediate)) + return out_y, out_kv, nnx.state(merged_layer) + + checkpointed_block_fn = jax.checkpoint( + pure_block_fn, + policy=self.get_remat_policy(), + prevent_cse=maxtext_utils.should_prevent_cse_in_remat(cfg), + ) + + graphdef, state = nnx.split(block) + y, updated_remainder_kv, new_state = checkpointed_block_fn(graphdef, state, y, remainder_kv) + nnx.update(block, new_state) + + if kv_caches is not None and updated_remainder_kv is not None: + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item return y diff --git a/tests/unit/deepseek_decode_consistency_test.py b/tests/unit/deepseek_decode_consistency_test.py index 181636d9a8..bdeeaa9c72 100644 --- a/tests/unit/deepseek_decode_consistency_test.py +++ b/tests/unit/deepseek_decode_consistency_test.py @@ -87,6 +87,32 @@ _DEEPSEEK_MOE = _DEEPSEEK | {"first_num_dense_layers": 1} _DEEPSEEK_DENSE = _DEEPSEEK | {"first_num_dense_layers": 3} +# Qwen3-Next interleaves three linear-attention (GatedDeltaNet) layers with one +# full-attention layer, and scanned it runs those as a stacked scan nested inside +# the block scan. Both attention kinds carry state across a decode step -- a +# recurrent/conv state and a KV cache -- so a block that mis-threads its +# sub-layers' state shows up here as a diverging token. +_QWEN3_NEXT = { + "model_name": "qwen3-next-80b-a3b", + "override_model_config": True, + "base_num_decoder_layers": 8, # two whole periods of inhomogeneous_layer_cycle_interval=4 + "head_dim": 32, # the real model's 256 would make this test needlessly slow + "base_moe_mlp_dim": 32, + "num_experts": 4, + "num_experts_per_tok": 2, + "gdn_num_key_heads": 2, + "gdn_num_value_heads": 4, + "gdn_key_head_dim": 16, + "gdn_value_head_dim": 16, + "gdn_chunk_size": 4, + "sparse_matmul": True, + "megablox": False, +} +_QWEN3_NEXT_SCANNED = _QWEN3_NEXT | {"scan_layers": True} +# 6 layers is one whole period plus a two-layer remainder, which the scanned path +# has to apply on its own rather than drop. +_QWEN3_NEXT_SCANNED_REMAINDER = _QWEN3_NEXT_SCANNED | {"base_num_decoder_layers": 6} + def _make_config(**overrides): return pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | overrides)) @@ -193,6 +219,9 @@ def _build(self, cfg): @parameterized.named_parameters( ("deepseek_moe", _DEEPSEEK_MOE), ("deepseek_dense", _DEEPSEEK_DENSE), + ("qwen3_next_unscanned", _QWEN3_NEXT), + ("qwen3_next_scanned", _QWEN3_NEXT_SCANNED), + ("qwen3_next_scanned_remainder", _QWEN3_NEXT_SCANNED_REMAINDER), ("generic", {}), ) def test_greedy_decode_matches_forward_pass(self, overrides): diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index ac4f411a7b..c89bee12e2 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -963,16 +963,9 @@ def test_rejects_block_whose_global_layer_is_not_last(self): class TestNNXDecoderQwen3Next(unittest.TestCase): """Tests the NNXDecoder-level wiring of the Qwen3-Next scanned blocks.""" - def test_decoder_regroups_flat_kv_caches_per_block(self): - """A flat per-layer kv cache list must be regrouped per block and written back in order. - - The scan runs over blocks, not layers, so passing the flat list straight - through would hand block i only ``kv_caches[i]``. Guards - ``_apply_qwen3_next_scanned_blocks``, which must also keep - ``skip_block_remat=True`` on this path rather than falling back to the - generic (block-rematerialized) branch. - """ - cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, "base_num_decoder_layers": 8, "scan_layers": True}) + def _build(self, num_decoder_layers): + """Builds a scanned Qwen3-Next decoder and its shared embedding.""" + cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, "base_num_decoder_layers": num_decoder_layers, "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)) shared_embedding = Embed( @@ -983,17 +976,16 @@ def test_decoder_regroups_flat_kv_caches_per_block(self): mesh=mesh, rngs=nnx.Rngs(params=0), ) + return cfg, decoder, shared_embedding + def _run(self, cfg, decoder, shared_embedding, kv_caches=None): + """Runs one TRAIN-mode forward pass and returns the logits.""" batch = cfg.global_batch_size_to_train_on seq = cfg.max_target_length ids = jax.random.randint(jax.random.PRNGKey(0), (batch, seq), 0, cfg.vocab_size) segment_ids = jnp.full((batch, seq), DECODING_ACTIVE_SEQUENCE_INDICATOR) positions = jnp.broadcast_to(jnp.arange(seq)[None], (batch, seq)) - - # Distinct sentinel per layer: regrouping errors show up as caches landing on - # the wrong layer, which the pass-through in TRAIN mode makes visible. - kv_caches = [jnp.full((batch, seq), float(i)) for i in range(cfg.num_decoder_layers)] - decoder( + logits, _, _ = decoder( shared_embedding, ids, decoder_positions=positions, @@ -1002,11 +994,92 @@ def test_decoder_regroups_flat_kv_caches_per_block(self): model_mode=MODEL_MODE_TRAIN, kv_caches=kv_caches, ) + return logits + + def test_decoder_regroups_flat_kv_caches_per_block(self): + """A flat per-layer kv cache list must be regrouped per block and written back in order. + + The scan runs over blocks, not layers, so passing the flat list straight + through would hand block i only ``kv_caches[i]``. Guards + ``_apply_qwen3_next_scanned_blocks``, which must also keep + ``skip_block_remat=True`` on this path rather than falling back to the + generic (block-rematerialized) branch. + """ + cfg, decoder, shared_embedding = self._build(8) + batch = cfg.global_batch_size_to_train_on + seq = cfg.max_target_length + + # Distinct sentinel per layer: regrouping errors show up as caches landing on + # the wrong layer, which the pass-through in TRAIN mode makes visible. + kv_caches = [jnp.full((batch, seq), float(i)) for i in range(cfg.num_decoder_layers)] + self._run(cfg, decoder, shared_embedding, kv_caches=kv_caches) self.assertEqual(len(kv_caches), cfg.num_decoder_layers) for i, cache in enumerate(kv_caches): np.testing.assert_array_equal(np.asarray(cache), np.full((batch, seq), float(i))) + def test_decoder_keeps_layers_past_the_last_whole_block(self): + """Layers left over by the block scan must still be built and applied. + + ``num_decoder_layers // inhomogeneous_layer_cycle_interval`` blocks cover + only a whole number of periods, so with 6 layers and a period of 4 the last + two would be silently dropped -- the model would quietly run 4 layers. They + go into ``layers_remainder`` instead; perturbing only that block's weights + has to move the output, which it cannot do if the block is never applied. + """ + cfg, decoder, shared_embedding = self._build(6) + self.assertEqual(decoder.layers_remainder.num_local, 2) + # The remainder starts on a period boundary, so it holds no full-attention layer. + self.assertEqual(decoder.layers_remainder.num_global, 0) + + before = self._run(cfg, decoder, shared_embedding) + _, params, rest = nnx.split(decoder.layers_remainder, nnx.Param, ...) + nnx.update(decoder.layers_remainder, jax.tree.map(lambda x: x + 0.1, params), rest) + after = self._run(cfg, decoder, shared_embedding) + + self.assertFalse( + np.allclose(np.asarray(before), np.asarray(after)), + "the remainder block's weights did not affect the output, so it was not applied", + ) + + def test_decoder_has_no_remainder_block_when_layers_divide_evenly(self): + """With a whole number of periods every layer is inside the scan, so no remainder block.""" + _, decoder, _ = self._build(8) + self.assertFalse(hasattr(decoder, "layers_remainder")) + + +class TestQwen3NextDecoderParity(unittest.TestCase): + """The Linen `Decoder` and the pure-NNX `NNXDecoder` must emit the same parameter tree. + + One checkpoint mapping (`QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING`) serves both + decoders, so a name or shape that differs between them silently breaks + conversion on whichever side the mapping was not written against. + """ + + def _param_tree(self, num_decoder_layers, pure_nnx_decoder): + """Returns {parameter key: shape} for the chosen decoder implementation.""" + # pylint: disable=import-outside-toplevel + from maxtext.checkpoint_conversion.to_maxtext import get_maxtext_model_info + + cfg = _make_config( + **{ + **_QWEN3_NEXT_CONFIG, + "base_num_decoder_layers": num_decoder_layers, + "scan_layers": True, + "pure_nnx_decoder": pure_nnx_decoder, + } + ) + model_info, _ = get_maxtext_model_info(cfg) + return {key: shape for key, (_, shape) in model_info.items()} + + def test_decoders_agree_on_whole_periods(self): + self.assertEqual(self._param_tree(8, True), self._param_tree(8, False)) + + def test_decoders_agree_with_a_remainder(self): + """6 layers is one whole period plus a two-layer remainder, which both decoders + have to put in a `layers_remainder` block rather than spell out layer by layer.""" + self.assertEqual(self._param_tree(6, True), self._param_tree(6, False)) + class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" From ceedbea8b780bd7ed2eae283cd61ef4b59595f41 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Mon, 24 Aug 2026 21:37:32 +0000 Subject: [PATCH 8/8] Drop full_attention_layer_offset The flag was added alongside the scanned block, but the block rejects every value it can take other than the default: the local layers are scanned before the global one, so the full-attention layer has to be last in the cycle. Any other offset raised, or -- on the unscanned path, which does not go through the block -- silently built a different architecture from the scanned path. Stock Qwen3-Next puts full attention last, no config set the flag, and nothing tested a non-default value working. Go back to deriving the position from the cycle, and keep the block's validation, which still earns its place: a block starting off a cycle boundary straddles two periods and would reorder the model. --- src/maxtext/configs/base.yml | 3 --- src/maxtext/configs/types.py | 7 ------- src/maxtext/models/qwen3.py | 13 +++++++------ tests/unit/nnx_decoders_test.py | 13 ++++++++++--- 4 files changed, 17 insertions(+), 19 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index b46cb9d58f..31c4f0d165 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -305,9 +305,6 @@ batch_split_factor: 1 # the factor by which to split the batch. Only used if use # inhomogeneous layers. E.g. maverick uses [dense+rope, moe+rope, dense+rope, moe+nope] # which can only be scanned together in one large block of inhomogeneous_layer_cycle_interval=4 layers. inhomogeneous_layer_cycle_interval: 1 -# Position within each inhomogeneous cycle that holds the full-attention layer, for -# hybrid stacks such as Qwen3-Next. -1 (the default) puts it last in the cycle. -full_attention_layer_offset: -1 # pipeline parallelism # The number of decoder layers is equal to the product of num_stages, num_layers_per_pipeline_stage and num_pipeline_repeats. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 31e9e4bbb8..d7d56469d1 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1116,13 +1116,6 @@ class HardwareAndMesh(BaseModel): ) shard_mode: ShardMode = Field("auto", description="can be either auto or explicit") inhomogeneous_layer_cycle_interval: int = Field(1, description="The interval of repeated inhomogeneous layer patterns.") - full_attention_layer_offset: int = Field( - -1, - description=( - "Position within each inhomogeneous cycle that holds the full-attention layer, for hybrid stacks such as " - "Qwen3-Next. -1 places it last in the cycle." - ), - ) scan_layers: bool = Field( True, description=( diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 548bba2681..0861b4b3a0 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1244,7 +1244,9 @@ def __init__( self.layer_idx_offset = layer_idx_offset cycle_interval = cfg.inhomogeneous_layer_cycle_interval - full_attention_offset = cfg.full_attention_layer_offset % cycle_interval + # Qwen3-Next puts the full-attention layer last in every cycle, which is what + # Qwen3NextDecoderLayer derives from layer_idx when it is not told explicitly. + full_attention_offset = cycle_interval - 1 positions = [(layer_idx_offset + i) % cycle_interval for i in range(num_of_layers)] self.num_local = sum(1 for p in positions if p != full_attention_offset) @@ -1258,9 +1260,9 @@ def __init__( # model's layer order when the full-attention layer is last in the period. if self.num_global == 1 and positions[-1] != full_attention_offset: raise ValueError( - f"Qwen3-Next scannable block expects the full-attention layer last in the block, but with " - f"full_attention_layer_offset={cfg.full_attention_layer_offset} and layer_idx_offset={layer_idx_offset} it " - f"lands at block position {positions.index(full_attention_offset)} of {num_of_layers}." + f"Qwen3-Next scannable block expects the full-attention layer last in the block, but a block of " + f"{num_of_layers} layers starting at layer_idx_offset={layer_idx_offset} lands it at block position " + f"{positions.index(full_attention_offset)}. Blocks must start on a cycle boundary." ) if self.num_local > 0: @@ -1493,8 +1495,7 @@ def __init__( # knows each sub-layer's role up front and passes it explicitly, because inside a # scan the layer's position is not recoverable from layer_idx. if is_full_attention_layer is None: - offset = cfg.full_attention_layer_offset % cfg.inhomogeneous_layer_cycle_interval - is_full_attention_layer = self.layer_idx % cfg.inhomogeneous_layer_cycle_interval == offset + is_full_attention_layer = (self.layer_idx + 1) % cfg.inhomogeneous_layer_cycle_interval == 0 self.is_full_attention_layer = is_full_attention_layer # Conditionally instantiate either the Linear Attention or Full Attention block. diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index c89bee12e2..3573b98aae 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -874,13 +874,15 @@ def test_restores_local_state_and_preserves_kv_order(self): class TestQwen3NextScannableBlock(unittest.TestCase): """Tests Qwen3-Next's nested local(scan)/global(length-1 scan) decoder block.""" - def _build(self, **overrides): + def _build(self, layer_idx_offset=0, **overrides): + """Builds one scannable block; overrides go to the config, layer_idx_offset to the block.""" cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, **overrides}) mesh = _make_mesh(cfg) block = qwen3.Qwen3NextScannableBlock( config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, + layer_idx_offset=layer_idx_offset, rngs=nnx.Rngs(0), ) return cfg, mesh, block @@ -955,9 +957,14 @@ def test_external_kv_cache_matches_scanned_path(self): np.testing.assert_allclose(np.asarray(y_external), np.asarray(y_scanned), rtol=1e-5, atol=1e-5) def test_rejects_block_whose_global_layer_is_not_last(self): - """The local scan runs before the global layer, so any other ordering must be refused.""" + """The local scan runs before the global layer, so any other ordering must be refused. + + A block starting off a cycle boundary straddles two periods -- with a cycle of + 4, layers 1..4 put the full-attention layer third of four. Applying it as + local-scan-then-global would silently reorder the model, so it is rejected. + """ with self.assertRaisesRegex(ValueError, "full-attention layer last"): - self._build(full_attention_layer_offset=0) + self._build(layer_idx_offset=1) class TestNNXDecoderQwen3Next(unittest.TestCase):