[NNX] Checkpoint restore: error on weights missing from the checkpoint

src/maxtext/common/checkpointing.py

@@ -194,12 +194,23 @@ 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.
 
   Paths present in `partial_concrete` take the restored value; paths absent from
-  it
-  (NNX-only state the Linen checkpoint never had) get `_default_for_sds`.
+  it, or ones Orbax left as an unmaterialized ShapeDtypeStruct (NNX-only state the
+  Linen checkpoint never had, or a weight the missing-param policy chose to
+  zero-fill), get `_default_for_sds`.
   """
   if isinstance(abstract_pure, dict):
     return {
@@ -209,26 +220,94 @@ def _populate_pure_dict_from_partial(abstract_pure, partial_concrete):
         )
         for k, v in abstract_pure.items()
     }
-  if partial_concrete is not None and not isinstance(partial_concrete, dict):
+  if partial_concrete is not None and not isinstance(partial_concrete, (dict, jax.ShapeDtypeStruct)):
     return partial_concrete
   return _default_for_sds(abstract_pure)
 
 
+def _all_materialized(tree):
+  """True if Orbax restored every leaf, i.e. none is left as a jax.ShapeDtypeStruct."""
+  return not any(isinstance(leaf, jax.ShapeDtypeStruct) for leaf in jax.tree_util.tree_leaves(tree))
+
+
+def _missing_param_policy(config):
+  """Reads how to handle weights present in the model but absent from the checkpoint."""
+  return getattr(config, "checkpoint_missing_param_policy", "error") if config is not None else "error"
+
+
+def _missing_weight_paths(want, have, path=()):
+  """Paths of abstract weights in `want` that `have` didn't materialize (absent or a ShapeDtypeStruct)."""
+  if isinstance(want, dict):
+    out = []
+    for k, v in want.items():
+      out.extend(_missing_weight_paths(v, have.get(k) if isinstance(have, dict) else None, path + (k,)))
+    return out
+  if have is None or isinstance(have, jax.ShapeDtypeStruct):
+    return [("/".join(str(p) for p in path), getattr(want, "shape", "?"), getattr(want, "dtype", "?"))]
+  return []
+
+
+def _enforce_missing_weight_policy(abstract_nnx_state, restored_linen, missing_param_policy="error"):
+  """Every nnx.Param the model expects must be materialized from the checkpoint.
+
+  Filters weights by Variable type (nnx.Param) rather than key name, so only real
+  weights are checked -- rngs/dropout/batch stats live in `nnx_aux` and are restored
+  separately. A weight the checkpoint didn't provide (absent, or left by Orbax as an
+  unmaterialized ShapeDtypeStruct) is reported by path under "error", or zero-filled
+  with a warning under "warn".
+  """
+  want = nnx.split_state(abstract_nnx_state, nnx.Param, ...)[0].to_pure_dict().get("model", {})
+  have = restored_linen.get("params", {}).get("params", {})
+  missing = _missing_weight_paths(want, have)
+  if not missing:
+    return
+  lines = "\n".join(f"  - 'model/{p}' (model expects {s} {d})" for p, s, d in missing)
+  detail = (
+      "Checkpoint is missing model weights:\n"
+      f"{lines}\n"
+      "These would be silently zero-filled or left unmaterialized (a cryptic compile failure in the "
+      "first train_step). Verify the checkpoint matches the model architecture (emb_dim, mlp_dim, num "
+      "layers, scan_layers), or set checkpoint_missing_param_policy=warn to zero-fill and continue."
+  )
+  if missing_param_policy == "warn":
+    max_logging.log(f"WARNING: {detail}")
+  else:
+    raise ValueError(detail)
+
+
+def _linen_items_to_nnx(restored_linen, nnx_abstract_pure):
+  """Reshapes a restored Linen-layout `items` dict back into the NNX pure-dict state.
+
+  Merges `nnx_aux` (rngs/dropout/batch stats) back onto the model when present and
+  fills anything absent with defaults. A checkpoint without `nnx_aux` (Linen-trained
+  or pre-persistence) leaves those leaves unmaterialized, so they get the defaults.
+  Callers run `_enforce_missing_weight_policy` first, so a genuinely-missing weight
+  has already errored (or been chosen for zero-fill) before this fills defaults.
+  """
+  partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored_linen)
+  aux = restored_linen.get("nnx_aux")
+  if aux and _all_materialized(aux):
+    partial_nnx = _deep_merge(partial_nnx, aux)
+  return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx)
+
+
 def _load_linen_checkpoint_into_nnx(
     path,
     abstract_nnx_state,
     checkpoint_storage_concurrent_gb,
     use_ocdbt,
     use_zarr3,
+    missing_param_policy="error",
 ):
   """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.
+  Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via
+  `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
+  present, else fall back to defaults; a genuinely-missing weight is handled per
+  `missing_param_policy`.
   """
   max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}")
-  nnx_abstract_pure = abstract_nnx_state.to_pure_dict()
-  linen_abstract = train_state_nnx.to_linen_checkpoint_dict(nnx_abstract_pure)
+  linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state)
   ckptr = ocp.Checkpointer(
       ocp.PyTreeCheckpointHandler(
           restore_concurrent_gb=checkpoint_storage_concurrent_gb,
@@ -240,29 +319,35 @@ def _load_linen_checkpoint_into_nnx(
   restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract)
   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)
-  return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx)
+  _enforce_missing_weight_policy(abstract_nnx_state, restored, missing_param_policy)
+  return _linen_items_to_nnx(restored, abstract_nnx_state.to_pure_dict())
 
 
 def _restore_emergency_linen_checkpoint_into_nnx(
     checkpoint_manager,
     step,
     abstract_nnx_state,
     map_to_pspec,
+    missing_param_policy="error",
 ):
-  """Restores an emergency Linen-layout checkpoint into an NNX state."""
+  """Restores an emergency Linen-layout checkpoint into an NNX state.
+
+  The `nnx_aux` subtree is stored inside `items`, so an emergency checkpoint
+  carries it too; it's restored when present and otherwise filled with
+  deterministic defaults. A genuinely-missing weight is handled per
+  `missing_param_policy`.
+  """
   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)
+  linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state)
   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)
+  _enforce_missing_weight_policy(abstract_nnx_state, restored, missing_param_policy)
+  return _linen_items_to_nnx(restored, abstract_nnx_state.to_pure_dict())
 
 
 def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights):
@@ -318,6 +403,7 @@ def _load_full_state_from_path(
     checkpoint_storage_concurrent_gb,
     use_ocdbt,
     use_zarr3,
+    missing_param_policy="error",
 ):
   """Load full state from checkpoint at specified path.
 
@@ -340,6 +426,11 @@ def _load_full_state_from_path(
 
   if enable_orbax_v1:
     if source_checkpoint_layout == "orbax":
+      # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
+      if isinstance(abstract_unboxed_pre_state, nnx.State):
+        return _load_linen_checkpoint_into_nnx(
+            path, abstract_unboxed_pre_state, checkpoint_storage_concurrent_gb, use_ocdbt, use_zarr3, missing_param_policy
+        )
       context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.ORBAX)
       with context:
         return ocp_v1.load_pytree(path, abstract_unboxed_pre_state)
@@ -368,6 +459,7 @@ def combine_sharding(sds, shardings):
           checkpoint_storage_concurrent_gb,
           use_ocdbt,
           use_zarr3,
+          missing_param_policy,
       )
 
     # Original v0 logic.
@@ -378,11 +470,8 @@ def combine_sharding(sds, shardings):
         use_ocdbt=use_ocdbt,
         use_zarr3=use_zarr3,
     )
-    # NNX checkpoints are saved as pure dicts (see maybe_save_checkpoint); the
-    # restore target must match — a boxed nnx.State wouldn't.
+    # Only Linen TrainState reaches here; nnx.State returned above.
     restore_target = abstract_unboxed_pre_state
-    if isinstance(abstract_unboxed_pre_state, nnx.State):
-      restore_target = abstract_unboxed_pre_state.to_pure_dict()
     # Provide sharding info to ensure restoration returns JAX arrays (not NumPy arrays).
     restore_args = jax.tree_util.tree_map(
         lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding),
@@ -792,51 +881,59 @@ def map_to_pspec(data):
         )
         ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True)
 
-      # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
+      # pure_nnx saves in the Linen on-disk layout; restore that layout (weights +
+      # opt_state + step + nnx_aux), restoring the grain iterator in place when
+      # present, then reshape it back into the NNX state.
       # (Emergency managers use their own restore path below.)
       if isinstance(abstract_unboxed_pre_state, nnx.State) and not isinstance(
           checkpoint_manager,
           (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager),
       ):
-        checkpoint_path = str(checkpoint_manager.directory / str(step) / "items")
-        restored_nnx = _load_linen_checkpoint_into_nnx(
-            checkpoint_path,
-            abstract_unboxed_pre_state,
-            checkpoint_storage_concurrent_gb,
-            use_ocdbt,
-            use_zarr3,
+        linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state)
+        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)
+        if (
+            dataset_type == "grain"
+            and data_iterator
+            and not isinstance(data_iterator, PlaceHolderDataIterator)
+            and (checkpoint_manager.directory / str(step) / "iter").exists()
+        ):
+          restored, _ = _restore_grain_iterator(
+              checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data
+          )
+        else:
+          restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args))
+        _enforce_missing_weight_policy(
+            abstract_unboxed_pre_state, restored["items"], _missing_param_policy(maxtext_config)
         )
+        restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state.to_pure_dict())
         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))
         restored = _restore_emergency_linen_checkpoint_into_nnx(
             checkpoint_manager,
             step,
             abstract_unboxed_pre_state,
             map_to_pspec,
+            _missing_param_policy(maxtext_config),
         )
         return (
             restored,
             None,
         )
 
-      # Convert nnx.State to pure dict to match how checkpoints are saved for NNX
+      # Only Linen TrainState reaches here; the NNX cases returned above.
       restore_target = abstract_unboxed_pre_state
-      if isinstance(abstract_unboxed_pre_state, nnx.State):
-        restore_target = abstract_unboxed_pre_state.to_pure_dict()
-
       restore_args = jax.tree_util.tree_map(map_to_pspec, restore_target)
       checkpoint_args = ocp.args.PyTreeRestore(
           item=restore_target,
           restore_args=restore_args,
           partial_restore=True,
       )
 
-      checkpoint_path = str(checkpoint_manager.directory / str(step))
       match (checkpoint_manager, dataset_type, data_iterator):
         # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager
         # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and
@@ -908,6 +1005,7 @@ def map_to_pspec(data):
         checkpoint_storage_concurrent_gb=checkpoint_storage_concurrent_gb,
         use_ocdbt=use_ocdbt,
         use_zarr3=use_zarr3,
+        missing_param_policy=_missing_param_policy(maxtext_config),
     )
     return {"items": restored_state}, None
   else:
@@ -1040,7 +1138,9 @@ 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())
+      # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout
+      # stream continues across resumes instead of resetting to a base key.
+      state = train_state_nnx.to_checkpoint_dict(state)
 
   # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic.
   # This occurs if this function was called:

src/maxtext/common/train_state_nnx.py

@@ -66,7 +66,9 @@ def apply_gradients(self, grads: Any, **kwargs):
 #   2. weights: model/... -> params/params/... (Linen `params` collection)
 #   3. opt_state: int-keyed dict (empty entries skipped) -> list with None for EmptyState,
 #      and mu/nu wrapped under the `params` collection
-# NNX-only rngs/dropout state is dropped (Linen never had it).
+# NNX-only state Linen never had (rngs/dropout and batch stats) is split out by
+# Variable type and stored under an `nnx_aux` subtree of `items` (see
+# `to_checkpoint_dict`), so it survives a resume instead of resetting.
 
 _NNX_RNG_STATE_KEYS = ("rngs", "dropout")
 
@@ -193,3 +195,20 @@ def from_linen_checkpoint_dict(linen_pure_dict):
   if optimizer:
     result["optimizer"] = optimizer
   return result
+
+
+def to_checkpoint_dict(state):
+  """Reshapes an nnx.State into the on-disk checkpoint layout, splitting by Variable type.
+
+  Weights (nnx.Param) and the optimizer state map to the Linen
+  params/opt_state/step layout, so pure_nnx and Linen checkpoints stay
+  interchangeable. rngs/dropout (nnx.RngState) and batch stats (nnx.BatchStat) go
+  under an `nnx_aux` subtree so they survive a resume. Works on a concrete state
+  (save) or an abstract state (restore target).
+  """
+  params, rng_state, batch_stats, rest = nnx.split_state(state, nnx.Param, nnx.RngState, nnx.BatchStat, ...)
+  linen = to_linen_checkpoint_dict(nnx.merge_state(params, rest).to_pure_dict())
+  aux = nnx.merge_state(rng_state, batch_stats).to_pure_dict()
+  if aux:
+    linen["nnx_aux"] = aux
+  return linen

src/maxtext/configs/base.yml

@@ -56,6 +56,9 @@ enable_checkpointing: true
 save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
+# On NNX restore, how to handle a weight present in the model but absent from the checkpoint:
+# "error" raises naming the parameter; "warn" logs a warning and zero-fills.
+checkpoint_missing_param_policy: "error"
 max_num_checkpoints_to_keep: None
 enable_continuous_checkpointing: false
 # enables one replica to read the ckpt then broadcast to the rest

src/maxtext/configs/types.py

@@ -328,6 +328,13 @@ class Checkpointing(BaseModel):
   load_checkpoint_only_once: bool = Field(False, description="If True, deep copy the reference model to the actor model.")
   async_checkpointing: bool = Field(True, description="If True, uses an asynchronous checkpointer for performance.")
   checkpoint_period: int = Field(10_000, description="The frequency (in steps) at which to save checkpoints.")
+  checkpoint_missing_param_policy: Literal["error", "warn"] = Field(
+      "error",
+      description=(
+          "How to handle a weight present in the model but absent from the checkpoint on NNX restore. "
+          "'error' raises naming the parameter; 'warn' logs a warning and zero-fills."
+      ),
+  )
   max_num_checkpoints_to_keep: int | None = Field(None, description="Maximum number of checkpoints to keep.")
   enable_single_replica_ckpt_restoring: bool = Field(
       False, description="One replica reads and broadcasts the checkpoint."