Fix bug where the first collective call was using MSCCL instead of MSCCL++ (#1260)

This commit is contained in:
corey-derochie-amd
2024-07-22 15:46:47 -06:00
committed by GitHub
parent cf311b71ee
commit 69135976d6
3 changed files with 19 additions and 0 deletions
+10
View File
@@ -470,6 +470,11 @@ ncclResult_t mscclEnqueueCheck(
case mscclNoGroup:
#ifdef ENABLE_MSCCLPP
if (comm->mscclppCompatible) {
if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
INFO(NCCL_COLL, "MSCCL++: reading capture status");
NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
}
/* check if one rank per GPU and graph mode is enabled */
if ((threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold && (nBytes & 31) == 0) {
@@ -501,6 +506,11 @@ ncclResult_t mscclEnqueueCheck(
case mscclGroupSupportedOp:
#ifdef ENABLE_MSCCLPP
if (comm->mscclppCompatible) {
if (threadLocalStatus.captureStatus == mscclUnknownCaptureStatus) {
INFO(NCCL_COLL, "MSCCL++: reading capture status");
NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
}
/* check if one rank per GPU and graph mode is enabled */
if ((threadLocalStatus.captureStatus != mscclNoCapture) && comm->mscclCompatible) {
if (func == mscclFuncAllReduce && nBytes <= comm->mscclpp_threshold && (nBytes & 31) == 0) {