Merge remote-tracking branch 'nccl/master' into develop

This commit is contained in:
BertanDogancay
2025-04-23 20:46:36 -07:00
118 changed files with 12789 additions and 3358 deletions
+15 -7
View File
@@ -344,7 +344,7 @@ static void groupCleanup(struct ncclComm** groupCommHeadPtr, struct ncclComm** g
/* reset everything */
while (!ncclIntruQueueEmpty(asyncJobsPtr)) {
struct ncclAsyncJob* job = ncclIntruQueueDequeue(asyncJobsPtr);
if (job->comm && !job->comm->config.blocking)
if (!job->destroyFlag && job->comm && !job->comm->config.blocking)
(void) ncclCommSetAsyncError(job->comm, error);
if (job->undo) job->undo(job);
if (job->destructor) job->destructor((void*)job);
@@ -413,7 +413,6 @@ fail:
}
static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInfo = NULL) {
int savedDev;
ncclResult_t ret = ncclSuccess;
struct ncclGroupJob *gjob = (struct ncclGroupJob*) job_;
struct ncclComm *groupCommHeadMain = *gjob->groupCommHeadPtr;
@@ -422,8 +421,6 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
bool *groupAbortFlag = gjob->abortFlagPtr;
CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
if (!simInfo && groupCommPreconnectHeadMain != nullptr) {
struct ncclComm* comm = groupCommPreconnectHeadMain;
do {
@@ -475,12 +472,19 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
}
comm = comm->groupNext;
} while (comm);
NCCLCHECKGOTO(asyncJobLaunch(&asyncCollJobs, groupAbortFlag), ret, fail);
while (!ncclIntruQueueEmpty(&asyncCollJobs)) {
struct ncclAsyncJob* job = ncclIntruQueueDequeue(&asyncCollJobs);
if (job->destructor) job->destructor((void*)job);
}
// done with all buffer allocation, start registration and enqueue
comm = groupCommHeadMain;
do {
CUDACHECKGOTO(cudaSetDevice(comm->cudaDev), ret, fail);
NCCLCHECKGOTO(ncclTasksRegAndEnqueue(comm), ret, fail);
comm = comm->groupNext;
} while (comm);
}
if ((!simInfo) && (groupCommHeadMain != nullptr)) {
@@ -497,6 +501,9 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
while (groupCommHeadMain != nullptr) {
struct ncclComm* comm = groupCommHeadMain;
struct ncclComm* next = comm->groupNext;
// Poll for callbacks sent to us from other threads. Typically these free
// resources from to our memory pools and UB
NCCLCHECKGOTO(ncclCommPollCallbacks(comm, /*waitSome=*/false), ret, fail);
(void) ncclGroupCommLeave(comm);
if (!comm->config.blocking) {
(void) ncclCommSetAsyncError(comm, ret);
@@ -504,8 +511,6 @@ static ncclResult_t groupLaunch(struct ncclAsyncJob *job_, ncclSimInfo_t* simInf
groupCommHeadMain = next;
}
CUDACHECK(cudaSetDevice(savedDev));
exit:
return ret;
fail:
@@ -588,7 +593,10 @@ ncclResult_t ncclGroupEndInternal(ncclSimInfo_t* simInfo) {
ret = ncclInProgress;
} else {
/* blocking group */
int savedDev;
CUDACHECKGOTO(cudaGetDevice(&savedDev), ret, fail);
NCCLCHECKGOTO(groupLaunch(&ncclGroupJobMainPtr->base, internalSimInfoPtr), ret, fail);
CUDACHECKGOTO(cudaSetDevice(savedDev), ret, fail);
if (simInfo) memcpy((void*)simInfo, (void*)internalSimInfoPtr, realSize);
groupResetJobState(ncclGroupJobMainPtr);
}