[Raiden Weight Sync 2/7] Support inhomogeneous layer cycle interval in raiden_unscan

src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py

@@ -36,6 +36,7 @@
 simpler, single-axis case directly instead of adapting that function.
 """
 
+import re
 from typing import Any
 
 import jax
@@ -48,6 +49,7 @@ def unscan_layers(
     num_layers: int,
     layer_container: str = "layers",
     scan_axis: int = 1,
+    cycle_interval: int = 1,
 ) -> Any:
   """Splits `state`'s scanned `layer_container` axis into per-layer entries.
 
@@ -58,6 +60,13 @@ def unscan_layers(
       loop.
     layer_container: The pytree key holding the scanned per-layer params (MaxText's decoder body uses "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.
 
   Returns:
     A nested dict with `layer_container` keys replaced by `f"{layer_container}_{i}"` for each layer `i`, each holding
@@ -75,6 +84,14 @@ def unscan_layers(
   else:
     return state
 
+  if cycle_interval <= 0:
+    raise ValueError(f"unscan_layers: cycle_interval must be >= 1 (got {cycle_interval}).")
+
+  if cycle_interval > 1 and num_layers % cycle_interval != 0:
+    raise ValueError(
+        f"unscan_layers: num_layers ({num_layers}) must be cleanly divisible by " f"cycle_interval ({cycle_interval})."
+    )
+
   flat = flatten_dict(pure)
   new_flat = {}
   unscanned_count = 0
@@ -94,6 +111,26 @@ 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
+    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}."
+        )
+      suffix = suffix[1:]
     arr = getattr(value, "value", value)
 
     if arr is None or not hasattr(arr, "shape") or getattr(arr, "ndim", 0) <= scan_axis:
@@ -106,15 +143,18 @@ def unscan_layers(
       continue
     del value
 
-    if arr.shape[scan_axis] != num_layers:
+    # 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
+    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" num_layers={num_layers}."
+          f" {expected} (num_layers={num_layers}, cycle_interval={cycle_interval})."
       )
 
-    for i in range(num_layers):
+    for i in range(expected):
       sliced = jax.lax.index_in_dim(arr, i, axis=scan_axis, keepdims=False)
-      new_key = prefix + (f"{layer_container}_{i}",) + suffix
+      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
     del arr
     unscanned_count += 1

tests/unit/raiden_unscan_test.py

@@ -0,0 +1,257 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for `raiden_unscan.unscan_layers`.
+
+This transform decides the *names* Raiden binds. The trainer runs scanned
+(`scan_layers=True`); the sampler loads its MaxText model unscanned, and Raiden matches
+tensors by `jax.tree_util.keystr` path. Nothing downstream cross-checks the two name sets
+-- `raiden_handler._validate_metadata` only checks a single manifest's internal
+consistency (mesh rank, duplicate variable/layer keys, sharding specs) -- so a naming
+error here surfaces as weights that silently never transfer, not as an exception.
+
+The transform is pure pytree manipulation, so all of this runs on CPU in milliseconds.
+"""
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from maxtext.integration.tunix.weight_mapping import raiden_unscan
+import numpy as np
+import pytest
+
+# This transform exists only for the Tunix/Raiden weight-sync path, so it is graded with
+# the rest of that work. `tests/unit` is in cpu-post-training-unit's path list, so the
+# marker alone routes it there -- no file move needed (unlike tests/ or tests/integration,
+# which are not in that list and would be collected by no job at all).
+pytestmark = [pytest.mark.post_training]
+
+
+_NUM_LAYERS = 3
+_IN, _OUT, _VOCAB = 4, 8, 10
+
+
+def _unwrap(leaf):
+  """Reads a leaf's array whether or not it is wrapped in an `nnx.Param`."""
+  if isinstance(leaf, nnx.Variable):
+    return leaf[...]
+  return leaf
+
+
+def _names(tree) -> list[str]:
+  """The names Raiden binds: exactly what `raiden_synchronizer.flatten_weights` computes."""
+  return sorted(jax.tree_util.keystr(p) for p, _ in jax.tree_util.tree_leaves_with_path(tree))
+
+
+class ScannedInner(nnx.Module):
+  """One scanned param: the layer axis lives at axis 1 of a single array."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    self.kernel = nnx.Param(jnp.arange(_IN * num_layers * _OUT, dtype=jnp.float32).reshape(_IN, num_layers, _OUT))
+    self.scale = nnx.Param(jnp.arange(_OUT * num_layers, dtype=jnp.float32).reshape(_OUT, num_layers))
+
+
+class ScannedModel(nnx.Module):
+  """A trainer-side model: scanned `layers`, plus non-layer params that must pass through."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    self.layers = ScannedInner(num_layers)
+    self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN)))
+
+
+class UnscannedInner(nnx.Module):
+
+  def __init__(self):
+    self.kernel = nnx.Param(jnp.zeros((_IN, _OUT)))
+    self.scale = nnx.Param(jnp.zeros((_OUT,)))
+
+
+class UnscannedModel(nnx.Module):
+  """A sampler-side model: one submodule per layer, named `layers_0..N-1`."""
+
+  def __init__(self, num_layers: int = _NUM_LAYERS):
+    for i in range(num_layers):
+      setattr(self, f"layers_{i}", UnscannedInner())
+    self.embed = nnx.Param(jnp.zeros((_VOCAB, _IN)))
+
+
+class UnscanLayersTest(absltest.TestCase):
+
+  def _scanned_state(self, num_layers: int = _NUM_LAYERS):
+    return nnx.state(ScannedModel(num_layers), nnx.Param)
+
+  def test_names_match_an_unscanned_model_exactly(self):
+    """The point of the transform: trainer names must equal sampler names.
+
+    Raiden binds by `keystr` path on both sides, and nothing validates that the two sets
+    agree, so this is the assertion that a silent no-transfer would violate.
+    """
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    sampler_side = nnx.state(UnscannedModel(), nnx.Param)
+    self.assertEqual(_names(unscanned), _names(sampler_side))
+
+  def test_plain_dict_and_nnx_state_produce_identical_names(self):
+    """`unscan_layers` returns a plain nested dict, the sampler binds an `nnx.State`.
+
+    `keystr` renders both identically only because the transform rewraps leaves in
+    `nnx.Param`. Dropping that rewrap would rename every tensor (`['k']` vs `['k'].value`)
+    and break every transfer, so pin it.
+    """
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    self.assertIsInstance(jax.tree_util.tree_leaves(unscanned, is_leaf=lambda x: isinstance(x, nnx.Param))[0], nnx.Param)
+    self.assertTrue(all(n.endswith(".value") for n in _names(unscanned)), _names(unscanned))
+
+  def test_slices_carry_the_right_values(self):
+    """Layer i must receive index i of the scan axis -- not a transpose or an off-by-one."""
+    state = self._scanned_state()
+    original = np.asarray(state.to_pure_dict()["layers"]["kernel"])
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+
+    for i in range(_NUM_LAYERS):
+      got = unscanned[f"layers_{i}"]["kernel"]
+      got = np.asarray(_unwrap(got))
+      self.assertEqual(got.shape, (_IN, _OUT))
+      np.testing.assert_array_equal(got, original[:, i, :])
+
+  def test_rank_two_param_is_also_unscanned(self):
+    """A rank-2 scanned param (e.g. a norm scale) slices down to rank 1."""
+    unscanned = raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS)
+    for i in range(_NUM_LAYERS):
+      scale = unscanned[f"layers_{i}"]["scale"]
+      self.assertEqual(np.asarray(_unwrap(scale)).shape, (_OUT,))
+
+  def test_non_layer_entries_pass_through_unchanged(self):
+    """Embeddings and final norms have no layer axis and must survive untouched."""
+    state = self._scanned_state()
+    embed_before = np.asarray(state.to_pure_dict()["embed"])
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS)
+
+    self.assertIn("embed", unscanned)
+    embed_after = unscanned["embed"]
+    np.testing.assert_array_equal(np.asarray(_unwrap(embed_after)), embed_before)
+    self.assertNotIn("layers", unscanned)
+
+  def test_layer_count_mismatch_raises(self):
+    """A wrong num_layers must fail loudly rather than bind truncated weights."""
+    with self.assertRaisesRegex(ValueError, r"expected axis 1 to be 99 \(num_layers=99, cycle_interval=1\)"):
+      raiden_unscan.unscan_layers(self._scanned_state(), num_layers=99)
+
+  def test_already_unscanned_state_raises(self):
+    """The anti-silent-no-op guard.
+
+    Without it an already-unscanned (or wrongly-keyed) state would return unchanged and
+    bind under scanned names, transferring nothing with no error anywhere.
+    """
+    with self.assertRaisesRegex(ValueError, "found no scanned 'layers' entries"):
+      raiden_unscan.unscan_layers(nnx.state(UnscannedModel(), nnx.Param), num_layers=_NUM_LAYERS)
+
+  def test_wrong_layer_container_raises(self):
+    with self.assertRaisesRegex(ValueError, "found no scanned 'blocks' entries"):
+      raiden_unscan.unscan_layers(self._scanned_state(), num_layers=_NUM_LAYERS, layer_container="blocks")
+
+  def test_custom_scan_axis(self):
+    """`param_scan_axis` is configurable; axis 0 must slice the leading dim."""
+    state = {"layers": {"kernel": jnp.arange(_NUM_LAYERS * _OUT, dtype=jnp.float32).reshape(_NUM_LAYERS, _OUT)}}
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=_NUM_LAYERS, scan_axis=0)
+    for i in range(_NUM_LAYERS):
+      got = unscanned[f"layers_{i}"]["kernel"]
+      np.testing.assert_array_equal(np.asarray(_unwrap(got)), np.arange(i * _OUT, (i + 1) * _OUT))
+
+  def test_inhomogeneous_cycle_interval_unscans_correctly(self):
+    """Inhomogeneous scanned blocks unroll into interleaved per-layer entries."""
+    num_layers = 4
+    cycle_interval = 2
+    repeats = num_layers // cycle_interval  # 2
+
+    # slot 0 array (shape [_IN, 2, _OUT])
+    k0 = jnp.arange(_IN * repeats * _OUT, dtype=jnp.float32).reshape(_IN, repeats, _OUT)
+    # slot 1 array (shape [_IN, 2, _OUT])
+    k1 = (jnp.arange(_IN * repeats * _OUT, dtype=jnp.float32) + 1000).reshape(_IN, repeats, _OUT)
+
+    state = {
+        "layers": {
+            "layer_0": {"kernel": k0},
+            "layer_1": {"kernel": k1},
+        },
+        "embed": jnp.zeros((_VOCAB, _IN)),
+    }
+
+    unscanned = raiden_unscan.unscan_layers(state, num_layers=num_layers, cycle_interval=cycle_interval)
+    # Compare the names Raiden binds, not just the top-level keys: the `layer_<slot>` level
+    # has to be gone from *inside* each layer, and the `nnx.Param` rewrap has to survive.
+    expected_names = sorted(["['embed'].value"] + [f"['layers_{i}']['kernel'].value" for i in range(num_layers)])
+    self.assertEqual(_names(unscanned), expected_names)
+
+    # Check mapping: repeat i * cycle_interval + slot
+    # layer 0: repeat 0, slot 0
+    np.testing.assert_array_equal(np.asarray(_unwrap(unscanned["layers_0"]["kernel"])), np.asarray(k0[:, 0, :]))
+    # layer 1: repeat 0, slot 1
+    np.testing.assert_array_equal(np.asarray(_unwrap(unscanned["layers_1"]["kernel"])), np.asarray(k1[:, 0, :]))
+    # layer 2: repeat 1, slot 0
+    np.testing.assert_array_equal(np.asarray(_unwrap(unscanned["layers_2"]["kernel"])), np.asarray(k0[:, 1, :]))
+    # layer 3: repeat 1, slot 1
+    np.testing.assert_array_equal(np.asarray(_unwrap(unscanned["layers_3"]["kernel"])), np.asarray(k1[:, 1, :]))
+
+  def test_inhomogeneous_config_mismatches_raise(self):
+    """Every way a cycle config can disagree with the tree must fail loudly.
+
+    Each of these is otherwise silent: an out-of-range slot (or a `cycle_interval` the layer
+    count does not divide) makes two slots claim the same physical layer, so the later write
+    wins in `new_flat` and the overwritten weights never transfer; a missing `layer_<slot>`
+    level means config and tree describe different models. No downstream check would notice.
+    """
+    two_slots = {
+        "layers": {
+            "layer_0": {"kernel": jnp.zeros((_IN, 2, _OUT))},
+            "layer_1": {"kernel": jnp.zeros((_IN, 2, _OUT))},
+        }
+    }
+    cases = [
+        (
+            "indivisible num_layers",
+            two_slots,
+            5,
+            2,
+            r"num_layers \(5\) must be cleanly divisible by cycle_interval \(2\)",
+        ),
+        (
+            "slot at cycle_interval",
+            {"layers": {"layer_2": {"kernel": jnp.zeros((_IN, 2, _OUT))}}},
+            4,
+            2,
+            r"slot 2 parsed from key 'layers\.layer_2\.kernel' must be less than cycle_interval 2",
+        ),
+        (
+            "missing layer_<slot> level",
+            {"layers": {"kernel": jnp.zeros((_IN, 2, _OUT))}},
+            4,
+            2,
+            r"does not have a 'layer_<slot>' cycle prefix",
+        ),
+        ("cycle_interval below 1", two_slots, 4, 0, r"cycle_interval must be >= 1"),
+    ]
+    for name, state, num_layers, cycle_interval, expected in cases:
+      with self.subTest(name):
+        with self.assertRaisesRegex(ValueError, expected):
+          raiden_unscan.unscan_layers(state, num_layers=num_layers, cycle_interval=cycle_interval)
+
+  # TODO: Re-add streaming unscan tests (test_streaming_piece_count_and_parity,
+  # test_streaming_keys_per_piece_batching, test_streaming_leaves_are_nnx_param,
+  # test_streaming_already_unscanned_state_raises) once unscan_layers_streaming
+  # is implemented.
+
+
+if __name__ == "__main__":
+  absltest.main()