refactor: unify and simplify create_process_group implementation.

Author: yingxudeng

xllm/core/common/global_flags.cpp

@@ -465,7 +465,6 @@ DEFINE_bool(enable_constrained_decoding,
             "that the output meets specific format or structural requirements "
             "through pre-defined rules.");
 
-
 #if defined(USE_NPU)
 DEFINE_string(
     npu_kernel_backend,

xllm/core/framework/parallel_state/collective_communicator.cpp

@@ -28,6 +28,7 @@ limitations under the License.
 #endif
 #include "common/global_flags.h"
 #include "parallel_args.h"
+#include "process_group.h"
 #include "util/net.h"
 
 namespace xllm {

xllm/core/framework/parallel_state/cuda_process_group.h

@@ -21,12 +21,12 @@ limitations under the License.
 
 namespace xllm {
 
-class ProcessGroupNccl : public ProcessGroup {
+class ProcessGroupImpl : public ProcessGroup {
  public:
-  ProcessGroupNccl(int global_rank,
-                   int world_size,
-                   int rank_size,
-                   int port,
+  ProcessGroupImpl(int32_t global_rank,
+                   int32_t world_size,
+                   int32_t rank_size,
+                   int32_t port,
                    bool trans,
                    const std::string& host,
                    const std::string& group_name,
@@ -38,7 +38,7 @@ class ProcessGroupNccl : public ProcessGroup {
     (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 7)
     pg_options->group_name = group_name;
 #endif
-    int rank = global_rank;
+    int32_t rank = global_rank;
     if (world_size != rank_size) {
       auto [local_rank, group_ranks] =
           get_group_rank(world_size, global_rank, rank_size, trans);
@@ -52,16 +52,4 @@ class ProcessGroupNccl : public ProcessGroup {
   }
 };
 
-std::unique_ptr<xllm::ProcessGroup> create_process_group(
-    int rank,
-    int world_size,
-    int rank_size,
-    int port,
-    bool trans,
-    const std::string& host,
-    const std::string& group_name,
-    const torch::Device& device) {
-  return std::make_unique<ProcessGroupNccl>(
-      rank, world_size, rank_size, port, trans, host, group_name, device);
-}
 }  // namespace xllm

xllm/core/framework/parallel_state/mlu_process_group.h

@@ -23,9 +23,9 @@ namespace xllm {
 
 constexpr int32_t local_device_count = 8;
 
-class ProcessGroupCncl : public ProcessGroup {
+class ProcessGroupImpl : public ProcessGroup {
  public:
-  ProcessGroupCncl(int32_t global_rank,
+  ProcessGroupImpl(int32_t global_rank,
                    int32_t world_size,
                    int32_t rank_size,
                    int32_t port,
@@ -57,17 +57,4 @@ class ProcessGroupCncl : public ProcessGroup {
   }
 };
 
-std::unique_ptr<xllm::ProcessGroup> create_process_group(
-    int32_t rank,
-    int32_t world_size,
-    int32_t rank_size,
-    int32_t port,
-    bool trans,
-    const std::string& host,
-    const std::string& group_name,
-    const torch::Device& device) {
-  return std::make_unique<ProcessGroupCncl>(
-      rank, world_size, rank_size, port, trans, host, group_name, device);
-}
-
 }  // namespace xllm

xllm/core/framework/parallel_state/npu_process_group.cpp

@@ -32,10 +32,10 @@ namespace {
 
 namespace xllm {
 
-ProcessGroupHCCL::ProcessGroupHCCL(int global_rank,
-                                   int world_size,
-                                   int rank_size,
-                                   int port,
+ProcessGroupImpl::ProcessGroupImpl(int32_t global_rank,
+                                   int32_t world_size,
+                                   int32_t rank_size,
+                                   int32_t port,
                                    bool trans,
                                    const std::string& host,
                                    const std::string& group_name,
@@ -47,7 +47,7 @@ ProcessGroupHCCL::ProcessGroupHCCL(int global_rank,
     (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 7)
   hccl_pg_options->group_name = group_name;
 #endif
-  int rank = global_rank;
+  int32_t rank = global_rank;
   if (world_size != rank_size) {
     auto [local_rank, group_ranks] =
         get_group_rank(world_size, global_rank, rank_size, trans);
@@ -65,31 +65,18 @@ ProcessGroupHCCL::ProcessGroupHCCL(int global_rank,
 }
 
 // Destructor.
-ProcessGroupHCCL::~ProcessGroupHCCL() {
+ProcessGroupImpl::~ProcessGroupImpl() {
   if (pg_) {
     pg_->shutdown();
   } else {
     HCCLCHECK(HcclCommDestroy(comm_));
   }
 }
 
-ProcessGroupHCCL::ProcessGroupHCCL(int rank,
+ProcessGroupImpl::ProcessGroupImpl(int rank,
                                    int world_size,
                                    const torch::Device& device,
                                    HcclComm comm)
     : ProcessGroup(device), comm_(comm) {}
 
-std::unique_ptr<xllm::ProcessGroup> create_process_group(
-    int rank,
-    int world_size,
-    int rank_size,
-    int port,
-    bool trans,
-    const std::string& host,
-    const std::string& group_name,
-    const torch::Device& device) {
-  return std::make_unique<ProcessGroupHCCL>(
-      rank, world_size, rank_size, port, trans, host, group_name, device);
-}
-
 }  // namespace xllm

xllm/core/framework/parallel_state/npu_process_group.h

@@ -20,15 +20,15 @@ limitations under the License.
 
 namespace xllm {
 
-class ProcessGroupHCCL : public ProcessGroup {
+class ProcessGroupImpl : public ProcessGroup {
  public:
   // Constructor.
-  ProcessGroupHCCL(int rank,
+  ProcessGroupImpl(int rank,
                    int world_size,
                    const torch::Device& device,
                    HcclComm comm);
 
-  ProcessGroupHCCL(int rank,
+  ProcessGroupImpl(int rank,
                    int world_size,
                    int rank_size,
                    int port,
@@ -38,20 +38,10 @@ class ProcessGroupHCCL : public ProcessGroup {
                    const torch::Device& device);
 
   // Destructor.
-  ~ProcessGroupHCCL() override;
+  ~ProcessGroupImpl() override;
 
  private:
   HcclComm comm_ = nullptr;
 };
 
-std::unique_ptr<xllm::ProcessGroup> create_process_group(
-    int rank,
-    int world_size,
-    int rank_size,
-    int port,
-    bool trans,
-    const std::string& host,
-    const std::string& group_name,
-    const torch::Device& device);
-
 }  // namespace xllm

xllm/core/framework/parallel_state/parallel_state.cpp

@@ -215,7 +215,7 @@ std::vector<std::unique_ptr<ProcessGroup>> create_npu_process_groups(
   std::vector<std::unique_ptr<ProcessGroup>> process_groups;
   process_groups.reserve(devices.size());
   for (int i = 0; i < world_size; ++i) {
-    process_groups.emplace_back(std::make_unique<ProcessGroupHCCL>(
+    process_groups.emplace_back(std::make_unique<ProcessGroupImpl>(
         /*rank=*/i, world_size, devices[i], comms[i]));
   }
 

xllm/core/framework/parallel_state/process_group.cpp

@@ -15,6 +15,14 @@ limitations under the License.
 
 #include "process_group.h"
 
+#if defined(USE_NPU)
+#include "npu_process_group.h"
+#elif defined(USE_MLU)
+#include "mlu_process_group.h"
+#elif defined(USE_CUDA)
+#include "cuda_process_group.h"
+#endif
+
 namespace {
 std::pair<int, std::vector<uint64_t>> get_trans_group_rank(int world_size,
                                                            int global_rank,
@@ -75,4 +83,18 @@ void ProcessGroup::allgather(const torch::Tensor& input,
   std::vector<std::vector<torch::Tensor>> output_tensors = {outputs};
   pg_->allgather(output_tensors, input_tensors)->wait();
 }
+
+std::unique_ptr<ProcessGroup> create_process_group(
+    int32_t rank,
+    int32_t world_size,
+    int32_t rank_size,
+    int32_t port,
+    bool trans,
+    const std::string& host,
+    const std::string& group_name,
+    const torch::Device& device) {
+  return std::unique_ptr<ProcessGroup>(new ProcessGroupImpl(
+      rank, world_size, rank_size, port, trans, host, group_name, device));
+}
+
 }  // namespace xllm

xllm/core/framework/parallel_state/process_group.h

@@ -25,6 +25,9 @@ limitations under the License.
 #endif
 
 namespace xllm {
+
+class ProcessGroupImpl;
+
 std::pair<int, std::vector<uint64_t>> get_group_rank(int world_size,
                                                      int global_rank,
                                                      int split_size,
@@ -77,4 +80,14 @@ class ProcessGroup {
 #endif
 };
 
+std::unique_ptr<xllm::ProcessGroup> create_process_group(
+    int32_t rank,
+    int32_t world_size,
+    int32_t rank_size,
+    int32_t port,
+    bool trans,
+    const std::string& host,
+    const std::string& group_name,
+    const torch::Device& device);
+
 }  // namespace xllm