Use explicit grug checkpoint search paths

Author: dlwh

experiments/grug/base/train.py

@@ -372,20 +372,21 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        checkpoint_path = trainer.load_checkpoint_path
-        if checkpoint_path is None and checkpointer is not None:
-            checkpoint_path = trainer.checkpointer.expanded_path(run_id)
-        additional_checkpoint_paths = []
-        temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
-        if temp_path is not None:
-            additional_checkpoint_paths.append(temp_path)
+        if trainer.load_checkpoint_path is not None:
+            checkpoint_search_paths = [trainer.load_checkpoint_path]
+        elif checkpointer is not None:
+            checkpoint_search_paths = [trainer.checkpointer.expanded_path(run_id)]
+            temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
+            if temp_path is not None:
+                checkpoint_search_paths.append(temp_path)
+        else:
+            checkpoint_search_paths = []
         state = restore_grug_state_from_checkpoint(
             state,
-            checkpoint_path=checkpoint_path,
+            checkpoint_search_paths=checkpoint_search_paths,
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,
-            additional_checkpoint_paths=additional_checkpoint_paths,
         )
 
         levanter.tracker.log_summary({"parameter_count": parameter_count(state.params)})

experiments/grug/checkpointing.py

@@ -7,7 +7,7 @@
 import logging
 import os
 import urllib.parse
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import TypeVar
 
 import fsspec
@@ -26,29 +26,20 @@ def _get_fs_and_plain_path(path: str) -> tuple[AbstractFileSystem, str]:
     return fs, plain_path
 
 
-def _checkpoint_candidates(checkpoint_path: str, *, additional_paths: list[str] | None = None) -> list[str]:
-    if _is_checkpoint_dir(checkpoint_path):
-        return [checkpoint_path]
-
-    all_roots = [checkpoint_path] + (additional_paths or [])
-
+def _checkpoint_candidates(checkpoint_search_paths: Sequence[str]) -> list[str]:
     candidates: list[tuple[int, str, str]] = []
-    for root in all_roots:
-        candidates.extend(_scan_checkpoint_root(root))
+    for search_path in checkpoint_search_paths:
+        candidates.extend(_scan_checkpoint_root(search_path))
 
     candidates.sort(key=lambda item: (item[0], item[1]), reverse=True)
     ordered_candidates = [candidate for _, _, candidate in candidates]
-    if checkpoint_path not in ordered_candidates:
-        ordered_candidates.append(checkpoint_path)
 
+    for search_path in checkpoint_search_paths:
+        if search_path not in ordered_candidates:
+            ordered_candidates.append(search_path)
     return ordered_candidates
 
 
-def _is_checkpoint_dir(checkpoint_path: str) -> bool:
-    fs, plain_path = _get_fs_and_plain_path(checkpoint_path)
-    return fs.exists(os.path.join(plain_path, "metadata.json"))
-
-
 def _scan_checkpoint_root(root_path: str) -> list[tuple[int, str, str]]:
     """Scan a single root path and return (step, timestamp, path) tuples."""
     fs, plain_path = _get_fs_and_plain_path(root_path)
@@ -91,22 +82,21 @@ def maybe_unstrip_protocol(path: str) -> str:
 def restore_grug_state_from_checkpoint(
     state: StateT,
     *,
-    checkpoint_path: str | None,
+    checkpoint_search_paths: Sequence[str],
     load_checkpoint_setting: bool | None,
     mesh: jax.sharding.Mesh | None,
     allow_partial: bool,
-    additional_checkpoint_paths: list[str] | None = None,
     _load_fn: Callable[..., StateT] = load_checkpoint,
 ) -> StateT:
-    if checkpoint_path is None:
+    if not checkpoint_search_paths:
         if load_checkpoint_setting:
-            raise FileNotFoundError("load_checkpoint=True but no checkpoint path is configured.")
+            raise FileNotFoundError("load_checkpoint=True but no checkpoint search paths are configured.")
         return state
 
     if load_checkpoint_setting is False:
         return state
 
-    candidates = _checkpoint_candidates(checkpoint_path, additional_paths=additional_checkpoint_paths or [])
+    candidates = _checkpoint_candidates(checkpoint_search_paths)
     last_error: FileNotFoundError | None = None
 
     for candidate in candidates:
@@ -118,8 +108,8 @@ def restore_grug_state_from_checkpoint(
                 allow_partial=allow_partial,
                 load_fn=_load_fn,
             )
-            if candidate != checkpoint_path:
-                logger.info("Loaded checkpoint %s from %s", checkpoint_path, candidate)
+            if candidate not in checkpoint_search_paths:
+                logger.info("Loaded checkpoint from %s while searching %s", candidate, checkpoint_search_paths)
             return loaded
         except FileNotFoundError as exc:
             last_error = exc
@@ -128,14 +118,15 @@ def restore_grug_state_from_checkpoint(
             )
 
     if load_checkpoint_setting is True:
+        search_path_summary = ", ".join(checkpoint_search_paths)
         attempted = ", ".join(candidates)
         if last_error is None:
-            raise FileNotFoundError(f"Could not find checkpoint at {checkpoint_path}")
+            raise FileNotFoundError(f"Could not find checkpoint under any of: {search_path_summary}")
         raise FileNotFoundError(
-            f"Could not load a checkpoint from {checkpoint_path}. Attempted: {attempted}"
+            f"Could not load a checkpoint from search paths {search_path_summary}. Attempted: {attempted}"
         ) from last_error
 
-    logger.info(f"Checkpoint not found at {checkpoint_path}. Starting from scratch.")
+    logger.info("Checkpoint not found under %s. Starting from scratch.", checkpoint_search_paths)
     return state
 
 

experiments/grug/modular_opt/train.py

@@ -372,20 +372,21 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        checkpoint_path = trainer.load_checkpoint_path
-        if checkpoint_path is None and checkpointer is not None:
-            checkpoint_path = trainer.checkpointer.expanded_path(run_id)
-        additional_checkpoint_paths = []
-        temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
-        if temp_path is not None:
-            additional_checkpoint_paths.append(temp_path)
+        if trainer.load_checkpoint_path is not None:
+            checkpoint_search_paths = [trainer.load_checkpoint_path]
+        elif checkpointer is not None:
+            checkpoint_search_paths = [trainer.checkpointer.expanded_path(run_id)]
+            temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
+            if temp_path is not None:
+                checkpoint_search_paths.append(temp_path)
+        else:
+            checkpoint_search_paths = []
         state = restore_grug_state_from_checkpoint(
             state,
-            checkpoint_path=checkpoint_path,
+            checkpoint_search_paths=checkpoint_search_paths,
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,
-            additional_checkpoint_paths=additional_checkpoint_paths,
         )
 
         levanter.tracker.log_summary({"parameter_count": parameter_count(state.params)})

experiments/grug/moe/train.py

@@ -410,20 +410,21 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        checkpoint_path = trainer.load_checkpoint_path
-        if checkpoint_path is None and checkpointer is not None:
-            checkpoint_path = trainer.checkpointer.expanded_path(run_id)
-        additional_checkpoint_paths = []
-        temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
-        if temp_path is not None:
-            additional_checkpoint_paths.append(temp_path)
+        if trainer.load_checkpoint_path is not None:
+            checkpoint_search_paths = [trainer.load_checkpoint_path]
+        elif checkpointer is not None:
+            checkpoint_search_paths = [trainer.checkpointer.expanded_path(run_id)]
+            temp_path = trainer.checkpointer.expanded_temporary_path(run_id)
+            if temp_path is not None:
+                checkpoint_search_paths.append(temp_path)
+        else:
+            checkpoint_search_paths = []
         state = restore_grug_state_from_checkpoint(
             state,
-            checkpoint_path=checkpoint_path,
+            checkpoint_search_paths=checkpoint_search_paths,
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,
-            additional_checkpoint_paths=additional_checkpoint_paths,
         )
 
         levanter.tracker.log_summary({"parameter_count": parameter_count(state.params)})

tests/test_grug_checkpointing.py

@@ -34,7 +34,7 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
 
     loaded = restore_grug_state_from_checkpoint(
         {"state": "init"},
-        checkpoint_path=str(checkpoint_root),
+        checkpoint_search_paths=[str(checkpoint_root)],
         load_checkpoint_setting=True,
         mesh=None,
         allow_partial=False,
@@ -61,7 +61,7 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
 
     loaded = restore_grug_state_from_checkpoint(
         {"state": "init"},
-        checkpoint_path=str(checkpoint_root),
+        checkpoint_search_paths=[str(checkpoint_root)],
         load_checkpoint_setting=None,
         mesh=None,
         allow_partial=False,
@@ -86,15 +86,15 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
     with pytest.raises(FileNotFoundError, match="Could not load a checkpoint"):
         restore_grug_state_from_checkpoint(
             {"state": "init"},
-            checkpoint_path=str(checkpoint_root),
+            checkpoint_search_paths=[str(checkpoint_root)],
             load_checkpoint_setting=True,
             mesh=None,
             allow_partial=False,
             _load_fn=fake_load,
         )
 
 
-def test_restore_discovers_candidates_across_additional_paths(tmp_path: Path):
+def test_restore_discovers_candidates_across_search_paths(tmp_path: Path):
     permanent_root = tmp_path / "checkpoints"
     temp_root = tmp_path / "checkpoints-temp"
 
@@ -109,20 +109,19 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
 
     loaded = restore_grug_state_from_checkpoint(
         {"state": "init"},
-        checkpoint_path=str(permanent_root),
+        checkpoint_search_paths=[str(permanent_root), str(temp_root)],
         load_checkpoint_setting=True,
         mesh=None,
         allow_partial=False,
-        additional_checkpoint_paths=[str(temp_root)],
         _load_fn=fake_load,
     )
 
-    # step-150 from temp root should be preferred (highest step)
+    # step-150 from temp root should be preferred (highest step).
     assert attempted == [str(temp_root / "step-150")]
     assert loaded == {"loaded_from": str(temp_root / "step-150")}
 
 
-def test_restore_respects_explicit_checkpoint_path_with_additional_paths(tmp_path: Path):
+def test_restore_respects_explicit_checkpoint_path_as_single_search_path(tmp_path: Path):
     permanent_root = tmp_path / "checkpoints"
     temp_root = tmp_path / "checkpoints-temp"
     explicit_checkpoint = permanent_root / "step-100"
@@ -138,11 +137,10 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
 
     loaded = restore_grug_state_from_checkpoint(
         {"state": "init"},
-        checkpoint_path=str(explicit_checkpoint),
+        checkpoint_search_paths=[str(explicit_checkpoint)],
         load_checkpoint_setting=True,
         mesh=None,
         allow_partial=False,
-        additional_checkpoint_paths=[str(temp_root)],
         _load_fn=fake_load,
     )
 
@@ -167,11 +165,10 @@ def fake_load(state, path, *, discover_latest, axis_mapping, mesh, allow_partial
 
     loaded = restore_grug_state_from_checkpoint(
         {"state": "init"},
-        checkpoint_path=str(permanent_root),
+        checkpoint_search_paths=[str(permanent_root), str(temp_root)],
         load_checkpoint_setting=None,
         mesh=None,
         allow_partial=False,
-        additional_checkpoint_paths=[str(temp_root)],
         _load_fn=fake_load,
     )
 
@@ -197,7 +194,7 @@ def test_restore_supports_legacy_wrapped_and_current_checkpoint_formats(tmp_path
 
     loaded_legacy = restore_grug_state_from_checkpoint(
         template_state,
-        checkpoint_path=str(checkpoint_root),
+        checkpoint_search_paths=[str(checkpoint_root)],
         load_checkpoint_setting=True,
         mesh=None,
         allow_partial=False,
@@ -210,7 +207,7 @@ def test_restore_supports_legacy_wrapped_and_current_checkpoint_formats(tmp_path
 
     loaded_current = restore_grug_state_from_checkpoint(
         template_state,
-        checkpoint_path=str(checkpoint_root),
+        checkpoint_search_paths=[str(checkpoint_root)],
         load_checkpoint_setting=True,
         mesh=None,
         allow_partial=False,