AI/rocm-systems
2
0
Forkuj 0
Kod Zgłoszenia Oczekujące zmiany Actions Packages Projekty Wydania Wiki Aktywność

Fix MSCCL work FIFO allocation with HIP graph enabled (#967)

Przeglądaj źródła 
[ROCm/rccl commit: 7fc891bc8d]
This commit is contained in:
Ziyue Yang
2023-11-16 08:43:28 +08:00
zatwierdzone przez GitHub
rodzic 9e8eb41337
commit 6ce074d92d
5 zmienionych plików z 80 dodań i 40 usunięć
Ściągnij plik aktualizacji Ściągnij plik porównania Expand all files Collapse all files
projects/rccl/src/collectives/msccl.cc
+1 -1
Wyświetl plik
@@ -27,7 +27,7 @@ ncclResult_t mscclLoadAlgo(const char *mscclAlgoFilePath, mscclAlgoHandle_t *msc
status.hostAlgos[*mscclAlgoHandle] = hostAlgo;
struct mscclAlgo* devAlgo;
NCCLCHECK(ncclCudaCalloc(&devAlgo, 1));
NCCLCHECK(ncclCudaMalloc(&devAlgo, 1));
CUDACHECK(hipMemcpy(devAlgo, hostAlgo, sizeof(struct mscclAlgo), hipMemcpyHostToDevice));
status.devAlgos[*mscclAlgoHandle] = devAlgo;
projects/rccl/src/include/msccl/msccl_setup.h
+4
Wyświetl plik
@@ -27,4 +27,8 @@ 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);
ncclResult_t mscclInitWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus);
ncclResult_t mscclDestroyWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus);
#endif
projects/rccl/src/include/msccl/msccl_struct.h
+13 -8
Wyświetl plik
@@ -36,8 +36,6 @@
#define MSCCL_LOCAL_COPY 6
#define MSCCL_REDUCE 7
#define MSCCL_WORK_FIFO_DEPTH (64 << 10)
struct mscclTransmission {
int16_t dependencePointer; // index to the first dependence
int16_t numDependencies; // dependencePointer+numDependencies indicate the last dependence
@@ -189,6 +187,17 @@ struct mscclThreadLocalStatus {
hipGraph_t graph;
};
struct mscclWorkFifoStatus {
uint64_t workFifoDepth;
struct mscclWork* workFifo;
uint32_t* workFifoDone;
uint32_t workFifoSent;
uint32_t workFifoSentPerThreadBlock[MSCCL_MAX_NUM_THREAD_BLOCKS];
uint32_t workFifoAckdMin;
};
typedef std::map<unsigned long long, mscclWorkFifoStatus> mscclSavedGraphWorkFifoStatus;
struct mscclStatus {
std::vector<mscclAlgoHandle_t> freeAlgoHandles;
std::map<mscclAlgoHandle_t, mscclAlgo *> hostAlgos;
@@ -214,12 +223,8 @@ struct mscclStatus {
bool graphEnabled;
bool graphFirstKernel;
bool needsProxy;
uint64_t workFifoDepth;
struct mscclWork* workFifo;
uint32_t* workFifoDone;
uint32_t workFifoSent;
uint32_t workFifoSentPerThreadBlock[MSCCL_MAX_NUM_THREAD_BLOCKS];
uint32_t workFifoAckdMin;
mscclWorkFifoStatus defaultWorkFifoStatus;
mscclSavedGraphWorkFifoStatus graphWorkFifoStatus;
};
#pragma pack(push)
projects/rccl/src/misc/msccl/msccl_lifecycle.cc
+5 -10
Wyświetl plik
@@ -190,14 +190,7 @@ ncclResult_t mscclInit(ncclComm_t comm) {
NCCLCHECK(ncclCudaCalloc(&status.syncFlags, MSCCL_MAX_NUM_THREAD_BLOCKS));
status.lastStream = nullptr;
status.needsProxy = false;
status.workFifoDepth = MSCCL_WORK_FIFO_DEPTH;
NCCLCHECK(ncclCudaCalloc(&status.workFifo, status.workFifoDepth, nullptr, true));
NCCLCHECK(ncclCudaHostCalloc(&status.workFifoDone, MSCCL_MAX_NUM_THREAD_BLOCKS));
status.workFifoSent = 0;
for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) {
status.workFifoSentPerThreadBlock[i] = 0;
}
status.workFifoAckdMin = 0;
NCCLCHECK(mscclInitWorkFifoStatus(&(status.defaultWorkFifoStatus)));
mscclSchedulerTriedLoadAlgo = false;
NCCLCHECK(mscclSchedulerInit());
@@ -521,8 +514,10 @@ ncclResult_t mscclTeardown() {
} else {
NCCLCHECK(mscclInternalSchedulerTeardown());
}
NCCLCHECK(ncclCudaFree(status.workFifo));
NCCLCHECK(ncclCudaHostFree(status.workFifoDone));
NCCLCHECK(mscclDestroyWorkFifoStatus(&(status.defaultWorkFifoStatus)));
for (auto &p : status.graphWorkFifoStatus) {
NCCLCHECK(mscclDestroyWorkFifoStatus(&(p.second)));
}
mscclInitialized.store(false, std::memory_order_release);
}
projects/rccl/src/misc/msccl/msccl_setup.cc
+57 -21
Wyświetl plik
@@ -20,7 +20,10 @@
RCCL_PARAM(MscclEnableDoneEvent, "MSCCL_ENABLE_DONE_EVENT", 1);
#endif
RCCL_PARAM(MscclWorkFifoDepth, "MSCCL_WORK_FIFO_DEPTH", 64<<10);
ncclResult_t mscclGetCaptureStatus(hipStream_t stream) {
mscclStatus& status = mscclGetStatus();
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
mscclSavedProxyArgs& savedProxyArgs = mscclGetSavedProxyArgs();
cudaStreamCaptureStatus captureStatus;
@@ -30,6 +33,7 @@ ncclResult_t mscclGetCaptureStatus(hipStream_t stream) {
if (savedProxyArgs.count(captureId) == 0) {
threadLocalStatus.captureStatus = mscclNewCapture;
savedProxyArgs[captureId] = std::vector<struct mscclProxyArg>();
NCCLCHECK(mscclInitWorkFifoStatus(&(status.graphWorkFifoStatus[captureId])));
} else {
INFO(NCCL_NET,"mscclGetCaptureStatus: captureId %llu is same with the previous one\n", captureId);
threadLocalStatus.captureStatus = mscclExistingCapture;
@@ -69,7 +73,7 @@ ncclResult_t mscclSetupScratch(struct mscclAlgo* hostAlgo, hipStream_t stream) {
size_t sizeNeeded = (status.nBytes * (size_t)(hostAlgo->nScratchChunks)) / (size_t)(hostAlgo->nChunksPerLoop);
if (sizeNeeded > status.scratchBufferSize){
NCCLCHECK(ncclCudaFree(status.scratchBuffer));
NCCLCHECK(ncclCudaCalloc((char**)&status.scratchBuffer, sizeNeeded));
NCCLCHECK(ncclCudaMalloc((char**)&status.scratchBuffer, sizeNeeded));
status.scratchBufferSize = sizeNeeded;
}
return ncclSuccess;
@@ -333,12 +337,11 @@ static inline uint32_t rollingMin32(uint32_t a, uint32_t b) {
return (b-a <= PositiveMax) ? a : b;
}
static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) {
mscclStatus& status = mscclGetStatus();
if (__builtin_expect(rollingLess32(status.workFifoAckdMin + status.workFifoDepth, desiredSent), false)) {
static void mscclWaitWorkFifoAvailable(uint32_t desiredSent, mscclWorkFifoStatus* status) {
if (__builtin_expect(rollingLess32(status->workFifoAckdMin + status->workFifoDepth, desiredSent), false)) {
while (1) {
// We have to poll for notifications from device.
uint32_t* doneLive = status.workFifoDone;
uint32_t* doneLive = status->workFifoDone;
uint32_t ackd[MSCCL_MAX_NUM_THREAD_BLOCKS];
for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
ackd[c] = __atomic_load_n(&doneLive[c], __ATOMIC_RELAXED);
@@ -346,10 +349,10 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) {
// Compiler-only fence to prevent fusion of loops to encourage dense loads.
__atomic_signal_fence(__ATOMIC_SEQ_CST);
uint32_t ackdAll = status.workFifoSent;
uint32_t ackdAll = status->workFifoSent;
for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
// ackdAll is min over all non-quiesced channels
if (ackd[c] != status.workFifoSentPerThreadBlock[c])
if (ackd[c] != status->workFifoSentPerThreadBlock[c])
ackdAll = rollingMin32(ackdAll, ackd[c]);
}
@@ -359,15 +362,15 @@ static void mscclWaitWorkFifoAvailable(uint32_t desiredSent) {
for (int c=0; c < MSCCL_MAX_NUM_THREAD_BLOCKS; c++) {
// Advance counter on quiesced channels so they don't lag behind
// too far where they could get lost in 32-bit wraparound.
if (ackd[c] == status.workFifoSentPerThreadBlock[c]) {
status.workFifoSentPerThreadBlock[c] = ackdAll;
if (ackd[c] == status->workFifoSentPerThreadBlock[c]) {
status->workFifoSentPerThreadBlock[c] = ackdAll;
__atomic_store_n(&doneLive[c], ackdAll, __ATOMIC_RELAXED);
}
}
status.workFifoAckdMin = ackdAll;
status->workFifoAckdMin = ackdAll;
// See if that was enough.
if (!rollingLess32(status.workFifoAckdMin + status.workFifoDepth, desiredSent)) break;
if (!rollingLess32(status->workFifoAckdMin + status->workFifoDepth, desiredSent)) break;
sched_yield();
}
}
@@ -379,6 +382,7 @@ 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();
mscclThreadLocalStatus& threadLocalStatus = mscclGetThreadLocalStatus();
bool enableDoneEvent =
#ifndef HIP_EVENT_DISABLE_FENCE
@@ -421,9 +425,22 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
work.redOpArgIsPtr = opFull.scalarArgIsPtr;
work.fnIndex = fnIndex;
INFO(NCCL_COLL, "MSCCL: typeMask %x fnIndex %d Setup Kernel finished", hostAlgo->typeMask, fnIndex);
uint32_t workFifoIdxMask = status.workFifoDepth - 1;
uint32_t workFifoSent = status.workFifoSent;
mscclWorkFifoStatus* workFifoStatus = nullptr;
if (threadLocalStatus.captureStatus == mscclNoCapture) {
workFifoStatus = &(status.defaultWorkFifoStatus);
} else {
workFifoStatus = &(status.graphWorkFifoStatus[threadLocalStatus.captureId]);
}
uint32_t workFifoIdxMask = workFifoStatus->workFifoDepth - 1;
uint32_t workFifoSent = workFifoStatus->workFifoSent;
if (threadLocalStatus.captureStatus != mscclNoCapture && workFifoSent + numBlocks > workFifoStatus->workFifoDepth) {
WARN("MSCCL: number of captured works (%u) > max limit (%lu)", workFifoSent + numBlocks, workFifoStatus->workFifoDepth);
return ncclInternalError;
}
// First work for a channel has to be at workHeap+blockIdx.x which means
// we cannot tolerate fifo wraparound. So round up to the wrap boundary
// if not doing so would incur crossing it.
@@ -432,18 +449,19 @@ ncclResult_t mscclSetupKernel(const void* sendBuff, void* recvBuff, size_t count
// Need to update workFifoSent so waitWorkFifoAvailable() knows we've
// skipped those elements. Consider if all the channels report quiesced,
// this way the skipped slots will be considered consumed as well.
status.workFifoSent = workFifoSent;
workFifoStatus->workFifoSent = workFifoSent;
}
mscclWaitWorkFifoAvailable(workFifoSent + numBlocks);
mscclWaitWorkFifoAvailable(workFifoSent + numBlocks, workFifoStatus);
for (int i = 0; i < numBlocks; i++) {
work.workFifoDoneAck = workFifoSent + i;
work.workFifoDone = status.workFifoDone + i;
status.workFifoSentPerThreadBlock[i] = workFifoSent + i;
status.workFifo[(workFifoSent + i) & workFifoIdxMask] = work;
work.workFifoDone = workFifoStatus->workFifoDone + i;
workFifoStatus->workFifoSentPerThreadBlock[i] = workFifoSent + i;
workFifoStatus->workFifo[(workFifoSent + i) & workFifoIdxMask] = work;
}
status.workFifoSent = workFifoSent + numBlocks;
struct mscclWork *workPtr = status.workFifo + (workFifoSent & workFifoIdxMask);
struct mscclWork *workPtr = workFifoStatus->workFifo + (workFifoSent & workFifoIdxMask);
workFifoStatus->workFifoSent = workFifoSent + numBlocks;
void *args[3] = {&comm->devComm, &devAlgo, &workPtr};
void *func = mscclKernelEntries[fnIndex];
if (enableDoneEvent) {
@@ -472,3 +490,21 @@ size_t mscclKernMaxLocalSize() {
error:
return (res != ncclSuccess) ? 0 : max;
}
ncclResult_t mscclInitWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus) {
workFifoStatus->workFifoDepth = rcclParamMscclWorkFifoDepth();
NCCLCHECK(ncclCudaMalloc(&(workFifoStatus->workFifo), workFifoStatus->workFifoDepth, true));
NCCLCHECK(ncclCudaHostCalloc(&(workFifoStatus->workFifoDone), MSCCL_MAX_NUM_THREAD_BLOCKS));
workFifoStatus->workFifoSent = 0;
for (int i = 0; i < MSCCL_MAX_NUM_THREAD_BLOCKS; i++) {
workFifoStatus->workFifoSentPerThreadBlock[i] = 0;
}
workFifoStatus->workFifoAckdMin = 0;
return ncclSuccess;
}
ncclResult_t mscclDestroyWorkFifoStatus(mscclWorkFifoStatus* workFifoStatus) {
NCCLCHECK(ncclCudaFree(workFifoStatus->workFifo));
NCCLCHECK(ncclCudaHostFree(workFifoStatus->workFifoDone));
return ncclSuccess;
}