[fluid_ops] partial_send process_group (#67915)

co63oc · web-flow · commit b3b9126c3aee · 2024-09-04T14:32:01.000+08:00
diff --git a/paddle/phi/kernels/gpu/partial_send_kernel.cu b/paddle/phi/kernels/gpu/partial_send_kernel.cu
@@ -17,36 +17,28 @@
 #include <vector>
 
 #include "glog/logging.h"
+#include "paddle/phi/core/distributed/utils.h"
 #include "paddle/phi/core/kernel_registry.h"
 
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
-#include "paddle/phi/core/distributed/collective/process_group.h"
 #include "paddle/phi/core/distributed/nccl_comm_context.h"
 #endif
 
-#include "paddle/phi/core/distributed/comm_context_manager.h"
-
 namespace phi {
 
 template <typename T, typename Context>
 void PartialSendCUDAKernel(const Context& dev_ctx,
                            const DenseTensor& x_in,
-                           int ring_id,
+                           int ring_id UNUSED,
                            int peer,
-                           bool use_calc_stream,
+                           bool use_calc_stream UNUSED,
                            int num,
                            int id) {
 #if (defined(PADDLE_WITH_RCCL) || defined(PADDLE_WITH_NCCL)) && \
     NCCL_VERSION_CODE >= 2703
   auto x = &x_in;
   int numel = x->numel();
-  int rid = ring_id;
 
-  PADDLE_ENFORCE_GE(
-      rid,
-      0,
-      common::errors::InvalidArgument(
-          "The ring_id (%d) for partial_send op must be non-negative.", rid));
   PADDLE_ENFORCE_GE(
       peer,
       0,
@@ -70,68 +62,38 @@ void PartialSendCUDAKernel(const Context& dev_ctx,
   int64_t send_numel = numel / num;
   int64_t offset = send_numel * id;
 
-  auto map = distributed::ProcessGroupMapFromGid::getInstance();
-  if (map->has(rid)) {
-    // Use ProcessGroup
-    distributed::ProcessGroup* pg = map->get(rid);
-    phi::DenseTensor tmp = *x;
-    auto task = pg->Send(tmp, peer, offset, send_numel, /*sync_op*/ true);
-    task->Wait();
-  } else {
-    gpuStream_t stream = nullptr;
-
-    phi::distributed::NCCLCommContext* comm_ctx = nullptr;
-    int nranks = 0;
-    int rank = 0;
-
-    const auto& comm_context_manager =
-        phi::distributed::CommContextManager::GetInstance();
-
-    PADDLE_ENFORCE_EQ(comm_context_manager.Has(std::to_string(rid)),
-                      true,
-                      common::errors::InvalidArgument(
-                          "You choose to use new communication library by "
-                          "setting environment "
-                          "variable FLAGS_dynamic_static_unified_comm True. "
-                          "But ring_id(%d) is "
-                          "not found in comm_context_manager.",
-                          std::to_string(rid)));
-    comm_ctx = static_cast<phi::distributed::NCCLCommContext*>(
-        comm_context_manager.Get(std::to_string(rid)));
-    PADDLE_ENFORCE_NE(comm_ctx,
-                      nullptr,
-                      common::errors::Unavailable(
-                          "NCCLCommContext is nullptr, collective op should "
-                          "has ring_id attr."));
-
-    stream = comm_ctx->GetStream();
-    nranks = comm_ctx->GetSize();
-    rank = comm_ctx->GetRank();
-
-    VLOG(3) << "new comm_context_manager has ring_id " << rid;
-
-    if (use_calc_stream) {
-      // should ExecutionContext for calc stream.
-      stream = dev_ctx.stream();
-    }
-
-    PADDLE_ENFORCE_LT(
-        peer,
-        nranks,
-        common::errors::InvalidArgument("The value of peer (%d) you set must "
-                                        "be less than ranks (%d).",
-                                        peer,
-                                        nranks));
-
-    ncclDataType_t dtype = phi::ToNCCLDataType(x->dtype());
-
-    auto send_buf = distributed::GetPartialTensor(*x, offset, send_numel);
-
-    comm_ctx->Send(send_buf, send_numel, peer, stream);
-
-    VLOG(3) << "rank " << rank << " send " << send_numel << " from offset["
-            << offset << "] to " << peer;
-  }
+  gpuStream_t stream = nullptr;
+
+  phi::distributed::NCCLCommContext* comm_ctx = nullptr;
+  int nranks = 0;
+  int rank = 0;
+
+  comm_ctx =
+      static_cast<phi::distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
+  PADDLE_ENFORCE_NE(comm_ctx,
+                    nullptr,
+                    common::errors::Unavailable(
+                        "NCCLCommContext is nullptr, collective op should "
+                        "has ring_id attr."));
+
+  stream = dev_ctx.stream();
+  nranks = comm_ctx->GetSize();
+  rank = comm_ctx->GetRank();
+
+  PADDLE_ENFORCE_LT(
+      peer,
+      nranks,
+      common::errors::InvalidArgument("The value of peer (%d) you set must "
+                                      "be less than ranks (%d).",
+                                      peer,
+                                      nranks));
+
+  auto send_buf = distributed::GetPartialTensor(*x, offset, send_numel);
+  comm_ctx->Send(send_buf, send_numel, peer, stream);
+
+  VLOG(3) << "rank " << rank << " send " << send_numel << " from offset["
+          << offset << "] to " << peer;
+
 #else
   PADDLE_THROW(
       common::errors::Unavailable("PaddlePaddle should be compiled with NCCL "
