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

fix pytorch compatibility issue !55

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