diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index c527427144..6ae357b4e6 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,27 @@ def _populate_pure_dict_from_partial(abstract_pure, partial_concrete):
   return _default_for_sds(abstract_pure)
 
 
+def _restore_nnx_aux_state(step_dir, aux_abstract, checkpoint_storage_concurrent_gb, use_ocdbt, use_zarr3):
+  """Restores the separately-saved rngs/dropout state written next to `items`.
+
+  Returns the restored aux state, or None when the checkpoint has no `nnx_aux`
+  subdir (a Linen-trained checkpoint, or one saved before aux persistence).
+  """
+  aux_path = epath.Path(step_dir) / "nnx_aux"
+  if not aux_abstract or not aux_path.exists():
+    return None
+  ckptr = ocp.Checkpointer(
+      ocp.PyTreeCheckpointHandler(
+          restore_concurrent_gb=checkpoint_storage_concurrent_gb,
+          save_concurrent_gb=checkpoint_storage_concurrent_gb,
+          use_ocdbt=use_ocdbt,
+          use_zarr3=use_zarr3,
+      )
+  )
+  restore_args = ocp.checkpoint_utils.construct_restore_args(aux_abstract)
+  return ckptr.restore(aux_path, args=ocp.args.PyTreeRestore(item=aux_abstract, restore_args=restore_args))
+
+
 def _load_linen_checkpoint_into_nnx(
     path,
     abstract_nnx_state,
@@ -224,7 +255,8 @@ 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.
+  `from_linen_checkpoint_dict`, then restores the separately-saved rngs/dropout
+  state when present (falling back to deterministic defaults otherwise).
   """
   max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}")
   nnx_abstract_pure = abstract_nnx_state.to_pure_dict()
@@ -241,6 +273,13 @@ def _load_linen_checkpoint_into_nnx(
   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)
+
+  aux_abstract = train_state_nnx.extract_rng_state(nnx_abstract_pure.get("model", {}))
+  aux = _restore_nnx_aux_state(
+      epath.Path(path).parent, aux_abstract, checkpoint_storage_concurrent_gb, use_ocdbt, use_zarr3
+  )
+  if aux:
+    partial_nnx = _deep_merge(partial_nnx, {"model": aux})
   return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx)
 
 
@@ -250,7 +289,11 @@ 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.
+
+  Emergency checkpoints don't carry the separate rngs/dropout `nnx_aux` item, so
+  those are refilled 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)
@@ -417,8 +460,10 @@ def create_orbax_checkpoint_manager(
 
   max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}")
 
-  # Base configuration for all dataset types
-  item_names = ("items",)
+  # Base configuration for all dataset types.
+  # `nnx_aux` carries the NNX-only rngs/dropout state that pure_nnx runs persist
+  # alongside the Linen-layout `items` (empty/unused for Linen runs).
+  item_names = ("items", "nnx_aux")
   # we need to use ocdbt and zarr3 to control max file size in the checkpoint
   item_handlers = {
       "items": PyTreeCheckpointHandler(
@@ -426,7 +471,13 @@ def create_orbax_checkpoint_manager(
           save_concurrent_gb=checkpoint_storage_concurrent_gb,
           use_ocdbt=use_ocdbt,
           use_zarr3=use_zarr3,
-      )
+      ),
+      "nnx_aux": PyTreeCheckpointHandler(
+          restore_concurrent_gb=checkpoint_storage_concurrent_gb,
+          save_concurrent_gb=checkpoint_storage_concurrent_gb,
+          use_ocdbt=use_ocdbt,
+          use_zarr3=use_zarr3,
+      ),
   }
 
   if dataset_type is not None and dataset_type == "grain":
@@ -1032,6 +1083,7 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
+  nnx_aux_state = None
   if config.pure_nnx:
     # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
     if config.enable_diloco:
@@ -1040,7 +1092,11 @@ 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())
+      nnx_pure = state.to_pure_dict()
+      # rngs/dropout are stripped from the Linen layout; persist them separately so
+      # the RNG/dropout stream continues across resumes instead of resetting to a base key.
+      nnx_aux_state = train_state_nnx.extract_rng_state(nnx_pure.get("model", {}))
+      state = train_state_nnx.to_linen_checkpoint_dict(nnx_pure)
 
   # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic.
   # This occurs if this function was called:
@@ -1050,7 +1106,9 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
   force_ckpt_save = step is None and actual_step != -1 and (actual_step % config.checkpoint_period != 0)
 
   try:
-    checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save)
+    checkpoint_saved = save_checkpoint(
+        checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save, nnx_aux_state
+    )
     if checkpoint_saved:
       print_save_message(actual_step, config.async_checkpointing)
       if config.elastic_enabled:
@@ -1079,8 +1137,13 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     raise exceptions.StopTraining("Job is preempted.")
 
 
-def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False):
-  """Wrapper for saving checkpoint."""
+def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False, nnx_aux_state=None):
+  """Wrapper for saving checkpoint.
+
+  `nnx_aux_state` holds the NNX-only rngs/dropout state stripped from the Linen
+  layout; when non-empty it's saved as a separate `nnx_aux` composite item so the
+  RNG stream survives a resume. It's ignored by emergency managers.
+  """
   if config and config.enable_checkpointing:
     if (
         force
@@ -1110,6 +1173,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
   )
   save_args_composite = {"items": checkpoint_args}
 
+  if nnx_aux_state:
+    save_args_composite["nnx_aux"] = ocp.args.PyTreeSave(
+        item=nnx_aux_state,
+        save_args=jax.tree.map(lambda _: ocp.SaveArgs(chunk_byte_size=chunk_byte_size), nnx_aux_state),
+        ocdbt_target_data_file_size=chunk_byte_size,
+    )
+
   if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator):
     if isinstance(data_iterator, RemoteIteratorWrapper):
       # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step
diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py
index a123956a73..a251c02f2c 100644
--- a/src/maxtext/common/train_state_nnx.py
+++ b/src/maxtext/common/train_state_nnx.py
@@ -102,6 +102,26 @@ def _strip_rng_state(tree):
   return out
 
 
+def extract_rng_state(tree):
+  """Keeps only the NNX-only 'rngs'/'dropout' subtrees (complement of `_strip_rng_state`).
+
+  Returns the rngs/dropout state as a nested dict so it can be checkpointed
+  separately and restored to keep RNG continuity across resumes. Returns an empty
+  dict when the model carries no such state (e.g. dropout disabled).
+  """
+  if not isinstance(tree, dict):
+    return {}
+  out = {}
+  for k, v in tree.items():
+    if k in _NNX_RNG_STATE_KEYS:
+      out[k] = v
+    else:
+      sub = extract_rng_state(v)
+      if sub:
+        out[k] = sub
+  return out
+
+
 def _wrap_mu_nu_with_params(state):
   """Wraps mu/nu under an inner 'params' key (the Linen collection)."""
   if not isinstance(state, dict):
diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py
index 04aa898bde..573bf4ae31 100644
--- a/tests/unit/checkpointing_nnx_load_test.py
+++ b/tests/unit/checkpointing_nnx_load_test.py
@@ -185,5 +185,134 @@ def test_linen_layout_params_restore_into_nnx_state(self):
     self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]))
 
 
+class TestExtractRngState(unittest.TestCase):
+  """train_state_nnx.extract_rng_state keeps only the NNX-only rngs/dropout subtrees."""
+
+  def test_keeps_only_rng_and_dropout(self):
+    tree = {
+        "linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.zeros((1,))},
+        "rngs": {"params": {"key": jnp.asarray(1), "count": jnp.asarray(2)}},
+        "block": {"dropout": {"count": jnp.asarray(3)}, "kernel": jnp.ones((2, 2))},
+    }
+    aux = train_state_nnx.extract_rng_state(tree)
+    self.assertEqual(set(aux.keys()), {"rngs", "block"})
+    self.assertEqual(set(aux["block"].keys()), {"dropout"})  # kernel dropped
+    self.assertNotIn("linear", aux)
+
+  def test_empty_when_no_rng_state(self):
+    self.assertEqual(train_state_nnx.extract_rng_state({"linear": {"kernel": jnp.ones((2, 1))}}), {})
+
+  def test_extract_is_complement_of_strip(self):
+    """Every leaf lands in exactly one of extract / strip — nothing dropped or duplicated."""
+    tree = {
+        "linear": {"kernel": jnp.ones((2, 1))},
+        "rngs": {"params": {"count": jnp.asarray(2)}},
+    }
+    stripped = train_state_nnx._strip_rng_state(tree)  # pylint: disable=protected-access
+    extracted = train_state_nnx.extract_rng_state(tree)
+    self.assertEqual(set(stripped.keys()), {"linear"})
+    self.assertEqual(set(extracted.keys()), {"rngs"})
+
+
+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 TestRngStateAuxPersistence(unittest.TestCase):
+  """rngs/dropout persisted as a separate `nnx_aux` item, restored on resume."""
+
+  def _abstract_with_dropout(self):
+    abstract = mock.Mock()
+    abstract.to_pure_dict.return_value = {
+        "model": {
+            "linear": {"kernel": jax.ShapeDtypeStruct((2, 1), jnp.float32)},
+            "dropout": {"count": jax.ShapeDtypeStruct((), jnp.uint32)},
+        },
+        "optimizer": {"step": jax.ShapeDtypeStruct((), jnp.uint32)},
+    }
+    return abstract
+
+  def _write_linen_items(self, step_dir):
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(step_dir) / "items",
+        {"params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, "step": jnp.asarray(3, jnp.int32)},
+        force=True,
+    )
+
+  def test_restores_saved_rng_state_instead_of_default(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      step_dir = os.path.join(d, "5")
+      self._write_linen_items(step_dir)
+      # Persisted rng/dropout state: count=42 (a resumed stream, not the base 0).
+      ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+          epath.Path(step_dir) / "nnx_aux",
+          {"dropout": {"count": jnp.asarray(42, jnp.uint32)}},
+          force=True,
+      )
+      restored = checkpointing._load_linen_checkpoint_into_nnx(  # pylint: disable=protected-access
+          os.path.join(step_dir, "items"), self._abstract_with_dropout(), 8, True, True
+      )
+    self.assertEqual(int(restored["model"]["dropout"]["count"]), 42)
+    self.assertTrue(jnp.array_equal(restored["model"]["linear"]["kernel"], jnp.ones((2, 1))))
+
+  def test_falls_back_to_default_when_no_aux_dir(self):
+    """A Linen-trained checkpoint (no nnx_aux) still loads; rng/dropout gets the base default."""
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      step_dir = os.path.join(d, "5")
+      self._write_linen_items(step_dir)  # no nnx_aux written
+      restored = checkpointing._load_linen_checkpoint_into_nnx(  # pylint: disable=protected-access
+          os.path.join(step_dir, "items"), self._abstract_with_dropout(), 8, True, True
+      )
+    self.assertEqual(int(restored["model"]["dropout"]["count"]), 0)  # _default_for_sds
+
+  def test_save_checkpoint_adds_nnx_aux_item_when_present(self):
+    manager = mock.Mock()
+    manager.save.return_value = True
+    config = mock.Mock(
+        enable_checkpointing=False,
+        dataset_type="tfds",
+        lora=None,
+        checkpoint_storage_target_data_file_size_bytes=1,
+    )
+    aux = {"dropout": {"count": jnp.asarray(7, jnp.uint32)}}
+    checkpointing.save_checkpoint(manager, 5, {"params": {}}, config, None, False, aux)
+    composite = manager.save.call_args.kwargs["args"]
+    self.assertIn("nnx_aux", composite.keys())
+
+  def test_save_checkpoint_omits_nnx_aux_when_empty(self):
+    manager = mock.Mock()
+    manager.save.return_value = True
+    config = mock.Mock(
+        enable_checkpointing=False,
+        dataset_type="tfds",
+        lora=None,
+        checkpoint_storage_target_data_file_size_bytes=1,
+    )
+    checkpointing.save_checkpoint(manager, 5, {"params": {}}, config, None, False, {})
+    composite = manager.save.call_args.kwargs["args"]
+    self.assertNotIn("nnx_aux", composite.keys())
+
+
 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..32495cf76a
--- /dev/null
+++ b/tests/unit/checkpointing_nnx_roundtrip_test.py
@@ -0,0 +1,148 @@
+# 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 composite `nnx_aux` item, the RNG/dropout persistence, and
+the missing-parameter guard together the way training actually hits them.
+"""
+
+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 gets stripped 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)
+
+
+_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()
+    # rngs/dropout is persisted as a separate on-disk item, not folded into `items`.
+    self.assertIn("nnx_aux", os.listdir(os.path.join(self._dir, "ckpt", "1")))
+    restored = self._restore(manager, _Model)
+    self.assertEqual(int(restored["model"]["dropout"]["rngs"]["count"]), orig_count)
+    self.assertEqual(int(restored["optimizer"]["step"]), 1)
+
+  def test_missing_aux_dir_falls_back_to_default_rng(self):
+    """A checkpoint without nnx_aux (Linen-trained / pre-fix) still loads; rng gets the base default."""
+    manager, _ = self._save_step()
+    shutil.rmtree(os.path.join(self._dir, "ckpt", "1", "nnx_aux"))
+    restored = self._restore(manager, _Model)
+    self.assertEqual(int(restored["model"]["dropout"]["rngs"]["count"]), 0)
+
+
+if __name__ == "__main__":
+  unittest.main()