[CustomDevice] add inference MP support, PART2

paddle/fluid/operators/collective/c_comm_init_op.cc

@@ -25,7 +25,7 @@ limitations under the License. */
 #include "paddle/fluid/framework/op_registry.h"
 
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL) || \
-    defined(PADDLE_WITH_XPU_BKCL)
+    defined(PADDLE_WITH_XPU_BKCL) || defined(PADDLE_WITH_CUSTOM_DEVICE)
 #include "paddle/fluid/platform/collective_helper.h"
 #endif
 
@@ -48,43 +48,67 @@ class CCommInitOp : public framework::OperatorBase {
 
   void RunImpl(const framework::Scope& scope,
                const platform::Place& place) const override {
+    if (platform::is_custom_place(place)) {
+#if defined(PADDLE_WITH_CUSTOM_DEVICE)
+      auto var = scope.FindVar(Input("X"));
+      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");
+
+      int device_id = place.device;
+      if (Attr<int>("device_id") >= 0) {
+        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);
+#else
+      PADDLE_THROW(platform::errors::PreconditionNotMet(
+          "PaddlePaddle should compile with custom device."));
+#endif
+    } else {
 // TODO(wangxi): Put this in the unified header file
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
-    using UniqueId = ncclUniqueId;
-    using CommContext = platform::NCCLCommContext;
+      using UniqueId = ncclUniqueId;
+      using CommContext = platform::NCCLCommContext;
 #elif defined(PADDLE_WITH_XPU_BKCL)
-    using UniqueId = BKCLUniqueId;
-    using CommContext = platform::BKCLCommContext;
+      using UniqueId = BKCLUniqueId;
+      using CommContext = platform::BKCLCommContext;
 #else
-    PADDLE_THROW(platform::errors::PreconditionNotMet(
-        "PaddlePaddle should be compiled with GPU or XPU."));
+      PADDLE_THROW(platform::errors::PreconditionNotMet(
+          "PaddlePaddle should be compiled with GPU or XPU."));
 #endif
 
-    PADDLE_ENFORCE_EQ(
-        platform::is_gpu_place(place) || platform::is_xpu_place(place),
-        true,
-        platform::errors::PreconditionNotMet(
-            "CCommInitOp can run on gpu or xpu place only."));
+      PADDLE_ENFORCE_EQ(
+          platform::is_gpu_place(place) || platform::is_xpu_place(place),
+          true,
+          platform::errors::PreconditionNotMet(
+              "CCommInitOp can run on gpu or xpu place only."));
 
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL) || \
     defined(PADDLE_WITH_XPU_BKCL)
-    auto var = scope.FindVar(Input("X"));
-    PADDLE_ENFORCE_NOT_NULL(
-        var, platform::errors::InvalidArgument("Input con not be empty."));
-
-    UniqueId* comm_id = var->GetMutable<UniqueId>();
-
-    int nranks = Attr<int>("nranks");
-    int rid = Attr<int>("ring_id");
-
-    int device_id = place.device;
-    if (Attr<int>("device_id") >= 0) {
-      device_id = Attr<int>("device_id");
-    }
-    int rank_id = Attr<int>("rank");
-    CommContext::Instance().CreateComm(
-        comm_id, nranks, rank_id, device_id, rid);
+      auto var = scope.FindVar(Input("X"));
+      PADDLE_ENFORCE_NOT_NULL(
+          var, platform::errors::InvalidArgument("Input con not be empty."));
+
+      UniqueId* comm_id = var->GetMutable<UniqueId>();
+
+      int nranks = Attr<int>("nranks");
+      int rid = Attr<int>("ring_id");
+
+      int device_id = place.device;
+      if (Attr<int>("device_id") >= 0) {
+        device_id = Attr<int>("device_id");
+      }
+      int rank_id = Attr<int>("rank");
+      CommContext::Instance().CreateComm(
+          comm_id, nranks, rank_id, device_id, rid);
 #endif
+    }
   }
 };
 

paddle/fluid/operators/collective/c_gen_xccl_id_op.cc

@@ -0,0 +1,132 @@
+/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#include <string>
+
+#include "paddle/fluid/framework/op_proto_maker.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/framework/var_type_traits.h"
+#include "paddle/fluid/platform/device_context.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/fluid/platform/gen_comm_id_helper.h"
+#include "paddle/fluid/platform/place.h"
+#include "paddle/phi/backends/device_manager.h"
+
+namespace paddle {
+namespace operators {
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+static void CopyXCCLIDToVar(const std::vector<phi::ccl::CCLRootId>& xccl_ids,
+                            std::function<std::string(size_t)> func,
+                            const framework::Scope& scope) {
+  for (size_t i = 0; i < xccl_ids.size(); ++i) {
+    std::string var_name = func(i);
+    auto var = scope.FindVar(var_name);
+    PADDLE_ENFORCE_NOT_NULL(
+        var,
+        platform::errors::NotFound("Variable with name %s is not found",
+                                   var_name.c_str()));
+    auto xccl_id = var->GetMutable<phi::ccl::CCLRootId>();
+    *xccl_id = xccl_ids[i];
+  }
+}
+
+class CGenXCCLIdOp : public framework::OperatorBase {
+ public:
+  CGenXCCLIdOp(const std::string& type,
+               const framework::VariableNameMap& inputs,
+               const framework::VariableNameMap& outputs,
+               const framework::AttributeMap& attrs)
+      : 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);
+  }
+};
+
+#else
+class CGenXCCLIdOp : public framework::OperatorBase {
+ public:
+  CGenXCCLIdOp(const std::string& type,
+               const framework::VariableNameMap& inputs,
+               const framework::VariableNameMap& outputs,
+               const framework::AttributeMap& attrs)
+      : OperatorBase(type, inputs, outputs, attrs) {}
+
+  void RunImpl(const framework::Scope& scope,
+               const platform::Place& dev_place) const override {}
+};
+
+#endif
+
+class CGenXCCLIdOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddOutput("Out", "Raw variable contains a XCCL UniqueId instaces.");
+    AddComment(R"DOC(
+CGenXCCLId operator
+
+For trainer 0: generate a new UniqueId and send it to all the other trainers.
+For trainer 1~n: start a gRPC server to get the UniqueId, once got, stop the server.
+)DOC");
+    AddAttr<std::string>("endpoint",
+                         "(string), e.g. 127.0.0.1:6175 "
+                         "current listen endpoint");
+    AddAttr<std::vector<std::string>>(
+        "other_endpoints",
+        "['trainer1_ip:port', 'trainer2_ip:port', ...] "
+        "list of other trainer endpoints")
+        .SetDefault({});
+    AddAttr<int>("rank",
+                 "(int default 0) "
+                 "The rank of the trainer in distributed training.")
+        .SetDefault(0);
+    AddAttr<int>("ring_id", "(int default 0) user specified ring id")
+        .SetDefault(0);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OPERATOR(c_gen_xccl_id, ops::CGenXCCLIdOp, ops::CGenXCCLIdOpMaker);