[WIP] Add DiLoCo

ghstack-source-id: 68e071e88b5b238d137e0ecdaa33d97b79370b22
Pull Request resolved: #76

Author: H-Huang

torchft/local_sgd.py

@@ -11,17 +11,23 @@
 This module implements a fault tolerant version of LocalSGD and related methods.
 """
 
-from typing import Any, Dict, List, Mapping, Optional
+import logging
+from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional
 
 import torch
 from torch import nn, optim
 
+from torch.nn.parameter import Parameter
+from torch.optim.optimizer import Optimizer
+
 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
@@ -71,8 +77,8 @@ def __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"
 
@@ -93,7 +99,30 @@ def __init__(
         # 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):
+        # Add optimizer hook which increments the local step counter and syncs if necessary
+        self._opt_hook = self._local_optimizer.register_step_post_hook(
+            self._step_post_hook
+        )
+
+        # Register a forward prehook to check for quorum
+        self._forward_pre_hook = self._model.register_forward_pre_hook(
+            self._forward_step_pre_hook
+        )
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # 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
+        self._opt_hook.remove()
+        self._forward_pre_hook.remove()
+
+        return False  # Propagate exceptions
 
     def _save_parameters(self) -> None:
         # TODO: consider running copy on a separate stream
@@ -105,71 +134,53 @@ def _restore_parameters(self) -> None:
         for name, p in self._model.named_parameters():
             p.data.copy_(self._backup_parameters[name], non_blocking=True)
 
-    # 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, _args):
         """
-        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 should_sync(self) -> bool:
+    #     """
+    #     Checks if the model should be synchronized.
+    #     """
+    #     if self._local_step >= self._sync_every:
+    #         return True
+    #     else:
+    #         return False
 
-    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._local_step += 1
+        self._local_step = 0
+        self._perform_sync()
 
-        if self._local_step >= self._sync_every:
-            self._local_step = 0
-            self._average()
+        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()
 
-            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 +193,63 @@ 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:
+        super().__init__(
+            manager, model, inner_optimizer, sync_every, backup_device, pin_memory
+        )
+        self._outer_optimizer = outer_optimizer
+
+    def _model_sync(self) -> None:
+        """
+        ensure model has the same weights
+        """
+        pass
+
+    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()
+
+        # Use the outer optimizer to update the model parameters
+        self._outer_optimizer.step()
+
+    def _average_grads(self) -> None:
+        """
+        Average the gradients across the diloco group.
+        """
+        works = []
+        for p in self._model.parameters():
+            # Perform allreduce on the pseudogradients
+            work = self._manager.allreduce(p.grad)
+            works.append(work)
+        # Wait for all allreduce operations to complete
+        for work in works:
+            work.wait()

torchft/local_sgd_test.py

@@ -11,7 +11,7 @@
 import torch
 from torch import nn, optim
 
-from torchft.local_sgd import LocalSGD
+from torchft.local_sgd import DiLoCo, DiLoCoOptimizer, LocalSGD
 from torchft.manager import Manager
 
 
@@ -40,57 +40,103 @@ def _copy_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Ten
 
 class LocalSGDTest(TestCase):
     def test_local_sgd_healthy(self) -> None:
-        base_m = SimpleModel()
-        optimizer = optim.SGD(base_m.parameters())
+        model = SimpleModel()
+        optimizer = optim.SGD(model.parameters())
         manager = create_autospec(Manager)
-
-        m = LocalSGD(manager, base_m, optimizer, sync_every=2)
-        self.assertEqual(m._local_step, 0)
-
-        torch.testing.assert_close(m._backup_parameters, _params_dict(base_m))
-
-        inp = torch.rand(2, 3)
-
-        loss = m(inp).mean()
-        loss.backward()
-        optimizer.step()
-
-        self.assertEqual(m._local_step, 1)
-        self.assertEqual(manager.start_quorum.call_count, 1)
-
-        loss = m(inp).mean()
-        loss.backward()
-        optimizer.step()
-
-        manager.should_commit.return_value = True
-        self.assertEqual(m._local_step, 0)
-
-        torch.testing.assert_close(m._backup_parameters, _params_dict(base_m))
-        self.assertEqual(manager.should_commit.call_count, 1)
-        self.assertEqual(manager.allreduce.call_count, 4)
+        with LocalSGD(manager, model, optimizer, sync_every=2) as local_sgd:
+            self.assertEqual(local_sgd._local_step, 0)
+            torch.testing.assert_close(
+                local_sgd._backup_parameters, _params_dict(model)
+            )
+            inp = torch.rand(2, 3)
+            loss = model(inp).mean()
+            loss.backward()
+            optimizer.step()
+
+            self.assertEqual(local_sgd._local_step, 1)
+            self.assertEqual(manager.start_quorum.call_count, 1)
+            loss = model(inp).mean()
+            loss.backward()
+            optimizer.step()
+
+            manager.should_commit.return_value = True
+            self.assertEqual(local_sgd._local_step, 0)
+            torch.testing.assert_close(
+                local_sgd._backup_parameters, _params_dict(model)
+            )
+            self.assertEqual(manager.should_commit.call_count, 1)
+            self.assertEqual(manager.allreduce.call_count, 4)
 
     def test_local_sgd_recovery(self) -> None:
-        base_m = SimpleModel()
-        optimizer = optim.SGD(base_m.parameters())
+        model = SimpleModel()
+        optimizer = optim.SGD(model.parameters())
         manager = create_autospec(Manager)
 
-        m = LocalSGD(manager, base_m, optimizer, sync_every=2)
+        with LocalSGD(manager, model, optimizer, sync_every=2) as local_sgd:
+            torch.testing.assert_close(
+                local_sgd._backup_parameters, _params_dict(model)
+            )
+            og_state_dict = _copy_state_dict(model.state_dict())
+
+            inp = torch.rand(2, 3)
+
+            loss = model(inp).mean()
+            loss.backward()
+            optimizer.step()
 
-        torch.testing.assert_close(m._backup_parameters, _params_dict(base_m))
-        og_state_dict = _copy_state_dict(base_m.state_dict())
+            # Check that the model's state dict has been updated
+            for name, param in model.state_dict().items():
+                # Ensure the parameter has changed
+                self.assertFalse(
+                    torch.equal(og_state_dict[name], param),
+                    f"Parameter {name} did not change.",
+                )
+            self.assertEqual(local_sgd._local_step, 1)
 
-        inp = torch.rand(2, 3)
+            local_sgd._restore_parameters()
+            torch.testing.assert_close(
+                local_sgd._backup_parameters, _params_dict(model)
+            )
 
-        loss = m(inp).mean()
-        loss.backward()
-        optimizer.step()
 
-        self.assertEqual(m._local_step, 1)
+class DiLoCoTest(TestCase):
+    def test_diloco_healt(self) -> None:
+        model = SimpleModel()
 
-        state_dict = m.state_dict()
-        torch.testing.assert_close(state_dict, m._backup_parameters)
-        torch.testing.assert_close(state_dict, og_state_dict)
+        # Setup optimizers
+        inner_optimizer = torch.optim.AdamW(
+            model.parameters(), lr=4e-4, weight_decay=0.1, betas=(0.9, 0.95)
+        )
+        outer_optimizer = torch.optim.SGD(
+            model.parameters(), lr=0.7, momentum=0.9, nesterov=True
+        )
 
-        m.load_state_dict(state_dict)
-        torch.testing.assert_close(_params_dict(base_m), state_dict)
-        torch.testing.assert_close(m._backup_parameters, _params_dict(base_m))
+        manager = create_autospec(Manager)
+        with DiLoCo(
+            manager, model, inner_optimizer, outer_optimizer, sync_every=2
+        ) as diloco:
+            parameter_count = len(list(model.parameters()))
+            initial_outer_opt_state = outer_optimizer.state_dict()
+            self.assertEqual(initial_outer_opt_state["state"], {})
+
+            self.assertEqual(diloco._local_step, 0)
+            torch.testing.assert_close(diloco._backup_parameters, _params_dict(model))
+            inp = torch.rand(2, 3)
+            loss = model(inp).mean()
+            loss.backward()
+            inner_optimizer.step()
+
+            self.assertEqual(diloco._local_step, 1)
+            self.assertEqual(manager.start_quorum.call_count, 1)
+            loss = model(inp).mean()
+            loss.backward()
+            inner_optimizer.step()
+
+            manager.should_commit.return_value = True
+            self.assertEqual(diloco._local_step, 0)
+            torch.testing.assert_close(diloco._backup_parameters, _params_dict(model))
+            self.assertEqual(manager.should_commit.call_count, 1)
+            self.assertEqual(manager.allreduce.call_count, parameter_count)
+
+            outer_opt_state = outer_optimizer.state_dict()
+            self.assertEqual(len(outer_opt_state["state"]), parameter_count)