feat: support multi-node pp training

Author: kilinchange

example/gpt2/main.cc

@@ -29,7 +29,6 @@
 
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
-#include "example/common/utils.h"
 #include "example/gpt2/net.h"
 
 // I/O
@@ -321,7 +320,7 @@ void Train(const nn::parallel::Rank &rank) {
         const double duration_us = std::chrono::duration<double, std::micro>(iter_end - iter_start).count();
         const double tps = FLAGS_total_batch_size / (duration_us / 1e6);
 
-        if (rank.thread_rank() == pp_world_size - 1) {
+        if (rank.GlobalRank() == pp_world_size - 1) {
             LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s, "
                                       "DP={}, TP={}, SP={}, PP={})",
                                       step + 1, FLAGS_num_iteration, lossf, FLAGS_learning_rate, duration_us / 1e3f,
@@ -340,6 +339,10 @@ void Train(const nn::parallel::Rank &rank) {
     Profiler::Instance().Report("gpt2.report", Profiler::SortBy::DeviceTimePercentage);
     Profiler::Instance().PrintRecords("gpt2.records.log");
 #endif
+
+    if (pp_world_size > 1 && rank.IsMainRank()) {
+        pp_pg->Barrier();
+    }
 }
 
 int main(int argc, char *argv[]) {

example/llama3/main.cc

@@ -28,7 +28,6 @@
 
 #include "example/common/tiny_shakespeare_dataset.h"
 #include "example/common/tokenizer.h"
-#include "example/common/utils.h"
 #include "example/llama3/net.h"
 
 // I/O
@@ -124,8 +123,8 @@ void Train(const nn::parallel::Rank &rank) {
         }
 
         if (pp_world_size > 1) {
-            pp_pg = ProcessGroupFactory::Instance()->GetOrCreate(
-                GetPipelineParallelProcessGroupName(rank.thread_rank()), GetPipelineParallelGroupRanks(pp_world_size));
+            pp_pg = ProcessGroupFactory::Instance()->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()),
+                                                                 GetPipelineParallelGroupRanks(rank.GlobalRank()));
             pp_rank = pp_pg->GetGroupRank(rank.thread_rank());
 
             nn::parallel::pp_rank = pp_rank;
@@ -299,7 +298,7 @@ void Train(const nn::parallel::Rank &rank) {
         const double duration_us = std::chrono::duration<double, std::micro>(iter_end - iter_start).count();
         const double tps = FLAGS_total_batch_size / (duration_us / 1e6);
 
-        if (rank.thread_rank() == pp_world_size - 1) {
+        if (rank.GlobalRank() == pp_world_size - 1) {
             LOG(ERROR) << std::format("step {:4d}/{} | train loss {:.6f} | lr {:.2e} | ({:.2f} ms | {:.0f} tok/s, "
                                       "DP={}, TP={}, SP={}, PP={})",
                                       step + 1, FLAGS_num_iteration, lossf, FLAGS_learning_rate, duration_us / 1e3f,
@@ -318,6 +317,10 @@ void Train(const nn::parallel::Rank &rank) {
     Profiler::Instance().Report("llama3.report", Profiler::SortBy::DeviceTimePercentage);
     Profiler::Instance().PrintRecords("llama3.records.log");
 #endif
+
+    if (pp_world_size > 1 && rank.IsMainRank()) {
+        pp_pg->Barrier();
+    }
 }
 
 int main(int argc, char *argv[]) {

infini_train/include/nn/parallel/process_group.h

@@ -55,6 +55,8 @@ class ProcessGroup {
 
     std::vector<std::shared_ptr<Tensor>> NcclRecv(std::vector<std::shared_ptr<Tensor>> tensors, int src_rank) const;
 
+    void Barrier() const;
+
 private:
     void InitSingleProcess(const std::vector<int> &ranks);
 

infini_train/include/nn/parallel/utils.h

@@ -14,5 +14,5 @@ std::vector<int> GetDataParallelGroupRanks(int rank);
 
 std::vector<int> GetTensorParallelGroupRanks(int rank);
 
-std::vector<int> GetPipelineParallelGroupRanks(int pp_world_size);
+std::vector<int> GetPipelineParallelGroupRanks(int rank);
 } // namespace infini_train::nn::parallel

infini_train/src/nn/parallel/global.cc

@@ -111,8 +111,7 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq
 
     layout_.sizes[DP] = data_parallel_size_;
     layout_.sizes[TP] = tensor_parallel_size_;
-    // FIXME(zbl): set PP size
-    layout_.sizes[PP] = 1;
+    layout_.sizes[PP] = pipeline_parallel_size_;
     layout_.InitStrides();
 
     initialized_ = true;

infini_train/src/nn/parallel/pp/pipeline_schedule.cc

@@ -2,12 +2,10 @@
 #include "infini_train/include/nn/parallel/pp/pipeline_schedule.h"
 
 #include <cstddef>
-#include <cstdint>
 #include <vector>
 
 #include "glog/logging.h"
 
-#include "infini_train/include/autograd/grad_mode.h"
 #include "infini_train/include/device.h"
 #include "infini_train/include/nn/init.h"
 #include "infini_train/include/nn/modules/module.h"
@@ -90,9 +88,11 @@ float ScheduleGPipe::StepMicroBatches(const std::vector<std::shared_ptr<Tensor>>
         for (int mb = 0; mb < n; ++mb) {
             auto out_tensor = outputs[mb][0];
 
-            auto gradient = std::make_shared<Tensor>(out_tensor->Dims(), out_tensor->Dtype(), out_tensor->GetDevice());
+            auto dummy_gradient
+                = std::make_shared<Tensor>(out_tensor->Dims(), out_tensor->Dtype(), out_tensor->GetDevice());
 
-            out_tensor->Backward(gradient);
+            out_tensor->Backward(dummy_gradient);
+            cudaStreamSynchronize(dynamic_cast<const CudaDevice *>(stage_->device())->Stream());
         }
     } else {
         for (int mb = 0; mb < n; ++mb) {

infini_train/src/nn/parallel/pp/pipeline_stage.cc

@@ -15,7 +15,8 @@ PipelineStage::PipelineStage(const std::shared_ptr<Module> &model, int stage_ind
       prev_rank_(stage_index > 0 ? stage_index - 1 : -1),
       next_rank_(stage_index < num_stages - 1 ? stage_index + 1 : -1), recv_shape_(recv_shape),
       optimizer_(std::move(optimizer)),
-      device_(DeviceManager::Instance()->GetAllAvailableDevices(DeviceType::kCUDA).at(stage_index)) {}
+      // FIXME(dcj): use correct device
+      device_(DeviceManager::Instance()->GetAllAvailableDevices(DeviceType::kCUDA).at(stage_index % 8)) {}
 
 std::vector<std::shared_ptr<Tensor>>
 PipelineStage::ForwardOneChunk(const std::vector<std::shared_ptr<Tensor>> &inputs) {

infini_train/src/nn/parallel/process_group.cc

@@ -411,6 +411,21 @@ std::vector<std::shared_ptr<Tensor>> ProcessGroup::NcclRecv(std::vector<std::sha
     }
     return tensors;
 }
+
+void ProcessGroup::Barrier() const {
+    // NOTE(dcj): use ncclAllreduce to barrier all processes before destroying the communicators
+    // FIXME(dcj): should only call by one rank
+    int dummy = 1;
+    std::vector<int> results(1, 0);
+
+    NCCL_CHECK(ncclGroupStart());
+    for (const auto &device : devices_) {
+        auto comm = device_comm_map_.at(device);
+        auto cuda_dev = dynamic_cast<const CudaDevice *>(device);
+        NCCL_CHECK(ncclAllReduce(&dummy, &dummy, 1, ncclInt, ncclSum, comm, cuda_dev->Stream()));
+    }
+    NCCL_CHECK(ncclGroupEnd());
+}
 #endif
 
 ProcessGroupFactory *ProcessGroupFactory::Instance() {

infini_train/src/nn/parallel/utils.cc

@@ -20,10 +20,7 @@ std::vector<int> GetDataParallelGroupRanks(int thread_rank) { return global::Get
 
 std::vector<int> GetTensorParallelGroupRanks(int thread_rank) { return global::GetGroupRanks(global::TP, thread_rank); }
 
-std::vector<int> GetPipelineParallelGroupRanks(int pp_world_size) {
-    std::vector<int> ranks;
-    ranks.reserve(pp_world_size);
-    for (int i = 0; i < pp_world_size; ++i) { ranks.push_back(i); }
-    return ranks;
+std::vector<int> GetPipelineParallelGroupRanks(int thread_rank) {
+    return global::GetGroupRanks(global::PP, thread_rank);
 }
 } // namespace infini_train::nn::parallel