wait for futures while syncing fragments

Summary:
- we current wait for pg work's future when preparing for a fragment
- if we use gloo, this blocks the cpu
- move the wait call to when we perform the actual sync of the fragment
- the manager allreduce also returns the work object so we can wait for that as well when performing the sync
- use http transport instead of pg transport -- pg transport fails to resolve address when running locally
- deep copy the state dict for sending checkpoint because if the replica moves to the next step, the state dict can change before the checkpoint is sent

Test Plan:
gloo overlaps now

<img width="1284" height="662" alt="image" src="https://github.com/user-attachments/assets/e9b88e52-8053-432b-83a3-e689bcc4f9d4" />

nccl still overlaps

<img width="1283" height="664" alt="image" src="https://github.com/user-attachments/assets/cbd0a352-1529-42f7-b8d9-d45bd0e84a97" />

Author: tushar00jain

torchft/collectives.py

@@ -18,6 +18,7 @@
     AllreduceOptions,
     AllToAllOptions,
     ReduceScatterOptions,
+    Work,
 )
 from torch.futures import Future
 
@@ -288,7 +289,7 @@ def allreduce_quantized(
     opts: AllreduceOptions | ReduceOp,
     process_group: "ProcessGroup",
     sync_stream: cuda.Stream | None = None,
-) -> Future[list[torch.Tensor]]:
+) -> Work:
     """
     Performs a quantized all-reduce operation on a list of tensors.
 
@@ -379,6 +380,14 @@ def allreduce_quantized(
             [torch.split(quantized_tensors_out.view(world_size, -1), 1)[rank]],
             _to_allgather_options(allreduce_opts),
         )
+
+        # NOTE: This is not supposed to be used with gloo, only with NCCL.
+        # So we setup the stream dependency here by calling work.wait(),
+        # which doesn't block the CPU.
+        #
+        # The future callback below will run after the work has been
+        # completed.
+
         work.wait()
         fut = work.get_future()
 
@@ -392,4 +401,4 @@ def callback(fut: Future[list[torch.Tensor]]) -> list[torch.Tensor]:
                 return tensors
 
         fut = fut.then(callback)
-        return fut
+        return work

torchft/collectives_test.py

@@ -94,8 +94,8 @@ def _run_all_reduce_collective(
                     )
                 ]
 
-                fut = allreduce_quantized(tensors, reduce_op, pg)
-                fut.wait()
+                work = allreduce_quantized(tensors, reduce_op, pg)
+                work.get_future().wait()
 
                 work = pg.allreduce([expected], reduce_op)
                 work.get_future().wait()

torchft/ddp.py

@@ -68,7 +68,8 @@ def __init__(self, manager: "Manager", module: nn.Module, **kwargs: object) -> N
     def _comm_hook(
         state: "Manager", bucket: dist.GradBucket
     ) -> torch.futures.Future[torch.Tensor]:
-        return state.allreduce(bucket.buffer())
+        work = state.allreduce(bucket.buffer())
+        return work.get_future()
 
 
 class PureDistributedDataParallel(nn.Module):

torchft/ddp_test.py

@@ -10,11 +10,13 @@
 import torch
 import torch.distributed as dist
 from torch import nn
+from torch.distributed.distributed_c10d import Work
 from torch.futures import Future
 
 from torchft.ddp import DistributedDataParallel, PureDistributedDataParallel
 from torchft.manager import Manager
 from torchft.process_group import ProcessGroupBabyGloo, ProcessGroupGloo
+from torchft.work import _DummyWork
 
 
 class TestDDP(TestCase):
@@ -39,14 +41,14 @@ def test_ddp(self) -> None:
 
         call_count = 0
 
-        def allreduce(tensor: torch.Tensor) -> Future[torch.Tensor]:
+        def allreduce(
+            tensor: torch.Tensor,
+        ) -> Work:
             nonlocal call_count
 
             call_count += 1
 
-            fut = Future()  # pyre-fixme[29]: not a function
-            fut.set_result(tensor)
-            return fut
+            return _DummyWork(tensor)
 
         manager.allreduce = allreduce
 

torchft/local_sgd.py

@@ -18,6 +18,7 @@
 import torch
 import torch.distributed as dist
 from torch import nn, optim
+from torch.distributed.distributed_c10d import Work
 from torch.distributed.tensor import DTensor
 from torch.nn.parameter import Parameter
 from torch.optim.optimizer import Optimizer
@@ -154,7 +155,8 @@ def _average(self) -> list[torch.Tensor]:
         for p in self._model.parameters():
             # Create a new tensor to store the averaged parameter
             avg_param = extract_local_tensor(p)
-            works.append(self._manager.allreduce(avg_param))
+            work = self._manager.allreduce(avg_param)
+            works.append(work)
             averaged_parameters.append(avg_param)
         for work in works:
             work.wait()
@@ -200,7 +202,7 @@ def __init__(
         self._outer_optimizer = outer_optimizer
 
         # Stores pending all reduce
-        self._allreduce_futures: list[torch.futures.Future[torch.Tensor]] = []
+        self._allreduce_work: list[Work] = []
         self._stream: Optional[torch.cuda.Stream] = (
             torch.cuda.Stream() if torch.cuda.is_available() else None
         )
@@ -368,15 +370,15 @@ def wait(self) -> None:
         """
         Waits for the previously scheduled allreduce to finish
         """
-        if len(self._allreduce_futures) == 0:
+        if len(self._allreduce_work) == 0:
             return
 
         if self._stream is not None:
             assert self._stop_event is not None
             self._stop_event.synchronize()
             self._stop_event = None
 
-        self._allreduce_futures = []
+        self._allreduce_work = []
 
     @torch.profiler.record_function("torchft::local_sgd::prepare_sync")
     def prepare_sync(self) -> None:
@@ -386,7 +388,7 @@ def prepare_sync(self) -> None:
         """
         self._save_grads()
 
-        assert len(self._allreduce_futures) == 0
+        assert len(self._allreduce_work) == 0
 
         # Make sure tensors are available to `_stream`
         if self._stream is not None:
@@ -399,21 +401,28 @@ def prepare_sync(self) -> None:
         ):
             self._average_grads()
 
-            for work in self._allreduce_futures:
-                work.wait()
-
-            if self._stream is not None:
-                self._stop_event = torch.cuda.Event()
-                self._stop_event.record()
-
     @torch.profiler.record_function("torchft::local_sgd::perform_sync")
     def perform_sync(self) -> bool:
         """
         Overrides the sync method to wait for the scheduled allreduce to finish and
         steps using the outer optimizer.
         """
         # Waiting for an allreduce before it has been sent is currently not supported.
-        assert len(self._allreduce_futures) > 0
+        assert len(self._allreduce_work) > 0
+
+        with (
+            torch.cuda.stream(self._stream)
+            if self._stream is not None
+            else nullcontext()
+        ):
+            for work in self._allreduce_work:
+                # TODO: Setup proper stream dependency on the future
+                # attached to this work
+                work.wait()
+
+            if self._stream is not None:
+                self._stop_event = torch.cuda.Event()
+                self._stop_event.record()
 
         self.wait()
 
@@ -467,7 +476,8 @@ def _allreduce_per_param(self) -> None:
             work = self._manager.allreduce(
                 self._grads[name], should_quantize=self.should_quantize
             )
-            self._allreduce_futures.append(work)
+
+            self._allreduce_work.append(work)
 
     def _bucketize_and_allreduce(
         self,
@@ -508,17 +518,25 @@ def _bucketize_and_allreduce(
                 pack_offset += numel
                 flat_index += 1
 
-            work = self._manager.allreduce(
-                flat_buffer, should_quantize=self.should_quantize
-            )
+            with torch.cuda.stream(self._stream) if self._stream else nullcontext():
+                work = self._manager.allreduce(
+                    flat_buffer, should_quantize=self.should_quantize
+                )
 
             def callback(fut: torch.futures.Future[torch.Tensor]) -> None:
-                nonlocal bucket_tensors, flat_buffer
-                for t, pack_offset, numel in bucket_tensors:
-                    t.copy_(flat_buffer[pack_offset : pack_offset + numel].view_as(t))
+                with torch.cuda.stream(self._stream) if self._stream else nullcontext():
+                    nonlocal bucket_tensors, flat_buffer
+                    for t, pack_offset, numel in bucket_tensors:
+                        t.copy_(
+                            flat_buffer[pack_offset : pack_offset + numel].view_as(t)
+                        )
+
+            fut = work.get_future()
+            # TODO(tushar00jain): We need to call work.wait() here to ensure callback
+            # runs after work has been completed
+            fut = fut.then(callback)
 
-            work = work.then(callback)
-            self._allreduce_futures.append(work)
+            self._allreduce_work.append(work)
 
             offset += chunk_size
 

torchft/local_sgd_test.py

@@ -11,10 +11,12 @@
 import torch
 from parameterized import parameterized
 from torch import Tensor, nn, optim
+from torch.distributed.distributed_c10d import Work
 from torch.distributed.tensor import DTensor
 
 from torchft.local_sgd import DiLoCo, LocalSGD, extract_local_tensor
 from torchft.manager import Manager
+from torchft.work import _DummyWork
 
 
 def create_manager() -> MagicMock:
@@ -26,6 +28,11 @@ def create_manager() -> MagicMock:
 
     manager.errored.return_value = None
 
+    def mock_allreduce(tensor: torch.Tensor, should_quantize: bool = False) -> Work:
+        return _DummyWork(tensor)
+
+    manager.allreduce.side_effect = mock_allreduce
+
     return manager
 
 
@@ -66,7 +73,7 @@ class LocalSGDTest(TestCase):
     def test_local_sgd_healthy(self) -> None:
         model = SimpleModel()
         optimizer = optim.SGD(model.parameters())
-        manager = create_autospec(Manager)
+        manager = create_manager()
         with LocalSGD(manager, model, optimizer, sync_every=2) as local_sgd:
             self.assertEqual(local_sgd._local_step, 0)
             inp = torch.rand(2, 3)
@@ -240,13 +247,9 @@ def test_bucketization_correctness(self) -> None:
         manager.should_commit.return_value = True
 
         # Define fake allreduce: multiplies buffer by 2
-        def fake_allreduce(
-            tensor: Tensor, should_quantize: bool
-        ) -> torch.futures.Future[Tensor]:
+        def fake_allreduce(tensor: Tensor, should_quantize: bool) -> Work:
             tensor.mul_(2)
-            fut = torch.futures.Future()  # pyre-fixme[29]: not a function
-            fut.set_result(tensor)
-            return fut
+            return _DummyWork(tensor)
 
         manager.allreduce.side_effect = fake_allreduce
 
@@ -284,13 +287,9 @@ def test_gradient_correctness(self) -> None:
         manager.should_commit.return_value = True
 
         # Define fake allreduce: multiplies buffer by 2
-        def fake_allreduce(
-            tensor: Tensor, should_quantize: bool
-        ) -> torch.futures.Future[Tensor]:
+        def fake_allreduce(tensor: Tensor, should_quantize: bool) -> Work:
             tensor.mul_(2)
-            fut = torch.futures.Future()  # pyre-fixme[29]: not a function
-            fut.set_result(tensor)
-            return fut
+            return _DummyWork(tensor)
 
         manager.allreduce.side_effect = fake_allreduce
 

torchft/manager.py

@@ -26,6 +26,7 @@
 """
 
 import concurrent.futures
+import copy
 import logging
 import os
 import socket
@@ -39,11 +40,12 @@
 
 import torch
 from torch.distributed import ReduceOp, TCPStore
-from torch.distributed.distributed_c10d import AllreduceOptions, ReduceOp
+from torch.distributed.distributed_c10d import AllreduceOptions, ReduceOp, Work
 
 from torchft._torchft import ManagerClient, ManagerServer
 from torchft.checkpointing import CheckpointTransport, HTTPTransport
 from torchft.futures import future_timeout
+from torchft.work import _DummyWork
 
 if TYPE_CHECKING:
     from torchft.process_group import ProcessGroup
@@ -343,9 +345,7 @@ def shutdown(self, wait: bool = True) -> None:
         self._executor.shutdown(wait=wait)
 
     @torch.profiler.record_function("torchft::manager::allreduce")
-    def allreduce(
-        self, tensor: torch.Tensor, should_quantize: bool = False
-    ) -> torch.futures.Future[torch.Tensor]:
+    def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work:
         """
         Fault tolerant allreduce the tensor and return a Future that will be completed when
         the tensor is ready.
@@ -367,7 +367,7 @@ def allreduce(
         if self.errored():
             fut = torch.futures.Future()  # pyre-fixme[29]: not a function
             fut.set_result(tensor)
-            return fut
+            return _DummyWork(tensor)
 
         self.wait_quorum()
         num_participants: int = self.num_participants()
@@ -380,13 +380,14 @@ def allreduce(
             # Run the allreduce async and save the work object so we can wait on
             # it later.
             if should_quantize and IS_TRITON_AVAILABLE:
-                fut = allreduce_quantized(
+                work = allreduce_quantized(
                     [tensor], ReduceOp.SUM, self._pg, torch.cuda.current_stream()
                 )
             else:
                 work = self._pg.allreduce([tensor], ReduceOp.SUM)
-                work.wait()
-                fut = work.get_future()
+
+            # TODO(tushar00jain): Set up the stream dependency correctly
+            fut = work.get_future()
 
             stream: Optional[torch.cuda.Stream] = (
                 torch.cuda.current_stream() if torch.cuda.is_available() else None
@@ -411,17 +412,15 @@ def callback(
             fut = fut.then(callback)
 
             fut = self.wrap_future(fut, tensor)
-            return fut
+            return work
 
         except Exception as e:
             self._logger.exception(
                 f"got exception in all reduce -- skipping remaining: {e}"
             )
             self.report_error(e)
 
-            fut = torch.futures.Future()  # pyre-fixme[29]: not a function
-            fut.set_result(tensor)
-            return fut
+            return _DummyWork(tensor)
 
     def report_error(self, e: Exception) -> None:
         """
@@ -646,7 +645,7 @@ def _async_quorum(
                             self._checkpoint_transport.send_checkpoint(
                                 dst_ranks=quorum.recover_dst_replica_ranks,
                                 step=max_step,
-                                state_dict=self._manager_state_dict(),
+                                state_dict=copy.deepcopy(self._manager_state_dict()),
                                 timeout=self._timeout,
                             )
 

torchft/manager_integ_test.py

@@ -634,7 +634,7 @@ def all_reduce_callback(
 
         manager.start_quorum()
         t1 = torch.ones((1, 3), device=device)
-        fut = manager.allreduce(t1)
-        fut.wait()
+        work = manager.allreduce(t1)
+        work.get_future().wait()
         return t1
     return None