feat: Pipeline parallelism divides the model into chunks during construction

JYMiracle305 · JYMiracle305 · commit 9c47bc25d915 · 2025-11-12T16:55:48.000+08:00
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
@@ -148,6 +148,8 @@ void Train(const nn::parallel::Rank &rank) {
             pp_pg = ProcessGroupFactory::Instance()->GetOrCreate(
                 GetPipelineParallelProcessGroupName(rank.thread_rank()), GetPipelineParallelGroupRanks(pp_world_size));
             pp_rank = pp_pg->GetGroupRank(rank.thread_rank());
+
+            nn::parallel::pp_rank = pp_rank;
         }
     } else {
         device = FLAGS_device == kDeviceCPU ? DeviceManager::Instance()->GetDefaultDevice()
@@ -243,8 +245,9 @@ void Train(const nn::parallel::Rank &rank) {
         CHECK_EQ((FLAGS_batch_size * pp_world_size) % FLAGS_num_microbatches, 0)
             << "FLAGS_batch_size (" << (FLAGS_batch_size * pp_world_size)
             << ") must be divisible by FLAGS_num_microbatches (" << FLAGS_num_microbatches << ")";
-        auto shapes = std::vector<std::vector<int64_t>>{{(FLAGS_batch_size * pp_world_size) / FLAGS_num_microbatches,
-                                                         FLAGS_sequence_length, model->GetConfig()["n_embd"]}};
+
+        auto shapes = std::vector<std::vector<int64_t>>{
+            {(FLAGS_batch_size * pp_world_size) / FLAGS_num_microbatches, FLAGS_sequence_length, model_config.n_embd}};
 
         model = std::make_shared<nn::parallel::PipelineParallel>(model, pp_world_size, FLAGS_num_microbatches, shapes,
                                                                  pp_rank, optimizer_factory);
@@ -298,9 +301,9 @@ void Train(const nn::parallel::Rank &rank) {
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 
+            // FIXME(jym): without gradient accumulation
             if (pp_world_size > 1) {
                 lossf = model->TrainStep({x}, {y}, loss_fn);
-
                 auto loss_tensor = std::make_shared<Tensor>(std::vector<int64_t>{}, DataType::kFLOAT32);
                 static_cast<float *>(loss_tensor->DataPtr())[0] = lossf;
                 auto loss_device_ptr = std::make_shared<Tensor>(loss_tensor->To(device));
diff --git a/example/gpt2/net.cc b/example/gpt2/net.cc
diff --git a/example/gpt2/net.h b/example/gpt2/net.h
@@ -89,9 +89,6 @@ class GPT2 : public infini_train::nn::CloneableModule<GPT2> {
 
     explicit GPT2(const GPT2Config &config);
 
-    std::unordered_map<std::string, int64_t> GetConfig() const override;
-    std::vector<std::shared_ptr<infini_train::nn::Module>> GetPipelineLayers() override;
-
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
@@ -130,6 +130,8 @@ void Train(const nn::parallel::Rank &rank) {
             pp_pg = ProcessGroupFactory::Instance()->GetOrCreate(
                 GetPipelineParallelProcessGroupName(rank.thread_rank()), GetPipelineParallelGroupRanks(pp_world_size));
             pp_rank = pp_pg->GetGroupRank(rank.thread_rank());
+
+            nn::parallel::pp_rank = pp_rank;
         }
     } else {
         device = FLAGS_device == kDeviceCPU ? DeviceManager::Instance()->GetDefaultDevice()
@@ -222,8 +224,8 @@ void Train(const nn::parallel::Rank &rank) {
             << "FLAGS_batch_size (" << (FLAGS_batch_size * pp_world_size)
             << ") must be divisible by FLAGS_num_microbatches (" << FLAGS_num_microbatches << ")";
 
-        auto shapes = std::vector<std::vector<int64_t>>{{FLAGS_batch_size * pp_world_size / FLAGS_num_microbatches,
-                                                         FLAGS_sequence_length, model->GetConfig()["n_embd"]}};
+        auto shapes = std::vector<std::vector<int64_t>>{
+            {FLAGS_batch_size * pp_world_size / FLAGS_num_microbatches, FLAGS_sequence_length, model_config.n_embd}};
 
         model = std::make_shared<nn::parallel::PipelineParallel>(model, pp_world_size, FLAGS_num_microbatches, shapes,
                                                                  pp_rank, optimizer_factory);
diff --git a/example/llama3/net.cc b/example/llama3/net.cc
diff --git a/example/llama3/net.h b/example/llama3/net.h
@@ -129,10 +129,6 @@ class LLaMA3 : public infini_train::nn::CloneableModule<LLaMA3> {
 
     explicit LLaMA3(const LLaMA3Config &config);
 
-    std::vector<std::shared_ptr<infini_train::nn::Module>> GetPipelineLayers() override;
-
-    std::unordered_map<std::string, int64_t> GetConfig() const;
-
     std::vector<std::shared_ptr<infini_train::Tensor>>
     Forward(const std::vector<std::shared_ptr<infini_train::Tensor>> &x) override;
 
diff --git a/infini_train/include/nn/modules/container.h b/infini_train/include/nn/modules/container.h
@@ -5,6 +5,7 @@
 #include <vector>
 
 #include "infini_train/include/nn/modules/module.h"
+#include <iostream>
 
 namespace infini_train {
 class Tensor;
@@ -53,4 +54,24 @@ class ModuleList : public CloneableModule<ModuleList> {
 private:
     std::vector<std::shared_ptr<Module>> module_list_;
 };
+
+class PipelineModuleList : public CloneableModule<PipelineModuleList> {
+public:
+    static constexpr char kType[] = "PipelineModuleList";
+
+    explicit PipelineModuleList(std::vector<std::shared_ptr<nn::Module>> &&modules, int64_t global_start_index);
+
+    std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const override;
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+
+    auto begin() { return local_modules_.begin(); }
+    auto end() { return local_modules_.end(); }
+    auto begin() const { return local_modules_.begin(); }
+    auto end() const { return local_modules_.end(); }
+
+private:
+    std::vector<std::shared_ptr<nn::Module>> local_modules_;
+    int64_t global_start_;
+};
 } // namespace infini_train::nn
diff --git a/infini_train/include/nn/modules/module.h b/infini_train/include/nn/modules/module.h
@@ -43,12 +43,10 @@ class Module : public std::enable_shared_from_this<Module> {
     std::shared_ptr<Module> mutable_module(const std::string &name);
     const Module &module(const std::string &name) const;
 
-    std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const;
+    virtual std::unordered_map<std::string, std::shared_ptr<Tensor>> StateDict() const;
 
     virtual std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors);
 
-    virtual std::vector<std::shared_ptr<Module>> GetPipelineLayers() { return {}; }
-
     virtual std::unordered_map<std::string, int64_t> GetConfig() const { return {}; }
 
     virtual float TrainStep(const std::vector<std::shared_ptr<Tensor>> &input_tensors,
diff --git a/infini_train/include/nn/parallel/pp/pipeline_parallel.h b/infini_train/include/nn/parallel/pp/pipeline_parallel.h
@@ -13,6 +13,8 @@
 
 namespace infini_train::nn::parallel {
 
+extern thread_local int pp_rank;
+
 using OptimizerFactory = std::function<std::shared_ptr<Optimizer>(const std::vector<std::shared_ptr<Tensor>> &params)>;
 
 class PipelineParallel : public Module {
@@ -26,22 +28,14 @@ class PipelineParallel : public Module {
 private:
     int num_stages_;
     int rank_;
-    std::vector<const Device *> devices_;
     std::shared_ptr<Module> original_model_;
     std::shared_ptr<PipelineStage> pipeline_stage_;
     std::shared_ptr<PipelineSchedule> schedule_;
 
-    std::vector<std::vector<std::shared_ptr<Module>>>
-    SplitLayersIntoStages(std::vector<std::shared_ptr<Module>> layers);
-
-    void SplitModel(const std::vector<std::vector<int64_t>> &recv_shape, OptimizerFactory optimizer_factory);
-
-    std::vector<std::shared_ptr<Optimizer>>
-    CreateOptimizers(const std::vector<std::vector<std::shared_ptr<Module>>> &stage_layers,
-                     OptimizerFactory optimizer_factory);
+    std::shared_ptr<Optimizer> CreateOptimizer(const std::shared_ptr<Module> &model,
+                                               OptimizerFactory optimizer_factory);
 
-    void BuildPipelineStage(const std::vector<std::vector<std::shared_ptr<Module>>> &stage_layers,
-                            const std::vector<std::shared_ptr<Optimizer>> &optimizers,
+    void BuildPipelineStage(const std::shared_ptr<Module> &model, const std::shared_ptr<Optimizer> &optimizers,
                             const std::vector<std::vector<int64_t>> &recv_shape);
 
     void SetupSchedule(int num_microbatches);
diff --git a/infini_train/include/nn/parallel/pp/pipeline_stage.h b/infini_train/include/nn/parallel/pp/pipeline_stage.h
@@ -12,7 +12,7 @@ namespace infini_train::nn::parallel {
 
 class PipelineStage {
 public:
-    PipelineStage(const std::vector<std::shared_ptr<Module>> &layers, int stage_index, int num_stages,
+    PipelineStage(const std::shared_ptr<Module> &model, int stage_index, int num_stages,
                   const std::vector<std::vector<int64_t>> &recvShape, std::shared_ptr<Optimizer> optim);
 
     std::vector<std::shared_ptr<Tensor>> ForwardOneChunk(const std::vector<std::shared_ptr<Tensor>> &inputs);
@@ -34,7 +34,7 @@ class PipelineStage {
     int next_rank_;
     const Device *device_ = nullptr;
     std::vector<std::vector<int64_t>> recv_shape_;
-    std::vector<std::shared_ptr<Module>> layers_;
+    std::shared_ptr<Module> model_;
     std::shared_ptr<Optimizer> optim_;
 };
 
diff --git a/infini_train/src/nn/modules/container.cc b/infini_train/src/nn/modules/container.cc
@@ -41,4 +41,29 @@ ModuleList::ModuleList(std::vector<std::shared_ptr<Module>> &&layers)
 std::vector<std::shared_ptr<Tensor>> ModuleList::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
     LOG(FATAL) << "Not implemented";
 }
+
+PipelineModuleList::PipelineModuleList(std::vector<std::shared_ptr<nn::Module>> &&modules, int64_t global_start_index)
+    : CloneableModule(kType), local_modules_(std::move(modules)), global_start_(global_start_index) {
+    int idx = 0;
+    for (auto &layer : local_modules_) {
+        modules_[std::to_string(idx)] = layer;
+        ++idx;
+    }
+}
+
+std::unordered_map<std::string, std::shared_ptr<Tensor>> PipelineModuleList::StateDict() const {
+    // std::cout << "PipelineModuleList::StateDict() called!" << std::endl;
+    std::unordered_map<std::string, std::shared_ptr<Tensor>> state;
+    for (size_t i = 0; i < local_modules_.size(); ++i) {
+        int64_t global_i = global_start_ + i;
+        auto sub_dict = local_modules_[i]->StateDict();
+        for (auto &[sub_name, tensor] : sub_dict) { state.emplace(std::to_string(global_i) + "." + sub_name, tensor); }
+    }
+    return state;
+}
+
+std::vector<std::shared_ptr<Tensor>>
+PipelineModuleList::Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) {
+    LOG(FATAL) << "Not implemented";
+}
 } // namespace infini_train::nn
diff --git a/infini_train/src/nn/parallel/pp/pipeline_parallel.cc b/infini_train/src/nn/parallel/pp/pipeline_parallel.cc
@@ -12,69 +12,17 @@
 
 namespace infini_train::nn::parallel {
 
-std::vector<std::vector<std::shared_ptr<Module>>>
-PipelineParallel::SplitLayersIntoStages(std::vector<std::shared_ptr<Module>> layers) {
-    const int total_layers = layers.size();
-    CHECK_GT(total_layers, 0) << "Model has no layers to split!";
-    CHECK_GE(num_stages_, 1) << "num_stages must be >= 1";
-    CHECK_LE(num_stages_, total_layers) << "num_stages (" << num_stages_ << ") cannot be greater than total layers ("
-                                        << total_layers << ")";
+thread_local int pp_rank = 0;
 
-    std::vector<std::vector<std::shared_ptr<Module>>> stages(num_stages_);
-    int base_layers_per_stage = total_layers / num_stages_;
-    int remainder = total_layers % num_stages_;
-    int layer_idx = 0;
-
-    for (int s = 0; s < num_stages_; ++s) {
-        int layers_in_this_stage = base_layers_per_stage + (s < remainder ? 1 : 0);
-        for (int i = 0; i < layers_in_this_stage; ++i) {
-            auto layer = layers[layer_idx];
-            stages[s].emplace_back(layer);
-            layer_idx++;
-        }
-    }
-
-    return stages;
-}
-
-std::vector<std::shared_ptr<Optimizer>>
-PipelineParallel::CreateOptimizers(const std::vector<std::vector<std::shared_ptr<Module>>> &stage_layers,
-                                   OptimizerFactory optimizer_factory) {
-    std::vector<std::shared_ptr<Optimizer>> optims;
-    optims.reserve(stage_layers.size());
-
-    for (int s = 0; s < num_stages_; ++s) {
-        std::vector<std::shared_ptr<Tensor>> params;
-        for (const auto &layer : stage_layers[s]) {
-            layer->To(devices_[s]);
-            auto layer_params = layer->Parameters();
-            params.insert(params.end(), layer_params.begin(), layer_params.end());
-        }
-
-        auto optim = optimizer_factory(params);
-        CHECK(optim != nullptr) << "Optimizer factory returned null optimizer for stage " << s;
-        optims.push_back(std::move(optim));
-    }
-    return optims;
+std::shared_ptr<Optimizer> PipelineParallel::CreateOptimizer(const std::shared_ptr<Module> &model,
+                                                             OptimizerFactory optimizer_factory) {
+    return optimizer_factory(model->Parameters());
 }
 
-void PipelineParallel::BuildPipelineStage(const std::vector<std::vector<std::shared_ptr<Module>>> &stage_layers,
-                                          const std::vector<std::shared_ptr<Optimizer>> &optimizers,
+void PipelineParallel::BuildPipelineStage(const std::shared_ptr<Module> &model,
+                                          const std::shared_ptr<Optimizer> &optimizers,
                                           const std::vector<std::vector<int64_t>> &recv_shape) {
-    pipeline_stage_
-        = std::make_shared<PipelineStage>(stage_layers[rank_], rank_, num_stages_, recv_shape, optimizers[rank_]);
-}
-
-void PipelineParallel::SplitModel(const std::vector<std::vector<int64_t>> &recv_shape,
-                                  OptimizerFactory optimizer_factory) {
-    auto layers = original_model_->GetPipelineLayers();
-    CHECK(!layers.empty()) << "SplitModel: GetPipelineLayers returned empty vector";
-
-    auto stage_layer = SplitLayersIntoStages(layers);
-
-    auto optimizer = CreateOptimizers(stage_layer, optimizer_factory);
-
-    BuildPipelineStage(stage_layer, optimizer, recv_shape);
+    pipeline_stage_ = std::make_shared<PipelineStage>(model, rank_, num_stages_, recv_shape, optimizers);
 }
 
 void PipelineParallel::SetupSchedule(int num_microbatches) {
@@ -96,11 +44,11 @@ float PipelineParallel::TrainStep(const std::vector<std::shared_ptr<Tensor>> &in
 PipelineParallel::PipelineParallel(const std::shared_ptr<Module> &model, int num_stages, int num_microbatches,
                                    const std::vector<std::vector<int64_t>> &recv_shape, int rank,
                                    OptimizerFactory optimizer_factory)
-    : original_model_(model), devices_(DeviceManager::Instance()->GetAllAvailableDevices(DeviceType::kCUDA)),
-      num_stages_(num_stages), rank_(rank) {
-    CHECK(!devices_.empty()) << "Devices list is empty";
+    : original_model_(model), num_stages_(num_stages), rank_(rank) {
+
+    auto optimizer = CreateOptimizer(original_model_, optimizer_factory);
 
-    SplitModel(recv_shape, optimizer_factory);
+    BuildPipelineStage(original_model_, optimizer, recv_shape);
 
     SetupSchedule(num_microbatches);
 }
diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc
@@ -89,15 +89,18 @@ std::vector<std::shared_ptr<Tensor>> PipelineSchedule::ReceiveFromPrev() {
             tensor->set_is_leaf(false);
             recv_tensors.push_back(tensor);
         }
+
         return IRecv(recv_tensors, stage_->device(), stage_->stage_index(), stage_->prev_rank());
     }
+
     return recv_tensors;
 }
 
 std::vector<std::shared_ptr<Tensor>> PipelineSchedule::SendToNext(const std::vector<std::shared_ptr<Tensor>> &tensors) {
     if (!stage_->IsLastStage()) {
         return ISend(tensors, stage_->device(), stage_->stage_index(), stage_->next_rank(), stage_->recv_shape());
     }
+
     return tensors;
 }
 
diff --git a/infini_train/src/nn/parallel/pp/pipeline_stage.cc b/infini_train/src/nn/parallel/pp/pipeline_stage.cc
@@ -5,28 +5,23 @@
 #include <memory>
 
 #include "infini_train/include/device.h"
-#include "infini_train/include/dispatcher.h"
-#include "infini_train/include/nn/init.h"
+// #include "infini_train/include/dispatcher.h"
+// #include "infini_train/include/nn/init.h"
+// #include "infini_train/include/nn/modules/container.h"
 
 namespace infini_train::nn::parallel {
 
-PipelineStage::PipelineStage(const std::vector<std::shared_ptr<Module>> &layers, int stage_index, int num_stages,
+PipelineStage::PipelineStage(const std::shared_ptr<Module> &model, int stage_index, int num_stages,
                              const std::vector<std::vector<int64_t>> &recvShape, std::shared_ptr<Optimizer> optim)
-    : stage_index_(stage_index), num_stages_(num_stages), layers_(layers),
+    : model_(model), stage_index_(stage_index), num_stages_(num_stages),
       prev_rank_(stage_index > 0 ? stage_index - 1 : -1),
       next_rank_(stage_index < num_stages - 1 ? stage_index + 1 : -1), recv_shape_(recvShape), optim_(std::move(optim)),
       device_(DeviceManager::Instance()->GetAllAvailableDevices(DeviceType::kCUDA).at(stage_index)) {}
 
 std::vector<std::shared_ptr<Tensor>>
 PipelineStage::ForwardOneChunk(const std::vector<std::shared_ptr<Tensor>> &inputs) {
-    std::vector<std::shared_ptr<Tensor>> current = inputs;
-    int i = 0;
-    for (const auto &layer : layers_) {
-        current = layer->Forward(current);
-        ++i;
-    }
-
-    return current;
+
+    return model_->Forward(inputs);
 }
 
 } // namespace infini_train::nn::parallel
diff --git a/infini_train/src/nn/parallel/pp/send_recv.cc b/infini_train/src/nn/parallel/pp/send_recv.cc
@@ -91,7 +91,6 @@ std::vector<std::shared_ptr<Tensor>> ISend::Backward(const std::vector<std::shar
 
 std::vector<std::shared_ptr<Tensor>> IRecv::Forward(const std::vector<std::shared_ptr<Tensor>> &recv_tensors) {
     CHECK_NE(src_device_, nullptr) << "src_device_ must be set";
-
     auto pp_group
         = ProcessGroupFactory::Instance()->Get(GetPipelineParallelProcessGroupName(src_device_->rank().thread_rank()));
     pp_group->NcclRecv(recv_tensors, peer_rank_);
diff --git a/infini_train/src/optimizer.cc b/infini_train/src/optimizer.cc
@@ -20,6 +20,11 @@ SGD::SGD(const std::vector<std::shared_ptr<Tensor>> &params, float learning_rate
 
 void SGD::Step() {
     for (auto param : params_) {
+        // FIXME(jym): skip parameters with empty gradients
+        if (!param->grad()) {
+            LOG(INFO) << "Skipping param with null grad.";
+            continue;
+        }
         auto device = param->GetDevice();
         device->SetDevice();
         auto kernel = Dispatcher::Instance().GetKernel({device->Type(), "AccumulateGrad"});
