Centralize checkpoint search paths

Author: dlwh

experiments/grug/base/train.py

@@ -372,18 +372,9 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        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_search_paths=checkpoint_search_paths,
+            checkpoint_search_paths=trainer.checkpoint_search_paths(run_id),
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,

experiments/grug/modular_opt/train.py

@@ -372,18 +372,9 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        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_search_paths=checkpoint_search_paths,
+            checkpoint_search_paths=trainer.checkpoint_search_paths(run_id),
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,

experiments/grug/moe/train.py

@@ -410,18 +410,9 @@ def _init_state(model_rng):
         state = _init_state(model_key)
 
         checkpointer = trainer.checkpointer.create(run_id)
-        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_search_paths=checkpoint_search_paths,
+            checkpoint_search_paths=trainer.checkpoint_search_paths(run_id),
             load_checkpoint_setting=trainer.load_checkpoint,
             mesh=mesh,
             allow_partial=trainer.allow_partial_checkpoint,

lib/levanter/src/levanter/trainer.py

@@ -854,6 +854,17 @@ def batch_axis_name(self) -> str | None:
     """if None (default), we'll load a checkpoint if it exists. If true, we must load a checkpoint"""
     load_checkpoint_path: Optional[str] = None
     """can be a parent (to find latest) or a specific checkpoint. if None, will set to checkpointer.base_path."""
+
+    def checkpoint_search_paths(self, run_id: str) -> list[str]:
+        if self.load_checkpoint_path is not None:
+            return [self.load_checkpoint_path]
+
+        paths = [self.checkpointer.expanded_path(run_id)]
+        temp_path = self.checkpointer.expanded_temporary_path(run_id)
+        if temp_path is not None:
+            paths.append(temp_path)
+        return paths
+
     initialize_from: Optional[str] = None  # Levanter trainer checkpoint to initialize from
     """Load and continue training from a checkpoint. If None, will initialize from model_init."""
     allow_partial_checkpoint: bool = False

lib/levanter/tests/test_checkpoint.py

@@ -39,6 +39,7 @@
     save_checkpoint,
     unregister_debug_checkpointer_state_provider,
 )
+from levanter.trainer import TrainerConfig
 from levanter.trainer_state import TrainerState
 
 
@@ -348,6 +349,21 @@ def test_checkpointer_config_no_temporary_base_path():
     assert config.expanded_temporary_path("run1") is None
 
 
+def test_trainer_config_checkpoint_search_paths():
+    config = dataclasses.replace(
+        TrainerConfig(),
+        checkpointer=CheckpointerConfig(
+            base_path="/tmp/test-perm",
+            temporary_base_path="/tmp/test-temp",
+            append_run_id_to_base_path=True,
+        ),
+    )
+    assert config.checkpoint_search_paths("run1") == ["/tmp/test-perm/run1", "/tmp/test-temp/run1"]
+
+    pinned_config = dataclasses.replace(config, load_checkpoint_path="/tmp/test-perm/run1/step-100")
+    assert pinned_config.checkpoint_search_paths("run1") == ["/tmp/test-perm/run1/step-100"]
+
+
 def test_checkpointer_config_propagates_debug_settings():
     config = CheckpointerConfig(
         base_path="/tmp/checkpoints",