[WIP] Add DiLoCo

ghstack-source-id: 588c13d633fc69574f8a87e46bbe8ae4069d4a3c
Pull Request resolved: #76

Author: H-Huang

torchft/local_sgd.py

@@ -3,25 +3,29 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-
 """
 LocalSGD
 =========
-
 This module implements a fault tolerant version of LocalSGD and related methods.
 """
-
-from typing import Any, Dict, List, Mapping, Optional
+import logging
+from types import TracebackType
+from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional, Type
 
 import torch
 from torch import nn, optim
+from torch.nn.parameter import Parameter
+from torch.optim.optimizer import Optimizer
+from torch.utils.hooks import RemovableHandle
 
 from torchft.manager import Manager
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
-class LocalSGD(nn.Module):
+class LocalSGD:
     """
-    LocalSGD is a model wrapper similar to DistributedDataParallel that
+    LocalSGD is a context manager that
     implements the algorithm described in https://arxiv.org/pdf/1805.09767
 
     This will synchronize the model parameters periodically in a fault tolerant
@@ -60,26 +64,22 @@ def __init__(
     ) -> None:
         """
         Args:
-            manager: The manager to use.
-            model: The model to wrap.
-            optimizer: The optimizer used by the model.
-            sync_every: How often to sync the model weights.
-            backup_device: The device to store the backup of the model parameters on. (default cpu)
-            pin_memory: Whether to pin the memory used for the backup of the model parameters.
+            manager (Manager): The manager to use.
+            model (nn.Module): The model to wrap.
+            optimizer (optim.Optimizer): The optimizer used by the model.
+            sync_every (int): How often to sync the model weights.
+            backup_device (Optional[torch.device]): The device to store the backup of the model parameters on. (default cpu)
+            pin_memory (bool): Whether to pin the memory used for the backup of the model parameters.
         """
         super().__init__()
-
         self._manager = manager
         self._model = model
+        self._local_optimizer = optimizer
         self._local_step = 0
-        self._started_step = False
         self._sync_every = sync_every
         assert sync_every >= 1, "sync_every must be greater than or equal to 1"
-
         device = backup_device or torch.device("cpu")
-
         self._backup_parameters: Dict[str, torch.Tensor] = {}
-
         for name, p in self._model.named_parameters():
             t = torch.empty(*tuple(p.shape), dtype=p.dtype, device=device)
             if (
@@ -90,86 +90,88 @@ def __init__(
                 t = t.pin_memory()
             self._backup_parameters[name] = t
 
+        self._hooks: List[RemovableHandle] = []
         # Need to copy the parameters to the host to be safe if we are on the first step.
         self._save_parameters()
 
-        optimizer.register_step_post_hook(self._step_post_hook)
+    def __enter__(self) -> "LocalSGD":
+        # Add optimizer hook which increments the local step counter and syncs if necessary
+        self._hooks.append(
+            self._local_optimizer.register_step_post_hook(self._step_post_hook)
+        )
+        # Register a forward prehook to check for quorum
+        self._hooks.append(
+            self._model.register_forward_pre_hook(self._forward_step_pre_hook)
+        )
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> bool:
+        # Handle any cleanup or error handling here
+        if exc_type is not None:
+            # If an exception occurred, restore parameters
+            self._restore_parameters()
+        # Clean up hooks
+        for hook in self._hooks:
+            hook.remove()
+        self._hooks.clear()
+
+        return False  # Propagate exceptions
 
     def _save_parameters(self) -> None:
-        # TODO: consider running copy on a separate stream
-        for name, p in self._model.named_parameters():
-            self._backup_parameters[name].copy_(p.data, non_blocking=True)
+        with torch.no_grad():
+            # TODO: consider running copy on a separate stream
+            for name, p in self._model.named_parameters():
+                self._backup_parameters[name].copy_(p.data, non_blocking=True)
 
     def _restore_parameters(self) -> None:
-        # TODO: consider running copy on a separate stream
-        for name, p in self._model.named_parameters():
-            p.data.copy_(self._backup_parameters[name], non_blocking=True)
+        with torch.no_grad():
+            # TODO: consider running copy on a separate stream
+            for name, p in self._model.named_parameters():
+                p.copy_(self._backup_parameters[name], non_blocking=False)
 
-    # pyre-fixme[14]: support state_dict args
-    def state_dict(self) -> Dict[str, object]:
-        """
-        state_dict returns the state_dict from the last time LocalSGD
-        synchronized and not the current weights.
-        """
-        state_dict = self._model.state_dict()
-        for name, p in self._backup_parameters.items():
-            assert name in state_dict
-            state_dict[name] = p
-        return state_dict
-
-    def load_state_dict(
-        self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False
+    def _step_post_hook(
+        self, _optim: optim.Optimizer, _args: List[object], _kwargs: Dict[str, object]
     ) -> None:
         """
-        Loads the state dict to the model and the backup parameters.
-
-        This must be called while the model weights aren't being modified to
-        avoid corrupting the backup weights.
+        This hook is registered on the optimizer and is called after the optimizer step.
         """
-        self._model.load_state_dict(state_dict, strict=strict, assign=assign)
-        self._save_parameters()
+        self._local_step += 1
+        if self._local_step >= self._sync_every:
+            self.sync()
 
-    def forward(self, *args: object, **kwargs: object) -> object:
+    def _forward_step_pre_hook(self, _module: nn.Module, _args: List[object]) -> None:
         """
-        Run the model parameters.
-
-        This should be called before the optimizer step.
-
-        This will start the quorum and save the parameters if this is the first step.
+        Start the quorum before each module forward.
         """
         if self._local_step == 0:
             self._manager.start_quorum()
 
-        self._started_step = True
-
-        return self._model.forward(*args, **kwargs)
-
-    def _step_post_hook(
-        self, _optim: optim.Optimizer, _args: List[object], _kwargs: Dict[str, object]
-    ) -> None:
+    def sync(self) -> None:
         """
-        This hook is registered on the optimizer and is called after the optimizer step.
-
-        This will call the allreduce on the model weights every sync_every steps.
-        If any errors occur it will restore to the weights from the previous sync.
-
-        ``forward`` must be called before this function.
+        Synchronizes and averages the model weights across the manager.
         """
-        assert self._started_step, "forward must be called before step"
-        self._started_step = False
+        self._perform_sync()
 
-        self._local_step += 1
+        if self._manager.should_commit():
+            self._save_parameters()
+        else:
+            # commit failed, restore from the backup parameters
+            self._restore_parameters()
 
-        if self._local_step >= self._sync_every:
-            self._local_step = 0
-            self._average()
+        self._local_step = 0
 
-            if self._manager.should_commit():
-                # save the parameters so we can restore from them later if necessary.
-                self._save_parameters()
-            else:
-                # commit failed, restore from the backup parameters
-                self._restore_parameters()
+    def _perform_sync(self) -> None:
+        """
+        Performs the synchronization of the model weights across the manager.
+        This method is intended to be overridden by subclasses to implement custom
+        synchronization logic.
+        """
+        self._average()
 
     def _average(self) -> None:
         # TODO: do we need to broadcast buffers like DDP does?
@@ -182,3 +184,67 @@ def _average(self) -> None:
 
         for work in works:
             work.wait()
+
+
+class DiLoCo(LocalSGD):
+    """
+    DiLoCo is a subclass of LocalSGD that overrides the synchronization
+    mechanism to average and synchronize the pseudogradients (delta of the previous global weight and current local weights).
+
+    diloco: https://arxiv.org/pdf/2311.08105
+    """
+
+    def __init__(
+        self,
+        manager: Manager,
+        model: nn.Module,
+        inner_optimizer: optim.Optimizer,
+        outer_optimizer: optim.Optimizer,
+        sync_every: int,
+        backup_device: Optional[torch.device] = None,
+        pin_memory: bool = True,
+    ) -> None:
+        if manager._use_async_quorum:
+            raise ValueError(
+                "Using DiLoCo require synchronous quorum to be enabled. "
+                "Ensure that the manager is initialized with use_async_quorum=False"
+            )
+        super().__init__(
+            manager, model, inner_optimizer, sync_every, backup_device, pin_memory
+        )
+        self._outer_optimizer = outer_optimizer
+
+    def _perform_sync(self) -> None:
+        """
+        Overrides the sync method to calculate the pseugradient, average them across the manager group, and
+        step using the outer optimizer.
+        """
+
+        # Set the .grad field of each parameter to its pseudogradient
+        for name, p in self._model.named_parameters():
+            assert name in self._backup_parameters
+            pseudogradient = p.data - self._backup_parameters[name]
+            p.grad = pseudogradient
+
+        self._average_grads()
+
+        # Restore the parameters back to the previous state
+        self._restore_parameters()
+
+        # Use the outer optimizer to update the model parameters
+        self._outer_optimizer.step()
+        self._outer_optimizer.zero_grad()
+
+    def _average_grads(self) -> None:
+        """
+        Average the gradients across the diloco group.
+        """
+        works = []
+        for p in self._model.parameters():
+            # Perform allreduce on the pseudogradients
+            assert p.grad is not None
+            work = self._manager.allreduce(p.grad)
+            works.append(work)
+        # Wait for all allreduce operations to complete
+        for work in works:
+            work.wait()