feat: add multi-node default ProcessGroup

kilinchange · kilinchange · commit 78047a7f8028 · 2025-11-10T09:33:15.000Z
diff --git a/infini_train/include/nn/parallel/global.h b/infini_train/include/nn/parallel/global.h
@@ -4,6 +4,10 @@
 #include <string>
 #include <vector>
 
+#ifdef USE_NCCL
+#include <nccl.h>
+#endif
+
 namespace infini_train::nn::parallel::global {
 
 enum Axis : uint8_t { DP = 0, TP = 1, PP = 2, AXIS_COUNT = 3 };
@@ -45,6 +49,9 @@ class GlobalEnv {
     int data_parallel_size() const;
 
     Layout layout() const;
+#ifdef USE_NCCL
+    ncclUniqueId nccl_id() const;
+#endif
 
 private:
     GlobalEnv() = default;
@@ -65,6 +72,10 @@ class GlobalEnv {
 
     int data_parallel_size_ = 1;
 
+#ifdef USE_NCCL
+    ncclUniqueId nccl_id_;
+#endif
+
     mutable std::mutex mutex_;
     bool initialized_ = false;
 
@@ -108,5 +119,8 @@ inline std::vector<int> GetGroupRanks(Axis target, int rank) {
 }
 
 std::string ProcessGroupOverview(const Layout &L = GlobalEnv::Instance().layout(), bool skip_trivial_axes = true);
+#ifdef USE_NCCL
+inline ncclUniqueId GetNcclId() { return GlobalEnv::Instance().nccl_id(); }
+#endif
 
 } // namespace infini_train::nn::parallel::global
diff --git a/infini_train/include/nn/parallel/process_group.h b/infini_train/include/nn/parallel/process_group.h
@@ -3,6 +3,7 @@
 #include <memory>
 #include <mutex>
 #include <string>
+#include <type_traits>
 #include <unordered_map>
 #include <vector>
 
@@ -28,6 +29,9 @@ class ProcessGroup {
 public:
     explicit ProcessGroup(const std::vector<int> &device_indices);
 
+    // support for multi-node distributed training
+    explicit ProcessGroup(const ncclUniqueId &nccl_id);
+
     int GetGroupRank(int thread_rank) const;
 
     void AllReduce(const std::shared_ptr<Tensor> &tensor, function::ReduceOpType reduce_op) const;
@@ -52,14 +56,17 @@ class ProcessGroup {
 
     std::vector<std::shared_ptr<Tensor>> NcclRecv(std::vector<std::shared_ptr<Tensor>> tensors, int src_rank) const;
 
+private:
+    void Init(const std::vector<int> &device_indices);
+
 private:
     std::vector<ncclComm_t> comms_;
     std::vector<const Device *> devices_;
 
     std::unordered_map<const Device *, ncclComm_t> device_comm_map_;
     std::unordered_map<int, int> thread_group_rank_map_; // thread_rank : group_rank
 
-    int comm_size_ = 0;
+    int world_size_ = 0;
 };
 #endif
 
@@ -73,12 +80,37 @@ class ProcessGroupFactory {
 
     const ProcessGroup *GetOrCreate(const std::string &name, const std::vector<int> &device_indices);
 
+#ifdef USE_NCCL
+    const ProcessGroup *GetOrCreate(const std::string &name, const ncclUniqueId &nccl_id);
+#endif
+
     const ProcessGroup *Get(const std::string &name) const;
 
     const ProcessGroup *GetDefaultProcessGroup() const;
 
 private:
     ProcessGroupFactory();
+
+    template <typename Creator, typename = std::enable_if_t<std::is_invocable_v<Creator>>>
+    const ProcessGroup *GetOrCreate(const std::string &name, Creator &&creator) {
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            auto it = name_to_group_.find(name);
+            if (it != name_to_group_.end()) {
+                return it->second.get();
+            }
+        }
+
+        auto new_group = creator();
+
+        {
+            std::lock_guard<std::mutex> lock(mutex_);
+            auto [it, inserted] = name_to_group_.emplace(name, std::move(new_group));
+            return it->second.get();
+        }
+    }
+
+private:
     // TODO(dcj): maybe RWLock later?
     mutable std::mutex mutex_;
     std::unordered_map<std::string, std::unique_ptr<ProcessGroup>> name_to_group_;
diff --git a/infini_train/src/nn/parallel/global.cc b/infini_train/src/nn/parallel/global.cc
@@ -14,6 +14,11 @@ int GetEnvAsInt(const std::string &name, int default_value) {
     return value ? std::atoi(value) : default_value;
 }
 
+std::string GetEnvAsStr(const std::string &name, const std::string &default_value) {
+    const char *value = std::getenv(name.c_str());
+    return value ? std::string(value) : default_value;
+}
+
 #ifdef USE_NCCL
 ncclUniqueId StringToNcclId(const std::string &str) {
     ncclUniqueId id;
@@ -120,6 +125,10 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq
     // FIXME(zbl): set PP size
     layout_.sizes[PP] = 1;
     layout_.InitStrides();
+    // FIXME(dcj): what if no nccl id?
+#ifdef USE_NCCL
+    nccl_id_ = StringToNcclId(GetEnvAsStr("NCCL_UNIQUE_ID", ""));
+#endif
 
     initialized_ = true;
 }
@@ -267,5 +276,11 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) {
     oss << "\n";
     return oss.str();
 }
+#ifdef USE_NCCL
+ncclUniqueId GlobalEnv::nccl_id() const {
+    CHECK(initialized_) << "GlobalEnv is not initialized!";
+    return nccl_id_;
+}
+#endif
 
 } // namespace infini_train::nn::parallel::global
diff --git a/infini_train/src/nn/parallel/process_group.cc b/infini_train/src/nn/parallel/process_group.cc
@@ -1,5 +1,6 @@
 #include "infini_train/include/nn/parallel/process_group.h"
 
+#include <memory>
 #include <numeric>
 #include <vector>
 
@@ -40,11 +41,25 @@ const std::unordered_map<ReduceOpType, ncclRedOp_t> kNcclReduceOpMap = {
 namespace infini_train::nn::parallel {
 
 #ifdef USE_NCCL
-ProcessGroup::ProcessGroup(const std::vector<int> &device_indices) : comm_size_(device_indices.size()) {
-    comms_.resize(comm_size_);
-    NCCL_CHECK(ncclCommInitAll(comms_.data(), comm_size_, device_indices.data()));
+ProcessGroup::ProcessGroup(const ncclUniqueId &nccl_id) : world_size_(global::GetWorldSize()) {
+    int local_comm_size = global::GetNthreadPerProc();
+    comms_.resize(local_comm_size);
+    std::vector<int> device_indices(local_comm_size);
 
-    for (int i = 0; i < comm_size_; ++i) {
+    NCCL_CHECK(ncclGroupStart());
+    for (int i = 0; i < local_comm_size; ++i) {
+        device_indices[i] = i;
+
+        int global_rank = global::GetGlobalProcRank() * global::GetNthreadPerProc() + i;
+        NCCL_CHECK(ncclCommInitRank(&comms_[i], world_size_, nccl_id, global_rank));
+    }
+    NCCL_CHECK(ncclGroupEnd());
+
+    Init(device_indices);
+}
+
+void ProcessGroup::Init(const std::vector<int> &device_indices) {
+    for (int i = 0; i < world_size_; ++i) {
         auto device = DeviceManager::Instance()->GetDevice(DeviceType::kCUDA, device_indices[i]);
         devices_.push_back(device);
         device_comm_map_[device] = comms_[i];
@@ -92,7 +107,9 @@ ProcessGroup::BroadCast(const std::vector<std::shared_ptr<Tensor>> &input_tensor
     std::vector<ncclComm_t> comms;
     std::vector<const Device *> devices;
 
-    for (size_t i = 0; i < comm_size_; ++i) {
+    CHECK_EQ(world_size_, comms_.size());
+
+    for (size_t i = 0; i < world_size_; ++i) {
         auto device = devices_[i];
         for (const auto &input_tensor : input_tensors) {
             outputs.push_back(std::make_shared<Tensor>(input_tensor->Dims(), input_tensor->Dtype(), device));
@@ -323,31 +340,20 @@ ProcessGroupFactory *ProcessGroupFactory::Instance() {
 }
 
 const ProcessGroup *ProcessGroupFactory::GetOrCreate(const std::string &name, int comm_size) {
-    std::vector<int> devices(comm_size);
-    std::iota(devices.begin(), devices.end(), 0);
-    const std::vector<int> &device_indices = devices;
-
-    return GetOrCreate(name, device_indices);
+    std::vector<int> device_indices(comm_size);
+    std::iota(device_indices.begin(), device_indices.end(), 0);
+    return GetOrCreate(name, [&]() { return std::make_unique<ProcessGroup>(device_indices); });
 }
 
 const ProcessGroup *ProcessGroupFactory::GetOrCreate(const std::string &name, const std::vector<int> &device_indices) {
-    {
-        std::lock_guard<std::mutex> lock(mutex_);
-        auto it = name_to_group_.find(name);
-        if (it != name_to_group_.end()) {
-            return it->second.get();
-        }
-    }
-
-    auto new_group = std::make_unique<ProcessGroup>(device_indices);
-
-    {
-        std::lock_guard<std::mutex> lock(mutex_);
+    return GetOrCreate(name, [&]() { return std::make_unique<ProcessGroup>(device_indices); });
+}
 
-        auto [it, inserted] = name_to_group_.emplace(name, std::move(new_group));
-        return it->second.get();
-    }
+#ifdef USE_NCCL
+const ProcessGroup *ProcessGroupFactory::GetOrCreate(const std::string &name, const ncclUniqueId &nccl_id) {
+    return GetOrCreate(name, [&]() { return std::make_unique<ProcessGroup>(nccl_id); });
 }
+#endif
 
 const ProcessGroup *ProcessGroupFactory::Get(const std::string &name) const {
     std::lock_guard<std::mutex> lock(mutex_);
@@ -358,5 +364,11 @@ const ProcessGroup *ProcessGroupFactory::GetDefaultProcessGroup() const {
     return name_to_group_.at(kDefaltProcessGroupName).get();
 }
 
-ProcessGroupFactory::ProcessGroupFactory() { GetOrCreate(kDefaltProcessGroupName, global::GetWorldSize()); }
+ProcessGroupFactory::ProcessGroupFactory() {
+#ifdef USE_NCCL
+    GetOrCreate(kDefaltProcessGroupName, global::GetNcclId());
+#else
+    GetOrCreate(kDefaltProcessGroupName, global::GetWorldSize());
+#endif
+}
 } // namespace infini_train::nn::parallel
diff --git a/tools/infini_run/CMakeLists.txt b/tools/infini_run/CMakeLists.txt
@@ -1,2 +1,5 @@
 add_executable(infini_run infini_run.cc)
 target_link_libraries(infini_run PRIVATE gflags glog)
+if (USE_NCCL)
+    target_link_libraries(infini_run PRIVATE nccl)
+endif()
diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc
@@ -25,7 +25,7 @@ DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint (host:port)
 std::string NcclIdToString(const ncclUniqueId& id) {
     std::ostringstream oss;
     for (int i = 0; i < NCCL_UNIQUE_ID_BYTES; ++i) {
-        oss << std::hex << std::uppercase << (int)(unsigned char)id.internal[i];
+        oss << std::hex << std::uppercase << std::setw(2) << std::setfill('0') << (int)(unsigned char)id.internal[i];
     }
     return oss.str();
 }
@@ -99,5 +99,11 @@ int main(int argc, char **argv) {
         wait(&status);
     }
 
+#ifdef USE_NCCL
+    if (FLAGS_node_rank == 0) {
+        std::remove(nccl_id_path);
+    }
+#endif
+
     return 0;
 }

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint (host:port)`
`25`	`25`	`std::string NcclIdToString(const ncclUniqueId& id) {`
`26`	`26`	`std::ostringstream oss;`
`27`	`27`	`for (int i = 0; i < NCCL_UNIQUE_ID_BYTES; ++i) {`
`28`		`- oss << std::hex << std::uppercase << (int)(unsigned char)id.internal[i];`
	`28`	`+ oss << std::hex << std::uppercase << std::setw(2) << std::setfill('0') << (int)(unsigned char)id.internal[i];`
`29`	`29`	`}`
`30`	`30`	`return oss.str();`
`31`	`31`	`}`
`@@ -99,5 +99,11 @@ int main(int argc, char **argv) {`
`99`	`99`	`wait(&status);`
`100`	`100`	`}`
`101`	`101`
	`102`	`+#ifdef USE_NCCL`
	`103`	`+ if (FLAGS_node_rank == 0) {`
	`104`	`+ std::remove(nccl_id_path);`
	`105`	`+ }`
	`106`	`+#endif`
	`107`	`+`
`102`	`108`	`return 0;`
`103`	`109`	`}`