Update Qwen3.5 scan-remat pattern by re-using Qwen3-Next's nested block scan

src/maxtext/checkpoint_conversion/utils/tensor_handling.py

@@ -63,6 +63,8 @@ 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)``.
   """
+  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)

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_<slot>`
-      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_<n>` 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_<slot>`
+      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_<n>` 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_<slot>` 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 "<empty>"
-        raise ValueError(
-            f"unscan_layers: cycle_interval={cycle_interval} > 1, but key {'.'.join(key)!r} "
-            f"does not have a 'layer_<slot>' 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 "<empty>"
+      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_<slot>' 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
 

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,47 @@ 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`."""
+_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."""
+
+  container_idx: int
+  slot_idx: Optional[int]
+  consumed: int
+  nested: bool
+  trailing: bool
+
+
+_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 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
@@ -343,19 +365,34 @@ 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:
+    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)}")
+  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 +406,23 @@ 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 :]
+    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 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(
       "MaxTextVllmSampler: unrolled %d Qwen tensor components across %d layers for direct MaxText weight sync.",