diff --git a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py b/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py index b0d0deed62..15163f88b8 100644 --- a/src/maxtext/integration/tunix/weight_mapping/raiden_unscan.py +++ b/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_` + 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. 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_` 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 "" + 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}." + ) + 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 diff --git a/tests/unit/raiden_unscan_test.py b/tests/unit/raiden_unscan_test.py new file mode 100644 index 0000000000..f6bef16bd2 --- /dev/null +++ b/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_` 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_` + 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_ level", + {"layers": {"kernel": jnp.zeros((_IN, 2, _OUT))}}, + 4, + 2, + r"does not have a 'layer_' 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()