Run local_sgd/diloco in titan

Author: H-Huang

torchtitan/components/ft.py

@@ -6,8 +6,9 @@
 
 import copy
 import importlib
+from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import Optional
+from typing import ContextManager, Optional, TYPE_CHECKING, Union
 
 import torch
 import torch.distributed as dist
@@ -22,6 +23,9 @@
 if importlib.util.find_spec("torchft") is not None:
     import torchft as ft
 
+    if TYPE_CHECKING:
+        from torchft import local_sgd
+
     has_torchft = True
 else:
     has_torchft = False
@@ -85,13 +89,16 @@ def init_ft_manager(job: JobConfig) -> FTManager:
 
     pg = ft.ProcessGroupNCCL()
 
+    # If the training method is specific, then the quorum should be synchronous
+    use_async_quorum = job.fault_tolerance.semi_sync_method is None
+
     return FTManager(
         ft.Manager(
             pg=pg,
             min_replica_size=job.fault_tolerance.min_replica_size,
             load_state_dict=None,
             state_dict=None,
-            use_async_quorum=True,
+            use_async_quorum=use_async_quorum,
             replica_id=f"torchtitan_ft_{job.fault_tolerance.replica_id}",
         ),
         group_size=job.fault_tolerance.group_size,
@@ -158,3 +165,50 @@ def ft_clip_grad_norm_util(total_norm: DTensor) -> torch.Tensor:
             return DTensor.from_local(local_tensor, mesh.mesh, placements)
 
     return total_norm
+
+
+def maybe_semi_sync_training(
+    config: JobConfig,
+    ft_manager: FTManager,
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    sync_every: int,
+) -> ContextManager[Union["local_sgd.DiLoCo", "local_sgd.LocalSGD", None]]:
+    """
+    If TorchFT is enabled and the config is set, use semi_sync_method
+    """
+    semi_sync_method = config.fault_tolerance.semi_sync_method
+    torchft_enabled = config.fault_tolerance.enable
+    if torchft_enabled and semi_sync_method is not None:
+        from torchft import local_sgd
+
+        assert (
+            ft_manager._manager is not None
+        ), "FTManager must be enabled to use semi-sync training."
+        if semi_sync_method.lower() == "diloco":
+            # Create the outer optimizer based on the inner optimizer parameters.
+            params = [group["params"] for group in optimizer.param_groups]
+            params = [param for sublist in params for param in sublist]
+            outer_optimizer = torch.optim.SGD(
+                params, lr=0.7, momentum=0.9, nesterov=True
+            )
+
+            return local_sgd.DiLoCo(
+                manager=ft_manager._manager,
+                model=model,
+                inner_optimizer=optimizer,
+                outer_optimizer=outer_optimizer,
+                sync_every=sync_every,
+            )
+        elif semi_sync_method.lower() == "local_sgd":
+            return local_sgd.LocalSGD(
+                manager=ft_manager._manager,
+                model=model,
+                optimizer=optimizer,
+                sync_every=sync_every,
+            )
+        else:
+            raise ValueError(
+                f"Unknown training method: {semi_sync_method}, only 'diloco' and 'local_sgd' are supported."
+            )
+    return nullcontext()

torchtitan/config_manager.py

@@ -502,6 +502,19 @@ class FaultTolerance:
     min_replica_size: int = 1
     """The minimum number of FT replica for each step."""
 
+    semi_sync_method: str | None = None
+    """
+    The algorithm to use for semi-sync training. Currently, only "local_sgd" and "diloco" from
+    torchft are supported
+    (https://github.com/pytorch/torchft/blob/360c5c534bdeac959507e9d238ba9f3902d3fda9/torchft/local_sgd.py#L41)
+    """
+
+    sync_steps: int = 5
+    """
+    Number of steps to wait before performing synchronization. This is only used when "semi_sync_method"
+    is set.
+    """
+
 
 @dataclass
 class Experimental:

torchtitan/train.py

@@ -399,7 +399,13 @@ def train(self):
             job_config, global_step=self.step
         ) as torch_profiler, maybe_enable_memory_snapshot(
             job_config, global_step=self.step
-        ) as memory_profiler:
+        ) as memory_profiler, ft.maybe_semi_sync_training(
+            job_config,
+            ft_manager=self.ft_manager,
+            model=self.model_parts[0],
+            optimizer=self.optimizers,
+            sync_every=job_config.fault_tolerance.sync_steps,
+        ):
             data_iterator = iter(self.dataloader)
             while self.step < job_config.training.steps:
                 self.step += 1