internal review

torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp

@@ -180,13 +180,12 @@ ncclRedOpRAII getNcclReduceOp(
                 NCCL_MINOR));
         break;
       case ReduceOp::BAND:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BAND with NCCL");
-        break;
       case ReduceOp::BOR:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BOR with NCCL");
-        break;
       case ReduceOp::BXOR:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BXOR with NCCL");
+        C10_THROW_ERROR(
+            ValueError,
+            "Cannot use ReduceOp." + reduce_op_to_string(reduceOp) +
+                " with NCCL");
         break;
       default:
         C10_THROW_ERROR(ValueError, "Unhandled ReduceOp");

torch/csrc/distributed/c10d/ProcessGroupXCCL.cpp

@@ -67,17 +67,15 @@ ccl::reduction getXcclReduceOp(const ReduceOp& reduceOp, at::Tensor& input) {
     return xcclOps.at(reduceOp);
   } catch (const std::out_of_range&) {
     switch (reduceOp) {
-      case ReduceOp::AVG:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp AVG with XCCL");
-        break;
-      case ReduceOp::BAND:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BAND with XCCL");
-        break;
-      case ReduceOp::BOR:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BOR with XCCL");
-        break;
-      case ReduceOp::BXOR:
-        C10_THROW_ERROR(ValueError, "Cannot use ReduceOp.BXOR with XCCL");
+      case c10d::ReduceOp::BAND:
+      case c10d::ReduceOp::BOR:
+      case c10d::ReduceOp::BXOR:
+      case c10d::ReduceOp::AVG:
+      case c10d::ReduceOp::PREMUL_SUM:
+        C10_THROW_ERROR(
+            ValueError,
+            "Cannot use ReduceOp." + reduce_op_to_string(reduceOp) +
+                " with XCCL");
         break;
       default:
         C10_THROW_ERROR(ValueError, "Unhandled ReduceOp");
@@ -133,7 +131,6 @@ void ProcessGroupXCCL::WorkXCCL::synchronizeInternal(
       auto currentTimepoint = std::chrono::steady_clock::now();
       auto timeElapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
           currentTimepoint - workStartTime_);
-      std::chrono::milliseconds opTimeout = std::chrono::milliseconds(60000);
       if (timeElapsed >= timeout) {
         std::string exceptionMsg = c10::str(
             "Work ran for ",

torch/csrc/distributed/c10d/ProcessGroupXCCL.hpp

@@ -262,6 +262,34 @@ class TORCH_API ProcessGroupXCCL : public Backend {
   c10::intrusive_ptr<Store> store_;
   std::mutex mutex_;
   bool blockingWait_ = false;
+
+ private:
+  XCCL_KVS kvs;
+  std::mutex kvs_mutex;
+  XCCL_KVS get_kvs(int rank, c10d::Store& store) {
+    std::lock_guard<std::mutex> lock(kvs_mutex);
+    if (kvs)
+      return kvs;
+    std::string storeKey = "xccl_kvs";
+    // Rank 0 broadcast the bootstrap network information to other ranks
+    if (rank == 0) {
+      kvs = ccl::create_main_kvs();
+      ccl::kvs::address_type main_addr = kvs->get_address();
+      auto ccl_kvs_addr =
+          std::vector<uint8_t>(main_addr.begin(), main_addr.end());
+      store.set(storeKey, ccl_kvs_addr);
+    } else {
+      auto ccl_kvs_addr = store.get(storeKey);
+      if (ccl_kvs_addr.size() != ccl::kvs::address_max_size) {
+        throw std::runtime_error("Unexpected ccl kvs addr from the store\n");
+      }
+      ccl::kvs::address_type main_addr;
+      std::copy_n(
+          ccl_kvs_addr.begin(), ccl::kvs::address_max_size, main_addr.begin());
+      kvs = ccl::create_kvs(main_addr);
+    }
+    return kvs;
+  }
 };
 
 namespace {
@@ -297,35 +325,6 @@ bool with_mpirun() {
       : false;
 }
 
-XCCL_KVS kvs;
-std::mutex kvs_mutex;
-XCCL_KVS get_kvs(int rank, c10d::Store& store) {
-  std::lock_guard<std::mutex> lock(kvs_mutex);
-  if (kvs)
-    return kvs;
-  std::string storeKey = "xccl_kvs";
-
-  // Rank 0 broadcast the bootstrap network information to other ranks
-  if (rank == 0) {
-    kvs = ccl::create_main_kvs();
-    ccl::kvs::address_type main_addr = kvs->get_address();
-    auto ccl_kvs_addr =
-        std::vector<uint8_t>(main_addr.begin(), main_addr.end());
-    store.set(storeKey, ccl_kvs_addr);
-  } else {
-    auto ccl_kvs_addr = store.get(storeKey);
-    if (ccl_kvs_addr.size() != ccl::kvs::address_max_size) {
-      throw std::runtime_error("Unexpected ccl kvs addr from the store\n");
-    }
-    ccl::kvs::address_type main_addr;
-    std::copy_n(
-        ccl_kvs_addr.begin(), ccl::kvs::address_max_size, main_addr.begin());
-    kvs = ccl::create_kvs(main_addr);
-  }
-
-  return kvs;
-}
-
 } // namespace
 } // namespace c10d
 

torch/csrc/distributed/c10d/Utils.hpp

@@ -557,6 +557,31 @@ size_t computeLengthsAndOffsets(
   return offset;
 }
 
+inline std::string reduce_op_to_string(c10d::ReduceOp op) {
+  switch (op) {
+    case c10d::ReduceOp::SUM:
+      return "SUM";
+    case c10d::ReduceOp::PRODUCT:
+      return "PRODUCT";
+    case c10d::ReduceOp::MIN:
+      return "MIN";
+    case c10d::ReduceOp::MAX:
+      return "MAX";
+    case c10d::ReduceOp::BAND:
+      return "BAND";
+    case c10d::ReduceOp::BOR:
+      return "BOR";
+    case c10d::ReduceOp::BXOR:
+      return "BXOR";
+    case c10d::ReduceOp::AVG:
+      return "AVG";
+    case c10d::ReduceOp::PREMUL_SUM:
+      return "PREMUL_SUM";
+    default:
+      return "UNKNOWN";
+  }
+}
+
 using RankType = uint32_t;
 using SizeType = uint64_t;