Fix MTC NNX checkpoint restore

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):

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__":