diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c527427144..43e331b8c6 100644 --- a/src/maxtext/common/checkpointing.py +++ b/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,7 +278,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 +286,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 +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: diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index a123956a73..2e30fa7d3b 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/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): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..3db375465a 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a8c1745..25c931d10a 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 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." diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 04aa898bde..1ce44a75ea 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/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() diff --git a/tests/unit/checkpointing_nnx_missing_param_test.py b/tests/unit/checkpointing_nnx_missing_param_test.py new file mode 100644 index 0000000000..b0b46ea3e1 --- /dev/null +++ b/tests/unit/checkpointing_nnx_missing_param_test.py @@ -0,0 +1,158 @@ +# Copyright 2025-2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end proof that a weight missing from the checkpoint fails loudly on NNX restore. + +Drives the real create_orbax_checkpoint_manager -> maybe_save_checkpoint -> +load_state_if_possible stack (no mocks). A weight present in the model but absent +from the checkpoint -- whether Orbax omits it or leaves it as an unmaterialized +ShapeDtypeStruct -- must raise a descriptive error naming the parameter, instead of +being silently zero-filled or surfacing later as a cryptic compile error in the +first train_step. `checkpoint_missing_param_policy=warn` opts back into zero-filling. +""" + +import os +import shutil +import tempfile +import unittest +from types import SimpleNamespace + +from flax import nnx +import jax +from maxtext.common import checkpointing +from maxtext.common import train_state_nnx +import optax + + +class _Model(nnx.Module): + """Linear + dropout, so the state carries rngs/dropout that gets stripped on save.""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 1, rngs=rngs) + self.dropout = nnx.Dropout(rate=0.5, rngs=rngs) + + def __call__(self, x, deterministic=False): + return self.dropout(self.linear(x), deterministic=deterministic) + + +class _ModelExtraLayer(nnx.Module): + """`_Model` plus a layer absent from a `_Model` checkpoint (a weight with no array on disk).""" + + def __init__(self, rngs: nnx.Rngs): + self.linear = nnx.Linear(2, 1, rngs=rngs) + self.dropout = nnx.Dropout(rate=0.5, rngs=rngs) + self.extra = nnx.Linear(1, 1, rngs=rngs) + + def __call__(self, x, deterministic=False): + return self.extra(self.dropout(self.linear(x), deterministic=deterministic)) + + +_TX = optax.adam(1e-3) + + +def _config(missing_param_policy="error"): + """Minimal config with the fields save/restore reads for a pure_nnx run.""" + return SimpleNamespace( + pure_nnx=True, + enable_diloco=False, + enable_checkpointing=True, + enable_continuous_checkpointing=False, + enable_emergency_checkpoint=False, + enable_autocheckpoint=False, + checkpoint_period=1, + local_checkpoint_period=0, + async_checkpointing=False, + dataset_type="tfds", + lora=None, + checkpoint_storage_target_data_file_size_bytes=checkpointing.DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE, + elastic_enabled=False, + checkpoint_missing_param_policy=missing_param_policy, + ) + + +def _abstract_state(model_cls): + """An abstract (ShapeDtypeStruct) nnx.State for `model_cls`, the restore blueprint.""" + mesh = jax.sharding.Mesh(jax.devices(), ("x",)) + sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) + + def make(): + model = model_cls(nnx.Rngs(9)) + return nnx.state(train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param))) + + abstract = nnx.eval_shape(make) + return jax.tree.map( + lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype, sharding=sharding) if hasattr(x, "shape") else x, + abstract, + ) + + +class TestMissingParamRoundTrip(unittest.TestCase): + """Real save->restore covering the missing-weight policy end to end.""" + + def setUp(self): + self._dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self._dir, ignore_errors=True) + + def _save_model(self): + """Saves a one-step `_Model` checkpoint and returns its manager.""" + manager = checkpointing.create_orbax_checkpoint_manager( + os.path.join(self._dir, "ckpt"), + enable_checkpointing=True, + use_async=False, + save_interval_steps=1, + dataset_type="tfds", + ) + model = _Model(nnx.Rngs(0)) + state = train_state_nnx.TrainStateNNX(model, nnx.Optimizer(model, _TX, wrt=nnx.Param)) + checkpointing.maybe_save_checkpoint(manager, nnx.state(state), _config(), data_iterator=None, step=1) + manager.wait_until_finished() + return manager + + def _restore(self, manager, model_cls, policy="error"): + """Restores the saved checkpoint into an abstract state for `model_cls` under `policy`.""" + full, _ = checkpointing.load_state_if_possible( + manager, + data_iterator=None, + load_parameters_from_path="", + load_full_state_from_path="", + checkpoint_storage_concurrent_gb=8, + abstract_unboxed_pre_state=_abstract_state(model_cls), + dataset_type="tfds", + maxtext_config=_config(policy), + ) + return full["items"] + + def test_missing_weight_raises_naming_the_path(self): + manager = self._save_model() + with self.assertRaises(ValueError) as ctx: + self._restore(manager, _ModelExtraLayer, policy="error") + msg = str(ctx.exception) + self.assertIn("model/extra", msg) # the exact offending parameter path + self.assertIn("missing parameter", msg) + + def test_missing_weight_zero_fills_under_warn(self): + """The warn policy zero-fills the absent weight and lets restore proceed.""" + manager = self._save_model() + restored = self._restore(manager, _ModelExtraLayer, policy="warn") + self.assertEqual(int(restored["model"]["extra"]["bias"].sum()), 0) + + def test_matching_config_restores_clean(self): + """Negative control: an exact-arch restore returns concrete arrays, no error.""" + manager = self._save_model() + restored = self._restore(manager, _Model) + self.assertNotIsInstance(restored["model"]["linear"]["kernel"], jax.ShapeDtypeStruct) + + +if __name__ == "__main__": + unittest.main()