fix: remove wrong async

PrimeIntellect-ai · Sep 26, 2024 · e31fffa · e31fffa
1 parent 867350d
commit e31fffa
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/zeroband/diloco.py b/src/zeroband/diloco.py
@@ -98,10 +98,8 @@ def sync_pseudo_gradient(self, model: nn.Module):
 
             # gloo does not support AVG
             param_offloaded.grad = param_offloaded.grad / self.elastic_device_mesh.global_pg.size()
-            dist.all_reduce(
-                param_offloaded.grad, op=dist.ReduceOp.SUM, group=self.elastic_device_mesh.global_pg, async_op=True
-            )
-            # todo async here
+            dist.all_reduce(param_offloaded.grad, op=dist.ReduceOp.SUM, group=self.elastic_device_mesh.global_pg)
+            # todo maybe do async here
 
     def sync_inner_model(self, model: nn.Module):
         """