do not use batch_isend_irecv for UCC (NVIDIA#1675)

mmarcinkiewicz · Jun 9, 2023 · 82ee367 · 82ee367
1 parent 05091d4
commit 82ee367
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/apex/transformer/pipeline_parallel/p2p_communication.py b/apex/transformer/pipeline_parallel/p2p_communication.py
@@ -60,7 +60,7 @@ def _run_p2pops(
     need_to_sync = p2p_group.name() != default_group.name()
     reqs = []
 
-    if batch_p2p_comm:
+    if batch_p2p_comm and p2p_group.name() == "nccl":
         ops = []
         if tensor_send_prev is not None:
             send_prev_op = torch.distributed.P2POp(
@@ -101,9 +101,9 @@ def _run_p2pops(
             reqs = torch.distributed.batch_isend_irecv(ops)
     else:
         # sync before communication if needed
-        if need_to_sync and any(
+        if need_to_sync and any([
             tensor_send_prev is not None, tensor_recv_prev is not None,
-            tensor_send_next is not None, tensor_recv_next is not None):
+            tensor_send_next is not None, tensor_recv_next is not None]):
             torch.cuda.synchronize()
 
         if tensor_send_prev is not None: