[PyTorch][NCCL PG][Resubmit D67193887] Change getNCCLCommDumpMap to use new ncclCommDumpAll API

jiayulu · facebook-github-bot · commit b3c444f33194 · 2025-05-15T10:59:12.000-07:00
Summary: see D67193887 Test Plan: https://www.internalfb.com/mlhub/pipelines/runs/mast/aps-dlrm_fsdp2_test-7711e496fc?job_attempt=0&version=0&tab=execution_details&env=PRODUCTION Differential Revision: D74820576
diff --git a/torch/csrc/distributed/c10d/NCCLUtils.hpp b/torch/csrc/distributed/c10d/NCCLUtils.hpp
@@ -280,6 +280,13 @@ class NCCLComm {
 
   ncclUniqueId getNcclId();
   at::DeviceIndex getDeviceIndex();
+#if defined(IS_NCCLX) && defined(NCCL_COMM_GET_UNIQUE_HASH)
+  uint64_t getNcclUniqueHash() {
+    uint64_t ncclUniqueHash = 0;
+    C10D_NCCL_CHECK(ncclCommGetUniqueHash(ncclComm_, &ncclUniqueHash), std::nullopt);
+    return ncclUniqueHash;
+  }
+#endif
 
   // Must not be copyable
   NCCLComm(const NCCLComm&) = delete;
diff --git a/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp
@@ -364,11 +364,20 @@ static void attachAllocatorHooks() {
 static std::
     unordered_map<std::string, std::unordered_map<std::string, std::string>>
     getNCCLCommDumpMap() {
-#if (defined(IS_NCCLX) || defined(USE_ROCM)) && defined(NCCL_COMM_DUMP)
+#if (defined(IS_NCCLX) || defined(USE_ROCM)) && defined(NCCL_COMM_DUMP) && \
+    defined(NCCL_COMM_GET_UNIQUE_HASH)
   std::unordered_map<
-      std::string /* ncclUniqueID */,
+      std::string /* CommHash */,
       std::unordered_map<std::string, std::string> /* dump from this comm */>
       ncclDumpMap;
+#ifdef NCCL_COMM_DUMP_ALL
+  auto res = ncclCommDumpAll(ncclDumpMap);
+  if (res == ncclSuccess) {
+    return ncclDumpMap;
+  }
+  // Fall back to dump from each comm if ncclCommDumpAll failed
+#endif // NCCL_COMM_DUMP_ALL
+
   // dump_nccl_trace is only called from the default PG (local_id_=0), but we
   // want to dump from all comms so we need to iterate over ncclCommMemPoolMap,
   // which is static
@@ -382,8 +391,11 @@ static std::
     }
   }
   for (auto& ncclComm : allNCCLComms) {
-    std::string ncclUniqueIDStr = buildNcclUniqueIdStr(ncclComm->getNcclId());
-    ncclDumpMap[ncclUniqueIDStr] = ncclComm->ncclCommDump();
+    std::stringstream ss;
+    ss << std::hex << ncclComm->getNcclUniqueHash();
+    std::string ncclUniqueHashStr = ss.str();
+
+    ncclDumpMap[ncclUniqueHashStr] = ncclComm->ncclCommDump();
   }
   return ncclDumpMap;
 #else
