From 0da53c3b8438082c0eca4acee6210d9a4ec1d4a1 Mon Sep 17 00:00:00 2001 From: Sujeeth Jinesh Date: Wed, 1 Jul 2026 12:43:08 -0700 Subject: [PATCH] Fix emergency NNX checkpoint restore --- src/maxtext/common/checkpointing.py | 39 +++++++ tests/unit/checkpointing_nnx_load_test.py | 119 ++++++++++++---------- 2 files changed, 103 insertions(+), 55 deletions(-) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c778f92bf9..a47c024910 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -245,6 +245,27 @@ def _load_linen_checkpoint_into_nnx( return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) +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.""" + 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) + 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 + partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) + return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + + def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights): """Fills each Variable in `abstract_nnx_state` with the matching restored array.""" leaves, treedef = jax.tree_util.tree_flatten(abstract_nnx_state, is_leaf=lambda x: isinstance(x, nnx.Variable)) @@ -850,6 +871,24 @@ def map_to_pspec(data): _assert_no_shaped_dtype_struct(restored_nnx) 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)) + with handle_checkpoint_mismatch("restore emergency NNX checkpoint", checkpoint_path): + restored = _restore_emergency_linen_checkpoint_into_nnx( + checkpoint_manager, + step, + abstract_unboxed_pre_state, + map_to_pspec, + ) + _assert_no_shaped_dtype_struct(restored) + return ( + restored, + None, + ) + # Convert nnx.State to pure dict to match how checkpoints are saved for NNX restore_target = abstract_unboxed_pre_state if isinstance(abstract_unboxed_pre_state, nnx.State): diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 69a64ec1eb..5af3f9b0b8 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/tests/unit/checkpointing_nnx_load_test.py @@ -46,6 +46,56 @@ def _abstract_nnx_state(): 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) + + 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, + ), + ) + + 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) + 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, + (), + ) + 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 nnx.split(model, nnx.Param, ...) to get the params and forward them to load_params_from_path.""" @@ -126,8 +176,7 @@ def test_load_state_if_possible_wraps_load_params_mismatch_exception(self): str(ctx.exception), ) self.assertIn( - "This is often caused by a mismatch in the 'scan_layers'" - " configuration", + "This is often caused by a mismatch in the 'scan_layers'" " configuration", str(ctx.exception), ) @@ -156,69 +205,33 @@ class TestCheckpointMismatchHandling(unittest.TestCase): def test_is_structural_or_shape_mismatch(self): """Verifies that is_structural_or_shape_mismatch matches only shape/tree mismatches in ValueError/TypeError.""" # Matches - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("PyTree structure mismatch") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - TypeError("shape mismatch in leaf") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("tree paths matched 143/145") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("invalid type shapedtypestruct") - ) - ) + self.assertTrue(checkpointing.is_structural_or_shape_mismatch(ValueError("PyTree structure mismatch"))) + self.assertTrue(checkpointing.is_structural_or_shape_mismatch(TypeError("shape mismatch in leaf"))) + self.assertTrue(checkpointing.is_structural_or_shape_mismatch(ValueError("tree paths matched 143/145"))) + self.assertTrue(checkpointing.is_structural_or_shape_mismatch(ValueError("invalid type shapedtypestruct"))) # Does not match - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - ValueError("checkpoint directory does not exist") - ) - ) - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - FileNotFoundError("file not found: checkpoint") - ) - ) - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - RuntimeError("something went wrong") - ) - ) + self.assertFalse(checkpointing.is_structural_or_shape_mismatch(ValueError("checkpoint directory does not exist"))) + self.assertFalse(checkpointing.is_structural_or_shape_mismatch(FileNotFoundError("file not found: checkpoint"))) + self.assertFalse(checkpointing.is_structural_or_shape_mismatch(RuntimeError("something went wrong"))) def test_handle_checkpoint_mismatch_intercepts_matching_exceptions(self): """Verifies that handle_checkpoint_mismatch intercepts and wraps structural errors.""" with self.assertRaises(ValueError) as ctx: - with checkpointing.handle_checkpoint_mismatch( - "load parameters", "gs://bucket/params" - ): + with checkpointing.handle_checkpoint_mismatch("load parameters", "gs://bucket/params"): raise ValueError("PyTree structure mismatch") - self.assertIn( - "Failed to load parameters from gs://bucket/params.", str(ctx.exception) - ) + self.assertIn("Failed to load parameters from gs://bucket/params.", str(ctx.exception)) self.assertIn( "This is often caused by a mismatch in the 'scan_layers' configuration", str(ctx.exception), ) - self.assertIn( - "Original error: PyTree structure mismatch", str(ctx.exception) - ) + self.assertIn("Original error: PyTree structure mismatch", str(ctx.exception)) def test_handle_checkpoint_mismatch_re_raises_non_matching_exceptions(self): """Verifies that handle_checkpoint_mismatch does not intercept non-structural errors.""" with self.assertRaises(FileNotFoundError): - with checkpointing.handle_checkpoint_mismatch( - "load parameters", "gs://bucket/params" - ): + with checkpointing.handle_checkpoint_mismatch("load parameters", "gs://bucket/params"): raise FileNotFoundError("file not found: checkpoint") @@ -248,12 +261,8 @@ def test_linen_layout_params_restore_into_nnx_state(self): self.assertIsInstance(restored, nnx.State) pure = restored.to_pure_dict() - self.assertTrue( - jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"]) - ) - self.assertTrue( - jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]) - ) + self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"])) + self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"])) if __name__ == "__main__":