Merge branch 'main_fix' into 'main'

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

See merge request ADLR/megatron-lm!230

megatron/training.py

@@ -351,6 +351,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