option 1 - use block_current to overlap compute/communication

tushar00jain · tushar00jain · commit 61e177cc0fd7 · 2025-07-28T21:18:18.000-07:00
diff --git a/torchft/manager.py b/torchft/manager.py
@@ -388,7 +388,7 @@ def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work
                 )
             else:
                 work = self._pg.allreduce([tensor], ReduceOp.SUM)
-                work.wait()
+                work.block_current_stream()
 
             fut = work.get_future()
 

Original file line number	Diff line number	Diff line change
`@@ -388,7 +388,7 @@ def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work`
`388`	`388`	`)`
`389`	`389`	`else:`
`390`	`390`	`work = self._pg.allreduce([tensor], ReduceOp.SUM)`
`391`		`- work.wait()`
	`391`	`+ work.block_current_stream()`
`392`	`392`
`393`	`393`	`fut = work.get_future()`
`394`	`394`