Take job config out of checkpoint manager (#1433)

This PR takes job_config out of the CheckpointManager class. Why?
JobConfig is a monolith -- it has knowledge of every part of a titan
training job. As a result, it is hard to actually use CheckpointManager
in a standalone fashion. In practice the job config is mostly only used
for its checkpoint config, plus two other usages as far as I can tell:

1) Getting the replica_id from the FTManager
2) Taking the dump_folder from the job field and joining it with the
checkpoint folder

For (1) we can just get this directly from FTManager without accessing
the JobConfig field. For (2) we can pass `job_config.job.dump_folder`
explicitly as a base folder, then join to `checkpoint_config.folder`.
Personally I would try to consolidate `job.dump_folder` and
`checkpoint.folder` (though I understand there are cases where only the
former is needed) under Checkpoint, but not sure if this is preferable
from titan's pov.

Author: ebsmothers

tests/unit_tests/test_checkpoint.py

@@ -175,7 +175,8 @@ def test_save_load_restores_state(self, mock_load, mock_save, mock_rank):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 
@@ -207,7 +208,8 @@ def test_save_and_purge_keeps_last_k_checkpoints(
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 
@@ -247,7 +249,8 @@ def test_nonzero_rank_does_not_purge_or_save(self, mock_load, mock_save, mock_ra
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         manager.save(curr_step=1)
@@ -269,7 +272,8 @@ def test_load_returns_false_when_no_checkpoint_folder(self):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         self.assertFalse(manager.load(step=-1))
@@ -292,7 +296,8 @@ def test_load_finds_latest_and_calls_dcp_load(self, mock_load, mock_rank):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         res = manager.load(step=-1)
@@ -321,7 +326,8 @@ def test_interval_respects_interval(self, mock_load, mock_save, mock_rank):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         manager.save(curr_step=1)
@@ -354,7 +360,8 @@ def test_last_save_model_only_and_initial_load_model_only(
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         manager1.save(curr_step=1, last_step=True)
@@ -373,7 +380,8 @@ def test_last_save_model_only_and_initial_load_model_only(
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
         r1 = manager2.load(step=1)
@@ -404,7 +412,8 @@ def test_async_save_calls_async_wait(self, mock_async_save, mock_new_group):
         """
         # Configure async mode
         job_config = DummyJobConfig(job=self.job_config.job)
-        job_config.checkpoint.async_mode = "async"
+        checkpoint_config = job_config.checkpoint
+        checkpoint_config.async_mode = "async"
         ft_manager = DummyFTManager()
         states = {"trainer": torch.tensor([0])}
         manager = CheckpointManager(
@@ -413,8 +422,9 @@ def test_async_save_calls_async_wait(self, mock_async_save, mock_new_group):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=states,
-            job_config=job_config,
-            ft_manager=ft_manager,
+            checkpoint_config=checkpoint_config,
+            base_folder=self.job_config.job.dump_folder,
+            ft_manager=self.ft_manager,
         )
 
         # First save schedules async
@@ -445,7 +455,8 @@ def test_ft_async_save_calls_async_wait(
         Test that with FT enabled, AsyncMode.ASYNC via FT triggers correct waits.
         """
         job_config = DummyJobConfig(job=self.job_config.job)
-        job_config.checkpoint.async_mode = "async"
+        checkpoint_config = job_config.checkpoint
+        checkpoint_config.async_mode = "async"
         ft_manager = mock.Mock()
         ft_manager.manager.return_value = mock.Mock()
         ft_manager.manager.participating_rank = mock.Mock(return_value=0)
@@ -456,8 +467,9 @@ def test_ft_async_save_calls_async_wait(
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=job_config,
-            ft_manager=ft_manager,
+            checkpoint_config=checkpoint_config,
+            base_folder=self.job_config.job.dump_folder,
+            ft_manager=self.ft_manager,
         )
 
         # Initially no future
@@ -491,7 +503,8 @@ def test_enable_first_step_checkpoint(self, mock_save, mock_rank):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 
@@ -516,7 +529,8 @@ def test_enable_first_step_checkpoint(self, mock_save, mock_rank):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 
@@ -561,7 +575,8 @@ def __init__(self):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 
@@ -610,7 +625,8 @@ def fake_load(state_dict: dict, checkpoint_id=None):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states=self.states,
-            job_config=self.job_config,
+            checkpoint_config=self.job_config.checkpoint,
+            base_folder=self.job_config.job.dump_folder,
             ft_manager=self.ft_manager,
         )
 

torchtitan/components/checkpoint.py

@@ -36,7 +36,7 @@
 from torchtitan.components.ft import FTManager
 from torchtitan.components.lr_scheduler import LRSchedulersContainer
 from torchtitan.components.optimizer import OptimizersContainer
-from torchtitan.config_manager import JobConfig, TORCH_DTYPE_MAP
+from torchtitan.config_manager import Checkpoint, TORCH_DTYPE_MAP
 from torchtitan.protocols.state_dict_adapter import StateDictAdapter
 from torchtitan.tools.logging import logger
 from torchtitan.tools.utils import GarbageCollection
@@ -174,10 +174,13 @@ class CheckpointManager:
         lr_schedulers (LRSchedulersContainer): The lr schedulers used to optimize the model.
         states (Dict[str, Any]): The states that need to be saved, other than the
             previous 4 components.
-        job_config (JobConfig): The job config used to configure the checkpointing.
+        checkpoint_config (Checkpoint): The config used to configure the checkpointing.
+        base_folder (str): The base folder to save the checkpoint. Will be concatenated
+            with checkpoint_config.folder
         sd_adapter (Optional[type[StateDictAdapter]]): The adapter used to convert model state
             dicts between native format and other formats.
         ft_manager (Optional[ft.Manager]): The FTManager from TorchFT.
+
     """
 
     def __init__(
@@ -187,13 +190,13 @@ def __init__(
         optimizers: OptimizersContainer,
         lr_schedulers: LRSchedulersContainer,
         states: dict[str, Any],
-        job_config: JobConfig,
+        checkpoint_config: Checkpoint,
+        base_folder: str,
         sd_adapter: type[StateDictAdapter] | None = None,
         ft_manager: FTManager | None = None,
     ) -> None:
-        ckpt_config = job_config.checkpoint
-        self.enable_checkpoint = ckpt_config.enable_checkpoint
-        self.last_save_in_hf = ckpt_config.last_save_in_hf
+        self.enable_checkpoint = checkpoint_config.enable_checkpoint
+        self.last_save_in_hf = checkpoint_config.last_save_in_hf
         if self.last_save_in_hf:
             assert (
                 sd_adapter is not None
@@ -224,9 +227,9 @@ def load_state_dict(state_dict):
                     self.states[k].load_state_dict(v)
 
             self.ft_manager.set_state_dict_fns(load_state_dict, state_dict)
-            self.ft_replica_id = job_config.fault_tolerance.replica_id
+            self.ft_replica_id = ft_manager.replica_id
 
-        async_mode = ckpt_config.async_mode.lower()
+        async_mode = checkpoint_config.async_mode.lower()
         self.enable_staging = (
             self.enable_checkpoint and async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM
         ) or self.ft_manager
@@ -251,27 +254,29 @@ def load_state_dict(state_dict):
         self.cpu_offload_state_dict = None
         self.stager = None
 
-        self.folder = os.path.join(job_config.job.dump_folder, ckpt_config.folder)
+        self.folder = os.path.join(base_folder, checkpoint_config.folder)
 
         # Checkpoint policy related fields.
-        self.initial_load_path = ckpt_config.initial_load_path
-        self.initial_load_model_only = ckpt_config.initial_load_model_only
-        self.last_save_model_only = ckpt_config.last_save_model_only
-        self.export_dtype = TORCH_DTYPE_MAP[ckpt_config.export_dtype]
-        self.exclude_from_loading = ckpt_config.exclude_from_loading
-        self.interval = ckpt_config.interval
-        self.enable_first_step_checkpoint = ckpt_config.enable_first_step_checkpoint
+        self.initial_load_path = checkpoint_config.initial_load_path
+        self.initial_load_model_only = checkpoint_config.initial_load_model_only
+        self.last_save_model_only = checkpoint_config.last_save_model_only
+        self.export_dtype = TORCH_DTYPE_MAP[checkpoint_config.export_dtype]
+        self.exclude_from_loading = checkpoint_config.exclude_from_loading
+        self.interval = checkpoint_config.interval
+        self.enable_first_step_checkpoint = (
+            checkpoint_config.enable_first_step_checkpoint
+        )
 
         # Async checkpoint related fields.
-        async_mode = ckpt_config.async_mode.lower()
+        async_mode = checkpoint_config.async_mode.lower()
         if (
             async_mode == AsyncMode.ASYNC
             or async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM
             or self.ft_manager
         ):
             self.pg = dist.new_group(backend="gloo")
 
-        self.keep_latest_k = ckpt_config.keep_latest_k
+        self.keep_latest_k = checkpoint_config.keep_latest_k
         if self.keep_latest_k > 0:
             if self.keep_latest_k == 1:
                 raise ValueError(
@@ -296,7 +301,9 @@ def load_state_dict(state_dict):
         elif async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
             self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
         else:
-            raise ValueError(f"Unkown checkpoint async_mode {ckpt_config.async_mode}")
+            raise ValueError(
+                f"Unkown checkpoint async_mode {checkpoint_config.async_mode}"
+            )
 
         logger.info(
             f"Checkpointing active. Checkpoints will be loaded from and saved to {self.folder}"

torchtitan/train.py

@@ -294,7 +294,8 @@ def __init__(self, job_config: JobConfig):
             optimizers=self.optimizers,
             lr_schedulers=self.lr_schedulers,
             states={"train_state": self},
-            job_config=job_config,
+            checkpoint_config=job_config.checkpoint,
+            base_folder=job_config.job.dump_folder,
             sd_adapter=self.train_spec.state_dict_adapter,
             ft_manager=self.ft_manager,
         )