update kvs key

Chao1Han · Nov 12, 2024 · b442419 · b442419
1 parent 20b60b1
commit b442419
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/torch/csrc/distributed/c10d/ProcessGroupXCCL.cpp b/torch/csrc/distributed/c10d/ProcessGroupXCCL.cpp
@@ -239,7 +239,7 @@ ProcessGroupXCCL::ProcessGroupXCCL(
     const c10::intrusive_ptr<Store>& store,
     int rank,
     int size)
-    : Backend(rank, size), store_(store) {
+    : Backend(rank, size), store_(store), xcclCommCounter_(0) {
   blockingWait_ = getCvarBool(TORCH_XCCL_BLOCKING_WAIT, false);
   init();
 }

diff --git a/torch/csrc/distributed/c10d/ProcessGroupXCCL.hpp b/torch/csrc/distributed/c10d/ProcessGroupXCCL.hpp
@@ -318,6 +318,7 @@ class TORCH_API ProcessGroupXCCL : public Backend {
   std::unordered_map<std::string, at::xpu::XPUEvent> xcclEventsMap_;
   std::unordered_map<std::string, std::shared_ptr<xcclComm_t>> devXCCLCommMap_;
   c10::intrusive_ptr<Store> store_;
+  uint64_t xcclCommCounter_{0};
   std::mutex mutex_;
   std::set<int> usedDeviceIdxs_;
   int coalescing_state_ = 0;
@@ -331,15 +332,19 @@ class TORCH_API ProcessGroupXCCL : public Backend {
  private:
   std::mutex kvs_mutex;
 
-  ccl::shared_ptr_class<ccl::kvs> get_kvs(int rank, c10d::Store& store,
-    bool singleP2POp = false, const std::string& p2pKey = "", int p2pRank = 0) {
+  ccl::shared_ptr_class<ccl::kvs> get_kvs(
+      int rank,
+      c10d::Store& store,
+      bool singleP2POp = false,
+      const std::string& p2pKey = "",
+      int p2pRank = 0) {
     std::lock_guard<std::mutex> lock(kvs_mutex);
     ccl::shared_ptr_class<ccl::kvs> kvs;
     std::string storeKey;
     if (!singleP2POp) {
-       storeKey = "xccl_kvs";
+      storeKey = std::to_string(xcclCommCounter_++);
     } else {
-       storeKey = p2pKey;
+      storeKey = p2pKey;
     }
     // Rank 0 broadcast the bootstrap network information to other ranks
     if (rank == 0 || (singleP2POp && p2pRank == 0)) {