perf(qwen3-next): nested scan over hybrid attention with per-layer remat

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -67,7 +67,7 @@
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS
 from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
-from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns
+from maxtext.checkpoint_conversion.utils.tensor_handling import apply_hook_fns, nesting_depth, slice_shape, stacked_axes
 from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys
 from maxtext.inference.inference_utils import str2bool
 from maxtext.layers import quantizations
@@ -344,66 +344,59 @@ def get_maxtext_model_info(config):
 
 
 def _build_multi_axis_stacked_tensor(
-    hf_source_keys: List[List[str]],
+    hf_source_keys: List[Any],
     tensor_getter_fn: Callable[[str], np.ndarray],
     hook_fns: Any,
     target_shape: tuple,
     config,
     mt_key: str = "",
 ) -> np.ndarray:
-  """Builds a MaxText tensor by stacking HF weights along two axes.
+  """Builds a MaxText tensor by stacking HF weights along several axes.
 
-  Two cases share this helper:
+  Three layouts share this helper, distinguished by how deeply ``hf_source_keys``
+  is nested and by whether ``mt_key`` names a nested block scan:
     * MoE expert stacking: outer=experts, inner=layers, placed at the LEADING two
       axes -> shape (num_experts, num_layers, ...).
-    * Gemma4 nested block scan (``scanned_blocks-local_layers``): the block's local
-      layers are an inner scan nested inside the outer block scan, so the two
-      stacked axes live at ``(param_scan_axis, param_scan_axis + 1)`` rather than
-      the leading axes, e.g. (emb, blocks, local, ...).
+    * Nested block scan (``...-local_layers-...``, gemma4 and qwen3-next): the
+      block's local layers are an inner scan nested inside the outer block scan,
+      so the two stacked axes live at ``(param_scan_axis, param_scan_axis + 1)``
+      rather than the leading axes, e.g. (emb, blocks, local, ...).
+    * Both at once (qwen3-next's routed experts, which are expert-stacked inside
+      a nested block scan): three axes at ``(0, param_scan_axis,
+      param_scan_axis + 1)``, e.g. (experts, blocks, local, ...).
 
   Args:
-      hf_source_keys: A nested (2D) list of Hugging Face parameter names.
-                      The outer list is the outer stack axis, the inner list the
-                      inner stack axis.
+      hf_source_keys: A nested list of Hugging Face parameter names, one level of
+        nesting per stacked axis, outermost axis first. A ``tuple`` of names is a
+        composite HF source that the hook fuses into a single leaf, not an axis.
       tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array).
       hook_fns: The hook function(s) to apply to each individual weight.
       target_shape: The final shape of the target MaxText tensor.
       config: The MaxText pyconfig object.
-      mt_key: The MaxText parameter key, used to detect the gemma4 nested-scan case.
+      mt_key: The MaxText parameter key, used to detect the nested-scan case.
 
   Returns:
       The final, assembled NumPy array for the MaxText parameter.
   """
-  # Gemma4's local layers are an inner scan nested inside the block scan, so the
-  # two stacked axes go at (param_scan_axis, param_scan_axis + 1); everything else
-  # (MoE experts/layers) stacks at the leading axes.
-  if isinstance(mt_key, str) and "scanned_blocks-local_layers" in mt_key:
-    outer_axis, inner_axis = config.param_scan_axis, config.param_scan_axis + 1
-  else:
-    outer_axis, inner_axis = 0, 1
-
+  depth = nesting_depth(hf_source_keys)
+  axes = stacked_axes(mt_key, config, depth)
   # The hook function needs the shape of an individual slice: target_shape with
-  # the two stacked axes removed.
-  stacked_axes = {outer_axis, inner_axis}
-  mt_slice_shape = tuple(d for i, d in enumerate(target_shape) if i not in stacked_axes)
-
-  all_outer_tensors = []
-  # Outer loop iterates the outer stack axis (experts, or blocks for gemma4 local)
-  for inner_keys in hf_source_keys:
-    inner_tensors = []
-    # Inner loop iterates the inner stack axis (layers, or local layers for gemma4)
-    for hf_key_single in inner_keys:
-      if isinstance(hf_key_single, (list, tuple)):
-        hf_tensor_numpy = tuple(tensor_getter_fn(k) for k in hf_key_single)
+  # the stacked axes removed.
+  mt_slice_shape = slice_shape(target_shape, axes)
+
+  def gather(keys, level):
+    if level == depth:
+      if isinstance(keys, (list, tuple)):
+        raw = tuple(tensor_getter_fn(k) for k in keys)
       else:
-        hf_tensor_numpy = tensor_getter_fn(hf_key_single)
-      inner_tensors.append(apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns))
-    all_outer_tensors.append(np.stack(inner_tensors, axis=0))
-  stacked = np.stack(all_outer_tensors, axis=0)  # (outer, inner, *slice)
-
-  # Move the (outer, inner) axes from leading positions to their targets.
-  if (outer_axis, inner_axis) != (0, 1):
-    stacked = np.moveaxis(stacked, (0, 1), (outer_axis, inner_axis))
+        raw = tensor_getter_fn(keys)
+      return apply_hook_fns(raw, mt_slice_shape, hook_fns)
+    # Stack with the axes leading, outermost first; they are moved into place below.
+    return np.stack([gather(sub, level + 1) for sub in keys], axis=0)
+
+  stacked = gather(hf_source_keys, 0)
+  if axes != tuple(range(depth)):
+    stacked = np.moveaxis(stacked, tuple(range(depth)), axes)
   return stacked
 
 

src/maxtext/checkpoint_conversion/utils/load_dynamic.py

@@ -231,6 +231,7 @@ def tensor_getter(key):
         hook_fn,
         target_leaf,
         maxtext_config,
+        mt_key,
     )
 
     # Execute transformation and assign to flat_restored

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -1402,100 +1402,81 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F
   }
 
   if scan_layers:
-    # 2. Scan over block cycles
-    for block_idx in range(layer_cycle_interval):
-      hf_indices = list(range(block_idx, num_main_layers, layer_cycle_interval))
-      prefix = f"params-decoder-layers-layer_{block_idx}"
+    # 2. Scanned blocks. One block covers a single period of the hybrid attention
+    # pattern: `layer_cycle_interval - 1` linear-attention (GatedDeltaNet) layers
+    # run as an inner scan, then one full-attention layer. The resulting params are:
+    #   layers-local_layers-*  -> nested [block][local]  (doubly scanned)
+    #   layers-global_layer-*  -> flat [block]           (block scan only; the
+    #     length-1 _scan_global_layer scan is a runtime memory boundary, not a
+    #     param stack)
+    # Routed-expert weights carry an additional leading expert axis, so they are
+    # nested one level deeper: [expert][block][local] and [expert][block].
+    # Qwen3NextScannableBlock requires the full-attention layer to be last in the
+    # period, so the local positions are 0..cycle-2 and the global position is
+    # cycle-1.
+    num_blocks = num_main_layers // layer_cycle_interval
+    local_positions = list(range(layer_cycle_interval - 1))
+    global_position = layer_cycle_interval - 1
+
+    def hf_layer(block_idx, position, suffix):
+      return f"model.layers.{block_idx * layer_cycle_interval + position}.{suffix}"
+
+    # (maxtext subkey, hf suffix) pairs shared by both the local and global layers.
+    shared_specs = [
+        ("input_layernorm-scale", "input_layernorm.weight"),
+        ("post_attention_layernorm-scale", "post_attention_layernorm.weight"),
+        ("mlp-routed_experts-gate-kernel", "mlp.gate.weight"),
+        ("mlp-shared_expert-wi_0-kernel", "mlp.shared_expert.gate_proj.weight"),
+        ("mlp-shared_expert-wi_1-kernel", "mlp.shared_expert.up_proj.weight"),
+        ("mlp-shared_expert-wo-kernel", "mlp.shared_expert.down_proj.weight"),
+        ("mlp-shared_expert_gate-kernel", "mlp.shared_expert_gate.weight"),
+    ]
+    # Linear (GatedDeltaNet) attention: only ever on the local layers.
+    local_specs = shared_specs + [
+        ("attention-in_proj_qkvz-kernel", "linear_attn.in_proj_qkvz.weight"),
+        ("attention-in_proj_ba-kernel", "linear_attn.in_proj_ba.weight"),
+        ("attention-conv1d-kernel", "linear_attn.conv1d.weight"),
+        ("attention-A_log", "linear_attn.A_log"),
+        ("attention-dt_bias", "linear_attn.dt_bias"),
+        ("attention-norm-rms_norm-scale", "linear_attn.norm.weight"),
+        ("attention-out_proj-kernel", "linear_attn.out_proj.weight"),
+    ]
+    # Full attention: only ever on the global layer.
+    global_specs = shared_specs + [
+        ("attention-attention-query-kernel", "self_attn.q_proj.weight"),
+        ("attention-attention-key-kernel", "self_attn.k_proj.weight"),
+        ("attention-attention-value-kernel", "self_attn.v_proj.weight"),
+        ("attention-attention-out-kernel", "self_attn.o_proj.weight"),
+        ("attention-attention-query_norm-scale", "self_attn.q_norm.weight"),
+        ("attention-attention-key_norm-scale", "self_attn.k_norm.weight"),
+    ]
+    expert_specs = [
+        ("mlp-routed_experts-wi_0", "gate_proj.weight"),
+        ("mlp-routed_experts-wi_1", "up_proj.weight"),
+        ("mlp-routed_experts-wo", "down_proj.weight"),
+    ]
 
-      # Layer norms
-      mapping[f"{prefix}-input_layernorm-scale"] = [  # pyrefly: ignore[bad-assignment]
-          f"model.layers.{i}.input_layernorm.weight" for i in hf_indices
-      ]  # pyrefly: ignore[bad-assignment]
-      mapping[f"{prefix}-post_attention_layernorm-scale"] = [  # pyrefly: ignore[bad-assignment]
-          f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
+    local_prefix = "params-decoder-layers-local_layers"
+    for subkey, suffix in local_specs:
+      mapping[f"{local_prefix}-{subkey}"] = [  # pyrefly: ignore[bad-assignment, no-matching-overload]
+          [hf_layer(b, p, suffix) for p in local_positions] for b in range(num_blocks)
+      ]
+    for subkey, suffix in expert_specs:
+      mapping[f"{local_prefix}-{subkey}"] = [  # pyrefly: ignore[bad-assignment, no-matching-overload]
+          [[hf_layer(b, p, f"mlp.experts.{e}.{suffix}") for p in local_positions] for b in range(num_blocks)]
+          for e in range(num_experts)
       ]
 
-      # Handle Interleaved Attention (Linear vs Full)
-      is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0
-
-      if is_full_attention_layer:
-        mapping.update(  # pyrefly: ignore[no-matching-overload]
-            {
-                f"{prefix}-attention-attention-query-kernel": [
-                    f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-attention-key-kernel": [
-                    f"model.layers.{i}.self_attn.k_proj.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-attention-value-kernel": [
-                    f"model.layers.{i}.self_attn.v_proj.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-attention-out-kernel": [
-                    f"model.layers.{i}.self_attn.o_proj.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-attention-query_norm-scale": [
-                    f"model.layers.{i}.self_attn.q_norm.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-attention-key_norm-scale": [
-                    f"model.layers.{i}.self_attn.k_norm.weight" for i in hf_indices
-                ],
-            }
-        )
-      else:
-        # Linear/Hybrid Attention Block
-        mapping.update(  # pyrefly: ignore[no-matching-overload]
-            {
-                f"{prefix}-attention-in_proj_qkvz-kernel": [
-                    f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-in_proj_ba-kernel": [
-                    f"model.layers.{i}.linear_attn.in_proj_ba.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-conv1d-kernel": [f"model.layers.{i}.linear_attn.conv1d.weight" for i in hf_indices],
-                f"{prefix}-attention-A_log": [f"model.layers.{i}.linear_attn.A_log" for i in hf_indices],
-                f"{prefix}-attention-dt_bias": [f"model.layers.{i}.linear_attn.dt_bias" for i in hf_indices],
-                f"{prefix}-attention-norm-rms_norm-scale": [
-                    f"model.layers.{i}.linear_attn.norm.weight" for i in hf_indices
-                ],
-                f"{prefix}-attention-out_proj-kernel": [
-                    f"model.layers.{i}.linear_attn.out_proj.weight" for i in hf_indices
-                ],
-            }
-        )
-
-      # 3. Handle MLP: Gates and Shared Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices],
-              f"{prefix}-mlp-shared_expert-wi_0-kernel": [
-                  f"model.layers.{i}.mlp.shared_expert.gate_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert-wi_1-kernel": [
-                  f"model.layers.{i}.mlp.shared_expert.up_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert-wo-kernel": [
-                  f"model.layers.{i}.mlp.shared_expert.down_proj.weight" for i in hf_indices
-              ],
-              f"{prefix}-mlp-shared_expert_gate-kernel": [
-                  f"model.layers.{i}.mlp.shared_expert_gate.weight" for i in hf_indices
-              ],
-          }
-      )
-
-      # 4. Handle MoE Routed Experts
-      mapping.update(  # pyrefly: ignore[no-matching-overload]
-          {
-              f"{prefix}-mlp-routed_experts-wi_0": [
-                  [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts)
-              ],
-              f"{prefix}-mlp-routed_experts-wi_1": [
-                  [f"model.layers.{i}.mlp.experts.{e}.up_proj.weight" for i in hf_indices] for e in range(num_experts)
-              ],
-              f"{prefix}-mlp-routed_experts-wo": [
-                  [f"model.layers.{i}.mlp.experts.{e}.down_proj.weight" for i in hf_indices] for e in range(num_experts)
-              ],
-          }
-      )
+    global_prefix = "params-decoder-layers-global_layer"
+    for subkey, suffix in global_specs:
+      mapping[f"{global_prefix}-{subkey}"] = [  # pyrefly: ignore[bad-assignment, no-matching-overload]
+          hf_layer(b, global_position, suffix) for b in range(num_blocks)
+      ]
+    for subkey, suffix in expert_specs:
+      mapping[f"{global_prefix}-{subkey}"] = [  # pyrefly: ignore[bad-assignment, no-matching-overload]
+          [hf_layer(b, global_position, f"mlp.experts.{e}.{suffix}") for b in range(num_blocks)]
+          for e in range(num_experts)
+      ]
   else:
     # Unscanned layer mapping
     for i in range(num_main_layers):
@@ -1589,17 +1570,21 @@ def permute_conv(input_tensor, target_shape=None):
 
   layer_cycle_interval = maxtext_config.inhomogeneous_layer_cycle_interval
   num_main_layers = config["num_hidden_layers"]
-  loop_indices = range(layer_cycle_interval) if scan_layers else range(num_main_layers)
-
-  for i in loop_indices:
-    if scan_layers:
-      prefix = f"params-decoder-layers-layer_{i}"
-      block_idx = i
-    else:
-      prefix = f"params-decoder-layers_{i}"
-      block_idx = i % layer_cycle_interval
-    is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0
+  # Scanned blocks expose two prefixes -- the stacked local (linear-attention) layers
+  # and the single global (full-attention) layer -- rather than one prefix per position
+  # in the cycle. Unscanned models keep one prefix per decoder layer.
+  if scan_layers:
+    layer_prefixes = [
+        ("params-decoder-layers-local_layers", False),
+        ("params-decoder-layers-global_layer", True),
+    ]
+  else:
+    layer_prefixes = [
+        (f"params-decoder-layers_{i}", (i % layer_cycle_interval + 1) % layer_cycle_interval == 0)
+        for i in range(num_main_layers)
+    ]
 
+  for prefix, is_full_attention_layer in layer_prefixes:
     if is_full_attention_layer:
       for key in ["query", "key", "value", "out"]:
         hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel  # pyrefly: ignore[bad-assignment]