Merge pull request #4304 from AI-Hypercomputer:jackyf/proactive-scan-layers

PiperOrigin-RevId: 941788193

Author: Google-ML-Automation

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(

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(

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

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