set pg names

torchtitan/components/checkpoint.py

@@ -189,21 +189,32 @@ def __init__(
         self.enable = checkpoint_config.enable
         self.load_only = checkpoint_config.load_only
 
+        self.states = states
+        self.states.update(
+            {
+                MODEL: ModelWrapper(model_parts),
+                OPTIMIZER: optimizers,
+                DATALOADER: dataloader,
+                LR_SCHEDULER: lr_schedulers,
+            }
+        )
+
+        ft_inner_manager = (
+            ft_manager.manager if ft_manager and ft_manager.enabled else None
+        )
         self.ft_manager = (
-            ft_manager.manager
-            if ft_manager
-            and ft_manager.enabled
-            and checkpoint_config.enable_ft_dataloader_checkpoints
+            ft_inner_manager
+            if ft_inner_manager and checkpoint_config.enable_ft_dataloader_checkpoints
             else None
         )
 
-        if ft_manager and ft_manager.enabled and not self.ft_manager:
+        if ft_inner_manager and not self.ft_manager:
             logger.warn(
                 "Fault tolerance is enabled but enable_ft_dataloader_checkpoints is False. "
                 "This means replicas can retrain over the same data multiple times, which can result in overfitting."
             )
 
-        if self.ft_manager:
+        if ft_inner_manager:
             optimizers.init_cache_state_dict()
 
             def state_dict():
@@ -223,7 +234,7 @@ def load_state_dict(state_dict):
                 for k, v in state_dict.items():
                     self.states[k].load_state_dict(v)
 
-            self.ft_manager.set_state_dict_fns(load_state_dict, state_dict)
+            ft_inner_manager.set_state_dict_fns(load_state_dict, state_dict)
             self.ft_replica_id = ft_manager.replica_id
 
         async_mode = checkpoint_config.async_mode.lower()
@@ -234,15 +245,6 @@ def load_state_dict(state_dict):
         if not self.enable and self.ft_manager is None:
             return
 
-        self.states = states
-        self.states.update(
-            {
-                MODEL: ModelWrapper(model_parts),
-                OPTIMIZER: optimizers,
-                DATALOADER: dataloader,
-                LR_SCHEDULER: lr_schedulers,
-            }
-        )
         self.ft_states = {DATALOADER: dataloader}
 
         self.staging = False

torchtitan/distributed/utils.py

@@ -239,7 +239,7 @@ def maybe_enable_amp(
 
 
 def init_distributed(
-    comm_config: CommConfig, enable_cpu_backend: bool = False, base_folder: str = ""
+    comm_config: CommConfig, enable_cpu_backend: bool = False, base_folder: str = "", ranks: list[int] = []
 ):
     def _warn_overwrite_env(env, val):
         if env in os.environ:
@@ -283,6 +283,7 @@ def _get_distributed_backend(enable_cpu_backend):
     torch.distributed.init_process_group(
         backend=_get_distributed_backend(enable_cpu_backend),
         timeout=timedelta(seconds=comm_config.init_timeout_seconds),
+        ranks=ranks,
     )
 
 

torchtitan/train.py

@@ -86,11 +86,21 @@ def __init__(self, job_config: JobConfig):
         # Device has to be set before creating TorchFT manager.
         device_module.set_device(self.device)
 
+        ranks = []
+        ft_config = job_config.fault_tolerance
+        if ft_config.enable:
+            group_size = ft_config.group_size
+            replica_id = ft_config.replica_id
+            first_rank = replica_id * group_size
+            last_rank = first_rank + group_size - 1
+            ranks = list(range(first_rank, last_rank + 1))
+
         # init distributed and build meshes
         dist_utils.init_distributed(
             job_config.comm,
             enable_cpu_backend=job_config.training.enable_cpu_offload,
             base_folder=job_config.job.dump_folder,
+            ranks=ranks,
         )
         world_size = int(os.environ["WORLD_SIZE"])
         parallelism_config = job_config.parallelism