Merge pull request #123 from laekov/torch-v1.12.0-compat

Fix nccl uid bcast for torch v1.12.0

Merge pull request #123 from laekov/torch-v1.12.0-compat
Fix nccl uid bcast for torch v1.12.0
5a5dfa18 · Rick Ho · GitHub · 9a01f8fa · a0b1e31a · 5a5dfa18
Unverified Commit 5a5dfa18 authored Jul 18, 2022 by Rick Ho Committed by GitHub Jul 18, 2022
Show 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,