From b55b6be0cb2a46a9e08c4a07e272333dd0c369b9 Mon Sep 17 00:00:00 2001 From: Wenkai Du <43822138+wenkaidu@users.noreply.github.com> Date: Fri, 11 Oct 2024 09:02:59 -0700 Subject: [PATCH] Fix crash when PXN is enabled on some platforms (#1369) --- src/graph/search.cc | 6 +++--- src/graph/xml.cc | 15 +++------------ src/include/comm.h | 4 +--- src/include/transport.h | 1 + src/init.cc | 25 ++++++++++++++++--------- src/transport/net.cc | 12 ++++-------- 6 files changed, 28 insertions(+), 35 deletions(-) diff --git a/src/graph/search.cc b/src/graph/search.cc index a85c378048..7b159950bd 100644 --- a/src/graph/search.cc +++ b/src/graph/search.cc @@ -1252,10 +1252,10 @@ ncclResult_t ncclTopoGetNetDev(struct ncclComm* comm, int rank, struct ncclTopoG int pxnLevel = ncclPxnDisable(comm) == 1 ? 0 : ncclParamP2pPxnLevel(); // See whether we can use the remote rank preferred device. if (ncclParamCrossNic() == 0 || (pxnLevel != 0)) { - // Find local NIC number close to local cudaDev - int cudaDev = comm->peerInfo[peerRank].cudaDev; + // Find local NIC number close to local nvmlDev + int nvmlDev = comm->peerInfo[peerRank].nvmlDev; int localRank; - if (ncclTopoDevToRank(comm->topo, cudaDev, &localRank) != ncclSuccess) return ncclSuccess; + if (ncclTopoDevToRank(comm->topo, nvmlDev, &localRank) != ncclSuccess) return ncclSuccess; NCCLCHECK(ncclTopoGetLocalNet(comm->topo, localRank, channelId, &netId, &netDev)); // Check that device exists on our node diff --git a/src/graph/xml.cc b/src/graph/xml.cc index 75888de0bd..de7fcb1253 100644 --- a/src/graph/xml.cc +++ b/src/graph/xml.cc @@ -722,14 +722,7 @@ ncclResult_t ncclTopoGetXmlFromGpu(struct ncclXmlNode* pciNode, uint32_t rocmDev int dev = -1; NCCLCHECK(xmlGetAttrIndex(gpuNode, "dev", &index)); if (index == -1) { - if (rocmDev == -1) { - const char* busId; - NCCLCHECK(xmlGetAttr(pciNode, "busid", &busId)); - if (busId == NULL || cudaDeviceGetByPCIBusId(&dev, busId) != cudaSuccess) dev = -1; - } else { - dev = rocmDev; - } - NCCLCHECK(xmlSetAttrInt(gpuNode, "dev", dev)); + NCCLCHECK(xmlSetAttrInt(gpuNode, "dev", rocmDev)); } NCCLCHECK(xmlGetAttrInt(gpuNode, "dev", &dev)); if (dev == -1) { *gpuNodeRet = NULL; return ncclSuccess; } @@ -934,11 +927,9 @@ ncclResult_t ncclTopoFillGpu(struct ncclXml* xml, const char* busId, struct nccl uint32_t devIndex = 0; static int rocmsmiInit = 0; if (rocmsmiInit == 0) { - rocmsmiInit = (rocm_smi_init() != ncclSuccess) ? 2 : 1; - } - if (rocmsmiInit == 1) { - if (rocm_smi_getDeviceIndexByPciBusId(busId, &devIndex) != ncclSuccess) devIndex = -1; + NCCLCHECK(rocm_smi_init()); } + NCCLCHECK(rocm_smi_getDeviceIndexByPciBusId(busId, &devIndex)); NCCLCHECK(ncclTopoGetXmlFromGpu(node, devIndex, xml, gpuNode)); #else nvmlDevice_t nvmlDev; diff --git a/src/include/comm.h b/src/include/comm.h index 6408fd4c3d..160439c048 100644 --- a/src/include/comm.h +++ b/src/include/comm.h @@ -61,8 +61,6 @@ struct ncclRecvMem { uint64_t tail; char pad1[CACHE_LINE_SIZE-sizeof(uint64_t)]; struct ncclConnFifo connFifo[NCCL_STEPS]; - int sizesFifo[NCCL_STEPS]; - int offsFifo[NCCL_STEPS]; int flush; // For GDRCopy-based flush }; char pad4[MEM_ALIGN]; @@ -257,7 +255,7 @@ struct ncclComm { int rank; // my rank in the communicator int nRanks; // number of GPUs in communicator int cudaDev; // my cuda device index - //int nvmlDev; // my nvml device index + int nvmlDev; // my nvml device index int compCap; // compute capability of the GPU int minCompCap, maxCompCap; // min/max compute capability in the communicator int64_t busId; // my PCI bus ID in int format diff --git a/src/include/transport.h b/src/include/transport.h index 24d22734bf..e7c2bd22e2 100644 --- a/src/include/transport.h +++ b/src/include/transport.h @@ -36,6 +36,7 @@ struct ncclComm; struct ncclPeerInfo { int rank; int cudaDev; + int nvmlDev; int gdrSupport; bool hasFineGrain; uint64_t hostHash; diff --git a/src/init.cc b/src/init.cc index 067ddf19c7..ed20796f51 100644 --- a/src/init.cc +++ b/src/init.cc @@ -48,6 +48,7 @@ #ifdef ENABLE_MSCCLPP #include "mscclpp/mscclpp_nccl.h" #endif +#include "rocm_smi_wrap.h" // [/RCCL] #include "msccl/msccl_lifecycle.h" @@ -548,6 +549,11 @@ static ncclResult_t commAlloc(struct ncclComm* comm, struct ncclComm* parent, in CUDACHECK(cudaGetDevice(&comm->cudaDev)); NCCLCHECK(getBusId(comm->cudaDev, &comm->busId)); + char busId[]="0000:00:00.0"; + NCCLCHECK(int64ToBusId(comm->busId, busId)); + NCCLCHECK(rocm_smi_init()); + NCCLCHECK(rocm_smi_getDeviceIndexByPciBusId(busId, (unsigned int*)&comm->nvmlDev)); + comm->compCap = ncclCudaCompCap(); TRACE(NCCL_INIT,"comm %p rank %d nranks %d cudaDev %d busId %lx compCap %d", comm, rank, ndev, comm->cudaDev, comm->busId, comm->compCap); @@ -750,7 +756,8 @@ static void showVersion() { static ncclResult_t fillInfo(struct ncclComm* comm, struct ncclPeerInfo* info, uint64_t commHash) { info->rank = comm->rank; - CUDACHECK(cudaGetDevice(&info->cudaDev)); + info->cudaDev = comm->cudaDev; + info->nvmlDev = comm->nvmlDev; info->hostHash=getHostHash()+commHash; info->pidHash=getPidHash()+commHash; @@ -1991,11 +1998,11 @@ static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) { comm->commHash = getHash(job->commId.internal, NCCL_UNIQUE_ID_BYTES); if (job->parent) { - INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d busId %lx parent %p color %d key %d commId 0x%llx - Init START", - comm, comm->rank, comm->nRanks, comm->cudaDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId)); + INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx parent %p color %d key %d commId 0x%llx - Init START", + comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId)); } else { - INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d busId %lx commId 0x%llx - Init START", - comm, comm->rank, comm->nRanks, comm->cudaDev, comm->busId, (unsigned long long)hashUniqueId(job->commId)); + INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx commId 0x%llx - Init START", + comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, (unsigned long long)hashUniqueId(job->commId)); } NCCLCHECKGOTO(initTransportsRank(comm, job->parent), res, fail); @@ -2073,11 +2080,11 @@ static ncclResult_t ncclCommInitRankFunc(struct ncclAsyncJob* job_) { } if (job->parent) { - INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d busId %lx parent %p color %d key %d commId 0x%llx localSize %zi used %ld bytes on core %d - Init COMPLETE", - comm, comm->rank, comm->nRanks, comm->cudaDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId), maxLocalSizeBytes, allocTracker[comm->cudaDev].totalAllocSize, sched_getcpu()); + INFO(NCCL_INIT,"ncclCommSplit comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx parent %p color %d key %d commId 0x%llx localSize %zi used %ld bytes on core %d - Init COMPLETE", + comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, job->parent, job->color, job->key, (unsigned long long)hashUniqueId(job->commId), maxLocalSizeBytes, allocTracker[comm->cudaDev].totalAllocSize, sched_getcpu()); } else { - INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d busId %lx commId 0x%llx localSize %zi used %ld bytes on core %d - Init COMPLETE", - comm, comm->rank, comm->nRanks, comm->cudaDev, comm->busId, (unsigned long long)hashUniqueId(job->commId), maxLocalSizeBytes, allocTracker[comm->cudaDev].totalAllocSize, sched_getcpu()); + INFO(NCCL_INIT,"ncclCommInitRank comm %p rank %d nranks %d cudaDev %d nvmlDev %d busId %lx commId 0x%llx localSize %zi used %ld bytes on core %d - Init COMPLETE", + comm, comm->rank, comm->nRanks, comm->cudaDev, comm->nvmlDev, comm->busId, (unsigned long long)hashUniqueId(job->commId), maxLocalSizeBytes, allocTracker[comm->cudaDev].totalAllocSize, sched_getcpu()); } exit: if (job->newcomm) { diff --git a/src/transport/net.cc b/src/transport/net.cc index dd5ffa5403..8abb47a379 100644 --- a/src/transport/net.cc +++ b/src/transport/net.cc @@ -378,10 +378,10 @@ static ncclResult_t sendConnect(struct ncclComm* comm, struct ncclConnect* conne send->conn.tail = &recvMem->tail; send->conn.stepSize = comm->buffSizes[NCCL_PROTO_SIMPLE]/NCCL_STEPS; send->conn.connFifo = recvMem->connFifo; - // Only fuse P2P buffers, continue to allocate dedicated buffers for ring/tree + struct ncclRecvMem *recvMemCpu = (struct ncclRecvMem*) NCCL_NET_MAP_GET_POINTER(map, cpu, recvMem); for (int i=0; iconn.connFifo[i].offset = -1; - recvMem->connFifo[i].mode = map->shared ? NCCL_MODE_OFFSET : NCCL_MODE_NORMAL; + recvMemCpu->connFifo[i].offset = -1; + recvMemCpu->connFifo[i].mode = map->shared ? NCCL_MODE_OFFSET : NCCL_MODE_NORMAL; } for (int p=0; ptransmitted < sub->posted && sub->transmitted < sub->done + NCCL_STEPS) { int buffSlot = (sub->base+sub->transmitted)%NCCL_STEPS; - volatile int* sizesFifo = resources->recvMem->sizesFifo; volatile uint64_t* recvTail = &resources->recvMem->tail; uint64_t tail = sub->base + (sub->reg ? 0 : sub->transmitted); if ((sub->reg || connFifo[buffSlot].size != -1) && ((*recvTail > tail) || p == NCCL_PROTO_LL)) { @@ -1243,10 +1242,7 @@ static ncclResult_t sendProxyProgress(struct ncclProxyState* proxyState, struct sub->timestamp[buffSlot] = 0; #endif - TRACE(NCCL_NET, "sendProxy [%ld/%d] Isend posted, req %p", sub->transmitted, buffSlot, sub->requests[buffSlot]); - sizesFifo[buffSlot] = -1; - // Make sure size is reset to zero before we update the head. - __sync_synchronize(); + TRACE(NCCL_NET, "sendProxy [%ld/%d] Isend posted, req %p, size %d, proto %d, myRank %d, channelId %d", sub->transmitted, buffSlot, sub->requests[buffSlot], size, p, proxyState->tpRank, sub->channelId); sub->transmitted += args->sliceSteps; for (uint64_t step=sub->transmitted-args->sliceSteps; steptransmitted; step++) ncclProfilingRecord(args, s, step, ncclProxyProfileSendWait); args->idle = 0;