AI/rocm-systems
2
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Pivot AllToAll algorithm for Rome model (#503)

Browse Source 
* add a2a pivot interface

* remove debug info

* address comments

* fix bug

* remove custom script

* address comments

* fix bug

[ROCm/rccl commit: b569c0a1db]
This commit is contained in:
Ziyue Yang authored and GitHub committed 2022-02-20 21:09:47 -08:00
1 parent 0f0388ba0b
commit dfa9b9e958
18 files changed
+187 -27
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

projects/rccl/src/enqueue.cc
+24 -5
View File
@@ -361,7 +361,7 @@ static inline ncclResult_t getCollNetSupport(struct ncclInfo* info, int* collNet
static ncclResult_t getAlgoInfo(struct ncclInfo* info, int collNetTypeSupport, int numPipeOps) {
struct ncclComm* comm = info->comm;
if (comm->nRanks == 1) {
if (comm->nRanks == 1 || info->coll == ncclFuncAllToAllPivot) {
info->algorithm = NCCL_ALGO_RING;
info->protocol = NCCL_PROTO_SIMPLE;
}
@@ -423,7 +423,12 @@ static ncclResult_t getAlgoInfo(struct ncclInfo* info, int collNetTypeSupport, i
if (info->algorithm == NCCL_ALGO_COLLNET) nt += 3*WARP_SIZE;
}
#endif
info->nChannels = nc;
if (info->coll == ncclFuncAllToAllPivot) {
int pivotA2ANumUniRings = comm->topo->pivotA2ANumBiRings * 2;
info->nChannels = comm->nChannels / pivotA2ANumUniRings * pivotA2ANumUniRings;
} else {
info->nChannels = nc;
}
info->nThreads = nt;
return ncclSuccess;
}
@@ -436,6 +441,7 @@ static ncclResult_t getPatternInfo(struct ncclInfo* info) {
info->pattern = info->algorithm == NCCL_ALGO_TREE ? ncclPatternTreeUp : ncclPatternPipelineTo; break;
case ncclFuncReduceScatter:
case ncclFuncAllGather:
case ncclFuncAllToAllPivot:
info->pattern = ncclPatternRing; break;
case ncclFuncAllReduce:
info->pattern = info->algorithm == NCCL_ALGO_COLLNET ? ncclPatternCollTreeUpDown : info->algorithm == NCCL_ALGO_TREE ? ncclPatternTreeUpDown : ncclPatternRingTwice; break;
@@ -497,10 +503,12 @@ comp_next:
// one-rank reduce index
work->funcIndex = FUNC_INDEX_P2P - ncclNumTypes + int(info->datatype);
return ncclSuccess;
} else if (info->coll == ncclFuncAllToAllPivot) {
work->funcIndex = FUNC_INDEX_ALLTOALL_PIVOT;
} else {
work->funcIndex = FUNC_INDEX(info->coll, info->opFull.op, info->datatype, info->algorithm, info->protocol);
}
work->funcIndex = FUNC_INDEX(info->coll, info->opFull.op, info->datatype, info->algorithm, info->protocol);
work->coll.connIndex = 0;
proxyArgs->connIndex = 0;
if (info->protocol == NCCL_PROTO_SIMPLE && info->algorithm == NCCL_ALGO_RING) {
@@ -599,6 +607,12 @@ comp_next:
TRACE(NCCL_COLL,"opCount %lx slicesteps %d spl %d cpl %d nbytes %zi -> protocol %d nchannels %d nthreads %d, nloops %d nsteps %d chunksize %d comm %p",
proxyArgs->opCount, sliceSteps, info->nstepsPerLoop, info->nchunksPerLoop, info->nBytes, info->protocol, info->nChannels, info->nThreads,
nLoops, proxyArgs->subs[0].nsteps, chunkSize, info->comm);
// For Pivot A2A, lastChunkSize is not needed, set pivotA2ANumBiRings instead
if (info->coll == ncclFuncAllToAllPivot) {
work->coll.pivotA2ANumBiRings = info->comm->topo->pivotA2ANumBiRings;
}
return ncclSuccess;
}
@@ -760,7 +774,12 @@ ncclResult_t ncclSetupAsyncKernels(ncclComm_t comm) {
int allCollNetSupport = comm->collNetSupport;
for (int c = 0; c < comm->asyncOpCount; c++) {
struct ncclInfo* info = comm->asyncOps+c;
info->nChannels = std::min(std::max(1, (int)DIVUP(info->nBytes, channelSize)), comm->nChannels); // assign number of channels
if (info->coll == ncclFuncAllToAllPivot) {
int pivotA2ANumUniRings = comm->topo->pivotA2ANumBiRings * 2;
info->nChannels = comm->nChannels / pivotA2ANumUniRings * pivotA2ANumUniRings;
} else {
info->nChannels = std::min(std::max(1, (int)DIVUP(info->nBytes, channelSize)), comm->nChannels); // assign number of channels
}
channelUsed += info->nChannels;
// We can use fast path if all collectives are the same
homogeneous &= info->coll == comm->asyncOps[0].coll &&