From 82ef7bf0a070014423170f110ff2a096e53eff32 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Wed, 22 Jul 2026 14:33:15 +0000 Subject: [PATCH] Unify the NNX and Linen restore paths in load_state_if_possible The NNX path had two early-return blocks that each duplicated the Linen restore flow, one for the standard manager and one for the emergency managers. Any change to restore logic had to be made in three places, and they had already drifted: the NNX grain branch dropped the iterator element that the Linen branch returns. pure_nnx saves in the Linen on-disk layout, so both paths restore the same tree. Convert the NNX abstract to that layout going in, run the existing match statement unchanged, and reshape what comes back on the way out. Linen behavior is unchanged. The weights-only branch's if/else folds into `_abstract_params`, which uses nnx.split_state rather than nnx.split -- the State-native API, without the unused GraphDef. That leaves the safetensors paths, which had the same duplication problem and three bugs behind it: The dynamic safetensors loader received the whole abstract state rather than the weights. For an NNX state that is fatal before any weight is read: the optimizer's opt_state is keyed by integer chain indices, and the loader flattens its target with a "." separator, so the load died in flatten_dict with `TypeError: sequence item 2: expected str instance, int found`. The loader now receives the weights, the way the weights-only Orbax load already does, and its result is unwrapped back into the NNX params state. A weight that no mapping covered came back as the unmaterialized leaf it went in as and reached the model as an untrained init value. Orbax does not flag this -- a shape conflict it rejects itself, but a weight the checkpoint simply does not carry it leaves as a placeholder. The weights-only Orbax load already makes that check; the dynamic load now makes it too, for both state types. `checkpoint_conversion_fn` arrives from config as a dotted string but was called directly, so it raised TypeError for any value that was set. It is now imported, and resolved before the checkpoint is read so a bad config fails immediately. Full-state safetensors also handles NNX now: the conversion fn produces MaxText's on-disk layout, which is what pure_nnx reads, so it goes through the same reshape as every other NNX restore. --- .../utils/load_dynamic.py | 11 +- src/maxtext/common/checkpointing.py | 160 +++++----- tests/unit/checkpointing_nnx_load_test.py | 282 +++++++++++++++++- tests/unit/checkpointing_test.py | 9 +- tests/unit/setup_initial_state_nnx_test.py | 2 +- 5 files changed, 369 insertions(+), 95 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py index 2032cf1c41..098bd403f1 100644 --- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py +++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py @@ -251,9 +251,12 @@ def tensor_getter(key): return {"params": restored_params} -def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config): +def load_safetensors_dynamic_state(path, abstract_params, maxtext_config): """Main entry point to dynamically build and load safetensors into MaxText format. + `abstract_params` is the weights of the target state -- Linen's `params` collection, or the + NNX params state -- not the full train state; the HF param mappings name weights only. + Splits execution into: 1. Deriving Mappings 2. Loading Sharded arrays directly to TPUs @@ -349,11 +352,7 @@ def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_con param_map_mt_to_hf, hook_fn_map_mt = get_hf_config_and_mappings(maxtext_config) max_logging.log(f"[1/3] Mappings derived in {time.time() - t_total:.2f}s") - target_tree = ( - abstract_unboxed_pre_state.to_pure_dict() - if isinstance(abstract_unboxed_pre_state, nnx.State) - else abstract_unboxed_pre_state.params - ) + target_tree = abstract_params.to_pure_dict() if isinstance(abstract_params, nnx.State) else abstract_params t1 = time.time() hf_state = load_sharded_hf_state(path) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index cf51083597..7860386e72 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -16,6 +16,7 @@ """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" import datetime +import importlib import time from typing import Any @@ -157,33 +158,58 @@ 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) - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored)) - return _linen_items_to_nnx(restored, abstract_nnx_state) + return _restored_linen_to_nnx(restored, abstract_nnx_state) -def _restore_emergency_linen_checkpoint_into_nnx( - checkpoint_manager, - step, - abstract_nnx_state, - map_to_pspec, -): - """Restores an emergency Linen-layout checkpoint into an NNX state. +def _restored_linen_to_nnx(restored_linen, abstract_nnx_state): + """Reshapes a restored Linen-layout tree into the NNX state. - The `nnx_aux` subtree is stored inside `items`, so an emergency checkpoint - carries it too; it's restored when present and otherwise kept at its fresh - init value. A genuinely-missing weight raises. + Raises if the checkpoint is missing a weight. Every NNX restore path ends here: the load + itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout. """ - max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}") - 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, - restore_args=restore_args, - partial_restore=True, - ) - restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored)) - return _linen_items_to_nnx(restored, abstract_nnx_state) + _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen)) + return _linen_items_to_nnx(restored_linen, abstract_nnx_state) + + +def _abstract_params(abstract_unboxed_pre_state): + """Returns the state's weights: the NNX Param subtree, or Linen's `params` collection.""" + if isinstance(abstract_unboxed_pre_state, nnx.State): + return nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0] + return abstract_unboxed_pre_state.params + + +def _bare_weights(tree): + """Strips the Flax `params` collection wrapper so weights compare at the same depth. + + A Linen params tree is the collection, an NNX one the bare weights; the dynamic + safetensors loader always returns the collection. + """ + return tree["params"] if isinstance(tree, dict) and len(tree) == 1 and "params" in tree else tree + + +def _resolve_conversion_fn(checkpoint_conversion_fn): + """Returns `checkpoint_conversion_fn` as a callable. + + Config carries it as a dotted string ("my_pkg.my_module.my_fn"), so it has to be imported + before it can be called. A callable is used as is. + """ + if checkpoint_conversion_fn is None: + raise ValueError( + "source_checkpoint_layout='safetensors' needs `checkpoint_conversion_fn` to map the " + "checkpoint's weights onto the model's, e.g. checkpoint_conversion_fn=my_pkg.my_module.my_fn." + ) + if callable(checkpoint_conversion_fn): + return checkpoint_conversion_fn + module_name, _, fn_name = str(checkpoint_conversion_fn).rpartition(".") + if not module_name: + raise ValueError(f"`checkpoint_conversion_fn` must be a dotted path to a function, got {checkpoint_conversion_fn!r}.") + try: + fn = getattr(importlib.import_module(module_name), fn_name, None) + except ImportError as e: + raise ValueError(f"Could not import `checkpoint_conversion_fn` {checkpoint_conversion_fn!r}: {e}") from e + if not callable(fn): + raise ValueError(f"`checkpoint_conversion_fn` {checkpoint_conversion_fn!r} is not a function.") + return fn def _load_full_state_from_path( @@ -226,6 +252,8 @@ def _load_full_state_from_path( with context: return ocp_v1.load_pytree(path, abstract_unboxed_pre_state) elif source_checkpoint_layout == "safetensors": + # Resolved first, so a bad config fails before the weights are read. + conversion_fn = _resolve_conversion_fn(checkpoint_conversion_fn) context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS) with context: metadata = ocp_v1.pytree_metadata(path) @@ -237,7 +265,11 @@ def combine_sharding(sds, shardings): sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings) pre_transformed_state = ocp_v1.load_pytree(path, sharded_abstract_state) - state = checkpoint_conversion_fn(pre_transformed_state) + state = conversion_fn(pre_transformed_state) + # The conversion fn returns MaxText's on-disk (Linen) layout, which is what pure_nnx reads, + # so NNX needs the same reshape as every other restore. An NNX state passes through. + if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(state, nnx.State): + state = _restored_linen_to_nnx(state, abstract_unboxed_pre_state) return state else: raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}") @@ -406,6 +438,11 @@ def load_state_if_possible( set. """ + # pure_nnx saves in the Linen on-disk layout, so every branch below loads the same tree Linen + # does: the NNX abstract is converted to that layout going in, and what comes back is reshaped + # into the NNX state on the way out. + is_nnx = isinstance(abstract_unboxed_pre_state, nnx.State) + if checkpoint_manager is not None: max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint") @@ -437,49 +474,9 @@ 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; 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), - ): - 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, _ = grain_utility.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)) - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_unboxed_pre_state, restored["items"])) - restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state) - return ({"items": restored_nnx}, None) - - if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance( - checkpoint_manager, - (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), - ): - restored = _restore_emergency_linen_checkpoint_into_nnx( - checkpoint_manager, - step, - abstract_unboxed_pre_state, - map_to_pspec, - ) - return ( - restored, - None, - ) - - # Only Linen TrainState reaches here; the NNX cases returned above. - restore_target = abstract_unboxed_pre_state + restore_target = ( + train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) if is_nnx else abstract_unboxed_pre_state + ) restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target) checkpoint_args = ocp.args.PyTreeRestore( item=restore_target, @@ -499,6 +496,8 @@ def map_to_pspec(data): ), ): restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state + if is_nnx: + restored = _restored_linen_to_nnx(restored, abstract_unboxed_pre_state) return ( restored, None, @@ -515,30 +514,43 @@ def map_to_pspec(data): and not isinstance(data_iterator, PlaceHolderDataIterator) and (checkpoint_manager.directory / str(step) / "iter").exists() ): - return grain_utility.restore_grain_iterator( + restored, iterator = grain_utility.restore_grain_iterator( checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data, ) + if is_nnx: + restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)} + return (restored, iterator) # Case 3: Default/Fallback case. # This case acts as a wildcard ('_') and matches if none of the preceding cases were met. case _: restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) + if is_nnx: + restored = {"items": _restored_linen_to_nnx(restored["items"], abstract_unboxed_pre_state)} return (restored, None) if source_checkpoint_layout == "safetensors_dynamic": path = load_parameters_from_path or load_full_state_from_path max_logging.log(f"Dynamic On-the-Fly Formatting: Loading SafeTensors from {path}") - return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config) + # Weights-only for both paths, so the loader gets the weights rather than the whole state: + # the HF param mappings name weights, and an NNX state hides them under `model`. + params = _abstract_params(abstract_unboxed_pre_state) + restored, restored_params = load_safetensors_dynamic_state(path, params, maxtext_config) + # A weight no HF mapping covered comes back unmaterialized and would reach the model as an + # untrained init value. Same check the Orbax weights-only load makes. + _raise_on_weight_mismatch(_bare_weights(params.to_pure_dict() if is_nnx else params), _bare_weights(restored_params)) + if is_nnx: + # The loader returns the Linen `params` collection; NNX holds bare weights, so unwrap it + # back into the params state, the shape load_params_from_path returns. + nnx.replace_by_pure_dict(params, restored_params["params"]) + return restored, params + return restored, restored_params elif load_parameters_from_path != "": - if isinstance(abstract_unboxed_pre_state, nnx.State): - _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...) - else: - params = abstract_unboxed_pre_state.params - + params = _abstract_params(abstract_unboxed_pre_state) restored_params = load_params_from_path( load_parameters_from_path, params, diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 163a5b1e72..0c2dae51ff 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/tests/unit/checkpointing_nnx_load_test.py @@ -21,6 +21,7 @@ from etils import epath from flax import nnx +from flax.training import train_state import jax import jax.numpy as jnp from maxtext.common import checkpointing @@ -83,13 +84,16 @@ def test_emergency_restore_recovers_nnx_aux(self): saved_linen = train_state_nnx.to_checkpoint_dict(concrete) self.assertIn("nnx_aux", saved_linen) - checkpoint_manager = mock.Mock() + checkpoint_manager = mock.Mock(spec=checkpointing.EmergencyCheckpointManager) 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_dropout_state(), - lambda leaf: ocp.type_handlers.ArrayRestoreArgs(global_shape=leaf.shape, dtype=leaf.dtype), + restored, _ = checkpointing.load_state_if_possible( + checkpoint_manager=checkpoint_manager, + data_iterator=None, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=_abstract_dropout_state(), + step=14, ) restored = restored.to_pure_dict() @@ -147,6 +151,74 @@ def test_load_parameters_from_path_uses_state_params_for_linen(self): forwarded_params = m.call_args[0][1] self.assertIs(forwarded_params, fake_state.params) + def test_safetensors_dynamic_hands_over_weights_and_returns_nnx_params(self): + """The loader gets the bare weights, and its result comes back as NNX params.""" + abstract = _abstract_nnx_state() + weights = {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}} + + with mock.patch.object( + checkpointing, "load_safetensors_dynamic_state", return_value=(None, {"params": weights}) + ) as m: + full, params = checkpointing.load_state_if_possible( + checkpoint_manager=None, + data_iterator=None, + load_parameters_from_path="gs://does-not-exist/hf", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=abstract, + source_checkpoint_layout="safetensors_dynamic", + maxtext_config=mock.Mock(), + ) + + self.assertIsNone(full) + # The loader sees the weights the way the HF mappings name them: no `model` prefix above them. + forwarded = m.call_args[0][1] + self.assertEqual(set(forwarded.to_pure_dict().keys()), {"linear"}) + # And the restored weights come back as an nnx.State, the shape a params-only load returns. + self.assertIsInstance(params, nnx.State) + self.assertTrue(jnp.array_equal(params.to_pure_dict()["linear"]["bias"], weights["linear"]["bias"])) + + def test_safetensors_dynamic_unmapped_weight_raises(self): + """A weight no HF mapping covered comes back unmaterialized, which must raise.""" + unmapped = { + "linear": {"kernel": jnp.ones((2, 1)), "bias": jax.ShapeDtypeStruct((1,), jnp.float32)}, + } + with mock.patch.object(checkpointing, "load_safetensors_dynamic_state", return_value=(None, {"params": unmapped})): + with self.assertRaises(ValueError) as ctx: + checkpointing.load_state_if_possible( + checkpoint_manager=None, + data_iterator=None, + load_parameters_from_path="gs://does-not-exist/hf", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=_abstract_nnx_state(), + source_checkpoint_layout="safetensors_dynamic", + maxtext_config=mock.Mock(), + ) + self.assertIn("linear/bias", str(ctx.exception)) + + def test_safetensors_dynamic_forwards_params_collection_for_linen(self): + """Linen hands over its `params` collection, not the whole TrainState, and gets it back as-is.""" + weights = {"params": {"linear": {"kernel": jnp.ones((2, 1))}}} + fake_state = mock.Mock(spec=["params"]) + fake_state.params = weights + + with mock.patch.object(checkpointing, "load_safetensors_dynamic_state", return_value=(None, weights)) as m: + full, params = checkpointing.load_state_if_possible( + checkpoint_manager=None, + data_iterator=None, + load_parameters_from_path="gs://does-not-exist/hf", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=fake_state, + source_checkpoint_layout="safetensors_dynamic", + maxtext_config=mock.Mock(), + ) + + self.assertIsNone(full) + self.assertIs(m.call_args[0][1], fake_state.params) + self.assertIs(params, weights) # Linen keeps the collection layout it passed in + def test_no_paths_returns_none_none(self): """Sanity: with no checkpoint manager and no load paths, the function returns (None, None).""" full, params = checkpointing.load_state_if_possible( @@ -161,6 +233,158 @@ def test_no_paths_returns_none_none(self): self.assertIsNone(params) +class TestManagerRestoreParity(unittest.TestCase): + """NNX and Linen go through the same manager restore; only the ends of the call differ. + + Cases come in pairs, one per state type, so a divergence in the shared flow fails one half. + """ + + def _manager(self, restored): + manager = mock.MagicMock(spec=ocp.CheckpointManager) + manager.restore.return_value = restored + return manager + + def _linen_abstract(self): + """A real Linen TrainState -- the manager restore paths tree_map over the whole state.""" + + def make(): + params = {"params": {"linear": {"kernel": jnp.zeros((2, 1))}}} + return train_state.TrainState.create(apply_fn=lambda x: x, params=params, tx=optax.identity()) + + return jax.eval_shape(make) # every leaf abstract, `step` included, as get_abstract_state gives it + + def _load(self, manager, abstract, **kwargs): + kwargs.setdefault("data_iterator", None) + return checkpointing.load_state_if_possible( + checkpoint_manager=manager, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=abstract, + step=7, + **kwargs, + ) + + def test_nnx_restores_the_linen_layout_and_returns_an_nnx_state(self): + abstract = _abstract_dropout_state() + saved = train_state_nnx.to_checkpoint_dict(_dropout_state(0)) + manager = self._manager({"items": saved}) + + restored, _ = self._load(manager, abstract) + + # Going in: the manager is asked for the Linen on-disk layout, not the NNX one. + item = manager.restore.call_args.kwargs["args"]["items"].item + self.assertEqual(set(item) & {"params", "step"}, {"params", "step"}) + self.assertNotIn("model", item) + # Coming out: reshaped back under `items`, the same key the Linen path returns. + self.assertIsInstance(restored["items"], nnx.State) + self.assertIn("model", restored["items"].to_pure_dict()) + + def test_linen_restore_target_and_return_are_untouched(self): + abstract = self._linen_abstract() + sentinel = {"items": {"params": abstract.params}} + manager = self._manager(sentinel) + + restored, _ = self._load(manager, abstract) + + self.assertIs(manager.restore.call_args.kwargs["args"]["items"].item, abstract) + self.assertIs(restored, sentinel) # returned exactly as the manager gave it + + def test_grain_case_converts_items_and_passes_the_iterator_through(self): + """The grain branch is shared too: NNX only reshapes `items`, leaving the iterator element alone.""" + abstract = _abstract_nnx_state() + saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}}}} + manager = self._manager(None) + + with mock.patch.object( + checkpointing.grain_utility, "restore_grain_iterator", return_value=({"items": saved}, None) + ) as m: + restored, iterator = self._load(manager, abstract, dataset_type="grain", data_iterator=mock.MagicMock()) + + m.assert_called_once() + self.assertIsNone(iterator) + self.assertIsInstance(restored["items"], nnx.State) + self.assertTrue(jnp.array_equal(restored["items"].to_pure_dict()["model"]["linear"]["bias"], jnp.array([5.0]))) + + def test_grain_case_is_untouched_for_linen(self): + abstract = self._linen_abstract() + sentinel = ({"items": {"params": abstract.params}}, None) + manager = self._manager(None) + + with mock.patch.object(checkpointing.grain_utility, "restore_grain_iterator", return_value=sentinel): + restored, iterator = self._load(manager, abstract, dataset_type="grain", data_iterator=mock.MagicMock()) + + self.assertIs(restored, sentinel[0]) + self.assertIsNone(iterator) + + def test_missing_weight_raises_on_the_standard_path(self): + """The missing-weight check guards every NNX restore, not just the weights-only load.""" + saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}} # no bias + manager = self._manager({"items": saved}) + + with self.assertRaises(ValueError) as ctx: + self._load(manager, _abstract_nnx_state()) + self.assertIn("linear/bias", str(ctx.exception)) + + def test_missing_weight_raises_on_the_emergency_path(self): + saved = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}} # no bias + manager = mock.Mock(spec=checkpointing.EmergencyCheckpointManager) + manager.restore.return_value = mock.Mock(state=saved) + + with self.assertRaises(ValueError) as ctx: + self._load(manager, _abstract_nnx_state()) + self.assertIn("linear/bias", str(ctx.exception)) + + def test_no_step_in_manager_falls_through_to_the_load_paths(self): + """An empty manager (latest_step() is None) must not restore -- it falls through, for both types.""" + manager = mock.MagicMock(spec=ocp.CheckpointManager) + manager.latest_step.return_value = None + + for abstract in (_abstract_nnx_state(), self._linen_abstract()): + restored, params = checkpointing.load_state_if_possible( + checkpoint_manager=manager, + data_iterator=None, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=abstract, + ) + self.assertIsNone(restored) + self.assertIsNone(params) + manager.restore.assert_not_called() + + +class TestResolveConversionFn(unittest.TestCase): + """`checkpoint_conversion_fn` arrives from config as a dotted string and has to be imported.""" + + def test_dotted_string_is_imported(self): + fn = checkpointing._resolve_conversion_fn("json.loads") # pylint: disable=protected-access + self.assertEqual(fn('{"a": 1}'), {"a": 1}) + + def test_callable_passes_through(self): + self.assertIs(checkpointing._resolve_conversion_fn(len), len) # pylint: disable=protected-access + + def test_missing_fn_names_the_config_flag(self): + with self.assertRaises(ValueError) as ctx: + checkpointing._resolve_conversion_fn(None) # pylint: disable=protected-access + self.assertIn("checkpoint_conversion_fn", str(ctx.exception)) + + def test_bare_name_without_module_raises(self): + with self.assertRaises(ValueError) as ctx: + checkpointing._resolve_conversion_fn("my_fn") # pylint: disable=protected-access + self.assertIn("dotted path", str(ctx.exception)) + + def test_unimportable_module_raises(self): + with self.assertRaises(ValueError) as ctx: + checkpointing._resolve_conversion_fn("no_such_module_xyz.my_fn") # pylint: disable=protected-access + self.assertIn("no_such_module_xyz", str(ctx.exception)) + + def test_non_callable_attribute_raises(self): + with self.assertRaises(ValueError) as ctx: + checkpointing._resolve_conversion_fn("json.__doc__") # pylint: disable=protected-access + self.assertIn("not a function", str(ctx.exception)) + + class TestLoadParamsIntoNNX(unittest.TestCase): """Weight-only load (load_parameters_path) of a Linen-layout checkpoint into NNX.""" @@ -214,6 +438,52 @@ def test_linen_params_dict_restores_as_the_params_collection(self): self.assertTrue(jnp.array_equal(restored["params"]["linear"]["bias"], weights["linear"]["bias"])) +class TestSafetensorsFullStateIntoNNX(unittest.TestCase): + """A full-state safetensors load reshapes the conversion fn's output into the NNX state.""" + + def _load(self, converted, abstract): + """Runs the v1 safetensors branch, stubbing the read so only the conversion is under test.""" + with ( + mock.patch.object(checkpointing, "ocp_v1") as v1, + mock.patch.object(checkpointing, "sharding_utils") as shardings, + ): + v1.pytree_metadata.return_value = mock.Mock(metadata={"w": jax.ShapeDtypeStruct((1,), jnp.float32)}) + shardings.construct_maximal_shardings.return_value = {"w": None} + return checkpointing._load_full_state_from_path( # pylint: disable=protected-access + path="gs://does-not-exist/hf", + abstract_unboxed_pre_state=abstract, + enable_orbax_v1=True, + checkpoint_conversion_fn=lambda _: converted, + source_checkpoint_layout="safetensors", + checkpoint_storage_concurrent_gb=8, + use_ocdbt=True, + use_zarr3=True, + ) + + def test_linen_layout_output_is_reshaped_into_nnx(self): + converted = { + "params": {"params": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.array([5.0])}}}, + "step": jnp.asarray(0, jnp.int32), + } + restored = self._load(converted, _abstract_nnx_state()) + + self.assertIsInstance(restored, nnx.State) + pure = restored.to_pure_dict() + self.assertTrue(jnp.array_equal(pure["model"]["linear"]["kernel"], jnp.ones((2, 1)))) + + def test_missing_weight_raises(self): + """A bad conversion fn hits the same missing-weight check as the other NNX loads.""" + converted = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}} # no bias + with self.assertRaises(ValueError) as ctx: + self._load(converted, _abstract_nnx_state()) + self.assertIn("linear/bias", str(ctx.exception)) + + def test_linen_state_is_untouched(self): + """For Linen the conversion fn's output is returned as-is.""" + converted = {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}} + self.assertIs(self._load(converted, mock.Mock(spec=["params"])), converted) + + class TestToCheckpointDict(unittest.TestCase): """train_state_nnx.to_checkpoint_dict splits the NNX state by Variable type.""" diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index f21f025175..421a6d150f 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -228,15 +228,8 @@ def __init__(self): config = MockConfig() - class DummyAbstractState: - - def __init__(self): - self.params = {} - - abstract_state = DummyAbstractState() - path = "repo/meta-llama" - dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, abstract_state, config) + dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, {}, config) self.assertIsNone(dummy_ret_val) self.assertEqual(loaded_vars, {"params": {}}) diff --git a/tests/unit/setup_initial_state_nnx_test.py b/tests/unit/setup_initial_state_nnx_test.py index 0f6477c9bb..8130f349fd 100644 --- a/tests/unit/setup_initial_state_nnx_test.py +++ b/tests/unit/setup_initial_state_nnx_test.py @@ -216,7 +216,7 @@ def test_emergency_manager_overlay_is_not_unwrapped(self): trained = _init_fn(_Model, 0)() grads = nnx.grad(lambda m: jnp.mean(m(_TRAIN_X, deterministic=False) ** 2))(trained.model) trained.apply_gradients(grads) - overlay = nnx.state(trained) # what _restore_emergency_linen_checkpoint_into_nnx returns + overlay = nnx.state(trained) # what an emergency restore returns: the NNX state, unwrapped manager = mock.Mock(spec=emergency_checkpoint_manager.CheckpointManager) with mock.patch.object(checkpointing, "load_state_if_possible", return_value=(overlay, None)):