option 1 - use block_current to overlap compute/communication

tushar00jain · tushar00jain · commit dfe09bff6ea4 · 2025-07-25T18:12:15.000-07:00
diff --git a/torchft/manager.py b/torchft/manager.py
@@ -385,8 +385,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 so it doesn't block cpu when using gloo
-                work.wait()
+                work.block_current_stream()
 
             fut = work.get_future()
 

Original file line number	Diff line number	Diff line change
`@@ -385,8 +385,7 @@ def allreduce(self, tensor: torch.Tensor, should_quantize: bool = False) -> Work`
`385`	`385`	`)`
`386`	`386`	`else:`
`387`	`387`	`work = self._pg.allreduce([tensor], ReduceOp.SUM)`
`388`		`- # TODO(tushar00jain): Set up the stream dependency correctly so it doesn't block cpu when using gloo`
`389`		`- work.wait()`
	`388`	`+ work.block_current_stream()`
`390`	`389`
`391`	`390`	`fut = work.get_future()`
`392`	`391`