fix it

PrimeIntellect-ai · Sep 24, 2024 · b3e2dd8 · b3e2dd8
1 parent 13d97b0
commit b3e2dd8
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/open_diloco/train_pure_fsdp.py b/open_diloco/train_pure_fsdp.py
@@ -243,14 +243,15 @@ def train(config: Config):
 
             mask = torch.rand_like(param_offloaded.grad) > 0.95
 
-            data_to_all_reduce = param_offloaded.grad * mask
+            data_to_send = param_offloaded.grad * mask
+            data_to_send_pre_reduce = data_to_send.clone()
 
             # gloo does not support AVG
-            data_to_all_reduce = data_to_all_reduce / global_pg.size()
-            dist.all_reduce(data_to_all_reduce, op=dist.ReduceOp.SUM, group=global_pg)
-
-            param_offloaded.grad += data_to_all_reduce
+            data_to_send = data_to_send / global_pg.size()
+            dist.all_reduce(data_to_send, op=dist.ReduceOp.SUM, group=global_pg)
 
+            param_offloaded.grad += data_to_send - data_to_send_pre_reduce # removing the 
+
         outer_optimizer.step()
         outer_optimizer.zero_grad()