AI/rocm-systems
2
0
Форк 0
Код Задачі Запити на злиття Дії Пакети Проєкти Релізи Вікі Активність

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

Переглянути джерело
Цей коміт міститься в:
BertanDogancay
2024-04-23 13:33:19 -07:00
джерело 220066197a 6dd51f15bf
коміт e1a835910e
76 змінених файлів з 5157 додано та 2296 видалено
Завантажити патч Завантажити файл різниці Розгорнути всі файли Згорнути всі файли
src/graph/connect.cc
+90 -37
Переглянути файл
@@ -22,6 +22,7 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm, struct ncclTopoGraph** graphs
int localRanks = comm->topo->nodes[GPU].count;
int nChannels = comm->nChannels;
topoRanks->nvlsHeadNum = 0;
for (int c=0; c<nChannels; c++) {
struct ncclChannel* channel = comm->channels+c;
channel->ring.prev = channel->ring.next = -1;
@@ -33,20 +34,20 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm, struct ncclTopoGraph** graphs
channel->collnetDirect.headRank = -1;
channel->collnetDirect.nHeads = 0;
channel->collnetDirect.shift = 0;
for (int i=0; i<NCCL_MAX_DIRECT_ARITY+1; i++) channel->collnetDirect.heads[i] = -1;
for (int i=0; i<NCCL_MAX_DIRECT_ARITY; i++) channel->collnetDirect.up[i] = -1;
for (int i=0; i<NCCL_MAX_DIRECT_ARITY; i++) channel->collnetDirect.down[i] = -1;
int* ringIntra = graphs[NCCL_ALGO_RING]->intra+c*localRanks;
int* treeIntra = graphs[NCCL_ALGO_TREE]->intra+c*localRanks;
int* collNetIntra = graphs[NCCL_ALGO_COLLNET_CHAIN]->intra+c*localRanks;
int* nvlsIntra = graphs[NCCL_ALGO_NVLS]->intra+c*localRanks;
for (int i=0; i<localRanks; i++) {
if (ringIntra[i] == rank) {
topoRanks->ringRecv[c] = ringIntra[0];
topoRanks->ringSend[c] = ringIntra[localRanks-1];
channel->ring.prev = (i == 0) ? -1 : ringIntra[i-1];
channel->ring.next = (i == localRanks-1) ? -1 : ringIntra[i+1];
topoRanks->ringPrev[c] = (i == 0) ? -1 : ringIntra[i-1];
topoRanks->ringNext[c] = (i == localRanks-1) ? -1 : ringIntra[i+1];
}
if (treeIntra[i] == rank) {
int parentIndex = 0;
@@ -64,14 +65,28 @@ ncclResult_t ncclTopoPreset(struct ncclComm* comm, struct ncclTopoGraph** graphs
channel->collnetChain.down[0] = i == localRanks-1 ? -1 : collNetIntra[i+1];
}
}
topoRanks->ringPrev[c] = channel->ring.prev;
topoRanks->ringNext[c] = channel->ring.next;
topoRanks->nvlsHeads[c] = nvlsIntra[0];
}
// Duplicate channels rings/trees
// Duplicate channels trees
struct ncclChannel* channel0 = comm->channels;
struct ncclChannel* channel1 = (nChannels > MAXCHANNELS/2) ? 0 : channel0+nChannels;
if (channel1) memcpy(channel1, channel0, nChannels*sizeof(struct ncclChannel));
// Get nvls heads and the number of heads. Duplicate head is not allowed.
for (int c = 0; c < graphs[NCCL_ALGO_NVLS]->nChannels; ++c) {
bool addHead = true;
int* nvlsIntra = graphs[NCCL_ALGO_NVLS]->intra + c * localRanks;
for (int dup = 0; dup < topoRanks->nvlsHeadNum; dup++) {
if (topoRanks->nvlsHeads[dup] == nvlsIntra[0]) {
addHead = false;
break;
}
}
if (addHead) {
topoRanks->nvlsHeads[topoRanks->nvlsHeadNum++] = nvlsIntra[0];
}
}
return ncclSuccess;
}
@@ -203,26 +218,14 @@ static ncclResult_t connectRings(struct ncclComm* comm, int* ringRecv, int* ring
int* send = ringSend+c*comm->nNodes;
int* prev = ringPrev+c*comm->nRanks;
int* next = ringNext+c*comm->nRanks;
struct ncclChannel* channel0 = comm->channels+c;
struct ncclChannel* channel1 = (nChannels > MAXCHANNELS/2) ? 0 : channel0+nChannels;
for (int n=0; n<nNodes; n++) {
int recvRank = recv[n];
int prevSendRank = send[(n-1+nNodes)%nNodes];
prev[recvRank] = prevSendRank;
if (comm->rank == recvRank) {
channel0->ring.prev = prevSendRank;
if (channel1) channel1->ring.prev = prevSendRank;
}
int sendRank = send[n];
int nextRecvRank = recv[(n+1)%nNodes];
next[sendRank] = nextRecvRank;
if (comm->rank == sendRank) {
channel0->ring.next = nextRecvRank;
if (channel1) channel1->ring.next = nextRecvRank;
}
}
TRACE(NCCL_GRAPH, "Ring %d : %d -> %d -> %d", c, channel0->ring.prev, comm->rank, channel0->ring.next);
if (channel1) TRACE(NCCL_GRAPH, "Ring %d : %d -> %d -> %d", c+nChannels, channel1->ring.prev, comm->rank, channel1->ring.next);
}
// [RCCL] Print off the recv/send local ranks per node, per channel
@@ -404,6 +407,15 @@ static ncclResult_t connectCollNet(struct ncclComm* comm, struct ncclTopoGraph*
channel->collnetDirect.up[nUp++] = heads[h];
sprintf(line+strlen(line), " %d ", heads[h]);
}
sprintf(line+strlen(line), "heads ");
{ // heads[] is the list of heads ordered in head order startubg with self
int h0 = (channel->collnetDirect.headRank == -1) ? 0 : channel->collnetDirect.headRank;
for (int h1=0; h1 < nHeads; h1++) {
int h = (h0+h1)%nHeads;
channel->collnetDirect.heads[h1] = heads[h];
sprintf(line+strlen(line), " %d ", heads[h]);
}
}
channel->collnetDirect.nHeads = nHeads;
channel->collnetDirect.shift = (rank%localRanks)%nHeads; // Shift by intraRank so that leaves don't send to same head simultaneously
channel->collnetDirect.depth = (nUp == 0 && nDown == 0) ? 1 : 2;
@@ -412,27 +424,22 @@ static ncclResult_t connectCollNet(struct ncclComm* comm, struct ncclTopoGraph*
INFO(NCCL_GRAPH, "%s", line);
channel->collnetChain.depth = comm->nRanks/comm->nNodes;
}
for (int c=0; c<comm->nvlsChannels; c++) {
struct ncclChannel* channel = comm->channels+c;
if (channel->nvls.headRank != -1) channel->nvls.out = comm->nRanks;
}
free(heads);
return ncclSuccess;
}
static ncclResult_t connectNvls(struct ncclComm* comm, int* nvlsHeads, struct ncclTopoGraph* nvlsGraph) {
int nHeads = nvlsGraph->nChannels;
static ncclResult_t connectNvls(struct ncclComm* comm, int* nvlsHeads, int nHeads) {
int headRank = -1;
for (int h=0; h<nHeads; h++) {
if (nvlsGraph->intra[h*comm->localRanks] == comm->rank) headRank = h;
}
if (nHeads == 0) {
comm->nvlsChannels = 0;
return ncclSuccess;
}
for (int c=0; c<comm->nvlsChannels; c++) {
for (int h = 0; h < nHeads; h++) {
if (nvlsHeads[h * comm->nNodes + comm->node] == comm->rank) headRank = h;
}
for (int c=0; c<comm->nChannels; c++) {
struct ncclChannel* channel = comm->channels+c;
channel->nvls.nHeads = nHeads;
for (int h=0; h<nHeads; h++) channel->nvls.up[h] = comm->nRanks+1+h;
@@ -443,8 +450,10 @@ static ncclResult_t connectNvls(struct ncclComm* comm, int* nvlsHeads, struct nc
channel->nvls.treeUp = channel->nvls.treeDown[0] = channel->nvls.treeDown[1] = channel->nvls.treeDown[2] = -1;
channel->nvls.node = comm->node;
channel->nvls.nNodes = comm->nNodes;
if (comm->collNetSupport && channel->nvls.headRank != -1) channel->nvls.out = comm->nRanks;
}
if (comm->nNodes == 1) return ncclSuccess;
// MNNVL: NVLS not yet supported
if (comm->nNodes == 1 || comm->MNNVL) return ncclSuccess;
// Connect Trees
int tree0Parent, tree0Child0, tree0Child1, tree1Parent, tree1Child0, tree1Child1;
@@ -485,7 +494,7 @@ static ncclResult_t connectNvls(struct ncclComm* comm, int* nvlsHeads, struct nc
}
// Set prev/next in all channels (NVLS compute channels work
// orthogonally to NVLS search channels).
for (int c=0; c<comm->nvlsChannels; c++) {
for (int c=0; c<comm->nChannels; c++) {
struct ncclChannel* channel = comm->channels+c;
channel->nvls.treeUp = treeUp[c%2];
channel->nvls.treeDown[0] = channel->nvls.down;
@@ -543,12 +552,19 @@ static int copyChannels(struct ncclComm* comm, int start, int end, int* ringPrev
return c;
}
void exchangeValues(int* v0, int* v1) {
int tmp = *v1;
*v1 = *v0;
*v0 = tmp;
}
ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePatterns, struct ncclTopoRanks** allTopoRanks, int* rings, struct ncclTopoGraph** graphs, int nc) {
// Gather data from all ranks
int *ringRecv, *ringSend, *ringPrev, *ringNext, *treeToParent, *treeToChild0, *treeToChild1, *nvlsHeads;
int nranks = comm->nRanks;
int nNodes = comm->nNodes;
int nChannels = comm->nChannels;
int minHeadNum = INT_MAX;
NCCLCHECK(ncclCalloc(&ringRecv, nNodes*MAXCHANNELS));
NCCLCHECK(ncclCalloc(&ringSend, nNodes*MAXCHANNELS));
NCCLCHECK(ncclCalloc(&ringPrev, nranks*MAXCHANNELS));
@@ -557,6 +573,22 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
NCCLCHECK(ncclCalloc(&treeToChild0, nNodes*MAXCHANNELS));
NCCLCHECK(ncclCalloc(&treeToChild1, nNodes*MAXCHANNELS));
NCCLCHECK(ncclCalloc(&nvlsHeads, nNodes*MAXCHANNELS));
// Alternate rings to avoid crossing rails
if (graphs[NCCL_ALGO_RING]->crossNic && (comm->nNodes % 2) == 0 && (nChannels % 2) == 0) {
for (int r=0; r<comm->nRanks; r++) {
if (comm->rankToNode[r] % 2 == 1) {
// Exchange rings
for (int c=0; c<nChannels; c+=2) {
exchangeValues(allTopoRanks[r]->ringRecv+c, allTopoRanks[r]->ringRecv+(c^1));
exchangeValues(allTopoRanks[r]->ringSend+c, allTopoRanks[r]->ringSend+(c^1));
exchangeValues(allTopoRanks[r]->ringPrev+c, allTopoRanks[r]->ringPrev+(c^1));
exchangeValues(allTopoRanks[r]->ringNext+c, allTopoRanks[r]->ringNext+(c^1));
}
}
}
}
for (int c=0; c<nChannels;c++) {
for (int n=0; n<nNodes; n++) {
int r = firstRanks[n];
@@ -571,17 +603,23 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
ringNext[c*nranks+r] = allTopoRanks[r]->ringNext[c];
}
}
for (int c=0; c<graphs[NCCL_ALGO_NVLS]->nChannels; c++) {
for (int n=0; n<nNodes; n++) {
for (int n = 0; n < nNodes; n++) {
int r = firstRanks[n];
if (minHeadNum > allTopoRanks[r]->nvlsHeadNum)
minHeadNum = allTopoRanks[r]->nvlsHeadNum;
}
for (int c = 0; c < minHeadNum; c++) {
for (int n = 0; n < nNodes; n++) {
int r = firstRanks[n];
nvlsHeads[c*nNodes+n] = allTopoRanks[r]->nvlsHeads[c];
nvlsHeads[c * nNodes + n] = allTopoRanks[r]->nvlsHeads[c];
}
}
// Connect rings and trees. This should also duplicate the channels.
NCCLCHECK(connectRings(comm, ringRecv, ringSend, ringPrev, ringNext));
NCCLCHECK(connectTrees(comm, treeToParent, treeToChild0, treeToChild1, treePatterns));
NCCLCHECK(connectNvls(comm, nvlsHeads, graphs[NCCL_ALGO_NVLS]));
// Only use full MAXCHANNELS for gfx94x
int maxChannels = IsArchMatch(comm->topo->nodes[GPU].nodes[0].gpu.gcn, "gfx94") ? MAXCHANNELS : (MAXCHANNELS/2);
@@ -595,6 +633,14 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
nc = std::min(maxNchannels/comm->nChannels, nc);
nc *= comm->nChannels;
// Set ring prev/next for my rank
for (int c=0; c<nChannels; c++) {
struct ncclChannel* channel0 = comm->channels+c;
struct ncclChannel* channel1 = channel0+nChannels;
channel0->ring.prev = channel1->ring.prev = ringPrev[c*nranks+comm->rank];
channel0->ring.next = channel1->ring.next = ringNext[c*nranks+comm->rank];
}
// Duplication should be complete now
nChannels = comm->nChannels = std::min(maxChannels, (nChannels <= maxChannels/2) ? nChannels*2 : nChannels);
@@ -633,6 +679,13 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
nChannels = comm->nChannels = copyChannels(comm, nChannels, std::max(minNchannels, std::max(nc, comm->config.minCTAs)), ringPrev, ringNext);
}
comm->collChannels = comm->nChannels;
// Support maximal channel usage for aggregation
if (comm->nChannels < comm->nvlsChannels) {
nChannels = comm->nChannels = copyChannels(comm, comm->nChannels, comm->nvlsChannels, ringPrev, ringNext);
}
NCCLCHECK(connectNvls(comm, nvlsHeads, minHeadNum));
// Create rings array and check all is fine
NCCLCHECK(ncclBuildRings(nChannels, rings, comm->rank, comm->nRanks, ringPrev, ringNext));
@@ -646,4 +699,4 @@ ncclResult_t ncclTopoPostset(struct ncclComm* comm, int* firstRanks, int* treePa
free(nvlsHeads);
return ncclSuccess;
}
}