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

src/maxtext/common/checkpointing.py

@@ -194,23 +194,58 @@ def _make():
   return jax.jit(_make, out_shardings=sharding)()
 
 
-def _populate_pure_dict_from_partial(abstract_pure, partial_concrete):
-  """Fills `abstract_pure` with values from `partial_concrete` (by path), defaulting the rest.
+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"
 
-  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`.
+
+def _populate_pure_dict_from_partial(
+    abstract_pure, partial_concrete, missing_param_policy="error", path=(), in_rng=False
+):
+  """Fills `abstract_pure` with values from `partial_concrete` (by path), handling the rest by policy.
+
+  Paths present in `partial_concrete` take the restored value. A path absent from
+  it, or one Orbax left as an unmaterialized ShapeDtypeStruct (no matching array on
+  disk), falls into two cases:
+    - NNX-only rngs/dropout state the Linen checkpoint never carried: silently
+      filled with `_default_for_sds` (a deterministic base RNG / zeros).
+    - a genuinely missing weight: governed by `missing_param_policy` -- "error"
+      raises naming the path and shape, "warn" logs a warning and zero-fills.
+      A zero-filled or unmaterialized weight is silent accuracy loss (or a cryptic
+      compile failure in the first train_step), so "error" is the default.
   """
   if isinstance(abstract_pure, dict):
     return {
         k: _populate_pure_dict_from_partial(
             v,
             partial_concrete.get(k) if isinstance(partial_concrete, dict) else None,
+            missing_param_policy,
+            path + (k,),
+            in_rng or k in train_state_nnx.NNX_RNG_STATE_KEYS,
         )
         for k, v in abstract_pure.items()
     }
-  if partial_concrete is not None and not isinstance(partial_concrete, dict):
+  # A present, materialized leaf takes the restored value. A surviving ShapeDtypeStruct
+  # means Orbax had no matching array, so treat it as absent (handled below).
+  if partial_concrete is not None and not isinstance(partial_concrete, (dict, jax.ShapeDtypeStruct)):
     return partial_concrete
+
+  # Absent (or unmaterialized) leaf. rngs/dropout are expected to be absent from a
+  # Linen-layout checkpoint; anything else is a real weight the checkpoint should have had.
+  if not in_rng:
+    key = "/".join(str(p) for p in path)
+    shape = getattr(abstract_pure, "shape", "?")
+    dtype = getattr(abstract_pure, "dtype", "?")
+    if missing_param_policy == "warn":
+      max_logging.log(f"WARNING: checkpoint missing parameter '{key}' ({shape} {dtype}); zero-filling.")
+    else:
+      raise ValueError(
+          f"Checkpoint is missing parameter '{key}' (model expects {shape} {dtype}).\n"
+          "This weight would otherwise 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."
+      )
   return _default_for_sds(abstract_pure)
 
 
@@ -220,11 +255,13 @@ def _load_linen_checkpoint_into_nnx(
     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.
+  `from_linen_checkpoint_dict`, fills NNX-only rngs/dropout with defaults, and
+  handles any genuinely-missing weight 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()
@@ -241,14 +278,15 @@ 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)
-  return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx)
+  return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx, missing_param_policy)
 
 
 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."""
   max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}")
@@ -262,7 +300,7 @@ def _restore_emergency_linen_checkpoint_into_nnx(
   )
   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)
+  return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx, missing_param_policy)
 
 
 def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights):
@@ -318,6 +356,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.
 
@@ -368,6 +407,7 @@ def combine_sharding(sds, shardings):
           checkpoint_storage_concurrent_gb,
           use_ocdbt,
           use_zarr3,
+          missing_param_policy,
       )
 
     # Original v0 logic.
@@ -805,6 +845,7 @@ def map_to_pspec(data):
             checkpoint_storage_concurrent_gb,
             use_ocdbt,
             use_zarr3,
+            _missing_param_policy(maxtext_config),
         )
         return ({"items": restored_nnx}, None)
 
@@ -818,6 +859,7 @@ def map_to_pspec(data):
             step,
             abstract_unboxed_pre_state,
             map_to_pspec,
+            _missing_param_policy(maxtext_config),
         )
         return (
             restored,
@@ -908,6 +950,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:

src/maxtext/common/train_state_nnx.py

@@ -69,6 +69,9 @@ def apply_gradients(self, grads: Any, **kwargs):
 # NNX-only rngs/dropout state is dropped (Linen never had it).
 
 _NNX_RNG_STATE_KEYS = ("rngs", "dropout")
+# Public alias: checkpoint restore uses this to tell NNX-only rng/dropout state
+# (legitimately absent from a Linen checkpoint) apart from a genuinely missing weight.
+NNX_RNG_STATE_KEYS = _NNX_RNG_STATE_KEYS
 
 
 def _cast_step(step, dtype):

src/maxtext/configs/base.yml

@@ -57,6 +57,9 @@ save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
 max_num_checkpoints_to_keep: None
+# 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"
 enable_continuous_checkpointing: false
 # enables one replica to read the ckpt then broadcast to the rest
 enable_single_replica_ckpt_restoring: false

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."

tests/unit/checkpointing_nnx_load_test.py

@@ -185,5 +185,71 @@ def test_linen_layout_params_restore_into_nnx_state(self):
     self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]))
 
 
+class TestMissingParamPolicy(unittest.TestCase):
+  """Weights absent from the checkpoint must not silently zero-fill under the default policy."""
+
+  def setUp(self):
+    self.kernel = jax.ShapeDtypeStruct((2, 1), jnp.float32)
+    self.bias = jax.ShapeDtypeStruct((1,), jnp.float32)
+    self.count = jax.ShapeDtypeStruct((), jnp.uint32)
+
+  def _abstract(self):
+    return {
+        "model": {
+            "linear": {"kernel": self.kernel, "bias": self.bias},
+            "dropout": {"count": self.count},
+        }
+    }
+
+  def test_present_weights_pass_through(self):
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.full((1,), 5.0)}}}
+    out = checkpointing._populate_pure_dict_from_partial(self._abstract(), partial)  # pylint: disable=protected-access
+    self.assertTrue(jnp.array_equal(out["model"]["linear"]["kernel"], jnp.ones((2, 1))))
+    self.assertTrue(jnp.array_equal(out["model"]["linear"]["bias"], jnp.full((1,), 5.0)))
+
+  def test_rng_dropout_absent_is_silently_defaulted(self):
+    """rngs/dropout are legitimately absent from a Linen checkpoint -- default, don't error."""
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.full((1,), 5.0)}}}
+    out = checkpointing._populate_pure_dict_from_partial(  # pylint: disable=protected-access
+        self._abstract(), partial, missing_param_policy="error"
+    )
+    self.assertEqual(out["model"]["dropout"]["count"].shape, ())
+
+  def test_missing_weight_raises_under_error_policy(self):
+    """A genuine missing weight (bias) must raise naming its path, not zero-fill."""
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1))}}}  # bias missing
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._populate_pure_dict_from_partial(  # pylint: disable=protected-access
+          self._abstract(), partial, missing_param_policy="error"
+      )
+    self.assertIn("model/linear/bias", str(ctx.exception))
+    self.assertIn("(1,)", str(ctx.exception))
+    self.assertIn("missing parameter", str(ctx.exception))
+
+  def test_surviving_shape_dtype_struct_treated_as_missing(self):
+    """A weight Orbax left as an unmaterialized SDS is treated as missing, not passed through."""
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jax.ShapeDtypeStruct((1,), jnp.float32)}}}
+    with self.assertRaises(ValueError) as ctx:
+      checkpointing._populate_pure_dict_from_partial(  # pylint: disable=protected-access
+          self._abstract(), partial, missing_param_policy="error"
+      )
+    self.assertIn("model/linear/bias", str(ctx.exception))
+
+  def test_missing_weight_zero_fills_under_warn_policy(self):
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1))}}}  # bias missing
+    with mock.patch.object(checkpointing.max_logging, "log") as logmock:
+      out = checkpointing._populate_pure_dict_from_partial(  # pylint: disable=protected-access
+          self._abstract(), partial, missing_param_policy="warn"
+      )
+    self.assertTrue(jnp.array_equal(out["model"]["linear"]["bias"], jnp.zeros((1,))))
+    self.assertTrue(any("missing parameter 'model/linear/bias'" in str(c) for c in logmock.call_args_list))
+
+  def test_default_policy_is_error(self):
+    """No explicit policy means error (the safe default)."""
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1))}}}  # bias missing
+    with self.assertRaises(ValueError):
+      checkpointing._populate_pure_dict_from_partial(self._abstract(), partial)  # pylint: disable=protected-access
+
+
 if __name__ == "__main__":
   unittest.main()