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

Added new unit tests for AllReduce with Bias API (#2036)

Browse Source 
* Added new unit tests for AllReduce with Bias API

* Address review comments
This commit is contained in:
Atul Kulkarni
2025-12-03 17:37:34 -06:00
committed by GitHub
parent 185e78a8f0
commit 7c12b0b76b
9 changed files with 586 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/common/TestBedChild.cpp
+31 -14
View File
@@ -598,15 +598,32 @@ namespace RcclUnitTesting
"ncclReduceScatter");
break;
case ncclCollAllReduce:
CHILD_NCCL_CALL_RANK(errCode, ncclAllReduce(
collArg.inputGpu.ptr,
collArg.outputGpu.ptr,
collArg.numInputElements,
collArg.dataType,
collArg.options.redOp,
this->comms[localRank],
this->streams[groupId][localRank][collArg.streamIdx]),
"ncclAllReduce");
// Use ncclAllReduceWithBias if bias is enabled
if (collArg.options.useBias)
{
CHILD_NCCL_CALL_RANK(errCode, ncclAllReduceWithBias(
collArg.inputGpu.ptr,
collArg.outputGpu.ptr,
collArg.numInputElements,
collArg.dataType,
collArg.options.redOp,
this->comms[localRank],
this->streams[groupId][localRank][collArg.streamIdx],
collArg.options.biasPtr),
"ncclAllReduceWithBias");
}
else
{
CHILD_NCCL_CALL_RANK(errCode, ncclAllReduce(
collArg.inputGpu.ptr,
collArg.outputGpu.ptr,
collArg.numInputElements,
collArg.dataType,
collArg.options.redOp,
this->comms[localRank],
this->streams[groupId][localRank][collArg.streamIdx]),
"ncclAllReduce");
}
break;
case ncclCollGather:
CHILD_NCCL_CALL_RANK(errCode, ncclGather(
@@ -762,7 +779,7 @@ namespace RcclUnitTesting
{
streamsToComplete.erase(streamsToComplete.begin() + i);
i--;
}
}
}
usElapsed = duration_cast<microseconds>(Clock::now() - start).count();
}
@@ -773,13 +790,13 @@ namespace RcclUnitTesting
if (this->verbose) INFO("Collective timed out, aborting\n");
for (int localRank : localRanksToExecute)
{
ncclCommAbort(this->comms[localRank]);
ncclCommAbort(this->comms[localRank]);
timedout = 1;
}
}
// extra sync to flush GPU cache for validation later
// TODO: remove this after figuring out & fixing the exact behavior
// TODO: remove this after figuring out & fixing the exact behavior
// of fencing between kernels and at hipStreamQuery
for (int localRank : localRanksToExecute)
{
@@ -970,7 +987,7 @@ namespace RcclUnitTesting
PIPE_READ(groupId);
// Release graphs
for (int localRank = 0; localRank < this->deviceIds.size(); ++localRank)
for (int localRank = 0; localRank < this->deviceIds.size(); ++localRank)
{
CHECK_HIP(hipSetDevice(this->deviceIds[localRank]));
for (int streamIdx = 0; streamIdx < this->numStreamsPerGroup[groupId]; ++streamIdx)
@@ -989,7 +1006,7 @@ namespace RcclUnitTesting
{
for (int i = 0; i < this->numStreamsPerGroup[groupId]; ++i)
CHECK_HIP(hipStreamSynchronize(this->streams[groupId][localRank][i]));
}
}
this->graphs[groupId].clear();
this->graphExecs[groupId].clear();