From ef6d75b3ee4d053c878995c7acbddf49937a37d9 Mon Sep 17 00:00:00 2001 From: Mustafa Abduljabbar Date: Wed, 25 Sep 2024 15:24:25 -0400 Subject: [PATCH] MSCCL Multithreaded regression root cause fix (#1347) * Make sure the target device is used for MSCCL * Enable single process mode by default to use MSCCL in MT * Create a per-rank state when GPUs share a thread [ROCm/rccl commit: 03a3ef3c3479d79e583144127d20f45643bcc376] --- projects/rccl/src/misc/msccl/msccl_lifecycle.cc | 2 +- projects/rccl/src/misc/msccl/msccl_status.cc | 10 ++++++---- projects/rccl/src/msccl.cc | 8 ++++++++ 3 files changed, 15 insertions(+), 5 deletions(-) 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));