Fix test_barrier hang by using static global rank in ProcessGroupXCCL (#2036)

Fixes #1978 

In ProcessGroupNCCL, `globalRank()` returns a static int globalRank,
which is first initialized by the ProcessGroup setup, so the globalRank
assigned to each thread matches the id of the CUDA device. However, we
were not using this same pattern for XCCL. Instead, we were just using
the assigned rank of the thread, which were not necessarily the same as
the globalRank.

The failing test `test_barrier` created two separate groups of 2 ranks
each, and then 4 threads called barrier, but all on the same 2-thread
group. Since initially the device id is not specified in this barrier
call, the thread attempts to "guess" the device index. In the previous
code, this guess would be 0 or 1, since the rank of each thread was not
actually the globalRank. In `barrier`, this guessed id was used to
initialize XCCLComm objects, and then call allreduce on a single element
tensor. However, this resulted in an allreduce call two times on each
device, which could result in a hang based on the execution order of the
4 threads.

With the update in this PR, PGXCCL now uses the static globalRank in the
same places as ProcessGroupNCCL, so the initialized XCCLComm objects are
for unique devices and allreduce doesn't call on the same device
multiple times.

Author: frost-intel

src/xccl/ProcessGroupXCCL.cpp

@@ -352,6 +352,11 @@ const std::string& ProcessGroupXCCL::logPrefix() const {
   return logPrefix_;
 }
 
+const int& ProcessGroupXCCL::globalRank() const {
+  static int globalRank = rank_;
+  return globalRank;
+}
+
 ProcessGroupXCCL::ProcessGroupXCCL(
     c10::intrusive_ptr<Store> store,
     int rank,
@@ -379,7 +384,7 @@ ProcessGroupXCCL::ProcessGroupXCCL(
   std::string torch_distributed_debug =
       getCvarString({"TORCH_DISTRIBUTED_DEBUG"}, OFF.c_str());
   LOG(INFO) << logPrefix() << "ProcessGroupXCCL initialization options: "
-            << "size: " << size << ", global rank: " << rank_
+            << "size: " << size << ", global rank: " << globalRank()
             << ", USE_HIGH_PRIORITY_STREAM: "
             << options_->is_high_priority_stream
             << ", PG Name: " << options_->group_name;
@@ -410,7 +415,7 @@ bool ProcessGroupXCCL::dumpDebuggingInfo(bool includeStackTrace /*=true*/) {
   if (traceBufferSize_ > 0) {
     // TODO: dump_xccl_trace
     auto xcclTrace = dump_xccl_trace(true, includeStackTrace, false);
-    DebugInfoWriter& writer = DebugInfoWriter::getWriter(rank_);
+    DebugInfoWriter& writer = DebugInfoWriter::getWriter(globalRank());
     LOG(INFO) << logPrefix() << "ProcessGroupXCCL dumping xccl trace to "
               << writer.getWriterTarget();
     writer.write(xcclTrace);
@@ -2021,7 +2026,7 @@ c10::DeviceIndex ProcessGroupXCCL::guessDeviceId() const {
     return *usedDeviceIdxs_.begin();
   }
   int devIdx =
-      static_cast<int16_t>(rank_ % at::detail::getXPUHooks().getNumGPUs());
+      static_cast<int16_t>(globalRank() % at::detail::getXPUHooks().getNumGPUs());
   LOG(WARNING)
       << logPrefix()
       << c10::str(

src/xccl/ProcessGroupXCCL.hpp

@@ -423,6 +423,7 @@ class TORCH_API ProcessGroupXCCL : public Backend {
   c10::DeviceIndex guessDeviceId() const;
 
   const std::vector<uint64_t>& groupRanks() const;
+  const int& globalRank() const;
   void setEnqueuedPgStatus(c10::intrusive_ptr<ProcessGroupXCCL::WorkXCCL> work);
   bool dumpDebuggingInfo(bool includeStackTrace = true);
 

src/xccl/ProcessGroupXCCLMonitor.cpp

@@ -39,7 +39,7 @@ void HeartbeatMonitorXCCL::runLoop() {
   // We only need to dump once per PG, so we use local_id_ == 0 for the first PG
   if (pg_->local_id_ == 0) {
     // DumpPipe is one per-trainer process
-    dumpPipe.emplace(pg_->getRank());
+    dumpPipe.emplace(pg_->globalRank());
     while (true) {
       std::unique_lock<std::mutex> lock(monitorMutex_);
       if (monitorWakeUpCV_.wait_for(