[NNX] Checkpoint: persist and restore rngs/dropout across resumes

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,16 +460,24 @@ 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(
           restore_concurrent_gb=checkpoint_storage_concurrent_gb,
           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

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

tests/unit/checkpointing_nnx_load_test.py

@@ -185,5 +185,125 @@ 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}})
+
+
+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()