From 4220d264b06a6970dbc41a413d303c857a509541 Mon Sep 17 00:00:00 2001 From: chengnuojin Date: Wed, 26 Aug 2026 03:07:31 +0000 Subject: [PATCH 1/5] perf(qwen3.5): reuse Qwen3-Next's nested block scan --- .../utils/param_mapping.py | 339 ++++++++---------- .../utils/tensor_handling.py | 5 + src/maxtext/layers/nnx_decoders.py | 51 ++- src/maxtext/models/qwen3.py | 28 +- src/maxtext/models/qwen3_5.py | 108 ++---- tests/unit/nnx_decoders_test.py | 190 ++++++---- tests/unit/param_mapping_test.py | 223 +++++++++--- 7 files changed, 526 insertions(+), 418 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index bea7b83843..77090b5ade 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -889,197 +889,142 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING( } 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) + # Qwen3.5 stores its routed experts fused, so unlike Qwen3-Next there is no + # extra per-expert axis to nest. FP8 checkpoints are the exception: they store + # one tensor per expert, which adds a leading expert axis exactly as for + # Qwen3-Next ([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 - # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"{hf_prefix}.layers.{i}.input_layernorm.weight" for i in hf_indices + def hf_layer(block_idx, position, suffix): + return f"{hf_prefix}.layers.{block_idx * layer_cycle_interval + position}.{suffix}" + + # (maxtext subkey, hf suffix) pairs shared by both the local and global layers. + # A tuple of suffixes means MaxText concatenates the HF tensors into one kernel. + shared_specs = [ + ("input_layernorm-scale", "input_layernorm.weight"), + ("post_attention_layernorm-scale", "post_attention_layernorm.weight"), + ] + # Routed experts stored one tensor per expert (FP8 checkpoints only). + expert_specs = [] + # Whether the routed experts' wi_0/wi_1 come from one fused HF gate_up_proj. + fused_experts = False + if num_experts > 1: + # 3. MoE: router gate and shared expert. + shared_specs += [ + ("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"), ] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] - f"{hf_prefix}.layers.{i}.post_attention_layernorm.weight" for i in hf_indices + # 4. MoE routed experts. + if is_quantized: + num_experts = config.get("text_config", config).get("num_experts", 256) + shared_specs += [ + ("mlp-shared_expert-wi_0-kernel_scale", "mlp.shared_expert.gate_proj.weight_scale_inv"), + ("mlp-shared_expert-wi_1-kernel_scale", "mlp.shared_expert.up_proj.weight_scale_inv"), + ("mlp-shared_expert-wo-kernel_scale", "mlp.shared_expert.down_proj.weight_scale_inv"), + ] + expert_specs = [ + ("mlp-routed_experts-wi_0", "gate_proj.weight"), + ("mlp-routed_experts-wi_0_scale", "gate_proj.weight_scale_inv"), + ("mlp-routed_experts-wi_1", "up_proj.weight"), + ("mlp-routed_experts-wi_1_scale", "up_proj.weight_scale_inv"), + ("mlp-routed_experts-wo", "down_proj.weight"), + ("mlp-routed_experts-wo_scale", "down_proj.weight_scale_inv"), + ] + else: + shared_specs.append(("mlp-routed_experts-wo", "mlp.experts.down_proj")) + fused_experts = True + else: + # 3. Dense MLP. + shared_specs += [ + ("mlp-wi_0-kernel", "mlp.gate_proj.weight"), + ("mlp-wi_1-kernel", "mlp.up_proj.weight"), + ("mlp-wo-kernel", "mlp.down_proj.weight"), + ] + # Linear (GatedDeltaNet) attention: only ever on the local layers. + local_specs = shared_specs + [ + ("attention-in_proj_qkvz-kernel", ("linear_attn.in_proj_qkv.weight", "linear_attn.in_proj_z.weight")), + ("attention-in_proj_ba-kernel", ("linear_attn.in_proj_b.weight", "linear_attn.in_proj_a.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"), + ] + if is_quantized: + local_specs += [ + ( + "attention-in_proj_qkvz-kernel_scale", + ("linear_attn.in_proj_qkv.weight_scale_inv", "linear_attn.in_proj_z.weight_scale_inv"), + ), + ("attention-out_proj-kernel_scale", "linear_attn.out_proj.weight_scale_inv"), + ] + global_specs += [ + ("attention-attention-query-kernel_scale", "self_attn.q_proj.weight_scale_inv"), + ("attention-attention-key-kernel_scale", "self_attn.k_proj.weight_scale_inv"), + ("attention-attention-value-kernel_scale", "self_attn.v_proj.weight_scale_inv"), + ("attention-attention-out-kernel_scale", "self_attn.o_proj.weight_scale_inv"), ] - # Handle Interleaved Attention (Linear vs Full) - is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 + def hf_entry(block_idx, position, suffix): + """One HF reference: a tuple of suffixes stays a tuple of HF keys to concatenate.""" + if isinstance(suffix, tuple): + return tuple(hf_layer(block_idx, position, s) for s in suffix) + return hf_layer(block_idx, position, suffix) - if is_full_attention_layer: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-attention-attention-query-kernel": [ - f"{hf_prefix}.layers.{i}.self_attn.q_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key-kernel": [ - f"{hf_prefix}.layers.{i}.self_attn.k_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-value-kernel": [ - f"{hf_prefix}.layers.{i}.self_attn.v_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-out-kernel": [ - f"{hf_prefix}.layers.{i}.self_attn.o_proj.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-query_norm-scale": [ - f"{hf_prefix}.layers.{i}.self_attn.q_norm.weight" for i in hf_indices - ], - f"{prefix}-attention-attention-key_norm-scale": [ - f"{hf_prefix}.layers.{i}.self_attn.k_norm.weight" for i in hf_indices - ], - } - ) - if is_quantized: - mapping.update( - { - f"{prefix}-attention-attention-query-kernel_scale": [ - f"{hf_prefix}.layers.{i}.self_attn.q_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-attention-attention-key-kernel_scale": [ - f"{hf_prefix}.layers.{i}.self_attn.k_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-attention-attention-value-kernel_scale": [ - f"{hf_prefix}.layers.{i}.self_attn.v_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-attention-attention-out-kernel_scale": [ - f"{hf_prefix}.layers.{i}.self_attn.o_proj.weight_scale_inv" for i in hf_indices - ], - } - ) - else: - # Linear/Hybrid Attention Block - mapping.update( # pyrefly: ignore[no-matching-overload] - { - # Provide a tuple of HF keys so MaxText concatenates them into qkvz - f"{prefix}-attention-in_proj_qkvz-kernel": [ - ( - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_qkv.weight", - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_z.weight", - ) - for i in hf_indices - ], - # Provide a tuple of HF keys so MaxText concatenates them into ba - f"{prefix}-attention-in_proj_ba-kernel": [ - ( - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_b.weight", - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_a.weight", - ) - for i in hf_indices - ], - f"{prefix}-attention-conv1d-kernel": [ - f"{hf_prefix}.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices - ], - f"{prefix}-attention-A_log": [f"{hf_prefix}.layers.{i}.linear_attn.A_log" for i in hf_indices], - f"{prefix}-attention-dt_bias": [f"{hf_prefix}.layers.{i}.linear_attn.dt_bias" for i in hf_indices], - f"{prefix}-attention-norm-rms_norm-scale": [ - f"{hf_prefix}.layers.{i}.linear_attn.norm.weight" for i in hf_indices - ], - f"{prefix}-attention-out_proj-kernel": [ - f"{hf_prefix}.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices - ], - } - ) - if is_quantized: - mapping.update( - { - f"{prefix}-attention-in_proj_qkvz-kernel_scale": [ - ( - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_qkv.weight_scale_inv", - f"{hf_prefix}.layers.{i}.linear_attn.in_proj_z.weight_scale_inv", - ) - for i in hf_indices - ], - f"{prefix}-attention-out_proj-kernel_scale": [ - f"{hf_prefix}.layers.{i}.linear_attn.out_proj.weight_scale_inv" for i in hf_indices - ], - } - ) - - # 3. Handle MLP: Gates and Shared Experts (MoE) or Dense MLP - if num_experts > 1: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-gate-kernel": [ - f"{hf_prefix}.layers.{i}.mlp.gate.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_0-kernel": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert_gate-kernel": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices - ], - } - ) - if is_quantized: - mapping.update( - { - f"{prefix}-mlp-shared_expert-wi_0-kernel_scale": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.gate_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wi_1-kernel_scale": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.up_proj.weight_scale_inv" for i in hf_indices - ], - f"{prefix}-mlp-shared_expert-wo-kernel_scale": [ - f"{hf_prefix}.layers.{i}.mlp.shared_expert.down_proj.weight_scale_inv" 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_entry(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) + ] + if fused_experts: + # Fused gate_up_proj feeds both wi_0 and wi_1, so it is keyed by the pair. + mapping[(f"{local_prefix}-mlp-routed_experts-wi_0", f"{local_prefix}-mlp-routed_experts-wi_1")] = [ + [hf_layer(b, p, "mlp.experts.gate_up_proj") for p in local_positions] for b in range(num_blocks) + ] - # 4. Handle MoE Routed Experts - if is_quantized: - num_experts = config.get("text_config", config).get("num_experts", 256) - mapping.update( - { - f"{prefix}-mlp-routed_experts-wi_0": [ - [f"{hf_prefix}.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_0_scale": [ - [f"{hf_prefix}.layers.{i}.mlp.experts.{e}.gate_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1": [ - [f"{hf_prefix}.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wi_1_scale": [ - [f"{hf_prefix}.layers.{i}.mlp.experts.{e}.up_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo": [ - [f"{hf_prefix}.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] - for e in range(num_experts) - ], - f"{prefix}-mlp-routed_experts-wo_scale": [ - [f"{hf_prefix}.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" for i in hf_indices] - for e in range(num_experts) - ], - } - ) - else: - mapping.update( # pyrefly: ignore[no-matching-overload] - { - f"{prefix}-mlp-routed_experts-wo": [ - f"{hf_prefix}.layers.{i}.mlp.experts.down_proj" for i in hf_indices - ], - (f"{prefix}-mlp-routed_experts-wi_0", f"{prefix}-mlp-routed_experts-wi_1"): [ - f"{hf_prefix}.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices - ], - } - ) - else: - mapping.update( - { - f"{prefix}-mlp-wi_0-kernel": [f"{hf_prefix}.layers.{i}.mlp.gate_proj.weight" for i in hf_indices], - f"{prefix}-mlp-wi_1-kernel": [f"{hf_prefix}.layers.{i}.mlp.up_proj.weight" for i in hf_indices], - f"{prefix}-mlp-wo-kernel": [f"{hf_prefix}.layers.{i}.mlp.down_proj.weight" for i in hf_indices], - } - ) + 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_entry(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) + ] + if fused_experts: + mapping[(f"{global_prefix}-mlp-routed_experts-wi_0", f"{global_prefix}-mlp-routed_experts-wi_1")] = [ + hf_layer(b, global_position, "mlp.experts.gate_up_proj") for b in range(num_blocks) + ] else: # Unscanned layer mapping for i in range(num_main_layers): @@ -1460,17 +1405,21 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval num_main_layers = getattr(maxtext_config, "num_decoder_layers", None) or config["text_config"]["num_hidden_layers"] num_experts = getattr(maxtext_config, "num_experts", None) or config.get("text_config", config).get("num_experts", 1) - 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 e5eea08084..a04beff2d5 100644 --- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py +++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py @@ -63,6 +63,11 @@ def stacked_axes(mt_key: str, config, depth: int) -> tuple: routed experts): the expert axis still leads, giving ``(0, param_scan_axis, param_scan_axis + 1)``. """ + # A tuple is the composite-MaxText-key convention (e.g. Qwen3.5's fused + # gate_up_proj feeding wi_0 and wi_1); every component shares the same prefix, + # so the first one decides the layout. + if isinstance(mt_key, tuple) and mt_key: + mt_key = mt_key[0] 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) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index c2456fd0a6..0a73559b3c 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -364,6 +364,13 @@ def reshape_forced_routed_experts_for_scan( return jnp.reshape(fre, (scan_length, layers_per_cycle) + fre.shape[1:]) +def _qwen3_hybrid_block_class(config): + """Returns the nested-scan block class for a Qwen3-Next-style hybrid decoder.""" + if config.decoder_block == DecoderBlockType.QWEN3_5: + return qwen3_5.Qwen3_5ScannableBlock + return qwen3.Qwen3NextScannableBlock + + def deepstack_process(hidden_states, bidirectional_mask, visual_embeds): """Process deepstack visual embeddings by adding them to hidden states at visual token positions. @@ -580,8 +587,9 @@ 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 - self.is_qwen3_5 = self.config.decoder_block == DecoderBlockType.QWEN3_5 + # Qwen3-Next and Qwen3.5 share the same hybrid attention period, so they share + # the same nested-scan block structure. + self.is_qwen3_hybrid = self.config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5) if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: self.hc_head = mhc.DeepSeek4HyperHead( @@ -692,8 +700,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, mesh) + elif self.is_qwen3_hybrid: + self._init_scanned_qwen3_hybrid(rngs, mesh) else: self._init_scanned_generic(decoder_block_classes, rngs) @@ -877,8 +885,8 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): rngs=rngs, ) - def _init_scanned_qwen3_next(self, rngs, mesh): - """Initializes scanned Qwen3-Next blocks with per-layer (rather than per-block) remat. + def _init_scanned_qwen3_hybrid(self, rngs, mesh): + """Initializes scanned Qwen3-Next/Qwen3.5 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 @@ -887,13 +895,14 @@ def _init_scanned_qwen3_next(self, rngs, mesh): keep producing the same parameter tree. """ config = self.config + block_cls = _qwen3_hybrid_block_class(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, + block_cls, length=scan_length, metadata_axis_name="layers", rngs=rngs, @@ -904,7 +913,7 @@ def _init_scanned_qwen3_next(self, rngs, mesh): 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( + self.layers_remainder = block_cls( config=config, mesh=mesh, quant=self.quant, @@ -2228,8 +2237,8 @@ def __call__( layer_kwargs, kv_caches=kv_caches, ) - elif self.is_qwen3_next: - y = self._apply_qwen3_next_scanned_blocks( + elif self.is_qwen3_hybrid: + y = self._apply_qwen3_hybrid_scanned_blocks( y, layer_args, layer_kwargs, @@ -2269,30 +2278,14 @@ def __call__( # Pass the kv_caches list directly to avoid copying in jnp.stack, # which breaks vLLM PagedAttention in-place memory updates. # The _apply_layers_sequentially function will handle it by statically unrolling. - # - # Qwen3.5 is the exception: one scan step is a Qwen3_5ScannableBlock - # spanning `cycle_interval` decoder layers, while vLLM hands us one - # cache per layer. Group them per block before the scan and write the - # returned per-block tuples back afterwards, as - # _apply_qwen3_next_scanned_blocks does for Qwen3-Next. - group_blocks = self.is_qwen3_5 and cycle_interval > 1 - grouped_kv_caches = ( - maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, cycle_interval, stack=False) - if group_blocks - else kv_caches - ) y, self.layers, _ = self._apply_layers_sequentially( self.layers, y, *layer_args, length=scan_length, - kv_caches_stacked=grouped_kv_caches, + kv_caches_stacked=kv_caches, **layer_kwargs, ) - if group_blocks: - maxtext_utils.update_kv_caches_after_scan( - kv_caches, grouped_kv_caches, scan_length, cycle_interval, stacked=False - ) # kv_caches list is updated in-place inside _apply_layers_sequentially else: y, self.layers, _ = self._apply_layers_sequentially( @@ -2654,8 +2647,8 @@ 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, kv_caches=None): - """Applies the Qwen3-Next scanned blocks. + def _apply_qwen3_hybrid_scanned_blocks(self, y, layer_args, layer_kwargs, kv_caches=None): + """Applies the Qwen3-Next / Qwen3.5 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 diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index f5749a026c..d3916bb501 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1551,11 +1551,7 @@ def __init__( 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, + lambda layer_rngs: self._make_decoder_layer( layer_idx=0, is_full_attention_layer=False, rngs=layer_rngs, @@ -1569,11 +1565,7 @@ def __init__( self.local_layers = None if self.num_global > 0: - self.global_layer = Qwen3NextDecoderLayer( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, + self.global_layer = self._make_decoder_layer( layer_idx=full_attention_offset, is_full_attention_layer=True, rngs=self.rngs, @@ -1598,6 +1590,22 @@ def __init__( else: self._maybe_shard_with_logical = lambda inputs, *args, **kwargs: inputs + def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): + """Builds one sub-layer of the block. + + Subclasses for architectures that share this block structure (e.g. Qwen3.5) + override this to swap in their own decoder layer class. + """ + return Qwen3NextDecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + layer_idx=layer_idx, + is_full_attention_layer=is_full_attention_layer, + rngs=rngs, + ) + 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) diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index a043aea640..48ebe66d70 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -29,6 +29,7 @@ from maxtext.models.qwen3 import ( Qwen3NextFullAttention, Qwen3NextGatedDeltaNet, + Qwen3NextScannableBlock, Qwen3NextSparseMoeBlock, ) from maxtext.utils import max_utils @@ -51,75 +52,26 @@ class Qwen3_5SparseMoEBlock(Qwen3NextSparseMoeBlock): """Shares same MoE code as Qwen3-Next""" -class Qwen3_5ScannableBlock(nnx.Module): - """Scanned Structure for Text-only Architecture, explicitly invoking Qwen3_5 layers.""" +class Qwen3_5ScannableBlock(Qwen3NextScannableBlock): + """Scanned Structure for Text-only Architecture, explicitly invoking Qwen3_5 layers. - def __init__(self, config: Config, mesh: Mesh, model_mode: str, quant=None, *, rngs: nnx.Rngs): - self.config = config - self.mesh = mesh - self.model_mode = model_mode - self.quant = quant - self.rngs = rngs - cfg = self.config - - # Explicitly instantiate Qwen3_5DecoderLayer here - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer_rngs = self.rngs.fork() - layer_name = f"layer_{i}" - layer = Qwen3_5DecoderLayer( - config=self.config, - mesh=self.mesh, - quant=self.quant, - model_mode=self.model_mode, - layer_idx=i, - rngs=layer_rngs, - ) - setattr(self, layer_name, layer) - - def __call__( - self, - carry: jnp.ndarray, - decoder_segment_ids: None | jnp.ndarray, - decoder_positions: None | jnp.ndarray, - deterministic: bool, - model_mode: str, - previous_chunk=None, - slot: None | int = None, - kv_cache: None | tuple[Any, ...] = None, - attention_metadata: None | dict[str, Any] = None, - forced_routed_experts: jnp.ndarray | None = None, - ) -> tuple[Array, None | tuple[Any, ...]]: - cfg = self.config - x = carry - # Inference with externally-managed (vLLM) caches hands this block one - # entry per sub-layer of the cycle; the decoder groups the flat per-layer - # list before the scan and ungroups the returned tuple afterwards. - updated_kv_caches = [] - - for i in range(cfg.inhomogeneous_layer_cycle_interval): - layer = getattr(self, f"layer_{i}") - # forced_routed_experts, when present, is shaped - # [inhomogeneous_layer_cycle_interval, batch, seq, top_k]: one slice per - # sub-layer in this cycle (see nnx_decoders.py's scan wiring). - layer_forced_routed_experts = forced_routed_experts[i] if forced_routed_experts is not None else None - layer_kv_cache = kv_cache[i] if (kv_cache is not None and i < len(kv_cache)) else None - x, new_kv_cache = layer( - x, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - previous_chunk, - slot, - kv_cache=layer_kv_cache, - attention_metadata=attention_metadata, - forced_routed_experts=layer_forced_routed_experts, - ) - updated_kv_caches.append(new_kv_cache) + Qwen3.5 repeats the same hybrid attention period as Qwen3-Next -- several + GatedDeltaNet layers followed by one full-attention layer -- so it reuses + Qwen3-Next's hierarchical nested scans (an inner scan over the homogeneous + linear-attention layers plus a trip-count-one scan over the full-attention + layer) and only swaps in the Qwen3.5 decoder layer. + """ - if kv_cache is None: - return x, None - return x, tuple(updated_kv_caches) + def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): + return Qwen3_5DecoderLayer( + config=self.config, + mesh=self.mesh, + model_mode=self.model_mode, + quant=self.quant, + layer_idx=layer_idx, + is_full_attention_layer=is_full_attention_layer, + rngs=rngs, + ) class Qwen3_5DecoderLayer(nnx.Module): @@ -137,7 +89,15 @@ class Qwen3_5DecoderLayer(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 @@ -180,8 +140,12 @@ 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: + 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. if is_full_attention_layer: @@ -305,8 +269,10 @@ def __call__( # We sow the load balancing loss so it can be collected and added to the total loss # during training. + # Assigned rather than sown: `sow` appends to a tuple, so the layer's Intermediate + # structure would grow on every call, which a `jax.lax.scan` body cannot express. if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: - self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss) + self.moe_lb_loss = nnx.Intermediate(load_balance_loss) else: mlp_output = self.mlp( hidden_states, diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index c494e2cef6..abe3f054fa 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -995,26 +995,29 @@ def test_restores_local_state_and_preserves_kv_order(self): } -def _build_qwen3_next_block(layer_idx_offset=0, **overrides): - """Builds one Qwen3-Next scannable block; overrides go to the config.""" - cfg = _make_config(**{**_QWEN3_NEXT_CONFIG, **overrides}) - block = qwen3.Qwen3NextScannableBlock( - config=cfg, - mesh=_make_mesh(cfg), - model_mode=MODEL_MODE_TRAIN, - layer_idx_offset=layer_idx_offset, - rngs=nnx.Rngs(0), - ) - return cfg, block - - class TestQwen3NextScannableBlock(unittest.TestCase): """Tests Qwen3-Next's nested local(scan)/global(length-1 scan) decoder block.""" + MODEL_CONFIG = _QWEN3_NEXT_CONFIG + BLOCK_CLS = qwen3.Qwen3NextScannableBlock + + @classmethod + def _build(cls, layer_idx_offset=0, **overrides): + """Builds one scannable block; overrides go to the config, layer_idx_offset to the block.""" + cfg = _make_config(**{**cls.MODEL_CONFIG, **overrides}) + block = cls.BLOCK_CLS( + config=cfg, + mesh=_make_mesh(cfg), + model_mode=MODEL_MODE_TRAIN, + layer_idx_offset=layer_idx_offset, + rngs=nnx.Rngs(0), + ) + return cfg, block + @classmethod def setUpClass(cls): """Builds the block once; it costs several seconds and no test here mutates it.""" - cls.cfg, cls.block = _build_qwen3_next_block() + cls.cfg, cls.block = cls._build() def _inputs(self, cfg): inputs = jax.random.normal(jax.random.PRNGKey(1), (1, cfg.max_target_length, cfg.emb_dim), dtype=jnp.float32) @@ -1073,16 +1076,18 @@ def test_rejects_block_whose_global_layer_is_not_last(self): local-scan-then-global would silently reorder the model, so it is rejected. """ with self.assertRaisesRegex(ValueError, "full-attention layer last"): - _build_qwen3_next_block(layer_idx_offset=1) + self._build(layer_idx_offset=1) class TestNNXDecoderQwen3Next(unittest.TestCase): """Tests the NNXDecoder-level wiring of the Qwen3-Next scanned blocks.""" + MODEL_CONFIG = _QWEN3_NEXT_CONFIG + def _build(self, num_decoder_layers, **overrides): """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, **overrides} + **{**self.MODEL_CONFIG, "base_num_decoder_layers": num_decoder_layers, "scan_layers": True, **overrides} ) mesh = _make_mesh(cfg) decoder = NNXDecoder(config=cfg, mesh=mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=1)) @@ -1119,7 +1124,7 @@ def test_decoder_regroups_flat_kv_caches_per_block(self): 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 + ``_apply_qwen3_hybrid_scanned_blocks``, which must also keep ``skip_block_remat=True`` on this path rather than falling back to the generic (block-rematerialized) branch. @@ -1167,6 +1172,36 @@ def test_decoder_keeps_layers_past_the_last_whole_block(self): ) +# Qwen3.5 repeats Qwen3-Next's hybrid attention period and reuses its scannable +# block, so it gets the same coverage over its own layer classes and config. +_QWEN3_5_CONFIG = { + **_QWEN3_NEXT_CONFIG, + "run_name": "qwen3_5_scannable_block_test", + "model_name": "qwen3.5-35b-a3b", + # Qwen3.5 uses MRoPE, whose sections must cover rotary_dim / 2 frequencies. The + # shipped [11, 11, 10] is sized for the real head_dim (256), not this tiny one + # (32 * partial_rotary_factor 0.25 / 2 = 4). + "mrope_section": [2, 1, 1], +} + + +class TestQwen3_5ScannableBlock(TestQwen3NextScannableBlock): + """Runs the Qwen3-Next scannable-block tests against Qwen3.5's block.""" + + MODEL_CONFIG = _QWEN3_5_CONFIG + BLOCK_CLS = qwen3_5.Qwen3_5ScannableBlock + + def test_block_uses_qwen3_5_layers(self): + """The block must build Qwen3.5 decoder layers, not Qwen3-Next ones.""" + self.assertIsInstance(self.block.global_layer, qwen3_5.Qwen3_5DecoderLayer) + + +class TestNNXDecoderQwen3_5(TestNNXDecoderQwen3Next): + """Runs the NNXDecoder wiring tests against Qwen3.5.""" + + MODEL_CONFIG = _QWEN3_5_CONFIG + + class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" @@ -1743,105 +1778,126 @@ def test_fp8_forward_pass_execution(self): self.assertTrue(jnp.all(jnp.isfinite(logits))) -# Unbound forward of the Qwen3.5 scannable block. It is invoked below against a -# stand-in `self` rather than a real instance, so it is bound here once instead -# of being called as a dunder at each call site. -_QWEN3_5_BLOCK_FORWARD = qwen3_5.Qwen3_5ScannableBlock.__call__ +class _StatefulQwen3_5DecoderLayer(nnx.Module): + """Stand-in Qwen3.5 sub-layer that exposes the cache and metadata it receives.""" + + def __init__(self, *, is_full_attention_layer, **unused_kwargs): + self.increment = 10 if is_full_attention_layer else 1 + self.call_count = nnx.Intermediate(jnp.array(0, dtype=jnp.int32)) + self.received_kv_cache = nnx.Intermediate(jnp.array(False)) + self.received_attention_metadata = nnx.Intermediate(jnp.array(False)) + + def __call__(self, inputs, *unused_args, kv_cache=None, attention_metadata=None, **unused_kwargs): + self.call_count.value += 1 + self.received_kv_cache.value = jnp.array(kv_cache is not None) + self.received_attention_metadata.value = jnp.array(attention_metadata is not None) + output = inputs + self.increment + if kv_cache is None: + return output + return output, kv_cache + self.increment class Qwen3_5ScannableBlockKVCacheTest(unittest.TestCase): """Qwen3.5 must thread vLLM's externally-managed KV caches through its block. - `Qwen3_5ScannableBlock.__call__` only touches `self.config` and its - `layer_{i}` attributes, so it is exercised here against a stand-in `self`. - That keeps the test on CPU and focused on the cache plumbing rather than on - building a real MoE block. + Qwen3_5ScannableBlock reuses Qwen3NextScannableBlock, whose external-cache path + unrolls the stacked local layers and then runs the global layer. The sub-layers + are swapped for stand-ins, which keeps the test on CPU and focused on the cache + plumbing rather than on building a real MoE block. """ CYCLE = 4 - def _fake_block(self): - """Builds a stand-in block whose sub-layers record the kwargs they receive.""" - calls = [] - - def make_layer(idx): - def layer(x, *args, **kwargs): - calls.append({"idx": idx, "args": args, "kwargs": kwargs}) - return x + 1, f"updated_kv_{idx}" - - return layer + def setUp(self): + super().setUp() + self.config = SimpleNamespace( + inhomogeneous_layer_cycle_interval=self.CYCLE, + param_scan_axis=1, + remat_policy="none", + scan_layers=True, + shard_mode=ShardMode.AUTO, + ) - block = SimpleNamespace(config=SimpleNamespace(inhomogeneous_layer_cycle_interval=self.CYCLE)) - for i in range(self.CYCLE): - setattr(block, f"layer_{i}", make_layer(i)) - return block, calls + def _make_block(self): + with mock.patch.object(qwen3_5, "Qwen3_5DecoderLayer", _StatefulQwen3_5DecoderLayer): + return qwen3_5.Qwen3_5ScannableBlock( + config=self.config, + mesh=None, + model_mode=MODEL_MODE_AUTOREGRESSIVE, + rngs=nnx.Rngs(0), + ) def _call(self, block, **kwargs): - return _QWEN3_5_BLOCK_FORWARD( - block, - 0, # carry - None, # decoder_segment_ids - None, # decoder_positions - True, # deterministic - MODEL_MODE_AUTOREGRESSIVE, + return block( + jnp.zeros((1, 1, 1)), + decoder_segment_ids=None, + decoder_positions=None, + deterministic=True, + model_mode=MODEL_MODE_AUTOREGRESSIVE, **kwargs, ) @pytest.mark.cpu_only def test_each_sublayer_gets_its_own_cache_and_updates_are_returned(self): - block, calls = self._fake_block() - kv_cache = tuple(f"kv_{i}" for i in range(self.CYCLE)) + block = self._make_block() + kv_cache = tuple(jnp.array(i) for i in range(self.CYCLE)) - carry, updated = self._call(block, kv_cache=kv_cache, attention_metadata={"meta": 1}) + output, updated = self._call(block, kv_cache=kv_cache, attention_metadata=object()) - self.assertEqual(carry, self.CYCLE) # every sub-layer ran exactly once - self.assertEqual(updated, tuple(f"updated_kv_{i}" for i in range(self.CYCLE))) - self.assertEqual([c["kwargs"]["kv_cache"] for c in calls], list(kv_cache)) + # Every sub-layer ran exactly once: three local layers (+1 each), then the global one (+10). + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) + np.testing.assert_array_equal(jnp.stack(updated), jnp.array([1, 2, 3, 13])) + np.testing.assert_array_equal(block.local_layers.call_count.value, jnp.ones(self.CYCLE - 1, dtype=jnp.int32)) + np.testing.assert_array_equal(block.global_layer.call_count.value, 1) # attention_metadata must reach the sub-layers; without it the attention # kernel cannot address vLLM's paged cache. - self.assertTrue(all(c["kwargs"]["attention_metadata"] == {"meta": 1} for c in calls)) + np.testing.assert_array_equal( + block.local_layers.received_attention_metadata.value, jnp.ones(self.CYCLE - 1, dtype=jnp.bool_) + ) + np.testing.assert_array_equal(block.global_layer.received_attention_metadata.value, True) @pytest.mark.cpu_only def test_training_path_returns_none_and_passes_no_cache(self): - block, calls = self._fake_block() + block = self._make_block() - carry, updated = self._call(block) + output, updated = self._call(block) - self.assertEqual(carry, self.CYCLE) + np.testing.assert_array_equal(output, jnp.full((1, 1, 1), 13)) self.assertIsNone(updated) - self.assertTrue(all(c["kwargs"]["kv_cache"] is None for c in calls)) + np.testing.assert_array_equal(block.local_layers.received_kv_cache.value, jnp.zeros(self.CYCLE - 1, dtype=jnp.bool_)) + np.testing.assert_array_equal(block.global_layer.received_kv_cache.value, False) @pytest.mark.cpu_only def test_shorter_cache_does_not_index_out_of_range(self): - block, calls = self._fake_block() + block = self._make_block() - _, updated = self._call(block, kv_cache=("kv_0",)) + _, updated = self._call(block, kv_cache=(jnp.array(0),)) - self.assertEqual([c["kwargs"]["kv_cache"] for c in calls], ["kv_0", None, None, None]) self.assertEqual(len(updated), self.CYCLE) + np.testing.assert_array_equal(updated[0], 1) + self.assertEqual(updated[1:], (None,) * (self.CYCLE - 1)) @pytest.mark.cpu_only def test_flat_per_layer_caches_survive_the_per_block_round_trip(self): """vLLM hands over one cache per layer; the scan runs over 4-layer blocks.""" num_layers = 8 scan_length = num_layers // self.CYCLE - kv_caches = [f"kv_{i}" for i in range(num_layers)] + kv_caches = [jnp.array(i) for i in range(num_layers)] grouped = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length, self.CYCLE, stack=False) - self.assertEqual(grouped, [("kv_0", "kv_1", "kv_2", "kv_3"), ("kv_4", "kv_5", "kv_6", "kv_7")]) + self.assertEqual([[int(kv) for kv in group] for group in grouped], [[0, 1, 2, 3], [4, 5, 6, 7]]) # Mimic _apply_layers_sequentially replacing each block entry with the # tuple the block returned. + block = self._make_block() for block_idx in range(scan_length): - fake_block, _ = self._fake_block() - _, updated = self._call(fake_block, kv_cache=grouped[block_idx]) - grouped[block_idx] = tuple(f"{kv}_b{block_idx}" for kv in updated) + _, grouped[block_idx] = self._call(block, kv_cache=grouped[block_idx]) maxtext_utils.update_kv_caches_after_scan(kv_caches, grouped, scan_length, self.CYCLE, stacked=False) self.assertEqual( - kv_caches, - [f"updated_kv_{i % self.CYCLE}_b{i // self.CYCLE}" for i in range(num_layers)], + [int(kv) for kv in kv_caches], + [i + (10 if i % self.CYCLE == self.CYCLE - 1 else 1) for i in range(num_layers)], ) diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index 97f2dc40bd..3113ed6467 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -28,6 +28,7 @@ ) from maxtext.checkpoint_conversion.utils import hf_shape from maxtext.checkpoint_conversion.utils import param_mapping +from maxtext.checkpoint_conversion.utils import tensor_handling from maxtext.checkpoint_conversion.utils import utils from maxtext.checkpoint_conversion.utils.utils import process_maxtext_param @@ -105,12 +106,24 @@ def test_qwen3_5_dense_mapping(self): for scanned in (False, True): with self.subTest(scan_layers=scanned): mapping = param_mapping.PARAM_MAPPING["qwen3.5-9b"](config, maxtext_config, scan_layers=scanned) - for index in range(4 if scanned else 8): - prefix = f"params-decoder-layers-layer_{index}" if scanned else f"params-decoder-layers_{index}" - indices = range(index, 8, 4) if scanned else [index] - for mt_name, hf_name in (("wi_0", "gate_proj"), ("wi_1", "up_proj"), ("wo", "down_proj")): - expected = [f"model.language_model.layers.{i}.mlp.{hf_name}.weight" for i in indices] - self.assertEqual(mapping[f"{prefix}-mlp-{mt_name}-kernel"], expected if scanned else expected[0]) + for mt_name, hf_name in (("wi_0", "gate_proj"), ("wi_1", "up_proj"), ("wo", "down_proj")): + + def hf_key(layer, hf_name=hf_name): + return f"model.language_model.layers.{layer}.mlp.{hf_name}.weight" + + if scanned: + # The nested block scan stacks local layers [block][local] and the global layer [block]. + self.assertEqual( + mapping[f"params-decoder-layers-local_layers-mlp-{mt_name}-kernel"], + [[hf_key(b * 4 + p) for p in range(3)] for b in range(2)], + ) + self.assertEqual( + mapping[f"params-decoder-layers-global_layer-mlp-{mt_name}-kernel"], + [hf_key(b * 4 + 3) for b in range(2)], + ) + else: + for index in range(8): + self.assertEqual(mapping[f"params-decoder-layers_{index}-mlp-{mt_name}-kernel"], hf_key(index)) self.assertFalse(any("routed_experts" in str(key) or "shared_expert" in str(key) for key in mapping)) def test_qwen3_next_mapping(self): @@ -164,6 +177,127 @@ def test_qwen3_next_mapping_scanned(self): self.assertEqual(len(global_experts), num_experts) self.assertEqual(len(global_experts[0]), num_blocks) + def test_qwen3_5_mapping_scanned(self): + """Qwen3.5 reuses Qwen3-Next's nested block scan, but stores its routed experts fused.""" + num_layers, cycle = 8, 4 + config = {"text_config": {"num_hidden_layers": num_layers}} + maxtext_config = mock.Mock(num_decoder_layers=num_layers, num_experts=4, weight_dtype="bfloat16") + maxtext_config.inhomogeneous_layer_cycle_interval = cycle + maxtext_config.use_multimodal = False + mapping = param_mapping.QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) + + num_blocks, num_local = num_layers // cycle, cycle - 1 + local_prefix = "params-decoder-layers-local_layers" + global_prefix = "params-decoder-layers-global_layer" + # Linear attention only exists on the local layers, full attention only on the global one. + self.assertIn(f"{local_prefix}-attention-in_proj_qkvz-kernel", mapping) + self.assertIn(f"{global_prefix}-attention-attention-query-kernel", mapping) + 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. + # A tuple of HF names is a composite source the hook fuses, not a stacking axis. + 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.language_model.layers.0.linear_attn.in_proj_qkv.weight", + "model.language_model.layers.0.linear_attn.in_proj_z.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.language_model.layers.{cycle - 1}.self_attn.q_proj.weight") + + # Fused experts: one HF tensor feeds both wi_0 and wi_1, so it is keyed by the pair + # and carries no extra per-expert axis (unlike Qwen3-Next). + fused_local = mapping[(f"{local_prefix}-mlp-routed_experts-wi_0", f"{local_prefix}-mlp-routed_experts-wi_1")] + self.assertEqual(len(fused_local), num_blocks) + self.assertEqual(len(fused_local[0]), num_local) + self.assertEqual(fused_local[0][0], "model.language_model.layers.0.mlp.experts.gate_up_proj") + fused_global = mapping[(f"{global_prefix}-mlp-routed_experts-wi_0", f"{global_prefix}-mlp-routed_experts-wi_1")] + self.assertEqual(len(fused_global), num_blocks) + + def test_qwen3_5_fp8_mapping_scanned(self): + """FP8 checkpoints store each routed expert separately, which nests one level deeper.""" + num_layers, cycle, num_experts = 8, 4, 3 + config = {"text_config": {"num_hidden_layers": num_layers, "num_experts": num_experts}} + maxtext_config = mock.Mock( + num_decoder_layers=num_layers, + num_experts=num_experts, + inhomogeneous_layer_cycle_interval=cycle, + weight_dtype="float8_e4m3fn", + use_multimodal=False, + ) + mapping = param_mapping.QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=True) + linear_attn_dims = { + "linear_num_key_heads": 2, + "linear_num_value_heads": 4, + "linear_key_head_dim": 8, + "linear_value_head_dim": 8, + } + hooks = param_mapping.QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN( + {"text_config": {**config["text_config"], **linear_attn_dims}}, maxtext_config, scan_layers=True + ) + + num_blocks, num_local = num_layers // cycle, cycle - 1 + local_prefix = "params-decoder-layers-local_layers" + global_prefix = "params-decoder-layers-global_layer" + self.assertFalse(any("layers-layer_" in str(key) for key in mapping)) + # [expert][block][local] on the local layers, [expert][block] on the global one. + local_wo_scale = mapping[f"{local_prefix}-mlp-routed_experts-wo_scale"] + self.assertEqual( + (len(local_wo_scale), len(local_wo_scale[0]), len(local_wo_scale[0][0])), (num_experts, num_blocks, num_local) + ) + self.assertEqual( + local_wo_scale[2][1][0], f"model.language_model.layers.{cycle}.mlp.experts.2.down_proj.weight_scale_inv" + ) + global_wi_0 = mapping[f"{global_prefix}-mlp-routed_experts-wi_0"] + self.assertEqual((len(global_wi_0), len(global_wi_0[0])), (num_experts, num_blocks)) + self.assertEqual(global_wi_0[1][1], f"model.language_model.layers.{2 * cycle - 1}.mlp.experts.1.gate_proj.weight") + # Per-expert checkpoints have no fused gate_up_proj. + self.assertNotIn((f"{local_prefix}-mlp-routed_experts-wi_0", f"{local_prefix}-mlp-routed_experts-wi_1"), mapping) + # Attention and shared-expert scales follow their kernels' layout. + self.assertEqual( + mapping[f"{local_prefix}-attention-in_proj_qkvz-kernel_scale"][0][num_local - 1], + ( + f"model.language_model.layers.{num_local - 1}.linear_attn.in_proj_qkv.weight_scale_inv", + f"model.language_model.layers.{num_local - 1}.linear_attn.in_proj_z.weight_scale_inv", + ), + ) + self.assertEqual( + mapping[f"{global_prefix}-attention-attention-out-kernel_scale"], + [f"model.language_model.layers.{b * cycle + cycle - 1}.self_attn.o_proj.weight_scale_inv" for b in range(2)], + ) + self.assertEqual(len(mapping[f"{global_prefix}-mlp-shared_expert-wo-kernel_scale"]), num_blocks) + # The hooks are registered under the same nested prefixes as the mapping. + for key in ( + f"{local_prefix}-attention-in_proj_qkvz-kernel_scale", + f"{local_prefix}-mlp-routed_experts-wo_scale", + f"{global_prefix}-attention-attention-out-kernel_scale", + f"{global_prefix}-mlp-routed_experts-wi_0", + ): + self.assertIn(key, mapping) + self.assertIn(key, hooks) + + def test_qwen3_5_composite_key_uses_the_nested_scan_axes(self): + """A composite (tuple) MaxText key must still be recognised as a nested block scan. + + Qwen3.5's fused gate_up_proj is keyed by the (wi_0, wi_1) pair, so a layout + check that only looks at `str` keys would place the (blocks, local) axes at + the leading positions and silently transpose the weights. + """ + cfg = mock.Mock() + cfg.param_scan_axis = 1 + local_key = "params-decoder-layers-local_layers-mlp-routed_experts-wi_0" + self.assertEqual( + tensor_handling.stacked_axes((local_key, local_key.replace("wi_0", "wi_1")), cfg, depth=2), + tensor_handling.stacked_axes(local_key, cfg, depth=2), + ) + @staticmethod def _indices_from_name(name): """Parses a synthetic HF key such as "e1_b0_l2" back into its stacked-axis indices.""" @@ -956,23 +1090,25 @@ def test_qwen3_5_regression_scanned_and_unscanned(self): ) if scan_layers: - gdn_prefix = "params-decoder-layers-layer_0" - attn_prefix = "params-decoder-layers-layer_3" + # Local layers stack [block][local], the global layer [block]. + local_prefix = "params-decoder-layers-local_layers" + global_prefix = "params-decoder-layers-global_layer" + local_blocks = ([0, 1, 2], [4, 5, 6]) self.assertEqual( - mapping[f"{gdn_prefix}-attention-in_proj_qkvz-kernel"], + mapping[f"{local_prefix}-attention-in_proj_qkvz-kernel"], [ - ( - "model.language_model.layers.0.linear_attn.in_proj_qkv.weight", - "model.language_model.layers.0.linear_attn.in_proj_z.weight", - ), - ( - "model.language_model.layers.4.linear_attn.in_proj_qkv.weight", - "model.language_model.layers.4.linear_attn.in_proj_z.weight", - ), + [ + ( + f"model.language_model.layers.{i}.linear_attn.in_proj_qkv.weight", + f"model.language_model.layers.{i}.linear_attn.in_proj_z.weight", + ) + for i in block + ] + for block in local_blocks ], ) self.assertEqual( - mapping[f"{attn_prefix}-attention-attention-query-kernel"], + mapping[f"{global_prefix}-attention-attention-query-kernel"], [ "model.language_model.layers.3.self_attn.q_proj.weight", "model.language_model.layers.7.self_attn.q_proj.weight", @@ -980,30 +1116,24 @@ def test_qwen3_5_regression_scanned_and_unscanned(self): ) if weight_dtype == "float8_e4m3fn": self.assertEqual( - mapping[f"{gdn_prefix}-mlp-routed_experts-wo_scale"], + mapping[f"{local_prefix}-mlp-routed_experts-wo_scale"], [ [ - "model.language_model.layers.0.mlp.experts.0.down_proj.weight_scale_inv", - "model.language_model.layers.4.mlp.experts.0.down_proj.weight_scale_inv", - ], - [ - "model.language_model.layers.0.mlp.experts.1.down_proj.weight_scale_inv", - "model.language_model.layers.4.mlp.experts.1.down_proj.weight_scale_inv", - ], + [f"model.language_model.layers.{i}.mlp.experts.{e}.down_proj.weight_scale_inv" for i in block] + for block in local_blocks + ] + for e in range(2) ], ) else: self.assertEqual( mapping[ ( - f"{gdn_prefix}-mlp-routed_experts-wi_0", - f"{gdn_prefix}-mlp-routed_experts-wi_1", + f"{local_prefix}-mlp-routed_experts-wi_0", + f"{local_prefix}-mlp-routed_experts-wi_1", ) ], - [ - "model.language_model.layers.0.mlp.experts.gate_up_proj", - "model.language_model.layers.4.mlp.experts.gate_up_proj", - ], + [[f"model.language_model.layers.{i}.mlp.experts.gate_up_proj" for i in block] for block in local_blocks], ) else: self.assertEqual( @@ -1124,17 +1254,21 @@ def test_qwen3_8_scanned_ordering(self): mapping = param_mapping.QWEN3_8_MAXTEXT_TO_HF_PARAM_MAPPING(text_cfg, mt_cfg, scan_layers=True) mapping_fp8 = param_mapping.QWEN3_8_MAXTEXT_TO_HF_PARAM_MAPPING(text_cfg, mt_cfg_fp8, scan_layers=True) - expected_layer_groups = { - 0: [0, 4], - 1: [1, 5], - 2: [2, 6], - 3: [3, 7], - } - for block_idx, hf_indices in expected_layer_groups.items(): - prefix = f"params-decoder-layers-layer_{block_idx}" + # Layers 0-2 of each cycle are the local layers, nested [block][local]; + # layer 3 is the global layer, flat over blocks. + def local_layers(suffix): + return [[f"model.layers.{i}.{suffix}" for i in block] for block in ([0, 1, 2], [4, 5, 6])] + + def global_layer(suffix): + return [f"model.layers.{i}.{suffix}" for i in (3, 7)] + + for prefix, stacked in ( + ("params-decoder-layers-local_layers", local_layers), + ("params-decoder-layers-global_layer", global_layer), + ): self.assertEqual( mapping[f"{prefix}-input_layernorm-scale"], - [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices], + stacked("input_layernorm.weight"), ) scanned_wi_key = ( f"{prefix}-mlp-routed_experts-wi_0", @@ -1142,14 +1276,11 @@ def test_qwen3_8_scanned_ordering(self): ) self.assertEqual( mapping[scanned_wi_key], - [f"model.layers.{i}.mlp.experts.gate_up_proj" for i in hf_indices], + stacked("mlp.experts.gate_up_proj"), ) self.assertEqual( mapping_fp8[f"{prefix}-mlp-routed_experts-wi_0"], - [ - [f"model.layers.{i}.mlp.experts.0.gate_proj.weight" for i in hf_indices], - [f"model.layers.{i}.mlp.experts.1.gate_proj.weight" for i in hf_indices], - ], + [stacked(f"mlp.experts.{e}.gate_proj.weight") for e in range(2)], ) def test_qwen3_5_and_qwen3_8_namespace_boundaries_and_config_forms(self): From 684c0f3c76e5bd7d20d79a38317323d99d845d3f Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 05:11:54 +0000 Subject: [PATCH 2/5] Share Qwen3-Next's block scan with Qwen3.5 and keep both paths whole Qwen3.5's decoder layer and scannable block were near-duplicates of Qwen3-Next's. Collapse them onto the Qwen3-Next implementations, which already nest a scan over the cycle's local layers inside a trip-count-one scan around the full-attention layer, so Qwen3.5 picks up per-sub-layer remat instead of block-level remat. Sharing is via three overridable factory hooks on Qwen3NextDecoderLayer (`_make_full_attention`, `_make_linear_attention`, `_make_mlp`). Qwen3.5 overrides only those, keeping the submodule attribute names -- and therefore the checkpoint parameter paths -- identical, so the whole forward pass, both norms and the residual structure are inherited rather than copied. qwen3_5.py drops from 266 to 127 lines. Two things had to be repaired to keep functionality identical: Router replay. `_apply_layers_sequentially` filters call kwargs against the callee signature, and Qwen3NextScannableBlock did not accept `forced_routed_experts`, so scanned Qwen3.5 lost it. Thread it through `nnx_scan.apply_scanned_layers` (new optional per-layer `xs`), the block (scan xs for the local layers, closure for the global one, per-layer slices on the external-KV unroll) and the hybrid scan entry point, which also slices the scan remainder. Combining forced routing with externally-managed KV caches now raises rather than silently dropping one of them. vLLM weight sync. The RL rollout converter and the direct-sync unroll both assumed one module per cycle slot (`layers.layer_{slot}`) with a single stacked axis. The nested block stores `local_layers` with the cycle slot on `param_scan_axis + 1` and `global_layer` with the block axis alone, so both now read that layout, selecting the slot once per plan group. Trainer states in the old layout still convert. This also fixes Qwen3-Next, which produced the nested layout already and was therefore never syncable. AOT compile, per_device_batch_size=1, max_target_length=1024, flash attention with tokamax splash, temp_size_in_bytes vs main: qwen3.5-35b-a3b v5p-128 16.16 GB -> 11.03 GB (-31.7%) qwen3.5-35b-a3b v5p-512 11.99 GB -> 6.74 GB (-43.8%) qwen3.5-397b-a17b v5p-512 75.62 GB -> 46.01 GB (-39.2%) --- .../tunix/weight_mapping/raiden_unscan.py | 77 ++++-- .../integration/vllm/maxtext_vllm_rollout.py | 109 ++++++-- .../integration/vllm/weight_converter.py | 154 ++++++++--- src/maxtext/layers/nnx_decoders.py | 83 ++++-- src/maxtext/layers/nnx_scan.py | 22 +- src/maxtext/models/qwen3.py | 135 ++++++--- src/maxtext/models/qwen3_5.py | 259 ++++-------------- .../unit/vllm_rollout_unroll_test.py | 59 ++++ .../unit/weight_converter_test.py | 84 ++++++ tests/unit/nnx_decoders_test.py | 69 +++++ tests/unit/raiden_unscan_test.py | 40 +++ tests/unit/router_replay_test.py | 25 ++ 12 files changed, 758 insertions(+), 358 deletions(-) diff --git a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py index 15163f88b8..1360484c46 100644 --- a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py +++ b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py @@ -62,11 +62,14 @@ def unscan_layers( scan_axis: The axis along which layers are scanned (default 1). cycle_interval: `config.inhomogeneous_layer_cycle_interval`. When > 1 the trainer scans a *block* of this many heterogeneous layers (qwen3.5's - GDN/attention cycle is 4), so the tree carries an extra `layer_` - level under `layer_container` and the scan axis is - `num_layers // cycle_interval` repeats rather than `num_layers`. Slot `j` - of repeat `i` is physical layer `i * cycle_interval + j`, which is the - flat `layers_` name the unscanned rollout model uses. + GDN/attention cycle is 4), so the tree carries an extra level under + `layer_container` and the scan axis is `num_layers // cycle_interval` + repeats rather than `num_layers`. That level is either one `layer_` + per cycle slot, or the nested `Qwen3NextScannableBlock` layout (Qwen3-Next + and Qwen3.5), where `local_layers` stacks slots `0..cycle_interval-2` on + `scan_axis + 1` and `global_layer` is the last slot. Slot `j` of repeat + `i` is physical layer `i * cycle_interval + j`, which is the flat + `layers_` name the unscanned rollout model uses. Returns: A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding @@ -111,25 +114,32 @@ def unscan_layers( idx = key.index(layer_container) prefix = key[:idx] suffix = key[idx + 1 :] - # A scanned inhomogeneous block nests one `layer_` level under - # `layers`; the rollout has no such level, so drop it and fold the slot - # into the physical layer number below. - slot = None + # A scanned inhomogeneous block nests one cycle-slot level under `layers`; + # the rollout has no such level, so drop it and fold the slot into the + # physical layer number below. `slots` pairs each cycle slot this leaf holds + # with its index along the nested `local_layers` slot axis (None if none). + slots = ((0, None),) if cycle_interval > 1: - m = re.fullmatch(r"layer_(\d+)", suffix[0]) if suffix else None - if not m: - prefix_str = suffix[0] if suffix else "" - raise ValueError( - f"unscan_layers: cycle_interval={cycle_interval} > 1, but key {'.'.join(key)!r} " - f"does not have a 'layer_' cycle prefix under {layer_container!r} " - f"(got {prefix_str!r})." - ) - slot = int(m.group(1)) - if slot >= cycle_interval: - raise ValueError( - f"unscan_layers: slot {slot} parsed from key {'.'.join(key)!r} " - f"must be less than cycle_interval {cycle_interval}." - ) + token = suffix[0] if suffix else "" + if token == "local_layers": + slots = tuple((j, j) for j in range(cycle_interval - 1)) + elif token == "global_layer": + slots = ((cycle_interval - 1, None),) + else: + m = re.fullmatch(r"layer_(\d+)", token) + if not m: + raise ValueError( + f"unscan_layers: cycle_interval={cycle_interval} > 1, but key {'.'.join(key)!r} " + f"does not have a 'layer_' cycle prefix (or a nested 'local_layers'/'global_layer' " + f"level) under {layer_container!r} (got {token!r})." + ) + slot = int(m.group(1)) + if slot >= cycle_interval: + raise ValueError( + f"unscan_layers: slot {slot} parsed from key {'.'.join(key)!r} " + f"must be less than cycle_interval {cycle_interval}." + ) + slots = ((slot, None),) suffix = suffix[1:] arr = getattr(value, "value", value) @@ -144,18 +154,29 @@ def unscan_layers( del value # Each slot is scanned over the repeats of the cycle, not over all layers. - expected = num_layers // cycle_interval if slot is not None else num_layers + expected = num_layers // cycle_interval if arr.shape[scan_axis] != expected: raise ValueError( f"unscan_layers: {'.'.join(key)!r} has shape {arr.shape}, expected axis {scan_axis} to be" f" {expected} (num_layers={num_layers}, cycle_interval={cycle_interval})." ) + nested = slots[0][1] is not None + if nested and (arr.ndim <= scan_axis + 1 or arr.shape[scan_axis + 1] != len(slots)): + raise ValueError( + f"unscan_layers: {'.'.join(key)!r} has shape {arr.shape}, expected axis {scan_axis + 1} to hold" + f" the {len(slots)} 'local_layers' cycle slots (cycle_interval={cycle_interval})." + ) for i in range(expected): - sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False) - layer_no = i * cycle_interval + slot if slot is not None else i - new_key = prefix + (f"{layer_container}_{layer_no}",) + suffix - new_flat[new_key] = sliced + block = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False) + for slot, local_index in slots: + sliced = block + if local_index is not None: + # Dropping the repeat axis moved the nested slot axis down to `scan_axis`. + sliced = jax.lax.index_in_dim(block, local_index, axis=scan_axis, keepdims=False) + new_key = prefix + (f"{layer_container}_{i * cycle_interval + slot}",) + suffix + new_flat[new_key] = sliced + del block del arr unscanned_count += 1 diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 4061025ad0..60320c38d0 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -31,7 +31,7 @@ import re import time import traceback -from typing import Any, Optional, Tuple +from typing import Any, NamedTuple, Optional, Tuple from flax import nnx from flax.traverse_util import flatten_dict, unflatten_dict @@ -298,25 +298,61 @@ def source_covers(target_key): return len(target_layer_keys) -def _find_qwen_scanned_layer_idx(key_tuple): - """Finds a Qwen scanned block path like `layers.layer_0` or `layers.moe_block`.""" +# `Qwen3NextScannableBlock` (shared by Qwen3-Next and Qwen3.5) nests two scans: +# the cycle's linear-attention layers share a single `local_layers` module, with +# blocks on `scan_axis` and the cycle slot on `scan_axis + 1`, while the trailing +# full-attention layer sits in `global_layer` with blocks on `scan_axis` alone. +_QWEN_NESTED_LOCAL = "local_layers" +_QWEN_NESTED_GLOBAL = "global_layer" + + +class _QwenScannedRef(NamedTuple): + """Where a scanned Qwen parameter sits, and how it maps onto cycle slots. + + `slot_idx` is the parameter's fixed position in the layer cycle. It is None + for the two nested containers: `local_layers` spans several slots, read off + `scan_axis + 1` rather than off the key, and `global_layer` always takes + whichever slot closes the cycle. + """ + + container_idx: int # Position of "layers" in the key tuple, or -1 for no match. + slot_idx: Optional[int] + consumed: int # Key elements the container occupies. + nested: bool # Slots are stacked on `scan_axis + 1`. + trailing: bool # Occupies the last slot of the cycle. + + +_NO_QWEN_SCAN = _QwenScannedRef(-1, -1, 0, False, False) + + +def _find_qwen_scanned_layer_idx(key_tuple) -> _QwenScannedRef: + """Finds a Qwen scanned block path like `layers.local_layers` or `layers.layer_0`.""" for i in range(len(key_tuple) - 1): if key_tuple[i] != "layers" or not isinstance(key_tuple[i + 1], str): continue - if key_tuple[i + 1].startswith("layers_"): + name = key_tuple[i + 1] + if name.startswith("layers_"): continue - match = re.fullmatch(r"layer_(\d+)", key_tuple[i + 1]) + if name == _QWEN_NESTED_LOCAL: + return _QwenScannedRef(i, None, 2, True, False) + if name == _QWEN_NESTED_GLOBAL: + return _QwenScannedRef(i, None, 2, False, True) + match = re.fullmatch(r"layer_(\d+)", name) if match: - return i, int(match.group(1)), 2 - return i, 0, 1 - return -1, -1, 0 + return _QwenScannedRef(i, int(match.group(1)), 2, False, False) + return _QwenScannedRef(i, 0, 1, False, False) + return _NO_QWEN_SCAN def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Optional[int] = None): """Unroll Qwen's heterogeneous or homogeneous scanned blocks for an unscanned MaxText target. - Qwen 3 Next/3.5 training stores a repeating layer cycle as - `decoder.layers.layer_{slot}`, with repetitions stacked on `scan_axis`. + Qwen 3 Next/3.5 training stores a repeating layer cycle under a nested scan: + `decoder.layers.local_layers` holds the cycle's linear-attention layers with + repetitions on `scan_axis` and the cycle slot on `scan_axis + 1`, and + `decoder.layers.global_layer` holds the trailing full-attention layer with + repetitions on `scan_axis` alone. Trainer states predating that layout, which + kept one `decoder.layers.layer_{slot}` module per slot, are still accepted. Qwen 3 base training stores homogeneous layers as `decoder.layers.*` stacked on `scan_axis`. The inference model stores every layer as a direct decoder attribute named `layers_{global_index}`. Tunix's generic direct-sync mapper cannot bridge @@ -343,19 +379,38 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt scanned_keys = [] slot_indices = set() scan_lengths = set() + local_widths = set() + has_trailing = False for key, value in flat_w.items(): - container_idx, slot_idx, consumed = _find_qwen_scanned_layer_idx(key) - if container_idx == -1 or "dropout" in key or "rngs" in key: + ref = _find_qwen_scanned_layer_idx(key) + if ref.container_idx == -1 or "dropout" in key or "rngs" in key: continue - if not hasattr(value, "shape") or len(value.shape) <= scan_axis: + # A nested `local_layers` parameter carries its cycle slots on the axis just + # past the block axis, so it needs one dimension more than a flat one. + min_ndim = scan_axis + 2 if ref.nested else scan_axis + 1 + if not hasattr(value, "shape") or len(value.shape) < min_ndim: raise ValueError(f"Qwen scanned parameter {'.'.join(map(str, key))} has no scan axis {scan_axis}: {value!r}") - scanned_keys.append((key, value, container_idx, slot_idx, consumed)) - slot_indices.add(slot_idx) + scanned_keys.append((key, value, ref)) scan_lengths.add(value.shape[scan_axis]) + if ref.nested: + local_widths.add(value.shape[scan_axis + 1]) + slot_indices.update(range(value.shape[scan_axis + 1])) + elif ref.trailing: + has_trailing = True + else: + slot_indices.add(ref.slot_idx) if not scanned_keys: return weights + if len(local_widths) > 1: + raise ValueError(f"Qwen nested `local_layers` parameters disagree on cycle width: {sorted(local_widths)}") + # `global_layer` closes the cycle, so its slot is whatever follows the local + # ones. Resolved here because the key alone does not carry the index. + trailing_slot = (max(slot_indices) + 1) if slot_indices else 0 + if has_trailing: + slot_indices.add(trailing_slot) + if pattern_length is None: expected_slots = set(range(max(slot_indices) + 1)) if slot_indices != expected_slots: @@ -369,16 +424,26 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt raise ValueError(f"Qwen scanned parameters disagree on scan length: {sorted(scan_lengths)}") scan_length = scan_lengths.pop() - scanned_key_paths = {key for key, _, _, _, _ in scanned_keys} + scanned_key_paths = {key for key, _, _ in scanned_keys} new_flat_w = {key: value for key, value in flat_w.items() if key not in scanned_key_paths} - for key, value, container_idx, slot_idx, consumed in scanned_keys: - prefix = key[:container_idx] - suffix = key[container_idx + consumed :] + for key, value, ref in scanned_keys: + prefix = key[: ref.container_idx] + suffix = key[ref.container_idx + ref.consumed :] + # `local_layers` supplies several cycle slots from one tensor; every other + # container supplies exactly one, at a slot fixed above. + if ref.nested: + slots = tuple(range(value.shape[scan_axis + 1])) + else: + slots = (trailing_slot if ref.trailing else ref.slot_idx,) for repetition in range(scan_length): - global_idx = repetition * pattern_length + slot_idx - new_key = prefix + (f"layers_{global_idx}",) + suffix - new_flat_w[new_key] = jnp.take(value, repetition, axis=scan_axis) + block = jnp.take(value, repetition, axis=scan_axis) + for local_idx, slot_idx in enumerate(slots): + global_idx = repetition * pattern_length + slot_idx + new_key = prefix + (f"layers_{global_idx}",) + suffix + # Taking the block collapses `scan_axis`, so the nested slot axis has + # shifted down into it. + new_flat_w[new_key] = jnp.take(block, local_idx, axis=scan_axis) if ref.nested else block logging.info( "MaxTextVllmSampler: unrolled %d Qwen tensor components across %d layers for direct MaxText weight sync.", diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py index b6fd54fdfe..ee5de8abf0 100644 --- a/src/maxtext/integration/vllm/weight_converter.py +++ b/src/maxtext/integration/vllm/weight_converter.py @@ -518,6 +518,9 @@ class _PlanEntry: # Index along the scan axis, or None when the source is already per-layer. slice_index: Optional[int] op: str # "identity" | "slice" | "fuse_moe" + # Index along the nested cycle-slot axis (`scan_axis + 1`), or None when the + # source stacks blocks alone. Only `local_layers` sources carry one. + local_index: Optional[int] = None @dataclasses.dataclass(frozen=True) @@ -542,21 +545,25 @@ class _PlanGroup: # Dot-joined source path, precomputed for diagnostics and for the # leaf-name check inside `_align_per_axis`. source_path: str + # Shared by every entry in the group: see `_PlanEntry.local_index`. + local_index: Optional[int] = None def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]: """Collapses per-leaf plan entries into one group per source parameter. - Entries are grouped by `(source_keys, op)`. Target shape is deliberately - *not* part of the key: every rollout layer reading a given scanned source - must share a shape, and if two did not, the bulk unstack would produce the - wrong shape for one of them -- which `convert()`'s post-execution shape - check catches loudly rather than silently writing incorrect weights. + Entries are grouped by `(source_keys, op, local_index)`. Target shape is + deliberately *not* part of the key: every rollout layer reading a given + scanned source must share a shape, and if two did not, the bulk unstack would + produce the wrong shape for one of them -- which `convert()`'s post-execution + shape check catches loudly rather than silently writing incorrect weights. + `local_index` *is* part of the key, because one nested `local_layers` source + feeds a different set of rollout layers at each cycle slot. """ grouped: Dict[Tuple[Any, ...], List[_PlanEntry]] = {} order: List[Tuple[Any, ...]] = [] for entry in plan: - key = (entry.source_keys, entry.op) + key = (entry.source_keys, entry.op, entry.local_index) if key not in grouped: grouped[key] = [] order.append(key) @@ -565,11 +572,12 @@ def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]: groups: List[_PlanGroup] = [] for key in order: entries = grouped[key] - source_keys, op = key + source_keys, op, local_index = key groups.append( _PlanGroup( source_keys=source_keys, op=op, + local_index=local_index, targets=tuple( (e.slice_index, e.target_key) # `identity` entries carry no index; -1 keeps the sort key @@ -586,6 +594,19 @@ def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]: return groups +def _last_group_reading(groups: List[_PlanGroup]) -> Dict[Tuple[Any, ...], int]: + """Maps each source key to the index of the last group that reads it. + + A nested `local_layers` source feeds one group per cycle slot, so it can only + be released once the last of those groups has run. + """ + last: Dict[Tuple[Any, ...], int] = {} + for idx, group in enumerate(groups): + for key in group.source_keys: + last[key] = idx + return last + + class ConversionPlanError(ValueError): """Raised when the source and target trees cannot be fully reconciled.""" @@ -768,18 +789,29 @@ def _resolve_target_dtype(self): # -------------------------------------------------------------- # def _scanned_candidates( self, key_tuple: Tuple[Any, ...], prefix_len: int, suffix_start: int, slot: int - ) -> List[Tuple[Any, ...]]: - """Source keys that could hold the scanned form of a target layer key.""" + ) -> List[Tuple[Tuple[Any, ...], Optional[int]]]: + """Source keys that could hold the scanned form of a target layer key. + + Each candidate is paired with the index to take along the nested cycle-slot + axis, or None when the source stacks blocks alone. `Qwen3NextScannableBlock` + nests two scans: the cycle's linear-attention layers share one + `local_layers` module whose slot sits on `scan_axis + 1`, while the trailing + full-attention layer sits in `global_layer` with no slot axis at all. + """ prefix, suffix = key_tuple[:prefix_len], key_tuple[suffix_start:] if self.cycle > 1: - # Inhomogeneous: `Qwen3_5ScannableBlock` holds `layer_0..layer_{C-1}`. - return [ - prefix + ("layers", f"layer_{slot}") + suffix, - prefix + ("layers", str(slot)) + suffix, - prefix + ("layers", slot) + suffix, + if slot == self.cycle - 1: + candidates = [(prefix + ("layers", "global_layer") + suffix, None)] + else: + candidates = [(prefix + ("layers", "local_layers") + suffix, slot)] + # Trainer states predating the nested scan kept one module per slot. + return candidates + [ + (prefix + ("layers", f"layer_{slot}") + suffix, None), + (prefix + ("layers", str(slot)) + suffix, None), + (prefix + ("layers", slot) + suffix, None), ] # Homogeneous: a single scanned `layers` container. - return [prefix + ("layers",) + suffix] + return [(prefix + ("layers",) + suffix, None)] def _build_target_free_plan( self, @@ -803,19 +835,26 @@ def _build_target_free_plan( prefix = src_key[:idx] rest = src_key[idx + 1 :] + # (cycle slot, index along the nested cycle-slot axis or None) pairs this source feeds. if self.cycle == 1: - slot = 0 + slot_refs = ((0, None),) suffix = rest else: - # Inhomogeneous hybrid cycle: ("decoder", "layers", "layer_0", "input_layernorm", "scale") + # Inhomogeneous hybrid cycle: ("decoder", "layers", "layer_0", "input_layernorm", "scale"), + # or the nested `Qwen3NextScannableBlock` layout, where `local_layers` stacks every slot + # but the last on `scan_axis + 1` and `global_layer` is the last slot. slot_token = rest[0] match = re.fullmatch(r"layer_(\d+)", slot_token) if isinstance(slot_token, str) else None - if match: - slot = int(match.group(1)) + if slot_token == "local_layers": + slot_refs = tuple((slot, slot) for slot in range(self.cycle - 1)) + elif slot_token == "global_layer": + slot_refs = ((self.cycle - 1, None),) + elif match: + slot_refs = ((int(match.group(1)), None),) elif isinstance(slot_token, str) and slot_token.isdigit(): - slot = int(slot_token) + slot_refs = ((int(slot_token), None),) elif isinstance(slot_token, int): - slot = slot_token + slot_refs = ((slot_token, None),) else: raise ConversionPlanError(f"Unexpected slot token {slot_token!r} in key {src_key}") suffix = rest[1:] @@ -826,17 +865,19 @@ def _build_target_free_plan( if fuse_moe: consumed_wi_1.add(wi_1_key) - for b in range(self.num_blocks): - global_idx = b * self.cycle + slot - tgt_key = prefix + (f"layers_{global_idx}",) + suffix[:-1] + ("wi",) - plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), b, "fuse_moe")) + for slot, local_index in slot_refs: + for b in range(self.num_blocks): + global_idx = b * self.cycle + slot + tgt_key = prefix + (f"layers_{global_idx}",) + suffix[:-1] + ("wi",) + plan.append(_PlanEntry(tgt_key, (src_key, wi_1_key), b, "fuse_moe", local_index)) elif self.prefuse_moe_weights and suffix and suffix[-1] == "wi_1" and (src_key[:-1] + ("wi_0",) in src_flat): continue else: - for b in range(self.num_blocks): - global_idx = b * self.cycle + slot - tgt_key = prefix + (f"layers_{global_idx}",) + suffix - plan.append(_PlanEntry(tgt_key, (src_key,), b, "slice")) + for slot, local_index in slot_refs: + for b in range(self.num_blocks): + global_idx = b * self.cycle + slot + tgt_key = prefix + (f"layers_{global_idx}",) + suffix + plan.append(_PlanEntry(tgt_key, (src_key,), b, "slice", local_index)) return plan @@ -876,18 +917,19 @@ def _build_plan( candidates = self._scanned_candidates(tgt_key, prefix_len, suffix_start, slot) # 2. Scanned source, sliced at this block index. - matched = next((c for c in candidates if c in src_flat), None) + matched = next(((c, li) for c, li in candidates if c in src_flat), None) if matched is not None: - plan.append(_PlanEntry(tgt_key, (matched,), block, "slice")) - consumed.add(matched) + matched_key, local_index = matched + plan.append(_PlanEntry(tgt_key, (matched_key,), block, "slice", local_index)) + consumed.add(matched_key) continue # 3. Rollout pre-fuses MoE `wi`; trainer still has `wi_0`/`wi_1`. if tgt_key and tgt_key[-1] == "wi": - for cand in candidates: + for cand, local_index in candidates: wi_0, wi_1 = cand[:-1] + ("wi_0",), cand[:-1] + ("wi_1",) if wi_0 in src_flat and wi_1 in src_flat: - plan.append(_PlanEntry(tgt_key, (wi_0, wi_1), block, "fuse_moe")) + plan.append(_PlanEntry(tgt_key, (wi_0, wi_1), block, "fuse_moe", local_index)) consumed.update((wi_0, wi_1)) break else: @@ -1086,8 +1128,8 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat): raw_0 = src_flat[group.source_keys[0]] raw_1 = src_flat[group.source_keys[1]] wi_0, wi_1 = ( - _cast_target_free_floating_leaf(raw_0, target_dtype, path), - _cast_target_free_floating_leaf(raw_1, target_dtype, path), + self._take_cycle_slot(_cast_target_free_floating_leaf(raw_0, target_dtype, path), group), + self._take_cycle_slot(_cast_target_free_floating_leaf(raw_1, target_dtype, path), group), ) self._check_scan_axis(wi_0, path) per_block = self._fuse_moe_bulk_target_free(wi_0, wi_1, path) @@ -1095,7 +1137,7 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat): # group.op == "slice" raw_val = src_flat[group.source_keys[0]] - val = _cast_target_free_floating_leaf(raw_val, target_dtype, path) + val = self._take_cycle_slot(_cast_target_free_floating_leaf(raw_val, target_dtype, path), group) self._check_scan_axis(val, path) per_block = self._slice_bulk_target_free(val, path) return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets] @@ -1111,16 +1153,40 @@ def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat): return [(tgt_key, out) for _, tgt_key in group.targets] if group.op == "fuse_moe": - wi_0, wi_1 = (_apply_dtype_cast(src_flat[k], first_tgt.dtype, path) for k in group.source_keys) + wi_0, wi_1 = ( + self._take_cycle_slot(_apply_dtype_cast(src_flat[k], first_tgt.dtype, path), group) for k in group.source_keys + ) self._check_scan_axis(wi_0, path) per_block = self._fuse_moe_bulk(wi_0, wi_1, first_tgt, path) else: # "slice" val = _apply_dtype_cast(src_flat[group.source_keys[0]], first_tgt.dtype, path) + val = self._take_cycle_slot(val, group) self._check_scan_axis(val, path) per_block = _bulk_align_and_unstack(val, self.scan_axis, first_tgt, path) return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets] + def _take_cycle_slot(self, val, group: _PlanGroup): + """Drops the nested cycle-slot axis from a `local_layers` source. + + Everything downstream expects blocks on `scan_axis` and no other stacked + axis, so the slot is selected here -- once per group, before the bulk + alignment and unstack, rather than once per target layer. + """ + if group.local_index is None: + return val + slot_axis = self.scan_axis + 1 + if val.ndim <= slot_axis or val.shape[slot_axis] <= group.local_index: + raise ConversionPlanError( + f"Expected a nested cycle-slot axis holding at least " + f"{group.local_index + 1} slots on axis {slot_axis} of " + f"{group.source_path}, found shape {val.shape}. Check " + "param_scan_axis and inhomogeneous_layer_cycle_interval." + ) + if isinstance(val, jax.ShapeDtypeStruct): + return jax.ShapeDtypeStruct(val.shape[:slot_axis] + val.shape[slot_axis + 1 :], val.dtype) + return jnp.take(val, group.local_index, axis=slot_axis) + def _check_scan_axis(self, val, path: str) -> None: if val.shape[self.scan_axis] != self.num_blocks: raise ConversionPlanError( @@ -1175,7 +1241,8 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: ) result: Dict[Tuple[Any, ...], Any] = {} - for group in self._groups: + last_reader = _last_group_reading(self._groups) + for group_idx, group in enumerate(self._groups): if self.debug: for k in group.source_keys: logging.info( @@ -1215,7 +1282,8 @@ def convert(self, src_pytree: Any, target_state: Any = None) -> Dict[str, Any]: outs = self._execute_group(group, src_flat, tgt_flat) for k in group.source_keys: - src_flat.pop(k, None) + if last_reader[k] == group_idx: + src_flat.pop(k, None) for tgt_key, out in outs: tgt_val = tgt_flat[tgt_key] @@ -1268,13 +1336,15 @@ def convert_streaming( self._groups = _group_plan(self._plan) groups_per_piece = max(1, groups_per_piece) + last_reader = _last_group_reading(self._groups) for i in range(0, len(self._groups), groups_per_piece): piece_groups = self._groups[i : i + groups_per_piece] piece_result: Dict[Tuple[Any, ...], Any] = {} - for group in piece_groups: + for group_idx, group in enumerate(piece_groups, start=i): outs = self._execute_group_target_free(group, src_flat) for k in group.source_keys: - src_flat.pop(k, None) + if last_reader[k] == group_idx: + src_flat.pop(k, None) for tgt_key, out in outs: piece_result[src_root + tgt_key] = out del outs diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 0a73559b3c..eb5b2ba15c 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -323,6 +323,43 @@ def __call__( return layer_output, kv_cache +def forced_routed_experts_per_layer(forced_routed_experts: jax.Array, num_layers: int) -> jax.Array: + """Normalizes a batch's forced_routed_experts to a per-decoder-layer stack. + + Every supported architecture is homogeneous-MoE, so the layer axis is simply + the decoder layer index. + + Args: + forced_routed_experts: `[batch, seq, num_layers, top_k]` (per-layer) or + `[batch, seq, top_k]` (broadcast to every layer). + num_layers: Total decoder layers. + + Returns: + Array shaped `[num_layers, batch, seq, top_k]`. + + Raises: + ValueError: If forced_routed_experts is not 3D or 4D, or its layer axis + does not match `num_layers`. + """ + fre = forced_routed_experts + if fre.ndim not in (3, 4): + raise ValueError( + "forced_routed_experts must be [batch, seq, top_k] (3D, broadcast to" + " every layer) or [batch, seq, num_layers, top_k] (4D, per-layer); got" + f" ndim={fre.ndim} with shape {fre.shape}." + ) + if fre.ndim == 4: + if fre.shape[2] != num_layers: + raise ValueError( + "forced_routed_experts layer axis must equal the number of decoder" + f" layers ({num_layers}); got {fre.shape[2]} (shape {fre.shape})." + ) + # [batch, seq, num_layers, top_k] -> [num_layers, batch, seq, top_k] + return jnp.moveaxis(fre, 2, 0) + # [batch, seq, top_k]: same forced routing for every layer. + return jnp.broadcast_to(fre[None], (num_layers,) + fre.shape) + + def reshape_forced_routed_experts_for_scan( forced_routed_experts: jax.Array, num_layers: int, @@ -332,8 +369,7 @@ def reshape_forced_routed_experts_for_scan( """Reshapes a batch's forced_routed_experts into jax.lax.scan's xs layout. The outer scan runs `scan_length` iterations, each covering a cycle of - `layers_per_cycle` decoder layers. Every supported architecture is - homogeneous-MoE, so the layer axis is simply the decoder layer index. + `layers_per_cycle` decoder layers. Args: forced_routed_experts: `[batch, seq, num_layers, top_k]` (per-layer) or @@ -348,19 +384,7 @@ def reshape_forced_routed_experts_for_scan( Raises: ValueError: If forced_routed_experts is not 3D or 4D. """ - fre = forced_routed_experts - if fre.ndim not in (3, 4): - raise ValueError( - "forced_routed_experts must be [batch, seq, top_k] (3D, broadcast to" - " every layer) or [batch, seq, num_layers, top_k] (4D, per-layer); got" - f" ndim={fre.ndim} with shape {fre.shape}." - ) - if fre.ndim == 4: - # [batch, seq, num_layers, top_k] -> [num_layers, batch, seq, top_k] - fre = jnp.moveaxis(fre, 2, 0) - else: - # [batch, seq, top_k]: same forced routing for every layer. - fre = jnp.broadcast_to(fre[None], (num_layers,) + fre.shape) + fre = forced_routed_experts_per_layer(forced_routed_experts, num_layers) return jnp.reshape(fre, (scan_length, layers_per_cycle) + fre.shape[1:]) @@ -2664,10 +2688,34 @@ def _apply_qwen3_hybrid_scanned_blocks(self, y, layer_args, layer_kwargs, kv_cac 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. + + Forced routing (router replay) is sliced per decoder layer here and handed + to each block as its own `[block_length, batch, seq, top_k]` stack, which + the block then splits between its local scan and its global layer. """ cfg = self.config block_length = cfg.inhomogeneous_layer_cycle_interval scan_length = cfg.num_decoder_layers // block_length + num_scanned_layers = scan_length * block_length + + layer_kwargs = dict(layer_kwargs) + forced_routed_experts = layer_kwargs.pop("forced_routed_experts", None) + forced_routed_experts_scanned = None + remainder_forced_routed_experts = None + if forced_routed_experts is not None: + if kv_caches is not None: + # The kv-cache path below unrolls into _apply_layers_sequentially, which + # threads either kv caches or scan xs, not both. + raise NotImplementedError( + "Forced routing is not supported together with externally-managed (vLLM) kv_caches in scanned layers." + ) + per_layer = forced_routed_experts_per_layer(forced_routed_experts, cfg.num_decoder_layers) + if scan_length > 0: + forced_routed_experts_scanned = jnp.reshape( + per_layer[:num_scanned_layers], (scan_length, block_length) + per_layer.shape[1:] + ) + remainder_forced_routed_experts = per_layer[num_scanned_layers:] + 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( @@ -2677,19 +2725,22 @@ def _apply_qwen3_hybrid_scanned_blocks(self, y, layer_args, layer_kwargs, kv_cac length=scan_length, kv_caches_stacked=grouped_kv_caches, skip_block_remat=True, + forced_routed_experts_scanned=forced_routed_experts_scanned, **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: + if remainder_forced_routed_experts is not None: + layer_kwargs["forced_routed_experts"] = remainder_forced_routed_experts y = self._apply_remainder_block( self.layers_remainder, y, layer_args, layer_kwargs, kv_caches=kv_caches, - start_idx=scan_length * block_length, + start_idx=num_scanned_layers, num_remaining_layers=num_remaining_layers, ) return y diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index 19d49502bd..45e221b921 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -98,7 +98,8 @@ def apply_scanned_layers( *, length: int, param_scan_axis: int, - apply_fn: Callable[[nnx.Module, Any], Any], + apply_fn: Callable[..., Any], + xs: Any | None = None, remat: bool = False, remat_policy: Callable[..., Any] | None = None, prevent_cse: bool = True, @@ -111,6 +112,11 @@ def apply_scanned_layers( ``remat`` is separate from ``remat_policy`` because ``None`` is JAX's full rematerialization policy, not an indication that rematerialization is off. + ``xs`` is an optional pytree of per-layer inputs whose leaves have a leading + axis of ``length``; each iteration's slice is passed to ``apply_fn`` as a + third argument. Callers that pass no ``xs`` keep the two-argument ``apply_fn`` + signature. + Externally managed per-layer state, such as KV caches, is not supported by this scan path. @@ -172,14 +178,19 @@ def _strip_scan_metadata(leaf): return leaf def scan_body(current_carry, scanned_state): - current_params, current_rest = scanned_state + if xs is None: + current_params, current_rest = scanned_state + apply_args = () + else: + current_params, current_rest, current_xs = scanned_state + apply_args = (current_xs,) 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_rest) - next_carry = apply_fn(current_layer, current_carry) + next_carry = apply_fn(current_layer, current_carry, *apply_args) # 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. @@ -190,9 +201,8 @@ def scan_body(current_carry, scanned_state): 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_new_params, scanned_rest) = jax.lax.scan( - scan_fn, carry, (params, rest), length=length, unroll=unroll - ) + scan_xs = (params, rest) if xs is None else (params, rest, xs) + final_carry, (scanned_new_params, scanned_rest) = jax.lax.scan(scan_fn, carry, scan_xs, 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) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index d3916bb501..fb5e06a6b6 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1606,9 +1606,9 @@ def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): rngs=rngs, ) - def _run_layer(self, layer, y, layer_kwargs, kv_cache=None): + def _run_layer(self, layer, y, layer_kwargs, kv_cache=None, forced_routed_experts=None): """Invokes one Qwen3NextDecoderLayer, returning (output, updated_kv_cache).""" - out = layer(y, **layer_kwargs, kv_cache=kv_cache) + out = layer(y, **layer_kwargs, kv_cache=kv_cache, forced_routed_experts=forced_routed_experts) return out if isinstance(out, tuple) else (out, None) @property @@ -1616,21 +1616,30 @@ 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): + def _scan_local_layers(self, y, layer_kwargs, forced_routed_experts=None): """Runs the local (linear attention / GatedDeltaNet) layers via a per-layer rematerialized jax.lax.scan.""" remat = self._remat_enabled + if forced_routed_experts is None: + apply_fn = lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0] + else: + # Router replay: one slice of forced routing per local layer, fed through + # the scan's xs so each iteration sees its own layer's expert indices. + def apply_fn(layer, carry, routing): + return self._run_layer(layer, carry, layer_kwargs, forced_routed_experts=routing)[0] + 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], + apply_fn=apply_fn, + xs=forced_routed_experts, 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): + def _scan_global_layer(self, y, layer_kwargs, forced_routed_experts=None): """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, ...) @@ -1639,7 +1648,9 @@ def _scan_global_layer(self, y, layer_kwargs): 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_hidden_states = self._run_layer( + layer, hidden_states, layer_kwargs, forced_routed_experts=forced_routed_experts + )[0] _, new_intermediate, new_other = nnx.split(layer, nnx.Intermediate, ...) return (new_hidden_states, new_other), new_intermediate @@ -1669,7 +1680,7 @@ def run_global_layer(carry, intermediate_slice): nnx.update(self.global_layer, final_other, intermediate_state) return y - def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): + def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs, forced_routed_experts=None): """Runs the block with externally-supplied per-layer kv caches. Inference KV caches are a Python list of per-layer entries, so this path @@ -1687,7 +1698,8 @@ def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): 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) + current_routing = None if forced_routed_experts is None else forced_routed_experts[i] + y, new_kv = self._run_layer(layer, y, layer_kwargs, current_kv, forced_routed_experts=current_routing) updated_kvs.append(new_kv) # 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 @@ -1699,7 +1711,8 @@ def _forward_with_external_kv_cache(self, y, kv_cache, layer_kwargs): 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) + global_routing = None if forced_routed_experts is None else forced_routed_experts[self.num_local] + y, new_kv = self._run_layer(self.global_layer, y, layer_kwargs, global_kv, forced_routed_experts=global_routing) updated_kvs.append(new_kv) return y, tuple(updated_kvs) @@ -1715,11 +1728,15 @@ def __call__( slot: None | int = None, kv_cache=None, attention_metadata=None, + forced_routed_experts=None, ) -> tuple[Array, None]: """Applies the block of decoder layers to the input carry. Args: carry: The input tensor from the previous scan iteration. + forced_routed_experts: Optional router replay indices for this block, + shaped `[num_of_layers, batch, seq, top_k]` -- one slice per sub-layer, + in block order (the local layers first, the full-attention layer last). # ... other arguments are broadcasted to each iteration. Returns: @@ -1740,14 +1757,22 @@ def __call__( "attention_metadata": attention_metadata, } + if forced_routed_experts is not None and forced_routed_experts.shape[0] != self.num_of_layers: + raise ValueError( + f"forced_routed_experts must carry one slice per sub-layer of the block: expected leading axis " + f"{self.num_of_layers}, got {forced_routed_experts.shape[0]} (shape {forced_routed_experts.shape})." + ) + if kv_cache is not None: - return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs) + return self._forward_with_external_kv_cache(inputs, kv_cache, layer_kwargs, forced_routed_experts) y = inputs if self.local_layers is not None: - y = self._scan_local_layers(y, layer_kwargs) + local_routing = None if forced_routed_experts is None else forced_routed_experts[: self.num_local] + y = self._scan_local_layers(y, layer_kwargs, local_routing) if self.global_layer is not None: - y = self._scan_global_layer(y, layer_kwargs) + global_routing = None if forced_routed_experts is None else forced_routed_experts[self.num_local] + y = self._scan_global_layer(y, layer_kwargs, global_routing) if cfg.scan_layers: return y, None @@ -1819,7 +1844,7 @@ def __init__( num_features=cfg.emb_dim, epsilon=cfg.normalization_layer_epsilon, dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, + weight_dtype=self._norm_weight_dtype(), shard_mode=cfg.shard_mode, rngs=rngs, ) @@ -1833,39 +1858,69 @@ def __init__( # Conditionally instantiate either the Linear Attention or Full Attention block. if is_full_attention_layer: - self.attention = Qwen3NextFullAttention( - config=cfg, - mesh=self.mesh, - quant=self.quant, - model_mode=model_mode, - layer_idx=self.layer_idx, - rngs=rngs, - ) + self.attention = self._make_full_attention(rngs=rngs) else: - batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config, model_mode) - dummy_inputs_shape = (batch_size, seq_len, config.emb_dim) - self.attention = Qwen3NextGatedDeltaNet( - config=cfg, - inputs_shape=dummy_inputs_shape, - mesh=self.mesh, - dtype=cfg.dtype, - model_mode=model_mode, - quant=self.quant, - rngs=rngs, - ) + self.attention = self._make_linear_attention(rngs=rngs) # Second LayerNorm, applied before the MoE block. self.post_attention_layernorm = Qwen3NextRMSNorm( num_features=cfg.emb_dim, epsilon=cfg.normalization_layer_epsilon, dtype=cfg.dtype, - weight_dtype=cfg.weight_dtype, + weight_dtype=self._norm_weight_dtype(), shard_mode=cfg.shard_mode, rngs=rngs, ) - # Instantiate our `Qwen3NextSparseMoeBlock`. - self.mlp = Qwen3NextSparseMoeBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + self.mlp = self._make_mlp(rngs=rngs) + + # The hooks below are the only places a sibling architecture that reuses this layer + # (Qwen3.5) has to diverge, so they are overridable rather than inline. `__init__` + # binds what the factories return to `attention` and `mlp`, which are checkpoint + # parameter paths, so an override has to keep the sub-module's own parameter names too. + def _norm_weight_dtype(self): + """Storage dtype of the two RMSNorm scales.""" + return self.config.weight_dtype + + def _make_full_attention(self, *, rngs: nnx.Rngs): + """Builds the full-attention sub-block used on the last layer of every cycle.""" + return Qwen3NextFullAttention( + config=self.config, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + layer_idx=self.layer_idx, + rngs=rngs, + ) + + def _make_linear_attention(self, *, rngs: nnx.Rngs): + """Builds the linear-attention (GatedDeltaNet) sub-block used on all other layers.""" + cfg = self.config + batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(cfg, self.model_mode) + dummy_inputs_shape = (batch_size, seq_len, cfg.emb_dim) + return Qwen3NextGatedDeltaNet( + config=cfg, + inputs_shape=dummy_inputs_shape, + mesh=self.mesh, + dtype=cfg.dtype, + model_mode=self.model_mode, + quant=self.quant, + rngs=rngs, + ) + + def _make_mlp(self, *, rngs: nnx.Rngs): + """Builds the sparse MoE block.""" + return Qwen3NextSparseMoeBlock(config=self.config, mesh=self.mesh, quant=self.quant, rngs=rngs) + + def _apply_mlp(self, hidden_states, deterministic, forced_routed_experts): + """Runs the block built by `_make_mlp`, returning `(output, load_balance_loss)`.""" + return self.mlp( + hidden_states, + deterministic=deterministic, + forced_routed_experts=forced_routed_experts, + intermediate_sharding=self.mlp_intermediate_sharding, + out_sharding=self.out_sharding, + ) def __call__( self, @@ -1878,6 +1933,7 @@ def __call__( slot: None | int = None, kv_cache: None | dict[str, Array] = None, attention_metadata: None | dict[str, Any] = None, + forced_routed_experts: jnp.ndarray | None = None, ): # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) if isinstance(inputs, tuple): @@ -1927,15 +1983,12 @@ def __call__( hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) # Instantiate and call our `Qwen3NextSparseMoeBlock`. - mlp_output, load_balance_loss = self.mlp( - hidden_states, - deterministic=deterministic, - intermediate_sharding=self.mlp_intermediate_sharding, - out_sharding=self.out_sharding, - ) + mlp_output, load_balance_loss = self._apply_mlp(hidden_states, deterministic, forced_routed_experts) # We sow the load balancing loss so it can be collected and added to the total loss # during training. + # Assigned rather than sown: `sow` appends to a tuple, so the layer's Intermediate + # structure would grow on every call, which a `jax.lax.scan` body cannot express. if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.moe_lb_loss = nnx.Intermediate(load_balance_loss) diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 48ebe66d70..6ba8ce0d37 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -16,24 +16,17 @@ # pylint: disable=arguments-differ # pylint: disable=no-name-in-module -import functools -from typing import Any, cast - from flax import nnx -import jax.numpy as jnp -from jax.sharding import Mesh -from maxtext.common.common_types import Array, Config, ShardMode, get_weight_dtype +from maxtext.common.common_types import get_weight_dtype from maxtext.layers.linears import MlpBlock -from maxtext.layers.normalizations import Qwen3NextRMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant from maxtext.models.qwen3 import ( + Qwen3NextDecoderLayer, Qwen3NextFullAttention, Qwen3NextGatedDeltaNet, Qwen3NextScannableBlock, Qwen3NextSparseMoeBlock, ) from maxtext.utils import max_utils -from maxtext.utils.sharding import create_sharding, get_logical_axis_rules, maybe_shard_with_logical # ----------------------------------------- # Qwen3.5 Layer Implementations @@ -74,11 +67,16 @@ def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): ) -class Qwen3_5DecoderLayer(nnx.Module): - """ - This layer is a hybrid, capable of functioning as either: - 1. A standard attention + MoE layer. - 2. A linear attention + MoE layer. +class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): + """Qwen3.5 hybrid decoder layer. + + Qwen3.5's decoder layer is structurally identical to Qwen3-Next's -- norm, + either full or linear attention, norm, MoE (or, for the dense models, an MLP), + with residuals around the attention and MLP halves -- so the whole forward pass + is inherited. Only the sub-block hooks are overridden, to build the Qwen3.5 + sub-classes. The attribute names (`input_layernorm`, `attention`, + `post_attention_layernorm`, `mlp`) are therefore shared, and so are the + checkpoint parameter paths. Attributes: config: The model configuration object. @@ -88,204 +86,59 @@ class Qwen3_5DecoderLayer(nnx.Module): quant: Optional quantization configuration. """ - def __init__( - 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 - self.model_mode = model_mode - self.layer_idx = layer_idx - self.quant = quant - cfg = self.config - self.activation_axis_names = ("activation_batch", "activation_norm_length", "activation_embed") - self.mlp_activation_axis_names = ( - "activation_batch", - "activation_norm_length", - "activation_mlp", - ) + def _norm_weight_dtype(self): + return get_weight_dtype(self.config, "norm") - # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In - # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout. - if cfg.shard_mode == ShardMode.EXPLICIT: - self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules()) - self.mlp_intermediate_sharding = create_sharding( - mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules() - ) - self._maybe_shard_with_logical = functools.partial( - maybe_shard_with_logical, - mesh=mesh, - shard_mode=cfg.shard_mode, - debug_sharding=cfg.debug_sharding, - ) - else: - self.out_sharding = None - self.mlp_intermediate_sharding = None - self._maybe_shard_with_logical = lambda inputs, *args, **kwargs: inputs - - # First LayerNorm, applied before the attention block. - self.input_layernorm = Qwen3NextRMSNorm( - num_features=cfg.emb_dim, - epsilon=cfg.normalization_layer_epsilon, - dtype=cfg.dtype, - weight_dtype=get_weight_dtype(cfg, "norm"), - shard_mode=cfg.shard_mode, + def _make_full_attention(self, *, rngs: nnx.Rngs): + return Qwen3_5FullAttention( + config=self.config, + mesh=self.mesh, + quant=self.quant, + model_mode=self.model_mode, + layer_idx=self.layer_idx, rngs=rngs, ) - # 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: - 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. - if is_full_attention_layer: - self.attention = Qwen3_5FullAttention( - config=cfg, - mesh=self.mesh, - quant=self.quant, - model_mode=model_mode, - layer_idx=self.layer_idx, - rngs=rngs, - ) - else: - batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(config, model_mode) - dummy_inputs_shape = (batch_size, seq_len, config.emb_dim) - self.attention = Qwen3_5GatedDeltaNet( - config=cfg, - inputs_shape=dummy_inputs_shape, - mesh=self.mesh, - dtype=cfg.dtype, - model_mode=model_mode, - quant=self.quant, - rngs=rngs, - ) - - # Second LayerNorm, applied before the MoE / MLP block. - self.post_attention_layernorm = Qwen3NextRMSNorm( - num_features=cfg.emb_dim, - epsilon=cfg.normalization_layer_epsilon, + def _make_linear_attention(self, *, rngs: nnx.Rngs): + cfg = self.config + batch_size, seq_len = max_utils.get_batch_seq_len_for_mode(cfg, self.model_mode) + dummy_inputs_shape = (batch_size, seq_len, cfg.emb_dim) + return Qwen3_5GatedDeltaNet( + config=cfg, + inputs_shape=dummy_inputs_shape, + mesh=self.mesh, dtype=cfg.dtype, - weight_dtype=get_weight_dtype(cfg, "norm"), - shard_mode=cfg.shard_mode, + model_mode=self.model_mode, + quant=self.quant, rngs=rngs, ) + def _make_mlp(self, *, rngs: nnx.Rngs): + cfg = self.config if getattr(cfg, "num_experts", 1) > 1: - self.mlp = Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) - else: - self.mlp = MlpBlock( - config=cfg, - mesh=self.mesh, - in_features=cfg.emb_dim, - intermediate_dim=cfg.mlp_dim, - activations=tuple(cfg.mlp_activations), - intermediate_dropout_rate=cfg.dropout_rate, - dtype=cfg.dtype, - weight_dtype=get_weight_dtype(cfg, "mlp"), - quant=self.quant, - model_mode=model_mode, - rngs=rngs, - ) - - def __call__( - self, - inputs: jnp.ndarray, - decoder_segment_ids: None | jnp.ndarray, - decoder_positions: None | jnp.ndarray, - deterministic: bool, - model_mode: str, - previous_chunk=None, - slot: None | int = None, - kv_cache: None | dict[str, Array] = None, - attention_metadata: None | dict[str, Any] = None, - forced_routed_experts: jnp.ndarray | None = None, - ): - # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache)) - if isinstance(inputs, tuple): - inputs = inputs[0] - inputs = self._maybe_shard_with_logical(inputs, self.activation_axis_names) - residual = inputs - - if isinstance(attention_metadata, dict): - attention_metadata = attention_metadata.get(f"layer.{self.layer_idx}", attention_metadata.get(self.layer_idx)) - - # First LayerNorm, applied before the attention block. - hidden_states = self.input_layernorm(inputs, out_sharding=self.out_sharding) - hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) - - # Conditionally apply either the Linear Attention or Full Attention block. - if isinstance(self.attention, Qwen3_5FullAttention): - attention_output, new_kv_cache = cast(Qwen3_5FullAttention, self.attention)( - hidden_states, - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - kv_cache=kv_cache, - attention_metadata=attention_metadata, - out_sharding=self.out_sharding, - ) - else: - attention_output, new_kv_cache = cast(Qwen3_5GatedDeltaNet, self.attention)( - hidden_states, - model_mode=model_mode, - kv_cache=kv_cache, - decoder_segment_ids=decoder_segment_ids, - attention_metadata=attention_metadata, - out_sharding=self.out_sharding, - ) - - # First residual connection after attention - attention_output = self._maybe_shard_with_logical(attention_output, self.activation_axis_names) - hidden_states = residual + attention_output - hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) - - # Prepare for the MoE / MLP block by capturing the new residual - residual = hidden_states - - # Second LayerNorm, applied before the MoE / MLP block. - hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding) - hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names) + return Qwen3_5SparseMoEBlock(config=cfg, mesh=self.mesh, quant=self.quant, rngs=rngs) + return MlpBlock( + config=cfg, + mesh=self.mesh, + in_features=cfg.emb_dim, + intermediate_dim=cfg.mlp_dim, + activations=tuple(cfg.mlp_activations), + intermediate_dropout_rate=cfg.dropout_rate, + dtype=cfg.dtype, + weight_dtype=get_weight_dtype(cfg, "mlp"), + quant=self.quant, + model_mode=self.model_mode, + rngs=rngs, + ) - # Instantiate and call our `Qwen3_5SparseMoEBlock` or `MlpBlock`. + def _apply_mlp(self, hidden_states, deterministic, forced_routed_experts): if getattr(self.config, "num_experts", 1) > 1: - mlp_output, load_balance_loss = self.mlp( - hidden_states, - deterministic=deterministic, - forced_routed_experts=forced_routed_experts, - intermediate_sharding=self.mlp_intermediate_sharding, - out_sharding=self.out_sharding, - ) - - # We sow the load balancing loss so it can be collected and added to the total loss - # during training. - # Assigned rather than sown: `sow` appends to a tuple, so the layer's Intermediate - # structure would grow on every call, which a `jax.lax.scan` body cannot express. - if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: - self.moe_lb_loss = nnx.Intermediate(load_balance_loss) - else: - mlp_output = self.mlp( - hidden_states, - deterministic=deterministic, - intermediate_sharding=self.mlp_intermediate_sharding, - out_sharding=self.out_sharding, - ) - - # Final residual connection (after the MoE block) - mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names) - layer_output = residual + mlp_output - layer_output = self._maybe_shard_with_logical( - layer_output, - self.activation_axis_names, + return super()._apply_mlp(hidden_states, deterministic, forced_routed_experts) + # A dense MLP routes nothing, so there is no routing to replay and no load-balance loss. + mlp_output = self.mlp( + hidden_states, + deterministic=deterministic, + intermediate_sharding=self.mlp_intermediate_sharding, + out_sharding=self.out_sharding, ) - return layer_output, new_kv_cache + return mlp_output, None diff --git a/tests/post_training/unit/vllm_rollout_unroll_test.py b/tests/post_training/unit/vllm_rollout_unroll_test.py index c08801b99c..a3169abba7 100644 --- a/tests/post_training/unit/vllm_rollout_unroll_test.py +++ b/tests/post_training/unit/vllm_rollout_unroll_test.py @@ -228,6 +228,65 @@ def test_interleaves_slots_and_repetitions(self): } self.assertEqual(validate_direct_sync_layer_coverage(unrolled, target), 4) + @pytest.mark.cpu_only + def test_unrolls_the_nested_local_and_global_block_scan(self): + """Qwen3-Next and Qwen3.5 nest the cycle's local layers into one module. + + `local_layers` stacks blocks on `scan_axis` and the cycle slot on + `scan_axis + 1`; `global_layer` holds the trailing full-attention layer on + the block axis alone. Together they must still unroll to a dense + `layers_{global_idx}` run, with each slot landing on its own layer. + """ + num_blocks, num_local = 2, 3 + cycle = num_local + 1 + local_probe = np.zeros((2, num_blocks, num_local, 1), dtype=np.float32) + global_probe = np.zeros((2, num_blocks, 1), dtype=np.float32) + for block in range(num_blocks): + for slot in range(num_local): + local_probe[:, block, slot, :] = block * cycle + slot + global_probe[:, block, :] = block * cycle + num_local + weights = MockWeights( + { + "base": { + "decoder": { + "layers": { + "local_layers": {"probe": local_probe}, + "global_layer": {"probe": global_probe}, + }, + "decoder_norm": {"scale": np.ones(2)}, + } + } + } + ) + + unrolled = unroll_qwen_scanned_weights(weights) + + decoder = unrolled["base"]["decoder"] + for layer_idx in range(num_blocks * cycle): + self.assertIn(f"layers_{layer_idx}", decoder) + np.testing.assert_array_equal( + decoder[f"layers_{layer_idx}"]["probe"], + np.full((2, 1), layer_idx, dtype=np.float32), + ) + np.testing.assert_array_equal(decoder["decoder_norm"]["scale"], np.ones(2)) + target = { + "model": { + "decoder": { + f"layers_{layer_idx}": {"probe": np.zeros((2, 1), dtype=np.float32)} + for layer_idx in range(num_blocks * cycle) + } + } + } + self.assertEqual(validate_direct_sync_layer_coverage(unrolled, target), num_blocks * cycle) + + @pytest.mark.cpu_only + def test_rejects_nested_local_layers_without_a_cycle_slot_axis(self): + """A `local_layers` tensor missing its slot axis must not be read as flat.""" + weights = MockWeights({"decoder": {"layers": {"local_layers": {"probe": np.ones((2, 2))}}}}) + + with self.assertRaisesRegex(ValueError, "has no scan axis"): + unroll_qwen_scanned_weights(weights) + @pytest.mark.cpu_only def test_correctly_unrolls_homogeneous_scanned_weights(self): """Verify homogeneous scanned layers (LLaMA, Qwen Base) unroll to layers_{global_idx}.""" diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py index 929b8a63c3..d243ba6c45 100644 --- a/tests/post_training/unit/weight_converter_test.py +++ b/tests/post_training/unit/weight_converter_test.py @@ -126,6 +126,24 @@ def _source_tree(fused_moe_on_target: bool): } +def _nested_source_tree(fused_moe_on_target: bool): + """Trainer state in the nested-scan layout `Qwen3NextScannableBlock` produces. + + Same values as `_source_tree`, restacked the way the shared Qwen3-Next block + stores them: slots `0..C-2` share one `local_layers` module with the cycle + slot on `SCAN_AXIS + 1`, and slot `C-1` becomes `global_layer` with the block + axis alone. Converting either tree must yield the same rollout weights. + """ + legacy = _source_tree(fused_moe_on_target)["base"]["decoder"]["layers"] + local_slots = [legacy[f"layer_{slot}"] for slot in range(CYCLE - 1)] + tree = _source_tree(fused_moe_on_target) + tree["base"]["decoder"]["layers"] = { + "local_layers": jax.tree.map(lambda *slots: jnp.stack(slots, axis=SCAN_AXIS + 1), *local_slots), + "global_layer": legacy[f"layer_{CYCLE - 1}"], + } + return tree + + def _target_tree(moe_intermediate=MOE_DIM, fused=True, wo_intermediate=MOE_DIM): """Rollout state: unrolled `layers_{0..7}`, optionally fused MoE. @@ -197,6 +215,34 @@ def test_non_layer_params_pass_through(self): np.asarray(_arr(16, EMB)), ) + def test_nested_scan_layout_converts_identically(self): + """Qwen3.5 moved onto Qwen3-Next's nested block scan. + + The cycle's linear-attention layers now live in one `local_layers` module + with the slot on `SCAN_AXIS + 1` instead of one module per slot, so the + converter has to read that layout too -- and read it to the same answer. + """ + target = _target_tree() + legacy = traverse_util.flatten_dict(MaxTextToMaxTextConverter(_config()).convert(_source_tree(True), target)) + nested = traverse_util.flatten_dict(MaxTextToMaxTextConverter(_config()).convert(_nested_source_tree(True), target)) + self.assertEqual(set(legacy), set(nested)) + for key, want in legacy.items(): + np.testing.assert_array_equal( + np.asarray(nested[key]), + np.asarray(want), + err_msg=f"{'.'.join(map(str, key))} differs between the per-slot and nested layouts", + ) + + def test_nested_local_layers_shorter_than_the_cycle_is_rejected(self): + """A `local_layers` axis that cannot cover every slot must not slice silently.""" + source = _nested_source_tree(True) + local = source["base"]["decoder"]["layers"]["local_layers"] + local["input_layernorm"]["scale"] = jnp.take( + local["input_layernorm"]["scale"], jnp.arange(CYCLE - 2), axis=SCAN_AXIS + 1 + ) + with self.assertRaisesRegex(ConversionPlanError, "nested cycle-slot axis"): + MaxTextToMaxTextConverter(_config()).convert(source, _target_tree()) + def test_plan_is_built_once_and_reused(self): converter = MaxTextToMaxTextConverter(_config()) converter.convert(_source_tree(True), _target_tree()) @@ -936,6 +982,44 @@ def test_case_12_target_free_dtype_matches_rollout_per_module_rule(self): leaf = getattr(flat_out[path], "value", flat_out[path]) self.assertEqual(leaf.dtype, jnp.dtype(want)) + def test_case_13_nested_scan_layout_target_free(self): + """Target-free conversion reads the nested `Qwen3NextScannableBlock` layout to the per-slot answer. + + One `local_layers` source feeds a plan group per cycle slot, so it must stay + available until the last of those groups has run. + """ + cfg = _config(padded_base_moe_mlp_dim=MOE_DIM + 2, prefuse_moe_weights=True) + + def convert(source, streaming): + converter = MaxTextToMaxTextConverter(cfg, prefuse_moe_weights=True) + if streaming: + flat = {} + for piece in converter.convert_streaming(source, target_state=None, groups_per_piece=1): + flat.update(traverse_util.flatten_dict(piece)) + return flat + return traverse_util.flatten_dict(converter.convert(source, target_state=None)) + + for streaming in (False, True): + with self.subTest(streaming=streaming): + legacy = convert(_source_tree(True), streaming) + nested = convert(_nested_source_tree(True), streaming) + self.assertEqual(set(legacy), set(nested)) + for key, want in legacy.items(): + np.testing.assert_array_equal( + np.asarray(getattr(nested[key], "value", nested[key])), + np.asarray(getattr(want, "value", want)), + err_msg=f"{'.'.join(map(str, key))} differs between the per-slot and nested layouts", + ) + + def to_struct(x): + return jax.ShapeDtypeStruct(x.shape, x.dtype) + + abstract = { + key: getattr(value, "value", value).shape + for key, value in convert(jax.tree_util.tree_map(to_struct, _nested_source_tree(True)), False).items() + } + self.assertEqual(abstract, {key: getattr(value, "value", value).shape for key, value in legacy.items()}) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index abe3f054fa..fec69287b2 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -37,6 +37,7 @@ import numpy as np from flax import linen as nn from flax import nnx +from flax import traverse_util from jax.sharding import Mesh, NamedSharding, PartitionSpec from maxtext.common.common_types import ( @@ -50,6 +51,7 @@ ShardMode, ) from maxtext.configs import pyconfig +from maxtext.integration.tunix.weight_mapping import raiden_unscan from maxtext.layers import fsdp_prefetch from maxtext.layers import linears from maxtext.layers import nnx_decoders @@ -1078,6 +1080,23 @@ def test_rejects_block_whose_global_layer_is_not_last(self): with self.assertRaisesRegex(ValueError, "full-attention layer last"): self._build(layer_idx_offset=1) + def test_rejects_forced_routing_that_does_not_cover_every_sub_layer(self): + """Router replay is sliced per sub-layer, so a mismatched leading axis must not be broadcast.""" + cfg, block = self.cfg, self.block + inputs, segment_ids, positions = self._inputs(cfg) + forced_routed_experts = jnp.zeros( + (block.num_of_layers + 1, 1, cfg.max_target_length, cfg.num_experts_per_tok), dtype=jnp.int32 + ) + with self.assertRaisesRegex(ValueError, "one slice per sub-layer"): + block( + inputs, + segment_ids, + positions, + True, + MODEL_MODE_TRAIN, + forced_routed_experts=forced_routed_experts, + ) + class TestNNXDecoderQwen3Next(unittest.TestCase): """Tests the NNXDecoder-level wiring of the Qwen3-Next scanned blocks.""" @@ -1201,6 +1220,56 @@ class TestNNXDecoderQwen3_5(TestNNXDecoderQwen3Next): MODEL_CONFIG = _QWEN3_5_CONFIG + def test_raiden_unscan_reproduces_the_unscanned_decoder(self): + """Raiden's unscan of the nested block state must load into an unscanned decoder unchanged. + + Raiden binds trainer and sampler tensors by name, and the sampler runs + ``scan_layers=False``, so every name and shape has to match and the loaded + decoder has to give the scanned decoder's logits. A cycle of 3 over 6 layers + puts two slots in ``local_layers`` and two repeats on the scan axis, so a + swapped slot or repeat index changes the output. + """ + cfg, decoder, shared_embedding = self._build(6, inhomogeneous_layer_cycle_interval=3) + unscanned_cfg, unscanned_decoder, _ = self._build(6, inhomogeneous_layer_cycle_interval=3, scan_layers=False) + + unscanned = raiden_unscan.unscan_layers( + nnx.state(decoder, nnx.Param), + num_layers=cfg.num_decoder_layers, + scan_axis=cfg.param_scan_axis, + cycle_interval=cfg.inhomogeneous_layer_cycle_interval, + ) + + def shapes(tree): + return {jax.tree_util.keystr(path): leaf.shape for path, leaf in jax.tree_util.tree_leaves_with_path(tree)} + + self.assertEqual(shapes(unscanned), shapes(nnx.state(unscanned_decoder, nnx.Param))) + nnx.update(unscanned_decoder, unscanned) + + # Every loaded leaf must be exactly its block/slot slice of the scanned state: + # slot j of repeat i is layer i * cycle + j, and the last slot is the global layer. + axis, cycle = cfg.param_scan_axis, cfg.inhomogeneous_layer_cycle_interval + scanned = traverse_util.flatten_dict(nnx.state(decoder, nnx.Param).to_pure_dict()) + loaded = traverse_util.flatten_dict(nnx.state(unscanned_decoder, nnx.Param).to_pure_dict()) + for key, value in loaded.items(): + idx = next((i for i, k in enumerate(key) if isinstance(k, str) and re.fullmatch(r"layers_\d+", k)), None) + if idx is None: + expected = scanned[key] + else: + repeat, slot = divmod(int(key[idx].split("_")[1]), cycle) + sub = "global_layer" if slot == cycle - 1 else "local_layers" + expected = np.take(np.asarray(scanned[key[:idx] + ("layers", sub) + key[idx + 1 :]]), repeat, axis=axis) + if sub == "local_layers": + expected = np.take(expected, slot, axis=axis) + np.testing.assert_array_equal(np.asarray(value), expected, err_msg=str(key)) + + # Scanned and unscanned forwards differ by float rounding that grows with the device count. + np.testing.assert_allclose( + np.asarray(self._run(unscanned_cfg, unscanned_decoder, shared_embedding)), + np.asarray(self._run(cfg, decoder, shared_embedding)), + rtol=1e-3, + atol=1e-3, + ) + class TestNNXDecoderDeepseekAndGemma4(unittest.TestCase): """Tests for Deepseek and Gemma4 specific decoder logic.""" diff --git a/tests/unit/raiden_unscan_test.py b/tests/unit/raiden_unscan_test.py index f6bef16bd2..561b44823b 100644 --- a/tests/unit/raiden_unscan_test.py +++ b/tests/unit/raiden_unscan_test.py @@ -204,6 +204,39 @@ def test_inhomogeneous_cycle_interval_unscans_correctly(self): # layer 3: repeat 1, slot 1 np.testing.assert_array_equal(np.asarray(_unwrap(unscanned["layers_3"]["kernel"])), np.asarray(k1[:, 1, :])) + def test_nested_block_layout_unscans_correctly(self): + """Qwen3-Next/Qwen3.5's nested scanned block unrolls into interleaved per-layer entries. + + `local_layers` stacks every cycle slot but the last on `scan_axis + 1`, and `global_layer` + is the last slot. A cycle of 3 puts two slots on the nested axis, so a swapped repeat or + slot index lands a slice on the wrong layer. + """ + num_layers = 6 + cycle_interval = 3 + repeats = num_layers // cycle_interval # 2 + num_local = cycle_interval - 1 # 2 + + local = jnp.arange(_IN * repeats * num_local * _OUT, dtype=jnp.float32).reshape(_IN, repeats, num_local, _OUT) + glob = (jnp.arange(_IN * repeats * _OUT, dtype=jnp.float32) + 1000).reshape(_IN, repeats, _OUT) + state = { + "layers": { + "local_layers": {"kernel": local}, + "global_layer": {"kernel": glob}, + }, + "embed": jnp.zeros((_VOCAB, _IN)), + } + + unscanned = raiden_unscan.unscan_layers(state, num_layers=num_layers, cycle_interval=cycle_interval) + expected_names = sorted(["['embed'].value"] + [f"['layers_{i}']['kernel'].value" for i in range(num_layers)]) + self.assertEqual(_names(unscanned), expected_names) + + for i in range(repeats): + for j in range(num_local): + got = np.asarray(_unwrap(unscanned[f"layers_{i * cycle_interval + j}"]["kernel"])) + np.testing.assert_array_equal(got, np.asarray(local[:, i, j, :])) + got = np.asarray(_unwrap(unscanned[f"layers_{i * cycle_interval + num_local}"]["kernel"])) + np.testing.assert_array_equal(got, np.asarray(glob[:, i, :])) + def test_inhomogeneous_config_mismatches_raise(self): """Every way a cycle config can disagree with the tree must fail loudly. @@ -240,6 +273,13 @@ def test_inhomogeneous_config_mismatches_raise(self): 2, r"does not have a 'layer_' cycle prefix", ), + ( + "local_layers without the nested slot axis", + {"layers": {"local_layers": {"kernel": jnp.zeros((_IN, 2, _OUT))}}}, + 4, + 2, + r"expected axis 2 to hold the 1 'local_layers' cycle slots", + ), ("cycle_interval below 1", two_slots, 4, 0, r"cycle_interval must be >= 1"), ] for name, state, num_layers, cycle_interval, expected in cases: diff --git a/tests/unit/router_replay_test.py b/tests/unit/router_replay_test.py index 7c0d86eff7..ce842396e0 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -647,6 +647,31 @@ def test_per_layer_routing_is_not_broadcast_scanned(self): "qwen3.5 scanned per-layer", ) + def test_per_layer_routing_reaches_the_scan_remainder_block(self): + """Layers past the last whole cycle live in a separate remainder block. + + With 3 layers and a cycle of 2 the scan covers layers 0-1 and layer 2 is + applied afterwards, so the remainder needs its own slice of the routing. + """ + seq_len, batch_size, top_k, num_experts = 8, 1, 2, 4 + cycle_interval, layers = 2, 3 + cfg = self._qwen35_cfg( + seq_len, + batch_size, + num_experts, + top_k, + "test_router_replay_scan_remainder", + base_num_decoder_layers=layers, + num_decoder_layers=layers, + inhomogeneous_layer_cycle_interval=cycle_interval, + ) + # Only the trailing (remainder) layer's routing differs between a and b, so + # the losses can only diverge if the remainder block replays its own slice. + shared = jnp.zeros((_batch_rows(cfg), seq_len, layers, top_k), dtype=jnp.int32) + a = shared.at[:, :, -1, :].set(1) + b = shared.at[:, :, -1, :].set(3) + self._assert_routing_is_load_bearing(cfg, seq_len, a, b, "qwen3.5 scan remainder") + def test_loss_fn_with_forced_routed_experts_gemma4(self): """Gemma4 is supported unscanned only (its scanned path raises), and it is From 749ccfd73ee65459ebdbb4e07b6cbdfb27d360ba Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 16:23:12 +0000 Subject: [PATCH 3/5] fix(moe): keep the FSDP expert-kernel gather inside a scanned decoder The routed-MoE kernel pspecs drop the fsdp axis so XLA emits the FSDP all-gather at the sharding constraint in sparse_matmul. Under a scanned decoder the kernels reaching that constraint are a per-layer dynamic-slice of a stacked param, and with nothing in between XLA rewrites all_gather(dynamic_slice(w)) into dynamic_slice(all_gather(w)). That gather is loop invariant, so it is hoisted out of the scan and every layer's gathered expert kernel is live at once -- num_layers * num_experts * embed * moe_mlp * dtype bytes of extra peak memory. The jnp.asarray(kernel, dtype) cast in __call__ normally separates the slice from the collective and blocks the rewrite. It is an identity when weight_dtype == dtype, which is why the blow-up only shows up there: qwen3-next-80b-a3b at tpu7x-128 compiles to 14.20 GB of temporaries with weight_dtype=float32 but 65.59 GB with weight_dtype=bfloat16, on code that is byte-identical to main's. Stand an optimization_barrier in for the cast in exactly that case, matching the barrier moe_fsdp_use_two_stage_all_gather already uses a few lines above. AOT temp_size_in_bytes per device: model topology before after qwen3-next-80b-a3b tpu7x-128 65.59 GB 12.11 GB qwen3.5-35b-a3b tpu7x-128 30.23 GB 6.94 GB qwen3.5-35b-a3b tpu7x-16 35.71 GB 17.04 GB qwen3-30b-a3b tpu7x-16 13.03 GB 12.42 GB Gating on the identity keeps every weight_dtype=float32 build byte-identical (qwen3.5-35b-a3b 25132249440 at tpu7x-16 and 8510193056 at tpu7x-128, qwen3-next-80b-a3b 14204451200 at tpu7x-128, before and after). Applying the barrier unconditionally costs 3-8% on those. --- src/maxtext/layers/moe.py | 13 ++++++ tests/unit/moe_test.py | 97 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 01e6ce4445..8e6e297c20 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -3991,6 +3991,19 @@ def _fsdp_all_gather(w, pspec): return w return self._maybe_shard_with_pspec(w, pspec) + if jnp.dtype(self.weight_dtype) == jnp.dtype(self.dtype): + # Keep the FSDP all-gather below inside the decoder loop. Under a scanned decoder the + # kernels arrive as a per-layer dynamic-slice of a stacked, fsdp-sharded param, and the + # pspecs above have the fsdp axis stripped so XLA inserts the all-gather right here. + # With nothing in between, XLA rewrites all_gather(dynamic_slice(w)) into + # dynamic_slice(all_gather(w)); that gather is loop invariant, so it is hoisted out of + # the scan and every layer's gathered expert kernel is live at once + # (num_layers * num_experts * embed * mlp * dtype bytes -- 51 GB for qwen3-next-80b). + # The `jnp.asarray(kernel, self.dtype)` cast in __call__ normally sits between the two + # and blocks the rewrite; when weight_dtype == dtype it is an identity, so pin the + # gather here instead. Only needed in that case, as the barrier costs a few percent of + # peak memory when XLA is otherwise free to schedule around the cast. + w0_kernel, w1_kernel, wo_kernel = jax.lax.optimization_barrier((w0_kernel, w1_kernel, wo_kernel)) w0_kernel = _fsdp_all_gather(w0_kernel, w0_pspec) w1_kernel = _fsdp_all_gather(w1_kernel, w1_pspec) wo_kernel = _fsdp_all_gather(wo_kernel, wo_pspec) diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 85c6d0f256..80245cf909 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -714,6 +714,8 @@ def test_sparse_matmul_repairs_batch_specs_only_without_expert_parallelism(exper get_expert_parallelism_size=lambda: expert_parallelism, _expert_parallelism_name="expert", per_expert_scale=None, + weight_dtype=jnp.float32, + dtype=jnp.bfloat16, ) original_batch_partition = "fsdp" if expert_parallelism == 1 else ("fsdp", "expert") fake_moe._logical_to_mesh_axes = lambda logical_axes: P( # pylint: disable=protected-access @@ -758,6 +760,101 @@ def fake_shard_map(function, *, mesh, in_specs, out_specs, check_vma): assert captured["out_specs"][3] == P(batch_partition) +@pytest.mark.parametrize( + ("weight_dtype", "dtype", "expect_barrier"), + ( + (jnp.bfloat16, jnp.bfloat16, True), + (jnp.float32, jnp.float32, True), + (jnp.float32, jnp.bfloat16, False), + ), +) +def test_sparse_matmul_pins_expert_kernel_all_gather_when_the_upcast_is_an_identity(weight_dtype, dtype, expect_barrier): + """Sparse MoE keeps the FSDP kernel gather from being hoisted out of a scanned decoder. + + The kernel pspecs drop the fsdp axis so XLA emits the all-gather at the sharding + constraint. Under a scanned decoder the kernels are a per-layer dynamic-slice of a + stacked param, and with nothing between the slice and the collective XLA turns + all_gather(dynamic_slice(w)) into the loop-invariant dynamic_slice(all_gather(w)) and + hoists it, making every layer's gathered expert kernel live at once. The + `jnp.asarray(kernel, dtype)` upcast usually separates them; when weight_dtype == dtype + it is an identity, so an optimization barrier has to stand in for it. + """ + fake_moe = SimpleNamespace( + config=SimpleNamespace( + shard_exp_on_fsdp=False, + use_2d_fsdp_sharding=False, + shard_embed_moe_on_fsdp=False, + model_name="qwen3.5-35b-a3b", + check_vma=False, + moe_fsdp_use_two_stage_all_gather=False, + moe_pin_sparse_core_all_gathers=False, + moe_flat_fsdp_weights=False, + moe_dropless_fallback=None, + load_balance_loss_weight=0.0, + ), + mesh=SimpleNamespace(axis_names=("diloco", "fsdp", "expert"), shape={"diloco": 2, "fsdp": 32, "expert": 1}), + rngs=object(), + get_expert_parallelism_size=lambda: 1, + _expert_parallelism_name="expert", + per_expert_scale=None, + weight_dtype=weight_dtype, + dtype=dtype, + ) + fake_moe._logical_to_mesh_axes = lambda logical_axes: P( # pylint: disable=protected-access + *("fsdp" if axis == "activation_batch" else None for axis in logical_axes) + ) + sharded = [] + + def record_shard(value, _pspec, **_kwargs): + sharded.append(value) + return value + + fake_moe._maybe_shard_with_pspec = record_shard # pylint: disable=protected-access + fake_moe.sow = lambda *_args, **_kwargs: None + + w0 = SimpleNamespace(shape=(256, 2048, 512)) + w1 = SimpleNamespace(shape=(256, 2048, 512)) + wo = SimpleNamespace(shape=(256, 512, 2048)) + barriered = tuple(SimpleNamespace(shape=w.shape) for w in (w0, w1, wo)) + + def fake_shard_map(function, *, mesh, in_specs, out_specs, check_vma): + del function, mesh, in_specs, out_specs, check_vma + return lambda x, *_args: (x, None, None, jnp.bool_(False), jnp.bool_(False), None) + + with ( + mock.patch.object(jax, "shard_map", side_effect=fake_shard_map), + mock.patch.object(jax.lax, "optimization_barrier", return_value=barriered) as barrier, + ): + moe.RoutedMoE.sparse_matmul( + fake_moe, + SimpleNamespace(shape=(4, 1024, 2048)), + SimpleNamespace(shape=(4, 1024, 256)), + None, + w0, + w1, + wo, + None, + None, + None, + ) + + # Identity, not equality: the stand-in kernels all share a shape, so SimpleNamespace + # equality cannot tell a barrier output from the raw kernel it replaced. + def was_sharded(value): + return any(value is candidate for candidate in sharded) + + if expect_barrier: + assert barrier.call_count == 1 + assert all(a is b for a, b in zip(barrier.call_args.args[0], (w0, w1, wo), strict=True)) + # The constraint must consume the barrier's outputs; a barrier whose results are + # dropped still lets XLA reorder the gather above the per-layer slice. + assert all(was_sharded(kernel) for kernel in barriered) + assert not any(was_sharded(kernel) for kernel in (w0, w1, wo)) + else: + barrier.assert_not_called() + assert all(was_sharded(kernel) for kernel in (w0, w1, wo)) + + # The TC ragged sort tests' moe_mlp_dim (256) is smaller than the default 1024 gmm_v2 mlp tiles. _GMM_V2_SMALL_MLP_TILES = {f"{w}_tile_{p}_mlp_dim": 256 for w in ("wi", "wo") for p in ("fwd", "dlhs", "drhs")} From 2b3425c580c469a4d41df6fb1df1f939cad84757 Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Thu, 3 Sep 2026 23:26:55 +0000 Subject: [PATCH 4/5] style(qwen3): replace an assigned lambda with a def pylint C3001 unnecessary-lambda-assignment on the no-router-replay branch of Qwen3NextScannableBlock._scan_local_layers. It was the only finding failing the Code Quality Check / Pre-commit Linters job. Mirrors the def already used in the sibling branch; behaviour is unchanged. --- src/maxtext/models/qwen3.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index fb5e06a6b6..27371e5d58 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1620,7 +1620,10 @@ def _scan_local_layers(self, y, layer_kwargs, forced_routed_experts=None): """Runs the local (linear attention / GatedDeltaNet) layers via a per-layer rematerialized jax.lax.scan.""" remat = self._remat_enabled if forced_routed_experts is None: - apply_fn = lambda layer, carry: self._run_layer(layer, carry, layer_kwargs)[0] + + def apply_fn(layer, carry): + return self._run_layer(layer, carry, layer_kwargs)[0] + else: # Router replay: one slice of forced routing per local layer, fed through # the scan's xs so each iteration sees its own layer's expert indices. From ee342ff0b64a7d073a516a4ef4c6eb0d7b859bae Mon Sep 17 00:00:00 2001 From: NuojCheng Date: Fri, 4 Sep 2026 00:52:32 +0000 Subject: [PATCH 5/5] style: trim the comments this PR adds Cut the explanatory comments and docstring prose added by this branch down to the few that record something the code cannot show. 180 lines removed, 27 added; no code changes (AST identical, docstrings excluded). --- .../utils/param_mapping.py | 30 ++------------- .../utils/tensor_handling.py | 3 -- .../integration/vllm/maxtext_vllm_rollout.py | 37 ++++--------------- .../integration/vllm/weight_converter.py | 18 +-------- src/maxtext/layers/moe.py | 13 +------ src/maxtext/layers/nnx_decoders.py | 9 +---- src/maxtext/layers/nnx_scan.py | 5 --- src/maxtext/models/qwen3.py | 16 ++------ src/maxtext/models/qwen3_5.py | 17 +-------- .../unit/vllm_rollout_unroll_test.py | 8 +--- .../unit/weight_converter_test.py | 15 +------- tests/unit/moe_test.py | 15 +------- tests/unit/nnx_decoders_test.py | 8 +--- tests/unit/param_mapping_test.py | 13 +------ tests/unit/router_replay_test.py | 8 +--- 15 files changed, 28 insertions(+), 187 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 77090b5ade..ee567045a8 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -889,20 +889,8 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING( } if scan_layers: - # 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) - # Qwen3.5 stores its routed experts fused, so unlike Qwen3-Next there is no - # extra per-expert axis to nest. FP8 checkpoints are the exception: they store - # one tensor per expert, which adds a leading expert axis exactly as for - # Qwen3-Next ([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. + # The length-1 _scan_global_layer scan is a memory boundary, not a param stack, so + # global_layer params stay flat [block] while local_layers nest [block][local]. num_blocks = num_main_layers // layer_cycle_interval local_positions = list(range(layer_cycle_interval - 1)) global_position = layer_cycle_interval - 1 @@ -910,18 +898,14 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING( def hf_layer(block_idx, position, suffix): return f"{hf_prefix}.layers.{block_idx * layer_cycle_interval + position}.{suffix}" - # (maxtext subkey, hf suffix) pairs shared by both the local and global layers. - # A tuple of suffixes means MaxText concatenates the HF tensors into one kernel. shared_specs = [ ("input_layernorm-scale", "input_layernorm.weight"), ("post_attention_layernorm-scale", "post_attention_layernorm.weight"), ] - # Routed experts stored one tensor per expert (FP8 checkpoints only). + # FP8 checkpoints store one HF tensor per routed expert, which nests one level deeper. expert_specs = [] - # Whether the routed experts' wi_0/wi_1 come from one fused HF gate_up_proj. fused_experts = False if num_experts > 1: - # 3. MoE: router gate and shared expert. shared_specs += [ ("mlp-routed_experts-gate-kernel", "mlp.gate.weight"), ("mlp-shared_expert-wi_0-kernel", "mlp.shared_expert.gate_proj.weight"), @@ -929,7 +913,6 @@ def hf_layer(block_idx, position, suffix): ("mlp-shared_expert-wo-kernel", "mlp.shared_expert.down_proj.weight"), ("mlp-shared_expert_gate-kernel", "mlp.shared_expert_gate.weight"), ] - # 4. MoE routed experts. if is_quantized: num_experts = config.get("text_config", config).get("num_experts", 256) shared_specs += [ @@ -949,13 +932,11 @@ def hf_layer(block_idx, position, suffix): shared_specs.append(("mlp-routed_experts-wo", "mlp.experts.down_proj")) fused_experts = True else: - # 3. Dense MLP. shared_specs += [ ("mlp-wi_0-kernel", "mlp.gate_proj.weight"), ("mlp-wi_1-kernel", "mlp.up_proj.weight"), ("mlp-wo-kernel", "mlp.down_proj.weight"), ] - # Linear (GatedDeltaNet) attention: only ever on the local layers. local_specs = shared_specs + [ ("attention-in_proj_qkvz-kernel", ("linear_attn.in_proj_qkv.weight", "linear_attn.in_proj_z.weight")), ("attention-in_proj_ba-kernel", ("linear_attn.in_proj_b.weight", "linear_attn.in_proj_a.weight")), @@ -965,7 +946,6 @@ def hf_layer(block_idx, position, suffix): ("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"), @@ -1006,7 +986,6 @@ def hf_entry(block_idx, position, suffix): for e in range(num_experts) ] if fused_experts: - # Fused gate_up_proj feeds both wi_0 and wi_1, so it is keyed by the pair. mapping[(f"{local_prefix}-mlp-routed_experts-wi_0", f"{local_prefix}-mlp-routed_experts-wi_1")] = [ [hf_layer(b, p, "mlp.experts.gate_up_proj") for p in local_positions] for b in range(num_blocks) ] @@ -1405,9 +1384,6 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval num_main_layers = getattr(maxtext_config, "num_decoder_layers", None) or config["text_config"]["num_hidden_layers"] num_experts = getattr(maxtext_config, "num_experts", None) or config.get("text_config", config).get("num_experts", 1) - # 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), diff --git a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py index a04beff2d5..b8341b3a31 100644 --- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py +++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py @@ -63,9 +63,6 @@ def stacked_axes(mt_key: str, config, depth: int) -> tuple: routed experts): the expert axis still leads, giving ``(0, param_scan_axis, param_scan_axis + 1)``. """ - # A tuple is the composite-MaxText-key convention (e.g. Qwen3.5's fused - # gate_up_proj feeding wi_0 and wi_1); every component shares the same prefix, - # so the first one decides the layout. if isinstance(mt_key, tuple) and mt_key: mt_key = mt_key[0] if isinstance(mt_key, str) and "-local_layers" in mt_key: diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 60320c38d0..1364f0c040 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -298,28 +298,18 @@ def source_covers(target_key): return len(target_layer_keys) -# `Qwen3NextScannableBlock` (shared by Qwen3-Next and Qwen3.5) nests two scans: -# the cycle's linear-attention layers share a single `local_layers` module, with -# blocks on `scan_axis` and the cycle slot on `scan_axis + 1`, while the trailing -# full-attention layer sits in `global_layer` with blocks on `scan_axis` alone. _QWEN_NESTED_LOCAL = "local_layers" _QWEN_NESTED_GLOBAL = "global_layer" class _QwenScannedRef(NamedTuple): - """Where a scanned Qwen parameter sits, and how it maps onto cycle slots. + """Where a scanned Qwen parameter sits, and how it maps onto cycle slots.""" - `slot_idx` is the parameter's fixed position in the layer cycle. It is None - for the two nested containers: `local_layers` spans several slots, read off - `scan_axis + 1` rather than off the key, and `global_layer` always takes - whichever slot closes the cycle. - """ - - container_idx: int # Position of "layers" in the key tuple, or -1 for no match. + container_idx: int slot_idx: Optional[int] - consumed: int # Key elements the container occupies. - nested: bool # Slots are stacked on `scan_axis + 1`. - trailing: bool # Occupies the last slot of the cycle. + consumed: int + nested: bool + trailing: bool _NO_QWEN_SCAN = _QwenScannedRef(-1, -1, 0, False, False) @@ -347,12 +337,8 @@ def _find_qwen_scanned_layer_idx(key_tuple) -> _QwenScannedRef: def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Optional[int] = None): """Unroll Qwen's heterogeneous or homogeneous scanned blocks for an unscanned MaxText target. - Qwen 3 Next/3.5 training stores a repeating layer cycle under a nested scan: - `decoder.layers.local_layers` holds the cycle's linear-attention layers with - repetitions on `scan_axis` and the cycle slot on `scan_axis + 1`, and - `decoder.layers.global_layer` holds the trailing full-attention layer with - repetitions on `scan_axis` alone. Trainer states predating that layout, which - kept one `decoder.layers.layer_{slot}` module per slot, are still accepted. + Qwen 3 Next/3.5 training nests two scans: `decoder.layers.local_layers` holds the cycle's linear-attention + layers with the slot on `scan_axis + 1`, and `decoder.layers.global_layer` holds the trailing full-attention layer. Qwen 3 base training stores homogeneous layers as `decoder.layers.*` stacked on `scan_axis`. The inference model stores every layer as a direct decoder attribute named `layers_{global_index}`. Tunix's generic direct-sync mapper cannot bridge @@ -385,8 +371,6 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt ref = _find_qwen_scanned_layer_idx(key) if ref.container_idx == -1 or "dropout" in key or "rngs" in key: continue - # A nested `local_layers` parameter carries its cycle slots on the axis just - # past the block axis, so it needs one dimension more than a flat one. min_ndim = scan_axis + 2 if ref.nested else scan_axis + 1 if not hasattr(value, "shape") or len(value.shape) < min_ndim: raise ValueError(f"Qwen scanned parameter {'.'.join(map(str, key))} has no scan axis {scan_axis}: {value!r}") @@ -405,8 +389,6 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt if len(local_widths) > 1: raise ValueError(f"Qwen nested `local_layers` parameters disagree on cycle width: {sorted(local_widths)}") - # `global_layer` closes the cycle, so its slot is whatever follows the local - # ones. Resolved here because the key alone does not carry the index. trailing_slot = (max(slot_indices) + 1) if slot_indices else 0 if has_trailing: slot_indices.add(trailing_slot) @@ -430,8 +412,6 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt for key, value, ref in scanned_keys: prefix = key[: ref.container_idx] suffix = key[ref.container_idx + ref.consumed :] - # `local_layers` supplies several cycle slots from one tensor; every other - # container supplies exactly one, at a slot fixed above. if ref.nested: slots = tuple(range(value.shape[scan_axis + 1])) else: @@ -441,8 +421,7 @@ def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Opt for local_idx, slot_idx in enumerate(slots): global_idx = repetition * pattern_length + slot_idx new_key = prefix + (f"layers_{global_idx}",) + suffix - # Taking the block collapses `scan_axis`, so the nested slot axis has - # shifted down into it. + # Taking the block collapsed `scan_axis`, so the nested slot axis has shifted down into it. new_flat_w[new_key] = jnp.take(block, local_idx, axis=scan_axis) if ref.nested else block logging.info( diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py index ee5de8abf0..9081e45848 100644 --- a/src/maxtext/integration/vllm/weight_converter.py +++ b/src/maxtext/integration/vllm/weight_converter.py @@ -518,8 +518,6 @@ class _PlanEntry: # Index along the scan axis, or None when the source is already per-layer. slice_index: Optional[int] op: str # "identity" | "slice" | "fuse_moe" - # Index along the nested cycle-slot axis (`scan_axis + 1`), or None when the - # source stacks blocks alone. Only `local_layers` sources carry one. local_index: Optional[int] = None @@ -545,7 +543,6 @@ class _PlanGroup: # Dot-joined source path, precomputed for diagnostics and for the # leaf-name check inside `_align_per_axis`. source_path: str - # Shared by every entry in the group: see `_PlanEntry.local_index`. local_index: Optional[int] = None @@ -557,8 +554,6 @@ def _group_plan(plan: List[_PlanEntry]) -> List[_PlanGroup]: scanned source must share a shape, and if two did not, the bulk unstack would produce the wrong shape for one of them -- which `convert()`'s post-execution shape check catches loudly rather than silently writing incorrect weights. - `local_index` *is* part of the key, because one nested `local_layers` source - feeds a different set of rollout layers at each cycle slot. """ grouped: Dict[Tuple[Any, ...], List[_PlanEntry]] = {} order: List[Tuple[Any, ...]] = [] @@ -792,11 +787,7 @@ def _scanned_candidates( ) -> List[Tuple[Tuple[Any, ...], Optional[int]]]: """Source keys that could hold the scanned form of a target layer key. - Each candidate is paired with the index to take along the nested cycle-slot - axis, or None when the source stacks blocks alone. `Qwen3NextScannableBlock` - nests two scans: the cycle's linear-attention layers share one - `local_layers` module whose slot sits on `scan_axis + 1`, while the trailing - full-attention layer sits in `global_layer` with no slot axis at all. + Each candidate is paired with the index to take along the nested cycle-slot axis, or None when there is none. """ prefix, suffix = key_tuple[:prefix_len], key_tuple[suffix_start:] if self.cycle > 1: @@ -1167,12 +1158,7 @@ def _execute_group(self, group: _PlanGroup, src_flat, tgt_flat): return [(tgt_key, per_block[idx]) for idx, tgt_key in group.targets] def _take_cycle_slot(self, val, group: _PlanGroup): - """Drops the nested cycle-slot axis from a `local_layers` source. - - Everything downstream expects blocks on `scan_axis` and no other stacked - axis, so the slot is selected here -- once per group, before the bulk - alignment and unstack, rather than once per target layer. - """ + """Drops the nested cycle-slot axis from a `local_layers` source.""" if group.local_index is None: return val slot_axis = self.scan_axis + 1 diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 8e6e297c20..fb8660d1cf 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -3992,17 +3992,8 @@ def _fsdp_all_gather(w, pspec): return self._maybe_shard_with_pspec(w, pspec) if jnp.dtype(self.weight_dtype) == jnp.dtype(self.dtype): - # Keep the FSDP all-gather below inside the decoder loop. Under a scanned decoder the - # kernels arrive as a per-layer dynamic-slice of a stacked, fsdp-sharded param, and the - # pspecs above have the fsdp axis stripped so XLA inserts the all-gather right here. - # With nothing in between, XLA rewrites all_gather(dynamic_slice(w)) into - # dynamic_slice(all_gather(w)); that gather is loop invariant, so it is hoisted out of - # the scan and every layer's gathered expert kernel is live at once - # (num_layers * num_experts * embed * mlp * dtype bytes -- 51 GB for qwen3-next-80b). - # The `jnp.asarray(kernel, self.dtype)` cast in __call__ normally sits between the two - # and blocks the rewrite; when weight_dtype == dtype it is an identity, so pin the - # gather here instead. Only needed in that case, as the barrier costs a few percent of - # peak memory when XLA is otherwise free to schedule around the cast. + # Without this, XLA rewrites all_gather(dynamic_slice(w)) into dynamic_slice(all_gather(w)) and hoists the + # gather out of the scanned decoder, keeping every layer's experts live; the dtype cast normally blocks that. w0_kernel, w1_kernel, wo_kernel = jax.lax.optimization_barrier((w0_kernel, w1_kernel, wo_kernel)) w0_kernel = _fsdp_all_gather(w0_kernel, w0_pspec) w1_kernel = _fsdp_all_gather(w1_kernel, w1_pspec) diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index eb5b2ba15c..564e764146 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -611,8 +611,6 @@ 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 - # Qwen3-Next and Qwen3.5 share the same hybrid attention period, so they share - # the same nested-scan block structure. self.is_qwen3_hybrid = self.config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5) if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4: @@ -2688,10 +2686,6 @@ def _apply_qwen3_hybrid_scanned_blocks(self, y, layer_args, layer_kwargs, kv_cac 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. - - Forced routing (router replay) is sliced per decoder layer here and handed - to each block as its own `[block_length, batch, seq, top_k]` stack, which - the block then splits between its local scan and its global layer. """ cfg = self.config block_length = cfg.inhomogeneous_layer_cycle_interval @@ -2704,8 +2698,7 @@ def _apply_qwen3_hybrid_scanned_blocks(self, y, layer_args, layer_kwargs, kv_cac remainder_forced_routed_experts = None if forced_routed_experts is not None: if kv_caches is not None: - # The kv-cache path below unrolls into _apply_layers_sequentially, which - # threads either kv caches or scan xs, not both. + # _apply_layers_sequentially threads either kv caches or scan xs, not both. raise NotImplementedError( "Forced routing is not supported together with externally-managed (vLLM) kv_caches in scanned layers." ) diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index 45e221b921..2eb9419d3c 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -112,11 +112,6 @@ def apply_scanned_layers( ``remat`` is separate from ``remat_policy`` because ``None`` is JAX's full rematerialization policy, not an indication that rematerialization is off. - ``xs`` is an optional pytree of per-layer inputs whose leaves have a leading - axis of ``length``; each iteration's slice is passed to ``apply_fn`` as a - third argument. Callers that pass no ``xs`` keep the two-argument ``apply_fn`` - signature. - Externally managed per-layer state, such as KV caches, is not supported by this scan path. diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 27371e5d58..a4068125c3 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1591,11 +1591,7 @@ def __init__( self._maybe_shard_with_logical = lambda inputs, *args, **kwargs: inputs def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): - """Builds one sub-layer of the block. - - Subclasses for architectures that share this block structure (e.g. Qwen3.5) - override this to swap in their own decoder layer class. - """ + """Builds one sub-layer of the block.""" return Qwen3NextDecoderLayer( config=self.config, mesh=self.mesh, @@ -1625,8 +1621,7 @@ def apply_fn(layer, carry): return self._run_layer(layer, carry, layer_kwargs)[0] else: - # Router replay: one slice of forced routing per local layer, fed through - # the scan's xs so each iteration sees its own layer's expert indices. + def apply_fn(layer, carry, routing): return self._run_layer(layer, carry, layer_kwargs, forced_routed_experts=routing)[0] @@ -1877,10 +1872,6 @@ def __init__( self.mlp = self._make_mlp(rngs=rngs) - # The hooks below are the only places a sibling architecture that reuses this layer - # (Qwen3.5) has to diverge, so they are overridable rather than inline. `__init__` - # binds what the factories return to `attention` and `mlp`, which are checkpoint - # parameter paths, so an override has to keep the sub-module's own parameter names too. def _norm_weight_dtype(self): """Storage dtype of the two RMSNorm scales.""" return self.config.weight_dtype @@ -1990,8 +1981,7 @@ def __call__( # We sow the load balancing loss so it can be collected and added to the total loss # during training. - # Assigned rather than sown: `sow` appends to a tuple, so the layer's Intermediate - # structure would grow on every call, which a `jax.lax.scan` body cannot express. + # Assigned, not sown: `sow` appends, so Intermediate would grow per call, which a scan body cannot express. if self.config.load_balance_loss_weight > 0.0 and load_balance_loss is not None: self.moe_lb_loss = nnx.Intermediate(load_balance_loss) diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 6ba8ce0d37..f9e0277bf9 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -46,14 +46,7 @@ class Qwen3_5SparseMoEBlock(Qwen3NextSparseMoeBlock): class Qwen3_5ScannableBlock(Qwen3NextScannableBlock): - """Scanned Structure for Text-only Architecture, explicitly invoking Qwen3_5 layers. - - Qwen3.5 repeats the same hybrid attention period as Qwen3-Next -- several - GatedDeltaNet layers followed by one full-attention layer -- so it reuses - Qwen3-Next's hierarchical nested scans (an inner scan over the homogeneous - linear-attention layers plus a trip-count-one scan over the full-attention - layer) and only swaps in the Qwen3.5 decoder layer. - """ + """Scanned Structure for Text-only Architecture, explicitly invoking Qwen3_5 layers.""" def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): return Qwen3_5DecoderLayer( @@ -70,14 +63,6 @@ def _make_decoder_layer(self, *, layer_idx, is_full_attention_layer, rngs): class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer): """Qwen3.5 hybrid decoder layer. - Qwen3.5's decoder layer is structurally identical to Qwen3-Next's -- norm, - either full or linear attention, norm, MoE (or, for the dense models, an MLP), - with residuals around the attention and MLP halves -- so the whole forward pass - is inherited. Only the sub-block hooks are overridden, to build the Qwen3.5 - sub-classes. The attribute names (`input_layernorm`, `attention`, - `post_attention_layernorm`, `mlp`) are therefore shared, and so are the - checkpoint parameter paths. - Attributes: config: The model configuration object. mesh: The device mesh for sharding. diff --git a/tests/post_training/unit/vllm_rollout_unroll_test.py b/tests/post_training/unit/vllm_rollout_unroll_test.py index a3169abba7..875041e7b9 100644 --- a/tests/post_training/unit/vllm_rollout_unroll_test.py +++ b/tests/post_training/unit/vllm_rollout_unroll_test.py @@ -230,13 +230,7 @@ def test_interleaves_slots_and_repetitions(self): @pytest.mark.cpu_only def test_unrolls_the_nested_local_and_global_block_scan(self): - """Qwen3-Next and Qwen3.5 nest the cycle's local layers into one module. - - `local_layers` stacks blocks on `scan_axis` and the cycle slot on - `scan_axis + 1`; `global_layer` holds the trailing full-attention layer on - the block axis alone. Together they must still unroll to a dense - `layers_{global_idx}` run, with each slot landing on its own layer. - """ + """Nested `local_layers` plus `global_layer` unroll to a dense `layers_{global_idx}` run.""" num_blocks, num_local = 2, 3 cycle = num_local + 1 local_probe = np.zeros((2, num_blocks, num_local, 1), dtype=np.float32) diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py index d243ba6c45..1feb80f456 100644 --- a/tests/post_training/unit/weight_converter_test.py +++ b/tests/post_training/unit/weight_converter_test.py @@ -127,13 +127,7 @@ def _source_tree(fused_moe_on_target: bool): def _nested_source_tree(fused_moe_on_target: bool): - """Trainer state in the nested-scan layout `Qwen3NextScannableBlock` produces. - - Same values as `_source_tree`, restacked the way the shared Qwen3-Next block - stores them: slots `0..C-2` share one `local_layers` module with the cycle - slot on `SCAN_AXIS + 1`, and slot `C-1` becomes `global_layer` with the block - axis alone. Converting either tree must yield the same rollout weights. - """ + """Trainer state in the nested-scan layout `Qwen3NextScannableBlock` produces.""" legacy = _source_tree(fused_moe_on_target)["base"]["decoder"]["layers"] local_slots = [legacy[f"layer_{slot}"] for slot in range(CYCLE - 1)] tree = _source_tree(fused_moe_on_target) @@ -216,12 +210,7 @@ def test_non_layer_params_pass_through(self): ) def test_nested_scan_layout_converts_identically(self): - """Qwen3.5 moved onto Qwen3-Next's nested block scan. - - The cycle's linear-attention layers now live in one `local_layers` module - with the slot on `SCAN_AXIS + 1` instead of one module per slot, so the - converter has to read that layout too -- and read it to the same answer. - """ + """The nested-scan layout converts to the same rollout weights as the per-slot layout.""" target = _target_tree() legacy = traverse_util.flatten_dict(MaxTextToMaxTextConverter(_config()).convert(_source_tree(True), target)) nested = traverse_util.flatten_dict(MaxTextToMaxTextConverter(_config()).convert(_nested_source_tree(True), target)) diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 80245cf909..420ac803f2 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -769,16 +769,7 @@ def fake_shard_map(function, *, mesh, in_specs, out_specs, check_vma): ), ) def test_sparse_matmul_pins_expert_kernel_all_gather_when_the_upcast_is_an_identity(weight_dtype, dtype, expect_barrier): - """Sparse MoE keeps the FSDP kernel gather from being hoisted out of a scanned decoder. - - The kernel pspecs drop the fsdp axis so XLA emits the all-gather at the sharding - constraint. Under a scanned decoder the kernels are a per-layer dynamic-slice of a - stacked param, and with nothing between the slice and the collective XLA turns - all_gather(dynamic_slice(w)) into the loop-invariant dynamic_slice(all_gather(w)) and - hoists it, making every layer's gathered expert kernel live at once. The - `jnp.asarray(kernel, dtype)` upcast usually separates them; when weight_dtype == dtype - it is an identity, so an optimization barrier has to stand in for it. - """ + """Sparse MoE keeps the FSDP kernel gather from being hoisted out of a scanned decoder.""" fake_moe = SimpleNamespace( config=SimpleNamespace( shard_exp_on_fsdp=False, @@ -838,16 +829,12 @@ def fake_shard_map(function, *, mesh, in_specs, out_specs, check_vma): None, ) - # Identity, not equality: the stand-in kernels all share a shape, so SimpleNamespace - # equality cannot tell a barrier output from the raw kernel it replaced. def was_sharded(value): return any(value is candidate for candidate in sharded) if expect_barrier: assert barrier.call_count == 1 assert all(a is b for a, b in zip(barrier.call_args.args[0], (w0, w1, wo), strict=True)) - # The constraint must consume the barrier's outputs; a barrier whose results are - # dropped still lets XLA reorder the gather above the per-layer slice. assert all(was_sharded(kernel) for kernel in barriered) assert not any(was_sharded(kernel) for kernel in (w0, w1, wo)) else: diff --git a/tests/unit/nnx_decoders_test.py b/tests/unit/nnx_decoders_test.py index fec69287b2..624bb67b74 100644 --- a/tests/unit/nnx_decoders_test.py +++ b/tests/unit/nnx_decoders_test.py @@ -1005,7 +1005,7 @@ class TestQwen3NextScannableBlock(unittest.TestCase): @classmethod def _build(cls, layer_idx_offset=0, **overrides): - """Builds one scannable block; overrides go to the config, layer_idx_offset to the block.""" + """Builds one scannable block; overrides go to the config.""" cfg = _make_config(**{**cls.MODEL_CONFIG, **overrides}) block = cls.BLOCK_CLS( config=cfg, @@ -1191,15 +1191,11 @@ def test_decoder_keeps_layers_past_the_last_whole_block(self): ) -# Qwen3.5 repeats Qwen3-Next's hybrid attention period and reuses its scannable -# block, so it gets the same coverage over its own layer classes and config. _QWEN3_5_CONFIG = { **_QWEN3_NEXT_CONFIG, "run_name": "qwen3_5_scannable_block_test", "model_name": "qwen3.5-35b-a3b", - # Qwen3.5 uses MRoPE, whose sections must cover rotary_dim / 2 frequencies. The - # shipped [11, 11, 10] is sized for the real head_dim (256), not this tiny one - # (32 * partial_rotary_factor 0.25 / 2 = 4). + # MRoPE sections must cover rotary_dim / 2, which is 4 here, not the shipped [11, 11, 10]. "mrope_section": [2, 1, 1], } diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index 3113ed6467..c8d80f89ee 100644 --- a/tests/unit/param_mapping_test.py +++ b/tests/unit/param_mapping_test.py @@ -189,14 +189,11 @@ def test_qwen3_5_mapping_scanned(self): num_blocks, num_local = num_layers // cycle, cycle - 1 local_prefix = "params-decoder-layers-local_layers" global_prefix = "params-decoder-layers-global_layer" - # Linear attention only exists on the local layers, full attention only on the global one. self.assertIn(f"{local_prefix}-attention-in_proj_qkvz-kernel", mapping) self.assertIn(f"{global_prefix}-attention-attention-query-kernel", mapping) 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. - # A tuple of HF names is a composite source the hook fuses, not a stacking axis. 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) @@ -209,11 +206,8 @@ def test_qwen3_5_mapping_scanned(self): ) 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.language_model.layers.{cycle - 1}.self_attn.q_proj.weight") - # Fused experts: one HF tensor feeds both wi_0 and wi_1, so it is keyed by the pair - # and carries no extra per-expert axis (unlike Qwen3-Next). fused_local = mapping[(f"{local_prefix}-mlp-routed_experts-wi_0", f"{local_prefix}-mlp-routed_experts-wi_1")] self.assertEqual(len(fused_local), num_blocks) self.assertEqual(len(fused_local[0]), num_local) @@ -284,12 +278,7 @@ def test_qwen3_5_fp8_mapping_scanned(self): self.assertIn(key, hooks) def test_qwen3_5_composite_key_uses_the_nested_scan_axes(self): - """A composite (tuple) MaxText key must still be recognised as a nested block scan. - - Qwen3.5's fused gate_up_proj is keyed by the (wi_0, wi_1) pair, so a layout - check that only looks at `str` keys would place the (blocks, local) axes at - the leading positions and silently transpose the weights. - """ + """A composite (tuple) MaxText key must still be recognised as a nested block scan.""" cfg = mock.Mock() cfg.param_scan_axis = 1 local_key = "params-decoder-layers-local_layers-mlp-routed_experts-wi_0" diff --git a/tests/unit/router_replay_test.py b/tests/unit/router_replay_test.py index ce842396e0..403654a827 100644 --- a/tests/unit/router_replay_test.py +++ b/tests/unit/router_replay_test.py @@ -648,11 +648,7 @@ def test_per_layer_routing_is_not_broadcast_scanned(self): ) def test_per_layer_routing_reaches_the_scan_remainder_block(self): - """Layers past the last whole cycle live in a separate remainder block. - - With 3 layers and a cycle of 2 the scan covers layers 0-1 and layer 2 is - applied afterwards, so the remainder needs its own slice of the routing. - """ + """Layers past the last whole cycle live in a separate remainder block.""" seq_len, batch_size, top_k, num_experts = 8, 1, 2, 4 cycle_interval, layers = 2, 3 cfg = self._qwen35_cfg( @@ -665,8 +661,6 @@ def test_per_layer_routing_reaches_the_scan_remainder_block(self): num_decoder_layers=layers, inhomogeneous_layer_cycle_interval=cycle_interval, ) - # Only the trailing (remainder) layer's routing differs between a and b, so - # the losses can only diverge if the remainder block replays its own slice. shared = jnp.zeros((_batch_rows(cfg), seq_len, layers, top_k), dtype=jnp.int32) a = shared.at[:, :, -1, :].set(1) b = shared.at[:, :, -1, :].set(3)