fix nccl uid bcast

a0b1e31a · Rick Ho · dd68fd78 · a0b1e31a
Commit a0b1e31a authored Jul 18, 2022 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 0 deletions

cuda/global_exchange.cpp cuda/global_exchange.cpp +6 -0

No files found.
--- a/cuda/global_exchange.cpp
+++ b/cuda/global_exchange.cpp
@@ -109,6 +109,12 @@ public:
            ncclGetUniqueId(&ncclID);
        }
 #if defined(TORCH_VERSION_MAJOR) && (TORCH_VERSION_MAJOR > 1 || \
+        (TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR >= 12))
+        broadcastUniqueNCCLID(&ncclID,
+                false,
+                "fastmoe_nccl_comm",
+                rank);
+#elif defined(TORCH_VERSION_MAJOR) && (TORCH_VERSION_MAJOR > 1 || \
        (TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR >= 8))
        broadcastUniqueNCCLID(&ncclID,
                c10d::OpType::SEND,