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

src/maxtext/common/checkpointing.py

@@ -194,23 +194,53 @@ 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. Paths absent from
+  it fall 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 (present in the model but not on disk): governed
+      by `missing_param_policy` -- "error" raises naming the path and shape,
+      "warn" logs a warning and zero-fills. Zero-filled weights are silent
+      accuracy loss, 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):
     return partial_concrete
+
+  # Absent 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. Verify the checkpoint matches "
+          "the model architecture, or set checkpoint_missing_param_policy=warn to zero-fill and continue."
+      )
   return _default_for_sds(abstract_pure)
 
 
@@ -220,11 +250,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`, then 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 +273,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 +295,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):
@@ -805,6 +838,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 +852,7 @@ def map_to_pspec(data):
             step,
             abstract_unboxed_pre_state,
             map_to_pspec,
+            _missing_param_policy(maxtext_config),
         )
         return (
             restored,

src/maxtext/common/train_state_nnx.py

@@ -69,6 +69,7 @@ def apply_gradients(self, grads: Any, **kwargs):
 # NNX-only rngs/dropout state is dropped (Linen never had it).
 
 _NNX_RNG_STATE_KEYS = ("rngs", "dropout")
+NNX_RNG_STATE_KEYS = _NNX_RNG_STATE_KEYS
 
 
 def _cast_step(step, dtype):

src/maxtext/configs/base.yml

@@ -56,6 +56,8 @@ enable_checkpointing: true
 save_checkpoint_on_completion: true
 async_checkpointing: true
 checkpoint_period: 10_000
+# How to handle a weight present in the model but missing from a restored checkpoint: "error" or "warn" (zero-fill).
+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 a restored checkpoint. "
+          "'error' raises naming the parameter; 'warn' logs a warning and zero-fills it."
+      ),
+  )
   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,68 @@ 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):
+  """Absent weights 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):
+    abstract = self._abstract()
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.full((1,), 5.0)}}}
+    out = checkpointing._populate_pure_dict_from_partial(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."""
+    abstract = self._abstract()
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1)), "bias": jnp.full((1,), 5.0)}}}
+    # policy="error" but dropout absence must not raise, because it is NNX-only state.
+    out = checkpointing._populate_pure_dict_from_partial(  # pylint: disable=protected-access
+        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."""
+    abstract = self._abstract()
+    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
+          abstract, partial, missing_param_policy="error"
+      )
+    self.assertIn("model/linear/bias", str(ctx.exception))
+    self.assertIn("(1,)", str(ctx.exception))  # the missing weight's expected shape
+    self.assertIn("missing parameter", str(ctx.exception))
+
+  def test_missing_weight_zero_fills_under_warn_policy(self):
+    abstract = self._abstract()
+    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
+          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)."""
+    abstract = self._abstract()
+    partial = {"model": {"linear": {"kernel": jnp.ones((2, 1))}}}  # bias missing
+    with self.assertRaises(ValueError):
+      checkpointing._populate_pure_dict_from_partial(abstract, partial)  # pylint: disable=protected-access
+
+
 if __name__ == "__main__":
   unittest.main()