Commit 17e2a5c5 authored by Rick Ho's avatar Rick Ho
Browse files

fix pytorch compatibility issue !55

parent 7d41fe88
...@@ -81,10 +81,15 @@ public: ...@@ -81,10 +81,15 @@ public:
if (rank == 0) { if (rank == 0) {
ncclGetUniqueId(&ncclID); ncclGetUniqueId(&ncclID);
} }
#if defined(TORCH_VERSION) && (TORCH_VERSION_MAJOR > 1 || \
(TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR >= 8))
broadcastUniqueNCCLID(&ncclID, broadcastUniqueNCCLID(&ncclID,
c10d::OpType::SEND, c10d::OpType::SEND,
"fastmoe_nccl_comm", "fastmoe_nccl_comm",
rank); rank);
#else
broadcastUniqueNCCLID(&ncclID);
#endif
ncclComm_t comm; ncclComm_t comm;
NCCL_SAFE_CALL(ncclCommInitRank(&comm, getSize(), ncclID, rank)); NCCL_SAFE_CALL(ncclCommInitRank(&comm, getSize(), ncclID, rank));
return comm; return comm;
......
...@@ -32,6 +32,7 @@ ...@@ -32,6 +32,7 @@
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include <cublas_v2.h> #include <cublas_v2.h>
#include <stdio.h> #include <stdio.h>
#include <stdlib.h>
#ifndef HELPER_CUDA_H #ifndef HELPER_CUDA_H
#define HELPER_CUDA_H #define HELPER_CUDA_H
......
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