From d6e313d2d0a1bbea14d95caa9e8cc046a89a4aa2 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Thu, 2 Jul 2026 18:49:06 +0000 Subject: [PATCH] [NNX] Checkpoint: persist and restore rngs/dropout across resumes The NNX-only rngs/dropout state is stripped from the Linen on-disk layout, so on resume it was refilled from a base key(0) -- every preemption reset the dropout RNG stream even though the dataset iterator was restored. Persist that state as a separate nnx_aux composite item (the items payload stays byte-identical to the Linen layout, preserving interchangeability) and restore it when present, falling back to deterministic defaults for Linen-trained or pre-fix checkpoints. Emergency managers keep the default-fill path. --- src/maxtext/common/checkpointing.py | 101 +++++++--- src/maxtext/common/train_state_nnx.py | 21 +- tests/unit/checkpointing_nnx_load_test.py | 189 ++++++++++++++---- .../unit/checkpointing_nnx_roundtrip_test.py | 165 +++++++++++++++ 4 files changed, 404 insertions(+), 72 deletions(-) create mode 100644 tests/unit/checkpointing_nnx_roundtrip_test.py diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c527427144..5af1467a88 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -194,6 +194,16 @@ def _make(): return jax.jit(_make, out_shardings=sharding)() +def _deep_merge(base, overlay): + """Recursively merges `overlay` into `base`, returning a new dict. Overlay wins on leaves.""" + if not (isinstance(base, dict) and isinstance(overlay, dict)): + return overlay if overlay is not None else base + out = dict(base) + for k, v in overlay.items(): + out[k] = _deep_merge(base.get(k), v) if k in base else v + return out + + def _populate_pure_dict_from_partial(abstract_pure, partial_concrete): """Fills `abstract_pure` with values from `partial_concrete` (by path), defaulting the rest. @@ -214,6 +224,25 @@ def _populate_pure_dict_from_partial(abstract_pure, partial_concrete): return _default_for_sds(abstract_pure) +def _all_materialized(tree): + """True if Orbax restored every leaf, i.e. none is left as a jax.ShapeDtypeStruct.""" + return not any(isinstance(leaf, jax.ShapeDtypeStruct) for leaf in jax.tree_util.tree_leaves(tree)) + + +def _linen_items_to_nnx(restored_linen, nnx_abstract_pure): + """Reshapes a restored Linen-layout `items` dict back into the NNX pure-dict state. + + Merges `nnx_aux` (rngs/dropout/batch stats) back onto the model when present and + fills anything absent with defaults. A checkpoint without `nnx_aux` (Linen-trained + or pre-persistence) leaves those leaves unmaterialized, so they get the defaults. + """ + partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored_linen) + aux = restored_linen.get("nnx_aux") + if aux and _all_materialized(aux): + partial_nnx = _deep_merge(partial_nnx, aux) + return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + + def _load_linen_checkpoint_into_nnx( path, abstract_nnx_state, @@ -223,12 +252,12 @@ def _load_linen_checkpoint_into_nnx( ): """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume). - Restores against a Linen-shape abstract, reshapes back via - `from_linen_checkpoint_dict`, then fills NNX-only rngs/dropout with defaults. + Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via + `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when + present, else fall back to defaults. """ max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}") - nnx_abstract_pure = abstract_nnx_state.to_pure_dict() - linen_abstract = train_state_nnx.to_linen_checkpoint_dict(nnx_abstract_pure) + linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state) ckptr = ocp.Checkpointer( ocp.PyTreeCheckpointHandler( restore_concurrent_gb=checkpoint_storage_concurrent_gb, @@ -240,8 +269,7 @@ def _load_linen_checkpoint_into_nnx( restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract) restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) restored = ckptr.restore(epath.Path(path), args=restored) - partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) - return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + return _linen_items_to_nnx(restored, abstract_nnx_state.to_pure_dict()) def _restore_emergency_linen_checkpoint_into_nnx( @@ -250,10 +278,14 @@ def _restore_emergency_linen_checkpoint_into_nnx( abstract_nnx_state, map_to_pspec, ): - """Restores an emergency Linen-layout checkpoint into an NNX state.""" + """Restores an emergency Linen-layout checkpoint into an NNX state. + + The `nnx_aux` subtree is stored inside `items`, so an emergency checkpoint + carries it too; it's restored when present and otherwise filled with + deterministic defaults. + """ max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}") - nnx_abstract_pure = abstract_nnx_state.to_pure_dict() - linen_abstract = train_state_nnx.to_linen_checkpoint_dict(nnx_abstract_pure) + linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state) restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract) checkpoint_args = ocp.args.PyTreeRestore( item=linen_abstract, @@ -261,8 +293,7 @@ def _restore_emergency_linen_checkpoint_into_nnx( partial_restore=True, ) restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state - partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) - return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + return _linen_items_to_nnx(restored, abstract_nnx_state.to_pure_dict()) def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights): @@ -340,6 +371,11 @@ def _load_full_state_from_path( if enable_orbax_v1: if source_checkpoint_layout == "orbax": + # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. + if isinstance(abstract_unboxed_pre_state, nnx.State): + return _load_linen_checkpoint_into_nnx( + path, abstract_unboxed_pre_state, checkpoint_storage_concurrent_gb, use_ocdbt, use_zarr3 + ) context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.ORBAX) with context: return ocp_v1.load_pytree(path, abstract_unboxed_pre_state) @@ -378,11 +414,8 @@ def combine_sharding(sds, shardings): use_ocdbt=use_ocdbt, use_zarr3=use_zarr3, ) - # NNX checkpoints are saved as pure dicts (see maybe_save_checkpoint); the - # restore target must match — a boxed nnx.State wouldn't. + # Only Linen TrainState reaches here; nnx.State returned above. restore_target = abstract_unboxed_pre_state - if isinstance(abstract_unboxed_pre_state, nnx.State): - restore_target = abstract_unboxed_pre_state.to_pure_dict() # Provide sharding info to ensure restoration returns JAX arrays (not NumPy arrays). restore_args = jax.tree_util.tree_map( lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding), @@ -792,27 +825,35 @@ def map_to_pspec(data): ) ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True) - # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. + # pure_nnx saves in the Linen on-disk layout; restore that layout (weights + + # opt_state + step + nnx_aux), restoring the grain iterator in place when + # present, then reshape it back into the NNX state. # (Emergency managers use their own restore path below.) if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance( checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), ): - checkpoint_path = str(checkpoint_manager.directory / str(step) / "items") - restored_nnx = _load_linen_checkpoint_into_nnx( - checkpoint_path, - abstract_unboxed_pre_state, - checkpoint_storage_concurrent_gb, - use_ocdbt, - use_zarr3, - ) + linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) + restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract) + checkpoint_args = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) + if ( + dataset_type == "grain" + and data_iterator + and not isinstance(data_iterator, PlaceHolderDataIterator) + and (checkpoint_manager.directory / str(step) / "iter").exists() + ): + restored, _ = _restore_grain_iterator( + checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data + ) + else: + restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) + restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state.to_pure_dict()) return ({"items": restored_nnx}, None) if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance( checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), ): - checkpoint_path = str(checkpoint_manager.directory / str(step)) restored = _restore_emergency_linen_checkpoint_into_nnx( checkpoint_manager, step, @@ -824,11 +865,8 @@ def map_to_pspec(data): None, ) - # Convert nnx.State to pure dict to match how checkpoints are saved for NNX + # Only Linen TrainState reaches here; the NNX cases returned above. restore_target = abstract_unboxed_pre_state - if isinstance(abstract_unboxed_pre_state, nnx.State): - restore_target = abstract_unboxed_pre_state.to_pure_dict() - restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target) checkpoint_args = ocp.args.PyTreeRestore( item=restore_target, @@ -836,7 +874,6 @@ def map_to_pspec(data): partial_restore=True, ) - checkpoint_path = str(checkpoint_manager.directory / str(step)) match (checkpoint_manager, dataset_type, data_iterator): # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and @@ -1040,7 +1077,9 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) else: - state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) + # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout + # stream continues across resumes instead of resetting to a base key. + state = train_state_nnx.to_checkpoint_dict(state) # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic. # This occurs if this function was called: diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index a123956a73..3bd0352b58 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -66,7 +66,9 @@ def apply_gradients(self, grads: Any, **kwargs): # 2. weights: model/... -> params/params/... (Linen `params` collection) # 3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState, # and mu/nu wrapped under the `params` collection -# NNX-only rngs/dropout state is dropped (Linen never had it). +# NNX-only state Linen never had (rngs/dropout and batch stats) is split out by +# Variable type and stored under an `nnx_aux` subtree of `items` (see +# `to_checkpoint_dict`), so it survives a resume instead of resetting. _NNX_RNG_STATE_KEYS = ("rngs", "dropout") @@ -193,3 +195,20 @@ def from_linen_checkpoint_dict(linen_pure_dict): if optimizer: result["optimizer"] = optimizer return result + + +def to_checkpoint_dict(state): + """Reshapes an nnx.State into the on-disk checkpoint layout, splitting by Variable type. + + Weights (nnx.Param) and the optimizer state map to the Linen + params/opt_state/step layout, so pure_nnx and Linen checkpoints stay + interchangeable. rngs/dropout (nnx.RngState) and batch stats (nnx.BatchStat) go + under an `nnx_aux` subtree so they survive a resume. Works on a concrete state + (save) or an abstract state (restore target). + """ + params, rng_state, batch_stats, rest = nnx.split_state(state, nnx.Param, nnx.RngState, nnx.BatchStat, ...) + linen = to_linen_checkpoint_dict(nnx.merge_state(params, rest).to_pure_dict()) + aux = nnx.merge_state(rng_state, batch_stats).to_pure_dict() + if aux: + linen["nnx_aux"] = aux + return linen diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 04aa898bde..d4f27e7621 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/tests/unit/checkpointing_nnx_load_test.py @@ -36,6 +36,17 @@ def __init__(self, rngs: nnx.Rngs): self.linear = nnx.Linear(2, 1, rngs=rngs) +class _ModelDropout(nnx.Module): + """Linear + dropout, so the state carries rngs that split out into nnx_aux.""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 1, rngs=rngs) + self.dropout = nnx.Dropout(rate=0.5, rngs=rngs) + + def __call__(self, x, deterministic=False): + return self.dropout(self.linear(x), deterministic=deterministic) + + def _abstract_nnx_state(): """Build an nnx.State from a TrainStateNNX — same shape that pre_train passes in.""" model = _Model(rngs=nnx.Rngs(0)) @@ -43,58 +54,52 @@ def _abstract_nnx_state(): return nnx.state(train_state_nnx.TrainStateNNX(model, optimizer)) +def _dropout_state(seed): + """A concrete TrainStateNNX state for `_ModelDropout` with its dropout stream advanced.""" + model = _ModelDropout(nnx.Rngs(seed)) + state = train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param)) + grads = nnx.grad(lambda m: jnp.mean(m(jnp.ones((4, 2)), deterministic=False) ** 2))(state.model) + state.apply_gradients(grads) # advances step + dropout rng count off the base 0 + return nnx.state(state) + + +def _abstract_dropout_state(): + def make(): + model = _ModelDropout(nnx.Rngs(9)) + return nnx.state(train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param))) + + return nnx.eval_shape(make) + + class TestLoadStateIfPossibleNNX(unittest.TestCase): """Cover the NNX branches in load_state_if_possible.""" - def test_emergency_linen_restore_converts_back_to_nnx(self): - kernel_abstract = jax.ShapeDtypeStruct((2, 1), jnp.float32) - step_abstract = jax.ShapeDtypeStruct((), jnp.uint32) - abstract_nnx_pure = { - "model": { - "linear": {"kernel": kernel_abstract}, - "dropout": {"rngs": {"params": {"count": step_abstract}}}, - }, - "optimizer": {"step": step_abstract}, - } - abstract_nnx_state = mock.Mock() - abstract_nnx_state.to_pure_dict.return_value = abstract_nnx_pure - restored_linen = { - "params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, - "step": jnp.asarray(7, dtype=jnp.int32), - } - checkpoint_manager = mock.Mock() - checkpoint_manager.restore.return_value = mock.Mock(state=restored_linen) + def test_emergency_restore_recovers_nnx_aux(self): + """Emergency restore reshapes back to NNX and recovers rngs/dropout from items/nnx_aux.""" + concrete = _dropout_state(0) + orig_count = int(concrete.to_pure_dict()["model"]["dropout"]["rngs"]["count"]) + self.assertGreater(orig_count, 0) + # The single state tree an emergency manager writes (weights + opt_state + step + nnx_aux). + saved_linen = train_state_nnx.to_checkpoint_dict(concrete) + self.assertIn("nnx_aux", saved_linen) + checkpoint_manager = mock.Mock() + checkpoint_manager.restore.return_value = mock.Mock(state=saved_linen) restored = checkpointing._restore_emergency_linen_checkpoint_into_nnx( # pylint: disable=protected-access checkpoint_manager, 14, - abstract_nnx_state, - lambda leaf: ocp.type_handlers.ArrayRestoreArgs( - global_shape=leaf.shape, - dtype=leaf.dtype, - ), + _abstract_dropout_state(), + lambda leaf: ocp.type_handlers.ArrayRestoreArgs(global_shape=leaf.shape, dtype=leaf.dtype), ) - checkpoint_manager.restore.assert_called_once() - restore_args = checkpoint_manager.restore.call_args.kwargs["args"].state - self.assertEqual( - set(restore_args.item.keys()), - {"params", "step"}, - ) - self.assertNotIn("model", restore_args.item) - self.assertNotIn("optimizer", restore_args.item) - self.assertTrue(restore_args.partial_restore) + # The restore target the manager was handed carries nnx_aux, so emergency checkpoints persist it. + restore_target = checkpoint_manager.restore.call_args.kwargs["args"].state.item + self.assertIn("nnx_aux", restore_target) + # Reshaped back to NNX, with the dropout stream recovered (not reset to 0). self.assertIn("model", restored) self.assertIn("optimizer", restored) self.assertNotIn("params", restored) - self.assertNotIn("opt_state", restored) - self.assertTrue(bool(jnp.array_equal(restored["model"]["linear"]["kernel"], jnp.ones((2, 1))))) - self.assertEqual(restored["optimizer"]["step"].dtype, jnp.uint32) - self.assertEqual(int(restored["optimizer"]["step"]), 7) - self.assertEqual( - restored["model"]["dropout"]["rngs"]["params"]["count"].shape, - (), - ) + self.assertEqual(int(restored["model"]["dropout"]["rngs"]["count"]), orig_count) def test_load_parameters_from_path_splits_nnx_state_for_param_view(self): """When abstract_unboxed_pre_state is an nnx.State, the function must call @@ -185,5 +190,109 @@ def test_linen_layout_params_restore_into_nnx_state(self): self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"])) +class TestToCheckpointDict(unittest.TestCase): + """train_state_nnx.to_checkpoint_dict splits the NNX state by Variable type.""" + + def test_splits_params_optimizer_and_nnx_aux(self): + model = _ModelDropout(nnx.Rngs(0)) + state = nnx.state(train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, optax.adam(1e-3), wrt=nnx.Param))) + ckpt = train_state_nnx.to_checkpoint_dict(state) + # Learnable weights land in the Linen params collection; optimizer + step alongside. + self.assertEqual(set(ckpt["params"]["params"].keys()), {"linear"}) + self.assertIn("opt_state", ckpt) + self.assertIn("step", ckpt) + # rngs/dropout land in nnx_aux, not in params. + self.assertIn("dropout", ckpt["nnx_aux"]["model"]) + self.assertNotIn("dropout", ckpt["params"]["params"]) + + def test_batch_stats_split_out_from_weights(self): + """BatchNorm scale/bias are weights; its running mean/var are batch stats -> nnx_aux.""" + + class _BN(nnx.Module): + + def __init__(self, rngs): + self.bn = nnx.BatchNorm(2, rngs=rngs) + + state = nnx.state(train_state_nnx.TrainStateNNX(_BN(nnx.Rngs(0)), None)) + ckpt = train_state_nnx.to_checkpoint_dict(state) + self.assertEqual(set(ckpt["params"]["params"]["bn"].keys()), {"scale", "bias"}) + self.assertEqual(set(ckpt["nnx_aux"]["model"]["bn"].keys()), {"mean", "var"}) + + def test_no_nnx_aux_when_state_has_none(self): + state = _abstract_nnx_state() # plain linear, no rngs/batch stats + self.assertNotIn("nnx_aux", train_state_nnx.to_checkpoint_dict(state)) + + +class TestDeepMerge(unittest.TestCase): + """checkpointing._deep_merge overlays leaves onto a base dict.""" + + def test_overlay_wins_and_bases_survive(self): + base = {"model": {"linear": {"kernel": 1}, "rngs": {"count": 0}}} + overlay = {"model": {"rngs": {"count": 99}}} + merged = checkpointing._deep_merge(base, overlay) # pylint: disable=protected-access + self.assertEqual(merged["model"]["linear"]["kernel"], 1) # untouched + self.assertEqual(merged["model"]["rngs"]["count"], 99) # overlaid + + def test_does_not_mutate_inputs(self): + base = {"a": {"b": 1}} + checkpointing._deep_merge(base, {"a": {"c": 2}}) # pylint: disable=protected-access + self.assertEqual(base, {"a": {"b": 1}}) + + def test_non_dict_leaves_prefer_overlay(self): + """Leaf vs leaf: overlay wins, but a None overlay keeps the base.""" + self.assertEqual(checkpointing._deep_merge(1, 2), 2) # pylint: disable=protected-access + self.assertEqual(checkpointing._deep_merge(1, None), 1) # pylint: disable=protected-access + + def test_adds_keys_absent_from_base(self): + merged = checkpointing._deep_merge({"a": 1}, {"b": 2}) # pylint: disable=protected-access + self.assertEqual(merged, {"a": 1, "b": 2}) + + +class TestLinenItemsToNnx(unittest.TestCase): + """checkpointing._linen_items_to_nnx reshapes restored items back to NNX, handling nnx_aux.""" + + def _abstract_pure(self): + return { + "model": { + "linear": {"kernel": jax.ShapeDtypeStruct((2, 1), jnp.float32)}, + "dropout": {"rngs": {"count": jax.ShapeDtypeStruct((), jnp.uint32)}}, + }, + "optimizer": {"step": jax.ShapeDtypeStruct((), jnp.uint32)}, + } + + def test_all_materialized_detects_surviving_sds(self): + # pylint: disable=protected-access + self.assertTrue(checkpointing._all_materialized({"a": jnp.ones((2,))})) + self.assertFalse(checkpointing._all_materialized({"a": jax.ShapeDtypeStruct((2,), jnp.float32)})) + + def test_materialized_aux_is_merged(self): + restored = { + "params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, + "step": jnp.asarray(3, jnp.int32), + "nnx_aux": {"model": {"dropout": {"rngs": {"count": jnp.asarray(42, jnp.uint32)}}}}, + } + out = checkpointing._linen_items_to_nnx(restored, self._abstract_pure()) # pylint: disable=protected-access + self.assertEqual(int(out["model"]["dropout"]["rngs"]["count"]), 42) # restored, not defaulted + self.assertTrue(jnp.array_equal(out["model"]["linear"]["kernel"], jnp.ones((2, 1)))) + + def test_unmaterialized_aux_falls_back_to_default(self): + """An old checkpoint returns nnx_aux as ShapeDtypeStruct -> treated as absent -> default.""" + restored = { + "params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, + "step": jnp.asarray(3, jnp.int32), + "nnx_aux": {"model": {"dropout": {"rngs": {"count": jax.ShapeDtypeStruct((), jnp.uint32)}}}}, + } + out = checkpointing._linen_items_to_nnx(restored, self._abstract_pure()) # pylint: disable=protected-access + self.assertEqual(int(out["model"]["dropout"]["rngs"]["count"]), 0) # _default_for_sds + + def test_no_aux_key_falls_back_to_default(self): + restored = { + "params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, + "step": jnp.asarray(3, jnp.int32), + } + out = checkpointing._linen_items_to_nnx(restored, self._abstract_pure()) # pylint: disable=protected-access + self.assertEqual(int(out["model"]["dropout"]["rngs"]["count"]), 0) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/checkpointing_nnx_roundtrip_test.py b/tests/unit/checkpointing_nnx_roundtrip_test.py new file mode 100644 index 0000000000..920a04a0c2 --- /dev/null +++ b/tests/unit/checkpointing_nnx_roundtrip_test.py @@ -0,0 +1,165 @@ +# Copyright 2025-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. + +"""End-to-end NNX checkpoint round trips through the real save/restore stack. + +Drives create_orbax_checkpoint_manager -> maybe_save_checkpoint -> +load_state_if_possible against a real on-disk Orbax checkpoint (no mocks), which +is what exercises the `items/nnx_aux` subtree and the RNG/dropout persistence the +way training actually hits them, including the fallback for a checkpoint that has +no nnx_aux (a Linen-trained or pre-persistence checkpoint). +""" + +import os +import shutil +import tempfile +import unittest +from types import SimpleNamespace + +from flax import nnx +import jax +import jax.numpy as jnp +from maxtext.common import checkpointing +from maxtext.common import train_state_nnx +import optax + + +class _Model(nnx.Module): + """Linear + dropout, so the state carries rngs/dropout that go under nnx_aux on save.""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 1, rngs=rngs) + self.dropout = nnx.Dropout(rate=0.5, rngs=rngs) + + def __call__(self, x, deterministic=False): + return self.dropout(self.linear(x), deterministic=deterministic) + + +class _PlainModel(nnx.Module): + """Linear only (no dropout), so its checkpoint carries no nnx_aux, like a Linen-trained one.""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 1, rngs=rngs) + + +_TX = optax.adam(1e-3) + + +def _config(): + """Minimal config with the fields save/restore reads for a pure_nnx run.""" + return SimpleNamespace( + pure_nnx=True, + enable_diloco=False, + enable_checkpointing=True, + enable_continuous_checkpointing=False, + enable_emergency_checkpoint=False, + enable_autocheckpoint=False, + checkpoint_period=1, + local_checkpoint_period=0, + async_checkpointing=False, + dataset_type="tfds", + lora=None, + checkpoint_storage_target_data_file_size_bytes=checkpointing.DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE, + elastic_enabled=False, + ) + + +def _abstract_state(model_cls): + """An abstract (ShapeDtypeStruct) nnx.State for `model_cls`, the restore blueprint. + + Mirrors what get_abstract_state hands load_state_if_possible: SDS leaves with a + replicated sharding so Orbax can build restore args in single- or multi-host CI. + """ + mesh = jax.sharding.Mesh(jax.devices(), ("x",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + + def make(): + model = model_cls(nnx.Rngs(9)) + return nnx.state(train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param))) + + abstract = nnx.eval_shape(make) + return jax.tree.map( + lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=sharding) if hasattr(x, "shape") else x, + abstract, + ) + + +class TestNNXCheckpointRoundTrip(unittest.TestCase): + """Real save->restore cycles covering rng persistence and fallback.""" + + def setUp(self): + self._dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self._dir, ignore_errors=True) + + def _save_step(self): + """Trains one step (advances step + dropout rng) and saves; returns (manager, orig_count).""" + manager = checkpointing.create_orbax_checkpoint_manager( + os.path.join(self._dir, "ckpt"), + enable_checkpointing=True, + use_async=False, + save_interval_steps=1, + dataset_type="tfds", + ) + model = _Model(nnx.Rngs(0)) + state = train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param)) + grads = nnx.grad(lambda m: jnp.mean(m(jnp.ones((4, 2)), deterministic=False) ** 2))(state.model) + state.apply_gradients(grads) + orig_count = int(nnx.state(state).to_pure_dict()["model"]["dropout"]["rngs"]["count"]) + self.assertGreater(orig_count, 0) # dropout actually advanced the stream + checkpointing.maybe_save_checkpoint(manager, nnx.state(state), _config(), data_iterator=None, step=1) + manager.wait_until_finished() + return manager, orig_count + + def _restore(self, manager, model_cls): + """Restores the saved checkpoint into an abstract state for `model_cls`.""" + full, _ = checkpointing.load_state_if_possible( + manager, + data_iterator=None, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=_abstract_state(model_cls), + dataset_type="tfds", + maxtext_config=_config(), + ) + return full["items"] + + def test_rng_dropout_state_survives_round_trip(self): + """The dropout rng stream continues across a save/restore instead of resetting to 0.""" + manager, orig_count = self._save_step() + restored = self._restore(manager, _Model) + # Recovered from items/nnx_aux; a lost stream would default back to 0. + self.assertEqual(int(restored["model"]["dropout"]["rngs"]["count"]), orig_count) + self.assertEqual(int(restored["optimizer"]["step"]), 1) + + def test_old_checkpoint_without_nnx_aux_still_loads(self): + """A checkpoint with no nnx_aux (Linen-trained / pre-persistence) loads; rng gets the base default.""" + manager = checkpointing.create_orbax_checkpoint_manager( + os.path.join(self._dir, "ckpt"), + enable_checkpointing=True, + use_async=False, + save_interval_steps=1, + dataset_type="tfds", + ) + model = _PlainModel(nnx.Rngs(0)) # no dropout -> to_checkpoint_dict writes no nnx_aux + state = train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param)) + checkpointing.maybe_save_checkpoint(manager, nnx.state(state), _config(), data_iterator=None, step=1) + manager.wait_until_finished() + # Restore into an abstract that DOES have dropout: the absent rng falls back to the base default. + restored = self._restore(manager, _Model) + self.assertEqual(int(restored["model"]["dropout"]["rngs"]["count"]), 0) + + +if __name__ == "__main__": + unittest.main()