Merge pull request #1158 from AtlantaPepsi/NPKit_fix

Prevent segfault from npkit-enabled rccl build [ROCm/rccl commit: cc39e91c6f]
2024-04-26 12:44:04 -04:00
parent 3c94f98688 8cf28704ce
commit afeaa17475
1 changed files with 7 additions and 0 deletions
@@ -1208,6 +1208,13 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, struct ncclComm* p
      ret = ncclInternalError;
      goto fail;
    }
+    #if defined(ENABLE_NPKIT)
+    if (intraProcRanks != 1) {
+      WARN("NPKit currently does not support more than 1 device per process");
+      ret = ncclInternalError;
+      goto fail;
+    }
+    #endif
    struct ncclComm* comm0 = comm->peerInfo[intraProcRank0].comm;
    assert(intraProcRank==0 ? comm==comm0 : true);
    comm->intraComm0 = comm0;