MSCCL: Improve executor and integrate scheduler (#694)
* MSCCL: improve executor and add scheduler for testing * Use external scheduler * Fix cmake error * Address comments * Fix thread safe issue * Make MSCCL lifecycle APIs thread safe * Make MSCCL internal scheduler aware of topology hint * Revise error message
This commit is contained in:
Ziyue Yang
GitHub
+318
-172
@@ -4,116 +4,192 @@
|
||||
************************************************************************/
|
||||
|
||||
#include <atomic>
|
||||
#include <map>
|
||||
#include <mutex>
|
||||
#include <set>
|
||||
|
||||
#include <dirent.h>
|
||||
#include <dlfcn.h>
|
||||
#include <error.h>
|
||||
#include <link.h>
|
||||
|
||||
#include "alloc.h"
|
||||
#include "checks.h"
|
||||
#include "graph/topo.h"
|
||||
|
||||
#include "msccl/msccl_lifecycle.h"
|
||||
#include "msccl/msccl_parser.h"
|
||||
#include "msccl/msccl_setup.h"
|
||||
#include "msccl/msccl_status.h"
|
||||
|
||||
RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 0);
|
||||
RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
|
||||
static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
|
||||
static std::atomic<bool> mscclInitialized;
|
||||
static bool mscclSchedulerTriedLoadAlgo = false;
|
||||
static std::mutex mscclLifecycleMutex;
|
||||
|
||||
bool mscclEnabled() {
|
||||
return rcclParamMscclEnabled();
|
||||
}
|
||||
|
||||
static bool mscclIsCallerFlag = false;
|
||||
|
||||
void mscclSetIsCallerFlag() {
|
||||
mscclIsCallerFlag = true;
|
||||
mscclGetThreadLocalStatus().mscclIsCallerFlag = true;
|
||||
}
|
||||
|
||||
void mscclClearIsCallerFlag() {
|
||||
mscclIsCallerFlag = false;
|
||||
mscclGetThreadLocalStatus().mscclIsCallerFlag = false;
|
||||
}
|
||||
|
||||
bool mscclIsCaller() {
|
||||
return mscclIsCallerFlag;
|
||||
return mscclGetThreadLocalStatus().mscclIsCallerFlag;
|
||||
}
|
||||
|
||||
bool mscclAvailable() {
|
||||
return mscclEnabled() && mscclInitialized.load(std::memory_order_acquire);
|
||||
}
|
||||
|
||||
ncclResult_t mscclInit(ncclComm_t comm) {
|
||||
if (comm->intraRanks > 1) {
|
||||
mscclInitialized.store(false, std::memory_order_release);
|
||||
INFO(NCCL_INIT, "MSCCL doesn't support multiple GPUs in one process and is not available");
|
||||
return ncclSuccess;
|
||||
static bool mscclCommCompatible(ncclComm_t comm) {
|
||||
std::map<uint64_t, std::set<uint64_t>> hostHashToPidHashes;
|
||||
for (int i = 0; i < comm->nRanks; i++) {
|
||||
uint64_t hostHash = comm->peerInfo[i].hostHash;
|
||||
uint64_t pidHash = comm->peerInfo[i].pidHash;
|
||||
if (hostHashToPidHashes.find(hostHash) != hostHashToPidHashes.end()) {
|
||||
auto& pidHashSet = hostHashToPidHashes[hostHash];
|
||||
if (pidHashSet.find(pidHash) != pidHashSet.end()) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
hostHashToPidHashes[hostHash].insert(pidHash);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static const char* mscclSchedulerPathEnv = "MSCCL_SCHEDULER";
|
||||
static const char* mscclSchedulerDefaultPath = "libmsccl-scheduler.so";
|
||||
static const char* mscclAlgoDirEnv = "MSCCL_ALGO_DIR";
|
||||
static const char* mscclAlgoDefaultDir = "msccl-algorithms";
|
||||
extern "C" bool mscclUnitTestMode() __attribute__((__weak__));
|
||||
static const char* mscclUnitTestAlgoDefaultDir = "msccl-unit-test-algorithms";
|
||||
|
||||
static ncclResult_t mscclInternalSchedulerInit() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
const char* mscclAlgoDir = getenv(mscclAlgoDirEnv);
|
||||
std::string mscclAlgoDirStr;
|
||||
if (mscclAlgoDir == nullptr) {
|
||||
// Try to find default algorithm directory based on librccl.so path
|
||||
Dl_info dl_info;
|
||||
struct link_map *link_map_ptr = nullptr;
|
||||
if (!dladdr1((void *)mscclInternalSchedulerInit, &dl_info, (void **)&link_map_ptr, RTLD_DL_LINKMAP)) {
|
||||
WARN("MSCCL Internal Scheduler: dladdr1 failed");
|
||||
return ncclInvalidUsage;
|
||||
}
|
||||
std::string selfLibPath = link_map_ptr->l_name;
|
||||
mscclAlgoDirStr = selfLibPath.substr(0, selfLibPath.find_last_of("/\\") + 1);
|
||||
mscclAlgoDirStr += (mscclUnitTestMode && mscclUnitTestMode()) ? mscclUnitTestAlgoDefaultDir : mscclAlgoDefaultDir;
|
||||
mscclAlgoDir = mscclAlgoDirStr.c_str();
|
||||
}
|
||||
struct dirent *entry = nullptr;
|
||||
DIR *dp = nullptr;
|
||||
dp = opendir(mscclAlgoDir);
|
||||
if (dp == nullptr) {
|
||||
WARN("MSCCL Internal Scheduler: open algorithm directory %s failed", mscclAlgoDir);
|
||||
return ncclInvalidUsage;
|
||||
}
|
||||
while ((entry = readdir(dp))) {
|
||||
if (entry->d_type != DT_LNK && entry->d_type != DT_REG) {
|
||||
continue;
|
||||
}
|
||||
status.algoMetas.emplace_back();
|
||||
std::string fullPath = mscclAlgoDir;
|
||||
fullPath += "/";
|
||||
fullPath += entry->d_name;
|
||||
NCCLCHECK(mscclGetAlgoMetaFromXmlFile(fullPath.c_str(), &(status.algoMetas.back())));
|
||||
}
|
||||
if (closedir(dp)) {
|
||||
WARN("MSCCL Internal Scheduler: closedir failed, error %d", errno);
|
||||
return ncclInvalidUsage;
|
||||
}
|
||||
status.rankToAlgoHandles.resize(status.algoMetas.size());
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclSchedulerInit() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
bool useInternalScheduler = false;
|
||||
|
||||
const char* mscclSchedulerPath = getenv(mscclSchedulerPathEnv);
|
||||
if (mscclSchedulerPath) {
|
||||
status.mscclSchedulerLib = dlopen(mscclSchedulerPath, RTLD_NOW | RTLD_LOCAL);
|
||||
} else {
|
||||
status.mscclSchedulerLib = dlopen(mscclSchedulerDefaultPath, RTLD_NOW | RTLD_LOCAL);
|
||||
}
|
||||
if (status.mscclSchedulerLib == nullptr) {
|
||||
INFO(NCCL_INIT, "MSCCL: No external scheduler found, using internal implementation");
|
||||
useInternalScheduler = true;
|
||||
} else {
|
||||
status.mscclSchedulerPtr = (mscclSchedulerInterface *)dlsym(status.mscclSchedulerLib, "mscclScheduler");
|
||||
if (status.mscclSchedulerPtr == nullptr) {
|
||||
INFO(NCCL_INIT, "MSCCL: Failed to find mscclScheduler symbol, using internal implementation");
|
||||
useInternalScheduler = true;
|
||||
}
|
||||
}
|
||||
if (useInternalScheduler) {
|
||||
NCCLCHECK(mscclInternalSchedulerInit());
|
||||
} else {
|
||||
NCCLCHECK(status.mscclSchedulerPtr->init());
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t mscclInit(ncclComm_t comm) {
|
||||
// Always initialize thread local status
|
||||
mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
threadLocalStatus.groupStatus = mscclNoGroup;
|
||||
threadLocalStatus.groupDepth = 0;
|
||||
comm->mscclCompatible = mscclCommCompatible(comm);
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
|
||||
|
||||
if (mscclInitialized.load(std::memory_order_acquire)) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
status.scratchBuffer = nullptr;
|
||||
status.scratchBufferSize = 0;
|
||||
status.workIndex = 1;
|
||||
status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
|
||||
for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
|
||||
status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
|
||||
}
|
||||
NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
|
||||
status.lastStream = nullptr;
|
||||
mscclSchedulerTriedLoadAlgo = false;
|
||||
|
||||
NCCLCHECK(mscclSchedulerInit());
|
||||
|
||||
mscclInitialized.store(true, std::memory_order_release);
|
||||
}
|
||||
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
status.scratchBuffer = nullptr;
|
||||
status.scratchBufferSize = 0;
|
||||
status.rank = comm->rank;
|
||||
status.workIndex = 1;
|
||||
status.freeAlgoHandles.resize(MSCCL_MAX_NUM_ALGOS);
|
||||
for (int i = 0; i < MSCCL_MAX_NUM_ALGOS; i++) {
|
||||
status.freeAlgoHandles[i] = MSCCL_MAX_NUM_ALGOS - i - 1;
|
||||
}
|
||||
NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
|
||||
status.groupStatus = mscclNoGroup;
|
||||
status.groupDepth = 0;
|
||||
mscclSchedulerTriedLoadAlgo = false;
|
||||
INFO(NCCL_INIT, "MSCCL: Initialization finished");
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t mscclGroupStart() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
status.groupDepth++;
|
||||
if (status.groupStatus == mscclNoGroup) {
|
||||
status.groupStatus = mscclGroupSupportedOp;
|
||||
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
threadLocalStatus.groupDepth++;
|
||||
if (threadLocalStatus.groupStatus == mscclNoGroup) {
|
||||
threadLocalStatus.groupStatus = mscclGroupSupportedOp;
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
|
||||
static bool algoAvailable = false;
|
||||
static mscclAlgoHandle_t loadedAlgoHandle;
|
||||
static mscclAlgo* loadedHostAlgo = nullptr;
|
||||
|
||||
static ncclResult_t mscclInternalSchedulerSelectAlgo(struct mscclSchedulerParam* param) {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
param->scheduled = false;
|
||||
|
||||
if (!mscclSchedulerTriedLoadAlgo) {
|
||||
mscclSchedulerTriedLoadAlgo = true;
|
||||
const char* mscclAlgoFilePath = getenv(mscclAlgoFilePathEnv);
|
||||
if (mscclAlgoFilePath != nullptr) {
|
||||
NCCLCHECK(mscclLoadAlgo(mscclAlgoFilePath, &loadedAlgoHandle));
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
loadedHostAlgo = status.hostAlgos[loadedAlgoHandle];
|
||||
algoAvailable = true;
|
||||
}
|
||||
}
|
||||
if (!algoAvailable) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
bool mscclAlgoFuncIsValid = loadedHostAlgo->func == param->func;
|
||||
if (!mscclAlgoFuncIsValid) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
bool numGpusIsValid = loadedHostAlgo->nRanks == param->comm->nRanks;
|
||||
if (!numGpusIsValid) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
size_t nBytes = param->count * ncclTypeSize(param->dataType) * loadedHostAlgo->sizeMultiplier;
|
||||
bool msgSizeIsValid =
|
||||
param->count > 0 && (param->count % loadedHostAlgo->nChunksPerLoop) == 0 &&
|
||||
nBytes >= loadedHostAlgo->minBytes &&
|
||||
(loadedHostAlgo->maxBytes == 0 || nBytes <= loadedHostAlgo->maxBytes);
|
||||
if (!msgSizeIsValid) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
// Whether the algorithm is in-place
|
||||
bool isInPlace = false;
|
||||
if (param->func == mscclFuncReduce ||
|
||||
param->func == mscclFuncBroadcast ||
|
||||
@@ -123,120 +199,153 @@ static ncclResult_t mscclScheduler(struct mscclSchedulerParam* param) {
|
||||
isInPlace = param->sendBuff == param->recvBuff;
|
||||
} else if (param->func == mscclFuncAllGather ||
|
||||
param->func == mscclFuncGather) {
|
||||
isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
|
||||
isInPlace = (char*)param->sendBuff == (char*)param->recvBuff + param->rank * param->count * ncclTypeSize(param->dataType);
|
||||
} else if (param->func == mscclFuncReduceScatter ||
|
||||
param->func == mscclFuncScatter) {
|
||||
isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->comm->rank * param->count * ncclTypeSize(param->dataType);
|
||||
}
|
||||
bool inPlaceOutOfPlaceIsValid = isInPlace ? loadedHostAlgo->inPlace : loadedHostAlgo->outOfPlace;
|
||||
if (!inPlaceOutOfPlaceIsValid) {
|
||||
return ncclSuccess;
|
||||
isInPlace = (char*)param->recvBuff == (char*)param->sendBuff + param->rank * param->count * ncclTypeSize(param->dataType);
|
||||
}
|
||||
|
||||
// Search suitable algorithms
|
||||
for (size_t i = 0; i < status.algoMetas.size(); i++) {
|
||||
auto &m = status.algoMetas[i];
|
||||
size_t nBytes = param->count * ncclTypeSize(param->dataType) * m.sizeMultiplier;
|
||||
bool msgSizeIsValid =
|
||||
param->count > 0 && (param->count % m.nChunksPerLoop) == 0 &&
|
||||
nBytes >= m.minBytes && (m.maxBytes == 0 || nBytes <= m.maxBytes);
|
||||
if (msgSizeIsValid &&
|
||||
m.nRanks == param->nRanks &&
|
||||
m.func == param->func &&
|
||||
(isInPlace ? m.inPlace : m.outOfPlace)) {
|
||||
// If not loaded for current rank, load it
|
||||
if (status.rankToAlgoHandles[i].find(param->rank) == status.rankToAlgoHandles[i].end()) {
|
||||
mscclAlgoHandle_t algoHandle;
|
||||
NCCLCHECK(mscclLoadAlgo(m.filePath.c_str(), &algoHandle, param->rank));
|
||||
status.rankToAlgoHandles[i][param->rank] = algoHandle;
|
||||
}
|
||||
param->handle = status.rankToAlgoHandles[i][param->rank];
|
||||
param->scheduled = true;
|
||||
return ncclSuccess;
|
||||
}
|
||||
}
|
||||
|
||||
param->handle = loadedAlgoHandle;
|
||||
param->scheduled = true;
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclSetSchedulerParam(
|
||||
static ncclResult_t mscclSchedulerSelectAlgo(struct mscclSavedSchedulerParam* param) {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
if (status.mscclSchedulerPtr) {
|
||||
NCCLCHECK(status.mscclSchedulerPtr->selectAlgo(&(param->p)));
|
||||
} else {
|
||||
if (param->comm->topo->mscclEnabled) {
|
||||
NCCLCHECK(mscclInternalSchedulerSelectAlgo(&(param->p)));
|
||||
} else {
|
||||
param->p.scheduled = false;
|
||||
}
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclSetSavedSchedulerParam(
|
||||
const void* sendBuff, const size_t sendCounts[], const size_t sDisPls[],
|
||||
void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
|
||||
size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
|
||||
mscclFunc_t func, ncclComm_t comm, hipStream_t stream,
|
||||
struct mscclSchedulerParam* param) {
|
||||
param->sendBuff = sendBuff;
|
||||
param->sendCounts = sendCounts;
|
||||
param->sDisPls = sDisPls;
|
||||
param->recvBuff = recvBuff;
|
||||
param->recvCounts = recvCounts;
|
||||
param->rDisPls = rDisPls;
|
||||
param->count = count;
|
||||
param->dataType = dataType;
|
||||
param->root = root;
|
||||
param->peer = peer;
|
||||
param->op = op;
|
||||
param->func = func;
|
||||
struct mscclSavedSchedulerParam* param) {
|
||||
param->p.sendBuff = sendBuff;
|
||||
param->p.sendCounts = sendCounts;
|
||||
param->p.sDisPls = sDisPls;
|
||||
param->p.recvBuff = recvBuff;
|
||||
param->p.recvCounts = recvCounts;
|
||||
param->p.rDisPls = rDisPls;
|
||||
param->p.count = count;
|
||||
param->p.dataType = dataType;
|
||||
param->p.root = root;
|
||||
param->p.peer = peer;
|
||||
param->p.op = op;
|
||||
param->p.func = func;
|
||||
param->p.rank = comm->rank;
|
||||
param->p.nRanks = comm->nRanks;
|
||||
param->comm = comm;
|
||||
param->stream = stream;
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclSaveCountsAndDispls(struct mscclSchedulerParam* param) {
|
||||
if (param->sendCounts) {
|
||||
param->savedSendCounts.assign(param->sendCounts, param->sendCounts + param->comm->nRanks);
|
||||
param->sendCounts = param->savedSendCounts.data();
|
||||
param->savedSDisPls.assign(param->sDisPls, param->sDisPls + param->comm->nRanks);
|
||||
param->sDisPls = param->savedSDisPls.data();
|
||||
param->savedRecvCounts.assign(param->recvCounts, param->recvCounts + param->comm->nRanks);
|
||||
param->recvCounts = param->savedRecvCounts.data();
|
||||
param->savedRDisPls.assign(param->rDisPls, param->rDisPls + param->comm->nRanks);
|
||||
param->rDisPls = param->savedRDisPls.data();
|
||||
static ncclResult_t mscclSaveCountsAndDispls(struct mscclSavedSchedulerParam* param) {
|
||||
if (param->p.sendCounts) {
|
||||
param->savedSendCounts.assign(param->p.sendCounts, param->p.sendCounts + param->p.nRanks);
|
||||
param->p.sendCounts = param->savedSendCounts.data();
|
||||
param->savedSDisPls.assign(param->p.sDisPls, param->p.sDisPls + param->p.nRanks);
|
||||
param->p.sDisPls = param->savedSDisPls.data();
|
||||
param->savedRecvCounts.assign(param->p.recvCounts, param->p.recvCounts + param->p.nRanks);
|
||||
param->p.recvCounts = param->savedRecvCounts.data();
|
||||
param->savedRDisPls.assign(param->p.rDisPls, param->p.rDisPls + param->p.nRanks);
|
||||
param->p.rDisPls = param->savedRDisPls.data();
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclRunSavedParams() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
for (auto& param : status.savedSchedulerParams) {
|
||||
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
for (auto& param : threadLocalStatus.savedSchedulerParams) {
|
||||
NCCLCHECK(mscclRunAlgo(
|
||||
param.sendBuff, param.sendCounts, param.sDisPls,
|
||||
param.recvBuff, param.recvCounts, param.rDisPls,
|
||||
param.count, param.dataType, param.root, param.peer, param.op, param.handle, param.comm, param.stream));
|
||||
param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
|
||||
param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
|
||||
param.p.count, param.p.dataType, param.p.root, param.p.peer, param.p.op, param.p.handle, param.comm, param.stream));
|
||||
}
|
||||
status.savedSchedulerParams.clear();
|
||||
threadLocalStatus.savedSchedulerParams.clear();
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
static ncclResult_t mscclFallBackSavedParams() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
mscclSetIsCallerFlag();
|
||||
for (auto& param : status.savedSchedulerParams) {
|
||||
switch (param.func) {
|
||||
for (auto& param : threadLocalStatus.savedSchedulerParams) {
|
||||
switch (param.p.func) {
|
||||
case mscclFuncReduce:
|
||||
NCCLCHECK(ncclReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.op, param.root, param.comm, param.stream));
|
||||
NCCLCHECK(ncclReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.op, param.p.root, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncBroadcast:
|
||||
NCCLCHECK(ncclBroadcast(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.root, param.comm, param.stream));
|
||||
NCCLCHECK(ncclBroadcast(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.root, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncAllReduce:
|
||||
NCCLCHECK(ncclAllReduce(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.op, param.comm, param.stream));
|
||||
NCCLCHECK(ncclAllReduce(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.op, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncReduceScatter:
|
||||
NCCLCHECK(ncclReduceScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.op, param.comm, param.stream));
|
||||
NCCLCHECK(ncclReduceScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.op, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncAllGather:
|
||||
NCCLCHECK(ncclAllGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
NCCLCHECK(ncclAllGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncSend:
|
||||
NCCLCHECK(ncclSend(param.sendBuff, param.count, param.dataType,
|
||||
param.peer, param.comm, param.stream));
|
||||
NCCLCHECK(ncclSend(param.p.sendBuff, param.p.count, param.p.dataType,
|
||||
param.p.peer, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncRecv:
|
||||
NCCLCHECK(ncclRecv(param.recvBuff, param.count, param.dataType,
|
||||
param.peer, param.comm, param.stream));
|
||||
NCCLCHECK(ncclRecv(param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.peer, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncGather:
|
||||
NCCLCHECK(ncclGather(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.root, param.comm, param.stream));
|
||||
NCCLCHECK(ncclGather(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.root, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncScatter:
|
||||
NCCLCHECK(ncclScatter(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
param.root, param.comm, param.stream));
|
||||
NCCLCHECK(ncclScatter(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.p.root, param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncAllToAll:
|
||||
NCCLCHECK(ncclAllToAll(param.sendBuff, param.recvBuff, param.count, param.dataType,
|
||||
NCCLCHECK(ncclAllToAll(param.p.sendBuff, param.p.recvBuff, param.p.count, param.p.dataType,
|
||||
param.comm, param.stream));
|
||||
break;
|
||||
case mscclFuncAllToAllv:
|
||||
NCCLCHECK(ncclAllToAllv(
|
||||
param.sendBuff, param.sendCounts, param.sDisPls,
|
||||
param.recvBuff, param.recvCounts, param.rDisPls,
|
||||
param.dataType, param.comm, param.stream));
|
||||
param.p.sendBuff, param.p.sendCounts, param.p.sDisPls,
|
||||
param.p.recvBuff, param.p.recvCounts, param.p.rDisPls,
|
||||
param.p.dataType, param.comm, param.stream));
|
||||
break;
|
||||
default:
|
||||
WARN("Invalid MSCCL function type in saved parameter");
|
||||
@@ -244,7 +353,7 @@ static ncclResult_t mscclFallBackSavedParams() {
|
||||
}
|
||||
}
|
||||
mscclClearIsCallerFlag();
|
||||
status.savedSchedulerParams.clear();
|
||||
threadLocalStatus.savedSchedulerParams.clear();
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
@@ -253,40 +362,43 @@ ncclResult_t mscclEnqueueCheck(
|
||||
void* recvBuff, const size_t recvCounts[], const size_t rDisPls[],
|
||||
size_t count, ncclDataType_t dataType, int root, int peer, ncclRedOp_t op,
|
||||
mscclFunc_t func, ncclComm_t comm, hipStream_t stream) {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
hipStreamCaptureStatus captureStatus;
|
||||
unsigned long long pid;
|
||||
|
||||
status.savedSchedulerParams.push_back({});
|
||||
NCCLCHECK(mscclSetSchedulerParam(
|
||||
threadLocalStatus.savedSchedulerParams.push_back({});
|
||||
NCCLCHECK(mscclSetSavedSchedulerParam(
|
||||
sendBuff, sendCounts, sDisPls, recvBuff, recvCounts, rDisPls,
|
||||
count, dataType, root, peer, op, func, comm, stream,
|
||||
&status.savedSchedulerParams.back()));
|
||||
&threadLocalStatus.savedSchedulerParams.back()));
|
||||
|
||||
switch (status.groupStatus) {
|
||||
switch (threadLocalStatus.groupStatus) {
|
||||
case mscclNoGroup:
|
||||
CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
|
||||
if (captureStatus == hipStreamCaptureStatusNone) {
|
||||
NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
|
||||
if (status.savedSchedulerParams.back().scheduled) {
|
||||
NCCLCHECK(mscclRunSavedParams());
|
||||
break;
|
||||
if (comm->mscclCompatible) {
|
||||
CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
|
||||
if (captureStatus == hipStreamCaptureStatusNone) {
|
||||
NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
|
||||
if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
|
||||
NCCLCHECK(mscclRunSavedParams());
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
NCCLCHECK(mscclFallBackSavedParams());
|
||||
break;
|
||||
case mscclGroupSupportedOp:
|
||||
CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
|
||||
if (captureStatus == hipStreamCaptureStatusNone) {
|
||||
NCCLCHECK(mscclScheduler(&status.savedSchedulerParams.back()));
|
||||
if (status.savedSchedulerParams.back().scheduled) {
|
||||
// Only save counts and displs when there is suitable MSCCL algorithm for this
|
||||
NCCLCHECK(mscclSaveCountsAndDispls(&status.savedSchedulerParams.back()));
|
||||
break;
|
||||
if (comm->mscclCompatible) {
|
||||
CUDACHECK(hipStreamGetCaptureInfo(stream, &captureStatus, &pid));
|
||||
if (captureStatus == hipStreamCaptureStatusNone) {
|
||||
NCCLCHECK(mscclSchedulerSelectAlgo(&threadLocalStatus.savedSchedulerParams.back()));
|
||||
if (threadLocalStatus.savedSchedulerParams.back().p.scheduled) {
|
||||
// Only save counts and displs when there is suitable MSCCL algorithm for this
|
||||
NCCLCHECK(mscclSaveCountsAndDispls(&threadLocalStatus.savedSchedulerParams.back()));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
NCCLCHECK(mscclFallBackSavedParams());
|
||||
break;
|
||||
threadLocalStatus.groupStatus = mscclGroupUnsupportedOp;
|
||||
case mscclGroupUnsupportedOp:
|
||||
NCCLCHECK(mscclFallBackSavedParams());
|
||||
break;
|
||||
@@ -297,37 +409,71 @@ ncclResult_t mscclEnqueueCheck(
|
||||
}
|
||||
|
||||
ncclResult_t mscclGroupEnd() {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
status.groupDepth--;
|
||||
if (status.groupDepth == 0) {
|
||||
if (status.groupStatus == mscclGroupSupportedOp) {
|
||||
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
threadLocalStatus.groupDepth--;
|
||||
if (threadLocalStatus.groupDepth == 0) {
|
||||
if (threadLocalStatus.groupStatus == mscclGroupSupportedOp) {
|
||||
NCCLCHECK(mscclRunSavedParams());
|
||||
}
|
||||
status.groupStatus = mscclNoGroup;
|
||||
threadLocalStatus.groupStatus = mscclNoGroup;
|
||||
}
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t mscclTeardown() {
|
||||
if (!mscclInitialized.load(std::memory_order_acquire)) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
static ncclResult_t mscclInternalSchedulerTeardown() {
|
||||
ncclResult_t ret = ncclSuccess, tmpRet = ncclSuccess;
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
for (auto &p : status.hostAlgos) {
|
||||
free(p.second);
|
||||
status.freeAlgoHandles.push_back(p.first);
|
||||
for (auto &m : status.rankToAlgoHandles) {
|
||||
for (auto &p : m) {
|
||||
tmpRet = mscclUnloadAlgo(p.second);
|
||||
if (ret == ncclSuccess) {
|
||||
ret = tmpRet;
|
||||
}
|
||||
}
|
||||
}
|
||||
for (auto &p : status.devAlgos) {
|
||||
CUDACHECK(hipFree(p.second));
|
||||
status.algoMetas.clear();
|
||||
status.rankToAlgoHandles.clear();
|
||||
return ret;
|
||||
}
|
||||
|
||||
ncclResult_t mscclTeardown() {
|
||||
// Always teardown thread local status
|
||||
mscclThreadLocalStatus threadLocalStatus = mscclGetThreadLocalStatus();
|
||||
threadLocalStatus.savedSchedulerParams.clear();
|
||||
|
||||
{
|
||||
std::lock_guard<std::mutex> lock(mscclLifecycleMutex);
|
||||
|
||||
if (!mscclInitialized.load(std::memory_order_acquire)) {
|
||||
return ncclSuccess;
|
||||
}
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
for (auto &p : status.hostAlgos) {
|
||||
free(p.second);
|
||||
status.freeAlgoHandles.push_back(p.first);
|
||||
}
|
||||
for (auto &p : status.devAlgos) {
|
||||
CUDACHECK(hipFree(p.second));
|
||||
}
|
||||
CUDACHECK(hipFree(status.scratchBuffer));
|
||||
CUDACHECK(hipFree(status.syncFlags));
|
||||
status.hostAlgos.clear();
|
||||
status.devAlgos.clear();
|
||||
status.freeAlgoHandles.clear();
|
||||
status.scratchBuffer = nullptr;
|
||||
status.scratchBufferSize = 0;
|
||||
status.connectedAlgos.clear();
|
||||
if (status.mscclSchedulerPtr) {
|
||||
NCCLCHECK(status.mscclSchedulerPtr->teardown());
|
||||
status.mscclSchedulerPtr = nullptr;
|
||||
dlclose(status.mscclSchedulerLib);
|
||||
status.mscclSchedulerLib = nullptr;
|
||||
} else {
|
||||
NCCLCHECK(mscclInternalSchedulerTeardown());
|
||||
}
|
||||
mscclInitialized.store(false, std::memory_order_release);
|
||||
}
|
||||
CUDACHECK(hipFree(status.scratchBuffer));
|
||||
CUDACHECK(hipFree(status.syncFlags));
|
||||
status.hostAlgos.clear();
|
||||
status.devAlgos.clear();
|
||||
status.freeAlgoHandles.clear();
|
||||
status.scratchBuffer = nullptr;
|
||||
status.scratchBufferSize = 0;
|
||||
status.workIndex = 1;
|
||||
mscclInitialized.store(false, std::memory_order_release);
|
||||
|
||||
INFO(NCCL_INIT, "MSCCL: Teardown finished");
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
@@ -701,3 +701,88 @@ ncclResult_t mscclGetAlgoFromXmlFile(const char* str, struct mscclAlgo* algo, in
|
||||
free(xml);
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t mscclXmlLoadSingleNode(FILE* file, struct mscclXmlNode* node) {
|
||||
memset(node, 0, sizeof(struct mscclXmlNode));
|
||||
return mscclXmlGetNode(file, node);
|
||||
}
|
||||
|
||||
ncclResult_t mscclAlgoMetaXmlLoad(const char* xmlFilePath, struct mscclXmlNode* node) {
|
||||
ncclResult_t ret = ncclSuccess;
|
||||
FILE* file = fopen(xmlFilePath, "r");
|
||||
if (file == NULL) {
|
||||
fprintf(stderr, "Could not open MSCCL XML algorithm file %s : %s", xmlFilePath, strerror(errno));
|
||||
return ncclSystemError;
|
||||
}
|
||||
NCCLCHECK(mscclXmlLoadSingleNode(file, node));
|
||||
fclose(file);
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
ncclResult_t mscclGetAlgoMetaFromXmlFile(const char* str, struct mscclAlgoMeta* algoMeta) {
|
||||
ncclResult_t ret = ncclSuccess;
|
||||
struct mscclXmlNode* node;
|
||||
node = (struct mscclXmlNode *)malloc(sizeof(struct mscclXmlNode));
|
||||
NCCLCHECK(mscclAlgoMetaXmlLoad(str, node));
|
||||
|
||||
algoMeta->filePath = str;
|
||||
|
||||
int nChunksPerLoop;
|
||||
NCCLCHECK(mscclXmlGetAttrInt(node, "nchunksperloop", &nChunksPerLoop));
|
||||
algoMeta->nChunksPerLoop = nChunksPerLoop;
|
||||
|
||||
int nGpus;
|
||||
NCCLCHECK(mscclXmlGetAttrInt(node, "ngpus", &nGpus));
|
||||
algoMeta->nRanks = nGpus;
|
||||
|
||||
const char* coll;
|
||||
NCCLCHECK(mscclXmlGetAttrStr(node, "coll", &coll));
|
||||
algoMeta->sizeMultiplier = 1;
|
||||
if (strcmp(coll, "reduce") == 0) {
|
||||
algoMeta->func = mscclFuncReduce;
|
||||
} else if (strcmp(coll, "broadcast") == 0) {
|
||||
algoMeta->func = mscclFuncBroadcast;
|
||||
} else if (strcmp(coll, "allreduce") == 0) {
|
||||
algoMeta->func = mscclFuncAllReduce;
|
||||
} else if (strcmp(coll, "reducescatter") == 0) {
|
||||
algoMeta->sizeMultiplier = nGpus;
|
||||
algoMeta->func = mscclFuncReduceScatter;
|
||||
} else if (strcmp(coll, "allgather") == 0) {
|
||||
algoMeta->sizeMultiplier = nGpus;
|
||||
algoMeta->func = mscclFuncAllGather;
|
||||
} else if (strcmp(coll, "send") == 0) {
|
||||
algoMeta->func = mscclFuncSend;
|
||||
} else if (strcmp(coll, "recv") == 0) {
|
||||
algoMeta->func = mscclFuncRecv;
|
||||
} else if (strcmp(coll, "gather") == 0) {
|
||||
algoMeta->func = mscclFuncGather;
|
||||
} else if (strcmp(coll, "scatter") == 0) {
|
||||
algoMeta->func = mscclFuncScatter;
|
||||
} else if (strcmp(coll, "alltoall") == 0) {
|
||||
algoMeta->sizeMultiplier = nGpus;
|
||||
algoMeta->func = mscclFuncAllToAll;
|
||||
} else if (strcmp(coll, "alltoallv") == 0) {
|
||||
algoMeta->func = mscclFuncAllToAllv;
|
||||
} else {
|
||||
return ncclInvalidUsage;
|
||||
}
|
||||
|
||||
int64_t minBytes;
|
||||
NCCLCHECK(mscclXmlGetAttrInt64(node, "minBytes", &minBytes));
|
||||
algoMeta->minBytes = minBytes;
|
||||
|
||||
int64_t maxBytes;
|
||||
NCCLCHECK(mscclXmlGetAttrInt64(node, "maxBytes", &maxBytes));
|
||||
algoMeta->maxBytes = maxBytes;
|
||||
|
||||
int inplace;
|
||||
NCCLCHECK(mscclXmlGetAttrInt(node, "inplace", &inplace));
|
||||
algoMeta->inPlace = (bool)inplace;
|
||||
|
||||
int outofplace;
|
||||
NCCLCHECK(mscclXmlGetAttrInt(node, "outofplace", &outofplace));
|
||||
algoMeta->outOfPlace = (bool)outofplace;
|
||||
|
||||
free(node);
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
@@ -258,6 +258,11 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
|
||||
ncclDataType_t dataType, ncclRedOp_t op, struct mscclAlgo* hostAlgo, struct mscclAlgo* devAlgo,
|
||||
ncclComm_t comm, hipStream_t stream) {
|
||||
mscclStatus& status = mscclGetStatus();
|
||||
|
||||
if (status.lastStream != stream && status.lastStream != nullptr) {
|
||||
CUDACHECK(hipStreamWaitEvent(stream, comm->doneEvent, 0));
|
||||
}
|
||||
|
||||
dim3 grid = {(uint32_t)hostAlgo->nBlocks, 1, 1};
|
||||
dim3 block = {NCCL_MAX_NTHREADS, 1, 1};
|
||||
ncclDevRedOpFull opFull;
|
||||
@@ -278,7 +283,8 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
|
||||
|
||||
void *args[3] = {&comm->devComm, &devAlgo, &work};
|
||||
void *func = mscclKernelEntries[(opFull.op * ncclNumTypes + dataType) * NCCL_NUM_PROTOCOLS + hostAlgo->protocol];
|
||||
CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, NULL,0));
|
||||
CUDACHECK(hipExtLaunchKernel(func, grid, block, args, 0, stream, NULL, comm->doneEvent, 0));
|
||||
status.workIndex++;
|
||||
status.lastStream = stream;
|
||||
return ncclSuccess;
|
||||
}
|
||||
|
||||
@@ -9,3 +9,8 @@ mscclStatus& mscclGetStatus() {
|
||||
static mscclStatus status;
|
||||
return status;
|
||||
}
|
||||
|
||||
mscclThreadLocalStatus& mscclGetThreadLocalStatus() {
|
||||
static thread_local mscclThreadLocalStatus threadLocalStatus;
|
||||
return threadLocalStatus;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user