allow disabling ft checkpoints

Summary:
Allows disabling the storage of checkpoints related to torchft.

Users don't really have to rely on any external storage. So it reduces set up time to get things up and running. Since we also don't really need model checkpoints when we have torchft. And if checkpoint storage has issues, this can work as a killswitch to completely disable the storage so it doesn't impact training.

Author: tushar00jain

torchtitan/components/checkpoint.py

@@ -193,7 +193,11 @@ def __init__(
         self.load_only = checkpoint_config.load_only
 
         self.ft_manager = (
-            ft_manager.manager if ft_manager and ft_manager.enabled else None
+            ft_manager.manager
+            if ft_manager
+            and ft_manager.enabled
+            and checkpoint_config.enable_ft_checkpoints
+            else None
         )
         if self.ft_manager:
             optimizers.init_cache_state_dict()

torchtitan/config/job_config.py

@@ -422,6 +422,21 @@ class Checkpoint:
     enable: bool = False
     """Whether to enable checkpoint"""
 
+    enable_ft_checkpoints: bool = True
+    """
+    Used to enable checkpointing state that's used for fault tolerant training with torchft.
+
+    Fault tolerant training stores data loader index in the checkpoints, so that training can resume
+    without going over the same batch twice.
+
+    If enabled, data loader state is checkpointed. Otherwise the data loader index will be infered
+    from the step count.
+
+    Note, if regular checkpointing is enabled, we also checkpoint the
+    data loader state. For non-fault tolerant training, we don't infer
+    data loader index if checkpointing is disabled.
+    """
+
     folder: str = "checkpoint"
     """
     The folder to store the checkpoints.

torchtitan/train.py

@@ -638,6 +638,12 @@ def state_dict(self) -> dict[str, Any]:
     def load_state_dict(self, state_dict: dict[str, Any]):
         self.step = state_dict["step"]
         self.ntokens_seen = state_dict["ntokens_seen"]
+        if (
+            self.job_config.fault_tolerance.enable
+            and not self.job_config.checkpoint.enable_ft_checkpoints
+        ):
+            # TODO: set index in the dataloader
+            pass
 
     def close(self) -> None:
         if self.checkpointer: