From f4076346cec5b301d0395303017d11ac2ebf3aab Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Tue, 30 Jun 2026 10:46:47 +0000 Subject: [PATCH] feat(scan_layers): verify scan_layers compatibility from checkpoint metadata --- src/maxtext/common/checkpointing.py | 11 ++-- src/maxtext/utils/model_creation_utils.py | 10 ++++ tests/post_training/unit/lora_utils_test.py | 2 +- tests/unit/model_creation_utils_test.py | 59 +++++++++++++++++++++ 4 files changed, 76 insertions(+), 6 deletions(-) diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index 0316cd29b5..15b16bb1aa 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -1201,11 +1201,12 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= grain_iters_to_save.append((data_iter.local_iterator, process_index, process_count_total)) save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save) - custom_metadata = None - if config and hasattr(config, "lora") and config.lora: - lora_rank = getattr(config.lora, "lora_rank", 0) - if lora_rank > 0 and hasattr(config.lora, "model_dump"): - custom_metadata = {"lora": config.lora.model_dump()} + custom_metadata = {} + if config: + if hasattr(config, "scan_layers"): + custom_metadata["scan_layers"] = config.scan_layers + if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0: + custom_metadata["lora"] = config.lora.model_dump() match (checkpoint_manager, config, data_iterator): case (checkpoint_manager, _, _) if isinstance( diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index c0248783eb..ef8e25feb0 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -44,6 +44,7 @@ import jax import jax.numpy as jnp from jax.sharding import Mesh +from maxtext.common import checkpointing from maxtext.common.checkpointing import handle_checkpoint_mismatch from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN from maxtext.configs import pyconfig @@ -864,6 +865,15 @@ def from_pretrained( } ) config = pyconfig.HyperParameters(new_config) + # Proactive verification of scan_layers from checkpoint metadata + if config.load_parameters_path: + custom_metadata = checkpointing.load_checkpoint_metadata(config.load_parameters_path) + saved_scan_layers = custom_metadata.get("scan_layers") + if isinstance(saved_scan_layers, bool) and saved_scan_layers != config.scan_layers: + raise ValueError( + f"Configuration mismatch: Your run specifies scan_layers={config.scan_layers}, " + f"but the checkpoint was saved with scan_layers={saved_scan_layers}." + ) if config.pure_nnx: _create_model, abstract_model = create_nnx_abstract_model( diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index 0adb1a2768..e910995542 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/tests/post_training/unit/lora_utils_test.py @@ -389,7 +389,7 @@ def test_save_checkpoint_passes_metadata(self): mock_manager.save.assert_called_once() _, kwargs = mock_manager.save.call_args self.assertIn("custom_metadata", kwargs) - self.assertEqual(kwargs["custom_metadata"], {"lora": cfg.lora.model_dump()}) + self.assertEqual(kwargs["custom_metadata"]["lora"], cfg.lora.model_dump()) def test_save_and_restore_metadata_integration(self): """Integration test checking that Orbax CheckpointManager writes and reads custom LoRA metadata.""" diff --git a/tests/unit/model_creation_utils_test.py b/tests/unit/model_creation_utils_test.py index 2568547944..b1d034e8b7 100644 --- a/tests/unit/model_creation_utils_test.py +++ b/tests/unit/model_creation_utils_test.py @@ -723,6 +723,65 @@ def test_checkpoint_load_error_propagates(self, mock_ocp): with self.assertRaises(RuntimeError): model_creation_utils.from_pretrained(cfg, self.mesh) + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_scan_layers_mismatch_raises_error(self, mock_load_meta): + """ValueError is raised if run specifies scan_layers=True but checkpoint specifies scan_layers=False.""" + mock_load_meta.return_value = {"scan_layers": False} + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_false_ckpt", scan_layers=True + ) + + with self.assertRaises(ValueError) as context: + model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIn( + "Configuration mismatch: Your run specifies scan_layers=True, " + "but the checkpoint was saved with scan_layers=False", + str(context.exception), + ) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.ocp") + def test_scan_layers_match_no_error(self, mock_ocp, mock_load_meta): + """If the run specifies scan_layers=True and the checkpoint matches, it proceeds without error.""" + mock_load_meta.return_value = {"scan_layers": True} + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_linen_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_true_ckpt", scan_layers=True + ) + + model = model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIsInstance(model, models.Transformer) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.ocp") + def test_scan_layers_missing_metadata_no_error(self, mock_ocp, mock_load_meta): + """Skip verification and proceed if custom_metadata lacks 'scan_layers'.""" + mock_load_meta.return_value = {} + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_linen_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_missing_ckpt", scan_layers=True + ) + + model = model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIsInstance(model, models.Transformer) + class TestSetupDecodeStateFromNnx(unittest.TestCase): """Tests for setup_decode_state_from_nnx()."""