option 1 - use block_current to overlap compute/communication

Author: tushar00jain

torchft/local_sgd.py

@@ -416,9 +416,7 @@ def perform_sync(self) -> bool:
             else nullcontext()
         ):
             for work in self._allreduce_work:
-                # TODO: Setup proper stream dependency on the future
-                # attached to this work
-                work.wait()
+                work.get_future().wait()
 
             if self._stream is not None:
                 self._stop_event = torch.cuda.Event()
@@ -522,6 +520,7 @@ def _bucketize_and_allreduce(
                 work = self._manager.allreduce(
                     flat_buffer, should_quantize=self.should_quantize
                 )
+                work.block_current_stream()
 
             def callback(fut: torch.futures.Future[torch.Tensor]) -> None:
                 with torch.cuda.stream(self._stream) if self._stream else nullcontext():
@@ -532,8 +531,6 @@ def callback(fut: torch.futures.Future[torch.Tensor]) -> None:
                         )
 
             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)
 
             self._allreduce_work.append(work)

torchft/manager.py

@@ -386,7 +386,7 @@ def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work
             else:
                 work = self._pg.allreduce([tensor], ReduceOp.SUM)
 
-            # TODO(tushar00jain): Set up the stream dependency correctly
+            work.block_current_stream()
             fut = work.get_future()
 
             stream: Optional[torch.cuda.Stream] = (

torchft/process_group.py

@@ -1072,7 +1072,7 @@ def wait(self, timeout: Optional[timedelta] = None) -> bool:
                 if timeout is not None:
                     self._work.wait(timeout)
                 else:
-                    self._work.wait()
+                    self._work.block_current_stream()
         except Exception as e:
             self._manager.report_error(e)