option 1 - use block_current to overlap compute/communication

tushar00jain · tushar00jain · commit bfe221fa1040 · 2025-07-26T10:48:38.000-07:00
diff --git a/torchft/manager.py b/torchft/manager.py
@@ -382,7 +382,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
`@@ -382,7 +382,7 @@ def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work`
`382`	`382`	`)`
`383`	`383`	`else:`
`384`	`384`	`work = self._pg.allreduce([tensor], ReduceOp.SUM)`
`385`		`- work.wait()`
	`385`	`+ work.block_current_stream()`
`386`	`386`
`387`	`387`	`fut = work.get_future()`
`388`	`388`