Use torch.cuda.synchronize() right after calling batch_isend_irecv() …

…communication API
singleheart · Feb 6, 2021 · 1b3dfa2 · 1b3dfa2
1 parent be473a5
commit 1b3dfa2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/megatron/training.py b/megatron/training.py
@@ -347,6 +347,8 @@ def communicate(tensor_send_next, tensor_send_prev, recv_forward, recv_backward)
     reqs = torch.distributed.batch_isend_irecv(ops)
     for req in reqs:
         req.wait()
+    # Temporary workaround for batch_isend_irecv() race condition.
+    torch.cuda.synchronize()
 
     return tensor_recv_prev, tensor_recv_next