diff --git a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc index 161eb93b7b..39c4bd25f5 100644 --- a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc +++ b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc @@ -28,7 +28,7 @@ RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1); RCCL_PARAM(MscclForceEnabled, "MSCCL_FORCE_ENABLE", 0); -RCCL_PARAM(MscclEnableSingleProcess, "MSCCL_ENABLE_SINGLE_PROCESS", 0); +RCCL_PARAM(MscclEnableSingleProcess, "MSCCL_ENABLE_SINGLE_PROCESS", 1); static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH"; bool mscclEnabled() { diff --git a/projects/rccl/src/misc/msccl/msccl_status.cc b/projects/rccl/src/misc/msccl/msccl_status.cc index 3812112dd6..f2b26663b7 100644 --- a/projects/rccl/src/misc/msccl/msccl_status.cc +++ b/projects/rccl/src/misc/msccl/msccl_status.cc @@ -27,9 +27,9 @@ static mutex rankStatesMutex; static unordered_map> rankStates; static inline mscclRankState& mscclGetRankState(int rank) { - static thread_local shared_ptr threadRankState = make_shared(); - + // In the unlikely case of negative rank, return a per-thread state if (rank < 0) { + static thread_local shared_ptr threadRankState(new mscclRankState()); return *threadRankState; } @@ -37,8 +37,10 @@ static inline mscclRankState& mscclGetRankState(int rank) { auto rankStateIt = rankStates.find(rank); if (rankStateIt == rankStates.end()) { - rankStateIt = rankStates.insert(make_pair(rank, make_shared(*threadRankState))).first; - rankStateIt->second->rank = rank; + // Create a per rank threadRankState rather than per thread + shared_ptr newthreadRankState(new mscclRankState()); + newthreadRankState->rank = rank; + rankStateIt = rankStates.insert(make_pair(rank, newthreadRankState)).first; } return *(rankStateIt->second); } diff --git a/projects/rccl/src/msccl.cc b/projects/rccl/src/msccl.cc index e912703d13..5814ac353a 100644 --- a/projects/rccl/src/msccl.cc +++ b/projects/rccl/src/msccl.cc @@ -61,6 +61,14 @@ ncclResult_t mscclRunAlgo_impl( struct mscclAlgo* hostAlgo = status.hostAlgos[mscclAlgoHandle]; struct mscclAlgo* devAlgo = status.devAlgos[mscclAlgoHandle]; + // NCCL adds a lot of guarantees that target device is getting used + // in its group management code, which we entirely skip when MSCCL is used + // Therefore, in single thread multiGPU mode + // setting the device is critical to be sure + // communication is done on the intended device + + CUDACHECK(hipSetDevice(comm->cudaDev)); + NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream)); NCCLCHECK(mscclSetupCount(hostAlgo, comm, count, dataType));