update header file processgroupnccl

c7e6a3db · Rick Ho · a762d33c · c7e6a3db · c7e6a3db
Commit c7e6a3db authored Dec 28, 2022 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 12 additions and 0 deletions

cuda/fmoe_cuda.cpp cuda/fmoe_cuda.cpp +7 -0

cuda/global_exchange.cpp cuda/global_exchange.cpp +5 -0

No files found.
--- a/cuda/fmoe_cuda.cpp
+++ b/cuda/fmoe_cuda.cpp
@@ -5,7 +5,14 @@

 // global_exchange
 #ifdef FMOE_USE_NCCL
+
+#if defined(TORCH_VERSION_MAJOR) && (TORCH_VERSION_MAJOR > 1 || \
+        (TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR >= 13))
+#include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
+#else
 #include <c10d/ProcessGroupNCCL.hpp>
+#endif
+
 torch::Tensor _expert_exchange(
        torch::Tensor local_expert_count,
        long n_expert, long n_workers);

--- a/cuda/global_exchange.cpp
+++ b/cuda/global_exchange.cpp
@@ -98,7 +98,12 @@ torch::Tensor _global_gather(
    return local_output_buf;
 }

+#if defined(TORCH_VERSION_MAJOR) && (TORCH_VERSION_MAJOR > 1 || \
+        (TORCH_VERSION_MAJOR == 1 && TORCH_VERSION_MINOR >= 13))
+#include <torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp>
+#else
 #include <c10d/ProcessGroupNCCL.hpp>
+#endif

 class HackNCCLGroup: public c10d::ProcessGroupNCCL {
 public: