Skip to content

GitLab

Commit 7ffea978 authored by Deepak Narayanan's avatar Deepak Narayanan
Browse files

Use torch.cuda.synchronize() right after calling batch_isend_irecv() communication API

parent 2096d356
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
megatron/training.py
View file @ 7ffea978
...@@ -351,6 +351,8 @@ def communicate(tensor_send_next, tensor_send_prev, recv_forward, recv_backward) ...@@ -351,6 +351,8 @@ def communicate(tensor_send_next, tensor_send_prev, recv_forward, recv_backward)
reqs = torch.distributed.batch_isend_irecv(ops) reqs = torch.distributed.batch_isend_irecv(ops)
for req in reqs: for req in reqs:
req.wait() req.wait()
# Temporary workaround for batch_isend_irecv() race condition.
torch.cuda.synchronize()
return tensor_recv_prev, tensor_recv_next return tensor_recv_prev, tensor_recv_next
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment