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,7 @@ 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

@@ -414,6 +414,17 @@ 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.
+    """
+
     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: