[NNX] Emergency checkpoint: build the manager with a Linen-layout abstract

pure_nnx runs save in the Linen on-disk layout (params/opt_state/step), but the
emergency checkpoint manager was constructed with the NNX-native abstract state
(model/optimizer + rngs/dropout). The manager bakes that abstract in and restores
against it, so on resume Orbax compared a Linen on-disk tree against an NNX target
and raised 'User-provided restore item and on-disk value metadata tree structures
do not match', failing the EMC resume-from-GCS job. The prior restore-side fix was
ineffective because the manager ignores the restore-time item and uses its
construction-time abstract.

Convert the abstract to the Linen layout in create_orbax_emergency_checkpoint_manager
so it matches what maybe_save_checkpoint writes; the restore path reshapes it back
to NNX. Adds a mock-based guard (asserts the manager gets a Linen abstract) plus an
end-to-end save/restore round trip through the real emergency manager.

Author: ecnal-cienet

src/maxtext/common/checkpointing.py

@@ -498,6 +498,13 @@ def create_orbax_emergency_checkpoint_manager(
 
   persistent_p = gcs_utils.mkdir_and_check_permissions(persistent_checkpoint_dir)
 
+  # pure_nnx passes an NNX-native abstract (model/optimizer), but pure_nnx runs save in the
+  # Linen on-disk layout (params/opt_state/step). The emergency manager bakes this abstract in
+  # and restores against it, so it must match the on-disk layout -- convert it to Linen here,
+  # matching what maybe_save_checkpoint writes. The restore path reshapes it back to NNX.
+  if isinstance(abstract_state, nnx.State):
+    abstract_state = train_state_nnx.to_linen_checkpoint_dict(abstract_state.to_pure_dict())
+
   manager = EmergencyCheckpointManager(
       local_checkpoint_dir,
       persistent_p,

tests/unit/checkpointing_emergency_nnx_test.py

@@ -0,0 +1,180 @@
+# 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.
+
+"""Emergency-checkpoint tests for the pure_nnx layout.
+
+Regression coverage for the EMC restore failure where the emergency manager was
+constructed with an NNX-native abstract (model/optimizer) while pure_nnx runs save
+in the Linen on-disk layout (params/opt_state/step). The manager bakes its abstract
+in and restores against it, so the two must agree -- otherwise Orbax raises
+"User-provided restore item and on-disk value metadata tree structures do not match".
+"""
+
+import os
+import tempfile
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+from etils import epath
+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 NNX state carries model, optimizer, and rngs."""
+
+  def __init__(self, rngs: nnx.Rngs):
+    self.linear = nnx.Linear(2, 3, 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)
+
+
+_TX = optax.adam(1e-3)
+
+
+def _nnx_abstract_state():
+  """An abstract NNX state (model/optimizer), the shape pure_nnx hands the manager."""
+
+  def make():
+    model = _Model(nnx.Rngs(0))
+    return nnx.state(train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param)))
+
+  return nnx.eval_shape(make)
+
+
+class TestEmergencyManagerAbstractLayout(unittest.TestCase):
+  """The regression guard: the emergency manager must be built with a Linen abstract.
+
+  Fast and environment-independent -- mocks the Orbax manager, so it catches a
+  reintroduction of the bug in CI without needing a real multi-host EMC run.
+  """
+
+  def _create(self, abstract_state):
+    """Runs create_orbax_emergency_checkpoint_manager with the Orbax manager mocked out."""
+    mesh = jax.sharding.Mesh(jax.devices(), ("x",))
+    with (
+        mock.patch.object(checkpointing, "EmergencyCheckpointManager") as manager_cls,
+        mock.patch.object(checkpointing.gcs_utils, "mkdir_and_check_permissions", side_effect=lambda p: epath.Path(p)),
+    ):
+      with tempfile.TemporaryDirectory() as d:
+        checkpointing.create_orbax_emergency_checkpoint_manager(
+            os.path.join(d, "local"),
+            os.path.join(d, "persist"),
+            mesh,
+            abstract_state,
+            local_save_interval_steps=1,
+            persistent_save_interval_steps=1,
+        )
+    return manager_cls.call_args.kwargs["abstract_state"]
+
+  def test_nnx_abstract_is_converted_to_linen_layout(self):
+    """An NNX-native abstract (model/optimizer) is reshaped to the Linen on-disk layout."""
+    passed = self._create(_nnx_abstract_state())
+    self.assertNotIsInstance(passed, nnx.State)
+    # Linen on-disk keys, not the NNX model/optimizer roots that caused the mismatch.
+    self.assertEqual(set(passed.keys()), {"params", "opt_state", "step"})
+    self.assertNotIn("model", passed)
+    self.assertNotIn("optimizer", passed)
+
+  def test_non_nnx_abstract_is_passed_through_unchanged(self):
+    """A Linen TrainState (not an nnx.State) is forwarded as-is, not double-converted."""
+    sentinel = SimpleNamespace(params={"a": 1}, opt_state=(), step=0)  # stand-in Linen state
+    passed = self._create(sentinel)
+    self.assertIs(passed, sentinel)
+
+
+class TestEmergencySaveRestoreRoundTrip(unittest.TestCase):
+  """End-to-end: save (Linen) then restore through the real emergency manager.
+
+  Reproduces the production DAG. Skips where a single-host EMC manager can't be
+  constructed (some CI backends); runs fully on TPU/GPU/CPU that support it.
+  """
+
+  def setUp(self):
+    self._dir = tempfile.mkdtemp()
+
+  def _config(self):
+    return SimpleNamespace(
+        pure_nnx=True,
+        enable_diloco=False,
+        enable_checkpointing=True,
+        enable_continuous_checkpointing=False,
+        enable_emergency_checkpoint=True,
+        enable_autocheckpoint=False,
+        enable_multi_tier_checkpointing=False,
+        checkpoint_period=1,
+        local_checkpoint_period=1,
+        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 test_emergency_round_trip_reshapes_back_to_nnx(self):
+    mesh = jax.sharding.Mesh(jax.devices(), ("x",))
+    sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+    abstract = jax.tree.map(
+        lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=sharding) if hasattr(x, "shape") else x,
+        _nnx_abstract_state(),
+    )
+
+    try:
+      manager = checkpointing.create_orbax_emergency_checkpoint_manager(
+          os.path.join(self._dir, "local"),
+          os.path.join(self._dir, "persist"),
+          mesh,
+          abstract,
+          local_save_interval_steps=1,
+          persistent_save_interval_steps=1,
+      )
+    except Exception as e:  # pylint: disable=broad-except
+      raise unittest.SkipTest(f"EMC manager unavailable in this environment: {e}")
+
+    # Train one step so weights/optimizer/step are non-trivial, then save (Linen layout).
+    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)
+    saved_kernel = jnp.asarray(nnx.state(state).to_pure_dict()["model"]["linear"]["kernel"])
+
+    checkpointing.maybe_save_checkpoint(manager, nnx.state(state), self._config(), data_iterator=None, step=1)
+    manager.wait_until_finished()
+
+    # Restore: must not raise the structural-mismatch ValueError, and must come back as NNX.
+    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,
+        dataset_type="tfds",
+        maxtext_config=self._config(),
+    )
+    self.assertIn("model", full)
+    self.assertIn("optimizer", full)
+    self.assertNotIn("params", full)
+    self.assertTrue(jnp.allclose(full["model"]["linear"]["kernel"], saved_kernel))
+
+
+if __name__ == "__main__":
+  unittest.main()