diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c527427144..961a5ba71b 100644 --- a/src/maxtext/common/checkpointing.py +++ b/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,7 +273,7 @@ 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( @@ -249,6 +281,7 @@ def _restore_emergency_linen_checkpoint_into_nnx( 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, diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index a123956a73..2620d8c23a 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/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): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..6537393ad2 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..9bdcbd79e0 100644 --- a/src/maxtext/configs/types.py +++ b/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." diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 04aa898bde..f3bc2f06bc 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/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()