MSCCL Multithreaded regression root cause fix (#1347)

* Make sure the target device is used for MSCCL

* Enable single process mode by default to use MSCCL in MT

* Create a per-rank state when GPUs share a thread

[ROCm/rccl commit: 03a3ef3c34]

projects/rccl/src/misc/msccl/msccl_lifecycle.cc

@@ -28,7 +28,7 @@
 
 RCCL_PARAM(MscclEnabled, "MSCCL_ENABLE", 1);
 RCCL_PARAM(MscclForceEnabled, "MSCCL_FORCE_ENABLE", 0);
-RCCL_PARAM(MscclEnableSingleProcess, "MSCCL_ENABLE_SINGLE_PROCESS", 0);
+RCCL_PARAM(MscclEnableSingleProcess, "MSCCL_ENABLE_SINGLE_PROCESS", 1);
 static const char* mscclAlgoFilePathEnv = "MSCCL_ALGO_FILE_PATH";
 
 bool mscclEnabled() {

projects/rccl/src/misc/msccl/msccl_status.cc

@@ -27,9 +27,9 @@ static mutex rankStatesMutex;
 static unordered_map<int, shared_ptr<mscclRankState>> rankStates;
 
 static inline mscclRankState& mscclGetRankState(int rank) {
-  static thread_local shared_ptr<mscclRankState> threadRankState = make_shared<mscclRankState>();
-
+  // In the unlikely case of negative rank, return a per-thread state
   if (rank < 0) {
+    static thread_local shared_ptr<mscclRankState> threadRankState(new mscclRankState());
     return *threadRankState;
   }
 
@@ -37,8 +37,10 @@ static inline mscclRankState& mscclGetRankState(int rank) {
 
   auto rankStateIt = rankStates.find(rank);
   if (rankStateIt == rankStates.end()) {
-    rankStateIt = rankStates.insert(make_pair(rank, make_shared<mscclRankState>(*threadRankState))).first;
-    rankStateIt->second->rank = rank;
+    // Create a per rank threadRankState rather than per thread
+    shared_ptr<mscclRankState> newthreadRankState(new mscclRankState());
+    newthreadRankState->rank = rank;
+    rankStateIt = rankStates.insert(make_pair(rank, newthreadRankState)).first;
   }
   return *(rankStateIt->second);
 }

projects/rccl/src/msccl.cc

@@ -61,6 +61,14 @@ ncclResult_t mscclRunAlgo_impl(
   struct mscclAlgo* hostAlgo = status.hostAlgos[mscclAlgoHandle];
   struct mscclAlgo* devAlgo = status.devAlgos[mscclAlgoHandle];
 
+  // NCCL adds a lot of guarantees that target device is getting used
+  // in its group management code, which we entirely skip when MSCCL is used
+  // Therefore, in single thread multiGPU mode
+  // setting the device is critical to be sure 
+  // communication is done on the intended device
+
+  CUDACHECK(hipSetDevice(comm->cudaDev)); 
+
   NCCLCHECK(mscclGetCaptureStatus(comm->rank, stream));
 
   NCCLCHECK(mscclSetupCount(hostAlgo, comm, count, dataType));