AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Integrated RCCL with MSCCL++ for small message sizes (#1231)

Browse Source
This commit is contained in:
corey-derochie-amd
2024-07-12 15:32:58 -06:00
committed by GitHub
parent c755b9cf93
commit 6dc47eecd7
15 changed files with 441 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/misc/msccl/msccl_lifecycle.cc
+45 -1
View File
@@ -22,6 +22,10 @@
#include "msccl/msccl_setup.h"
#include "msccl/msccl_status.h"
#ifdef ENABLE_MSCCLPP
#include "mscclpp/mscclpp_nccl.h"
#endif
RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
RCCL_PARAM(MscclForceEnabled, "MSCCL_FORCE_ENABLE", 0);
static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
@@ -448,18 +452,58 @@ ncclResult_t mscclEnqueueCheck(
count, dataType, root, peer, op, func, comm, stream,
&threadLocalStatus.savedSchedulerParams.back()));
size_t nBytes = count * ncclTypeSize(dataType);
switch (threadLocalStatus.groupStatus) {
case mscclNoGroup:
#ifdef ENABLE_MSCCLPP
if (comm->mscclppCompatible) {
/* check if one rank per GPU and graph mode is enabled */
if ((nBytes >= 32) && (threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold) {
INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllReduce (groupStatus=mscclNoGroup)");
NCCLCHECK(mscclpp_ncclAllReduce(sendBuff, recvBuff, count, dataType, op, comm->mscclpp_comm, stream));
threadLocalStatus.savedSchedulerParams.clear();
break;
}
else if (func == mscclFuncAllGather && nBytes * comm->nRanks <= comm->mscclpp_threshold) {
INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllGather (groupStatus=mscclNoGroup)");
NCCLCHECK(mscclpp_ncclAllGather(sendBuff, recvBuff, count, dataType, comm->mscclpp_comm, stream));
threadLocalStatus.savedSchedulerParams.clear();
break;
}
}
}
#endif
if (comm->mscclCompatible) {
NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
NCCLCHECK(mscclRunSavedParams());
break;
}
}
}
NCCLCHECK(mscclFallBackSavedParams());
break;
case mscclGroupSupportedOp:
#ifdef ENABLE_MSCCLPP
if (comm->mscclppCompatible) {
/* check if one rank per GPU and graph mode is enabled */
if ((nBytes >= 32) && (threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold) {
INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllReduce (groupStatus=mscclGroupSupportedOp)");
NCCLCHECK(mscclpp_ncclAllReduce(sendBuff, recvBuff, count, dataType, op, comm->mscclpp_comm, stream));
threadLocalStatus.savedSchedulerParams.clear();
break;
}
else if (func == mscclFuncAllGather && nBytes * comm->nRanks <= comm->mscclpp_threshold) {
INFO(NCCL_INIT, "MSCCL++: mscclpp_ncclAllGather (groupStatus=mscclGroupSupportedOp)");
NCCLCHECK(mscclpp_ncclAllGather(sendBuff, recvBuff, count, dataType, comm->mscclpp_comm, stream));
threadLocalStatus.savedSchedulerParams.clear();
break;
}
}
}
#endif
if (comm->mscclCompatible) {
NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {