diff --git a/projects/rccl/src/graph/connect.cc b/projects/rccl/src/graph/connect.cc index 37ebc6f3da..3637e0e3a6 100644 --- a/projects/rccl/src/graph/connect.cc +++ b/projects/rccl/src/graph/connect.cc @@ -700,7 +700,7 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa } if (mscclEnabled() && (comm->topo->mscclEnabled || mscclForceEnabled())) { - int mscclNumChannelsRequired = 0; + int mscclNumChannelsRequired = maxNchannels; mscclSchedulerInit(comm, &mscclNumChannelsRequired); minNchannels = std::max(minNchannels, mscclNumChannelsRequired); } diff --git a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc index d4ae1f2cdd..d9a54ee4f5 100644 --- a/projects/rccl/src/misc/msccl/msccl_lifecycle.cc +++ b/projects/rccl/src/misc/msccl/msccl_lifecycle.cc @@ -67,6 +67,20 @@ static bool mscclCommCompatible(ncclComm_t comm) { return true; } +const char *mscclFuncNames[] = { + "mscclFuncReduce", + "mscclFuncBroadcast", + "mscclFuncAllReduce", + "mscclFuncReduceScatter", + "mscclFuncAllGather", + "mscclFuncSend", + "mscclFuncRecv", + "mscclFuncGather", + "mscclFuncScatter", + "mscclFuncAllToAll", + "mscclFuncAllToAllv", + }; + static const char* mscclSchedulerPathEnv = "MSCCL_SCHEDULER"; static const char* mscclSchedulerDefaultPath = "libmsccl-scheduler.so"; static const char* mscclAlgoDirEnv = "MSCCL_ALGO_DIR"; @@ -80,12 +94,18 @@ static ncclResult_t mscclInternalSchedulerInit(ncclComm_t comm, int* numChannels static thread_local bool mscclAlgoMetaLoaded = false; mscclStatus& status = mscclGetStatus(comm->rank); + int maxNchannels = *numChannelsRequired; *numChannelsRequired = 0; // Query numChannelsRequired from loaded algorithm metas if (mscclAlgoMetaLoaded) { for (auto& m : status.algoMetas) { if (comm->nRanks == m.nRanks) { - *numChannelsRequired = std::max(*numChannelsRequired, m.nChannels); + if(m.nChannels <= maxNchannels) { + *numChannelsRequired = std::max(*numChannelsRequired, m.nChannels); + } else { + WARN("NCCL_MAX_NCHANNELS:%d is lesser than number of channels required by MSCCL:%d, so disabling MSCCL for %s between minBytes:%ld and maxBytes:%ld from file: %s", \ + maxNchannels, m.nChannels, mscclFuncNames[m.func], m.minBytes, m.maxBytes, m.filePath.c_str()); + } } } return ncclSuccess; @@ -138,7 +158,14 @@ static ncclResult_t mscclInternalSchedulerInit(ncclComm_t comm, int* numChannels fullPath += entry->d_name; NCCLCHECK(mscclGetAlgoMetaFromXmlFile(fullPath.c_str(), &(status.algoMetas.back()))); if (status.algoMetas.back().nRanks == comm->nRanks) { - *numChannelsRequired = std::max(*numChannelsRequired, status.algoMetas.back().nChannels); + if(status.algoMetas.back().nChannels <= maxNchannels) { + *numChannelsRequired = std::max(*numChannelsRequired, status.algoMetas.back().nChannels); + } else { + WARN("NCCL_MAX_NCHANNELS:%d is lesser than number of channels required by MSCCL:%d, so disabling MSCCL for %s between minBytes:%ld and maxBytes:%ld from file: %s", \ + maxNchannels, status.algoMetas.back().nChannels, mscclFuncNames[status.algoMetas.back().func], status.algoMetas.back().minBytes, status.algoMetas.back().maxBytes, \ + status.algoMetas.back().filePath.c_str()); + status.algoMetas.pop_back(); + } } } if (closedir(dp)) { @@ -151,7 +178,6 @@ static ncclResult_t mscclInternalSchedulerInit(ncclComm_t comm, int* numChannels } ncclResult_t mscclSchedulerInit(ncclComm_t comm, int* numChannelsRequired) { - *numChannelsRequired = 0; comm->mscclCompatible = mscclCommCompatible(comm); if (!comm->mscclCompatible) { return ncclSuccess; @@ -348,20 +374,6 @@ static ncclResult_t mscclSaveCountsAndDispls(struct mscclSavedSchedulerParam* pa return ncclSuccess; } -const char *mscclFuncNames[] = { - "mscclFuncReduce", - "mscclFuncBroadcast", - "mscclFuncAllReduce", - "mscclFuncReduceScatter", - "mscclFuncAllGather", - "mscclFuncSend", - "mscclFuncRecv", - "mscclFuncGather", - "mscclFuncScatter", - "mscclFuncAllToAll", - "mscclFuncAllToAllv", - }; - static ncclResult_t mscclRunSavedParams() { mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus(); for (auto& param : threadLocalStatus.savedSchedulerParams) {