[CustomDevice] use CommContextManager to create xccl comm (#57957)

ronny1996 · web-flow · commit 56d447639ee7 · 2023-10-10T17:54:08.000+08:00
diff --git a/paddle/fluid/distributed/collective/process_group_custom.cc b/paddle/fluid/distributed/collective/process_group_custom.cc
@@ -494,7 +494,7 @@ void ProcessGroupCustom::CreateXCCLEnvCache(const Place& place,
           << ", place: " << place_key;
 
   phi::distributed::CommContextManager::CreateXCCLCommContext(
-      store_, std::to_string(gid_), place.GetDeviceType(), rank_, size_);
+      store_, std::to_string(gid_), place, rank_, size_);
 
   auto* calc_ctx = static_cast<phi::CustomContext*>(
       platform::DeviceContextPool::Instance().Get(place));
diff --git a/paddle/fluid/operators/collective/c_comm_init_op.cc b/paddle/fluid/operators/collective/c_comm_init_op.cc
@@ -64,8 +64,6 @@ class CCommInitOp : public framework::OperatorBase {
       PADDLE_ENFORCE_NOT_NULL(
           var, platform::errors::InvalidArgument("Input con not be empty."));
 
-      phi::ccl::CCLRootId* comm_id = var->GetMutable<phi::ccl::CCLRootId>();
-
       int nranks = Attr<int>("nranks");
       int rid = Attr<int>("ring_id");
 
@@ -74,8 +72,17 @@ class CCommInitOp : public framework::OperatorBase {
         device_id = Attr<int>("device_id");
       }
       int rank_id = Attr<int>("rank");
-      platform::XCCLCommContext::Instance(place.GetDeviceType())
-          .CreateComm(comm_id, nranks, rank_id, device_id, rid);
+      auto store = phi::distributed::CreateOrGetGlobalTCPStore();
+      if (!phi::distributed::CommContextManager::GetInstance().Has(
+              std::to_string(rid))) {
+        phi::distributed::CommContextManager::CreateXCCLCommContext(
+            store,
+            std::to_string(rid),
+            phi::CustomPlace(place.GetDeviceType(), device_id),
+            rank_id,
+            nranks,
+            "c_comm_init_op");
+      }
 #else
       PADDLE_THROW(platform::errors::PreconditionNotMet(
           "PaddlePaddle should compile with custom device."));
diff --git a/paddle/fluid/operators/collective/c_gen_xccl_id_op.cc b/paddle/fluid/operators/collective/c_gen_xccl_id_op.cc
@@ -52,34 +52,7 @@ class CGenXCCLIdOp : public framework::OperatorBase {
       : OperatorBase(type, inputs, outputs, attrs) {}
 
   void RunImpl(const framework::Scope& scope,
-               const platform::Place& dev_place) const override {
-    int rank = Attr<int>("rank");
-    int ring_id = Attr<int>("ring_id");
-
-    std::function<std::string(size_t)> func = [&](size_t i) -> std::string {
-      return Output("Out");
-    };
-
-    std::string endpoint = Attr<std::string>("endpoint");
-    int server_fd = platform::SocketServer::GetInstance(endpoint).socket();
-
-    std::vector<phi::ccl::CCLRootId> xccl_ids;
-    xccl_ids.resize(1);
-
-    if (rank == 0) {
-      for (size_t i = 0; i < xccl_ids.size(); ++i) {
-        phi::DeviceManager::CCLGetUniqueId(dev_place.GetDeviceType(),
-                                           &xccl_ids[i]);
-      }
-      std::vector<std::string> endpoint_list =
-          Attr<std::vector<std::string>>("other_endpoints");
-      platform::SendBroadCastCommID(endpoint_list, &xccl_ids, ring_id);
-    } else {
-      platform::RecvBroadCastCommID(server_fd, endpoint, &xccl_ids, ring_id);
-    }
-
-    CopyXCCLIDToVar(xccl_ids, func, scope);
-  }
+               const platform::Place& dev_place) const override {}
 };
 
 #else
diff --git a/paddle/fluid/operators/custom_device_common_op_registry.cc b/paddle/fluid/operators/custom_device_common_op_registry.cc
@@ -22,6 +22,8 @@ limitations under the License. */
 #include "paddle/phi/api/backward/backward_api.h"
 #include "paddle/phi/api/include/api.h"
 #include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+#include "paddle/phi/core/distributed/xccl_comm_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/axis_utils.h"
 
@@ -99,13 +101,14 @@ class CConcatOpCustomDeviceKernel : public framework::OpKernel<T> {
       auto task = pg->AllGather(in_tensor, out_tensor);
       task->Wait();
     } else {
-      auto comm = platform::XCCLCommContext::Instance(place.GetDeviceType())
-                      .Get(rid, place);
+      auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+          phi::distributed::CommContextManager::GetInstance().Get(
+              std::to_string(rid)));
       PADDLE_ENFORCE_EQ(
           nranks,
-          comm->nranks(),
+          comm->GetSize(),
           platform::errors::InvalidArgument(
-              "nranks: %s should equal to %s", nranks, comm->nranks()));
+              "nranks: %s should equal to %s", nranks, comm->GetSize()));
 
       int64_t send_numel = x->numel();
       const T* send_buff = x->data<T>();
@@ -118,7 +121,7 @@ class CConcatOpCustomDeviceKernel : public framework::OpKernel<T> {
           recv_buff,
           send_numel,
           phi::ccl::ToCCLDataType(x->dtype()),
-          comm->comm(),
+          comm->GetXcclComm(),
           stream);
     }
     std::vector<phi::DenseTensor> inputs;
@@ -600,25 +603,25 @@ class CAllReduceOpCustomDeviceKernel : public framework::OpKernel<T> {
       return;
     }
 
-    auto comm =
-        paddle::platform::XCCLCommContext::Instance(place.GetDeviceType())
-            .Get(rid, place);
+    auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+        phi::distributed::CommContextManager::GetInstance().Get(
+            std::to_string(rid)));
 
     std::shared_ptr<phi::stream::Stream> stream;
     if (ctx.Attr<bool>("use_calc_stream")) {
       auto dev_ctx = paddle::platform::DeviceContextPool::Instance().Get(place);
       stream = static_cast<paddle::platform::CustomDeviceContext*>(dev_ctx)
                    ->GetStream();
     } else {
-      stream = comm->stream();
+      stream = comm->GetStream();
     }
     phi::DeviceManager::CCLAllReduce(place.GetDeviceType(),
                                      const_cast<void*>(sendbuff),
                                      recvbuff,
                                      numel,
                                      dtype,
                                      red_type,
-                                     comm->comm(),
+                                     comm->GetXcclComm(),
                                      *stream);
   }
 };
@@ -634,22 +637,30 @@ class CBroadcastOpCustomDeviceKernel : public framework::OpKernel<T> {
     int root = ctx.Attr<int>("root");
     int rid = ctx.Attr<int>("ring_id");
 
-    auto stream = static_cast<const phi::CustomContext&>(ctx.device_context())
-                      .GetStream();
+    auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+        phi::distributed::CommContextManager::GetInstance().Get(
+            std::to_string(rid)));
+
+    std::shared_ptr<phi::stream::Stream> stream;
+    if (ctx.Attr<bool>("use_calc_stream")) {
+      auto dev_ctx = paddle::platform::DeviceContextPool::Instance().Get(place);
+      stream = static_cast<paddle::platform::CustomDeviceContext*>(dev_ctx)
+                   ->GetStream();
+    } else {
+      stream = comm->GetStream();
+    }
 
     int numel = x->numel();
     auto dtype = phi::ccl::ToCCLDataType(x->dtype());
-    auto comm = platform::XCCLCommContext::Instance(place.GetDeviceType())
-                    .Get(rid, place);
-    if (root == comm->rank()) {
+    if (root == comm->GetRank()) {
       phi::DeviceManager::CCLBroadcast(place.GetDeviceType(),
                                        const_cast<void*>(x->data()),
                                        numel,
                                        dtype,
                                        root,
-                                       comm->comm(),
+                                       comm->GetXcclComm(),
                                        *stream);
-      VLOG(3) << "rank " << comm->rank() << " invoke Bcast. sent "
+      VLOG(3) << "rank " << comm->GetRank() << " invoke Bcast. sent "
               << x->numel();
       if (out != x) {
         framework::TensorCopy(
@@ -664,9 +675,9 @@ class CBroadcastOpCustomDeviceKernel : public framework::OpKernel<T> {
                                        numel,
                                        dtype,
                                        root,
-                                       comm->comm(),
+                                       comm->GetXcclComm(),
                                        *stream);
-      VLOG(3) << "rank " << comm->rank() << " invoke Bcast. received "
+      VLOG(3) << "rank " << comm->GetRank() << " invoke Bcast. received "
               << phi::product(out->dims());
     }
     out->set_lod(x->lod());
@@ -684,16 +695,27 @@ class BarrierOpCustomDeviceKernel : public framework::OpKernel<T> {
     const void* sendbuff = in->data();
     void* recvbuff = ctx.device_context().Alloc<T>(out);
     int rid = ctx.Attr<int>("ring_id");
-    auto comm = platform::XCCLCommContext::Instance(place.GetDeviceType())
-                    .Get(rid, place);
+
+    auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+        phi::distributed::CommContextManager::GetInstance().Get(
+            std::to_string(rid)));
+
+    std::shared_ptr<phi::stream::Stream> stream;
+    if (ctx.Attr<bool>("use_calc_stream")) {
+      auto dev_ctx = paddle::platform::DeviceContextPool::Instance().Get(place);
+      stream = static_cast<paddle::platform::CustomDeviceContext*>(dev_ctx)
+                   ->GetStream();
+    } else {
+      stream = comm->GetStream();
+    }
     phi::DeviceManager::CCLAllReduce(place.GetDeviceType(),
                                      const_cast<void*>(sendbuff),
                                      recvbuff,
                                      numel,
                                      phi::ccl::ToCCLDataType(in->dtype()),
                                      phi::ccl::CCLReduceOp::SUM,
-                                     comm->comm(),
-                                     *(comm->stream()));
+                                     comm->GetXcclComm(),
+                                     *stream);
   }
 };
 
@@ -993,16 +1015,22 @@ class GlobalScatterOpCustomDeviceKernel : public framework::OpKernel<T> {
         }
       }
     } else {
-      auto comm = platform::XCCLCommContext::Instance(place.GetDeviceType())
-                      .Get(rid, place);
+      auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+          phi::distributed::CommContextManager::GetInstance().Get(
+              std::to_string(rid)));
+
       std::shared_ptr<phi::stream::Stream> stream;
       if (ctx.Attr<bool>("use_calc_stream")) {
-        stream = dev_ctx.GetStream();
+        auto dev_ctx =
+            paddle::platform::DeviceContextPool::Instance().Get(place);
+        stream = static_cast<paddle::platform::CustomDeviceContext*>(dev_ctx)
+                     ->GetStream();
       } else {
-        stream = comm->stream();
+        stream = comm->GetStream();
       }
-      int nranks = comm->nranks();
-      int rank = comm->rank();
+
+      int nranks = comm->GetSize();
+      int rank = comm->GetRank();
       auto in_feat = x->dims()[1];
       auto n_expert = local_count->dims()[0] / nranks;
       int64_t fwd_count = 0;
@@ -1033,7 +1061,7 @@ class GlobalScatterOpCustomDeviceKernel : public framework::OpKernel<T> {
                 cpu_global_count_data[idx] * in_feat,
                 phi::ccl::ToCCLDataType(x->dtype()),
                 j,
-                comm->comm(),
+                comm->GetXcclComm(),
                 *stream);
             recv_ptr += cpu_global_count_data[idx];
           }
@@ -1049,7 +1077,7 @@ class GlobalScatterOpCustomDeviceKernel : public framework::OpKernel<T> {
                   cpu_local_count_data[idx] * in_feat,
                   phi::ccl::ToCCLDataType(x->dtype()),
                   j,
-                  comm->comm(),
+                  comm->GetXcclComm(),
                   *stream);
             }
           }
@@ -1072,7 +1100,7 @@ class GlobalScatterOpCustomDeviceKernel : public framework::OpKernel<T> {
                 cpu_global_count_data[idx] * in_feat,
                 phi::ccl::ToCCLDataType(x->dtype()),
                 j,
-                comm->comm(),
+                comm->GetXcclComm(),
                 *stream);
             recv_ptr += cpu_global_count_data[idx];
           }
@@ -1199,16 +1227,21 @@ class GlobalGatherOpCustomDeviceKernel : public framework::OpKernel<T> {
         }
       }
     } else {
-      auto comm = platform::XCCLCommContext::Instance(place.GetDeviceType())
-                      .Get(rid, place);
+      auto comm = reinterpret_cast<phi::distributed::XCCLCommContext*>(
+          phi::distributed::CommContextManager::GetInstance().Get(
+              std::to_string(rid)));
+
       std::shared_ptr<phi::stream::Stream> stream;
       if (ctx.Attr<bool>("use_calc_stream")) {
-        stream = dev_ctx.GetStream();
+        auto dev_ctx =
+            paddle::platform::DeviceContextPool::Instance().Get(place);
+        stream = static_cast<paddle::platform::CustomDeviceContext*>(dev_ctx)
+                     ->GetStream();
       } else {
-        stream = comm->stream();
+        stream = comm->GetStream();
       }
-      int nranks = comm->nranks();
-      int rank = comm->rank();
+      int nranks = comm->GetSize();
+      int rank = comm->GetRank();
       auto in_feat = x->dims()[1];
       auto n_expert = local_count->dims()[0] / nranks;
 
@@ -1238,7 +1271,7 @@ class GlobalGatherOpCustomDeviceKernel : public framework::OpKernel<T> {
                                         cpu_local_count_data[idx] * in_feat,
                                         phi::ccl::ToCCLDataType(x->dtype()),
                                         j,
-                                        comm->comm(),
+                                        comm->GetXcclComm(),
                                         *stream);
           }
         }
@@ -1253,7 +1286,7 @@ class GlobalGatherOpCustomDeviceKernel : public framework::OpKernel<T> {
                   cpu_global_count_data[idx] * in_feat,
                   phi::ccl::ToCCLDataType(x->dtype()),
                   j,
-                  comm->comm(),
+                  comm->GetXcclComm(),
                   *stream);
             } else {
               phi::DeviceManager::GetDeviceWithPlace(place)->MemoryCopyD2D(
@@ -1274,7 +1307,7 @@ class GlobalGatherOpCustomDeviceKernel : public framework::OpKernel<T> {
                                         cpu_local_count_data[idx] * in_feat,
                                         phi::ccl::ToCCLDataType(x->dtype()),
                                         j,
-                                        comm->comm(),
+                                        comm->GetXcclComm(),
                                         *stream);
           }
         }
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard_utils.cc b/paddle/phi/core/distributed/auto_parallel/reshard_utils.cc
@@ -95,11 +95,7 @@ CommContext* CreateOrGetCommContext(const DeviceContext& dev_ctx,
     } else if (phi::CustomContext::classof(&dev_ctx)) {
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
       CommContextManager::CreateXCCLCommContext(
-          store,
-          unique_comm_key,
-          dev_ctx.GetPlace().GetDeviceType(),
-          rank,
-          world_size);
+          store, unique_comm_key, dev_ctx.GetPlace(), rank, world_size);
 #endif
     } else {
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
diff --git a/paddle/phi/core/distributed/comm_context_manager.cc b/paddle/phi/core/distributed/comm_context_manager.cc
@@ -130,15 +130,19 @@ void CommContextManager::CreateGlooCommContext(
 void CommContextManager::CreateXCCLCommContext(
     const std::shared_ptr<Store>& store,
     const std::string& unique_comm_key,
-    const std::string& device_type,
+    const phi::Place& place,
     int rank,
-    int size) {
+    int size,
+    const std::string& hash_key) {
   phi::ccl::CCLRootId xccl_root_id;
   if (rank == 0) {
-    phi::DeviceManager::CCLGetUniqueId(device_type, &xccl_root_id);
+    phi::DeviceManager::CCLGetUniqueId(place.GetDeviceType(), &xccl_root_id);
   }
 
   std::string unique_key = "XCCLCommContext/" + unique_comm_key;
+  if (!hash_key.empty()) {
+    unique_key += "/" + hash_key;
+  }
   if (rank == 0) {
     store->set(unique_key, xccl_root_id);
   } else {
@@ -148,7 +152,7 @@ void CommContextManager::CreateXCCLCommContext(
           << ", unique_comm_key: " << unique_comm_key << ", xccl uniqueid: "
           << phi::ccl::SerializeXCCLUniqueId(xccl_root_id);
   auto xccl_comm_context =
-      std::make_unique<XCCLCommContext>(device_type, rank, size, xccl_root_id);
+      std::make_unique<XCCLCommContext>(place, rank, size, xccl_root_id);
   auto& comm_context_manager = CommContextManager::GetInstance();
   comm_context_manager.SetStore(store);
   comm_context_manager.Emplace(unique_comm_key, std::move(xccl_comm_context));
diff --git a/paddle/phi/core/distributed/comm_context_manager.h b/paddle/phi/core/distributed/comm_context_manager.h
diff --git a/paddle/phi/core/distributed/xccl_comm_context.cc b/paddle/phi/core/distributed/xccl_comm_context.cc
diff --git a/paddle/phi/core/distributed/xccl_comm_context.h b/paddle/phi/core/distributed/xccl_comm_context.h
