fix: fix grad accumulation/overlap under ZeRO-2

example/gpt2/main.cc

@@ -290,11 +290,20 @@ void Train(const nn::parallel::Rank &rank) {
                                                                  pp_rank, device, model_config.GetChunkSize());
         if (ddp_world_size > 1) {
             auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage};
-            auto *mutable_chunks = dynamic_cast<nn::parallel::PipelineParallel *>(model.get())->mutable_chunks();
+            auto *pipeline_model = dynamic_cast<nn::parallel::PipelineParallel *>(model.get());
+            auto *mutable_chunks = pipeline_model->mutable_chunks();
             for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) {
                 (*mutable_chunks)[chunk_id]
                     = std::make_shared<DistributedDataParallel>(mutable_chunks->at(chunk_id), rank, ddp_config);
             }
+            if (FLAGS_zero_stage >= 1) {
+                pipeline_model->SetNoSyncFunc([mutable_chunks] {
+                    std::vector<std::unique_ptr<nn::NoSyncGuard>> guards;
+                    guards.reserve(mutable_chunks->size());
+                    for (const auto &chunk : *mutable_chunks) { guards.push_back(chunk->no_sync()); }
+                    return guards;
+                });
+            }
         }
     } else if (ddp_world_size > 1) {
         // NOTE(dcj): Complete all device (.to(device)) and dtype (.to(dtype)) conversions
@@ -486,6 +495,10 @@ void Train(const nn::parallel::Rank &rank) {
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward";
 
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward";
+                std::unique_ptr<nn::NoSyncGuard> no_sync_guard;
+                if (ddp_world_size > 1 && FLAGS_zero_stage >= 1 && micro_step != grad_accum_steps - 1) {
+                    no_sync_guard = model->no_sync();
+                }
                 loss->Backward();
                 // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA
                 // between forward and backward.

example/llama3/main.cc

@@ -263,11 +263,20 @@ void Train(const nn::parallel::Rank &rank) {
                                                                  pp_rank, device, model_config.GetChunkSize());
         if (ddp_world_size > 1) {
             auto ddp_config = DistributedDataParallelConfig{.zero_stage = FLAGS_zero_stage};
-            auto *mutable_chunks = dynamic_cast<nn::parallel::PipelineParallel *>(model.get())->mutable_chunks();
+            auto *pipeline_model = dynamic_cast<nn::parallel::PipelineParallel *>(model.get());
+            auto *mutable_chunks = pipeline_model->mutable_chunks();
             for (int chunk_id = 0; chunk_id < mutable_chunks->size(); ++chunk_id) {
                 (*mutable_chunks)[chunk_id]
                     = std::make_shared<DistributedDataParallel>(mutable_chunks->at(chunk_id), rank, ddp_config);
             }
+            if (FLAGS_zero_stage >= 1) {
+                pipeline_model->SetNoSyncFunc([mutable_chunks] {
+                    std::vector<std::unique_ptr<nn::NoSyncGuard>> guards;
+                    guards.reserve(mutable_chunks->size());
+                    for (const auto &chunk : *mutable_chunks) { guards.push_back(chunk->no_sync()); }
+                    return guards;
+                });
+            }
         }
     } else if (ddp_world_size > 1) {
         // NOTE(dcj): Complete all device (.to(device)) and dtype (.to(dtype)) conversions
@@ -465,6 +474,10 @@ void Train(const nn::parallel::Rank &rank) {
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": finish loss forward";
 
                 LOG(INFO) << "Rank " << rank.GlobalRank() << ": start backward";
+                std::unique_ptr<nn::NoSyncGuard> no_sync_guard;
+                if (ddp_world_size > 1 && FLAGS_zero_stage >= 1 && micro_step != grad_accum_steps - 1) {
+                    no_sync_guard = model->no_sync();
+                }
                 loss->Backward();
                 // Defer the loss D2H copy until after backward; reading it earlier would synchronize CUDA
                 // between forward and backward.

infini_train/include/nn/modules/module.h

@@ -20,6 +20,22 @@ template <typename HookType> class HookHandleImpl;
 namespace infini_train::nn {
 class Module;
 
+class NoSyncGuard {
+public:
+    explicit NoSyncGuard(std::function<void()> exit_func) : exit_func_(std::move(exit_func)) {}
+    ~NoSyncGuard() {
+        if (exit_func_) {
+            exit_func_();
+        }
+    }
+
+    NoSyncGuard(const NoSyncGuard &) = delete;
+    NoSyncGuard &operator=(const NoSyncGuard &) = delete;
+
+private:
+    std::function<void()> exit_func_;
+};
+
 namespace parallel::function {
 std::vector<std::shared_ptr<Module>> Replicate(const std::shared_ptr<Module> &network,
                                                const std::vector<Device> &devices);
@@ -82,6 +98,10 @@ class Module : public std::enable_shared_from_this<Module> {
         return 0.0f;
     };
 
+    virtual std::unique_ptr<NoSyncGuard> no_sync() {
+        return std::make_unique<NoSyncGuard>([] {});
+    }
+
     virtual void To(Device device);
 
     virtual void To(DataType dtype);

infini_train/include/nn/parallel/ddp/distributed_data_parallel.h

@@ -31,6 +31,8 @@ class DistributedDataParallel : public nn::Module {
 
     std::shared_ptr<nn::Module> module() const;
 
+    std::unique_ptr<nn::NoSyncGuard> no_sync() override;
+
     DistributedDataParallelConfig ddp_config() const { return ddp_config_; }
 
     const std::vector<std::shared_ptr<ParamAndGradBuffer>> &param_grad_buffers() const { return param_grad_buffers_; }
@@ -41,6 +43,7 @@ class DistributedDataParallel : public nn::Module {
     void BuildParamAndGradBuffers();
     void RegisterBackwardHooks();
     void OnGradReady(const std::shared_ptr<Tensor> &param);
+    void SetIsLastMicrobatch(bool is_last_microbatch);
 
 private:
     std::shared_ptr<Reducer> reducer_ = nullptr;

infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h

@@ -97,6 +97,8 @@ class ParamAndGradBucketGroup {
     // When all params in a bucket group are ready, will call StartGradSync()
     void RegisterGradReady(const std::shared_ptr<Tensor> &parameter);
 
+    void SetIsLastMicrobatch(bool is_last_microbatch) { is_last_microbatch_ = is_last_microbatch; }
+
     // Start grad reduce
     void StartGradSync();
 

infini_train/include/nn/parallel/pp/pipeline_parallel.h

@@ -1,6 +1,7 @@
 // pipeline_parallel.h
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <vector>
 
@@ -40,6 +41,8 @@ class PipelineParallel : public Module {
 
     std::vector<std::shared_ptr<Module>> *mutable_chunks();
 
+    void SetNoSyncFunc(std::function<std::vector<std::unique_ptr<nn::NoSyncGuard>>()> func);
+
 private:
     void BuildPipelineStage(const std::vector<std::vector<int64_t>> &recv_shape, Device device,
                             std::vector<std::shared_ptr<Module>> &&chunks);

infini_train/include/nn/parallel/pp/pipeline_schedule.h

@@ -1,9 +1,11 @@
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <vector>
 
 #include "infini_train/include/datatype.h"
+#include "infini_train/include/nn/modules/module.h"
 
 namespace infini_train {
 class Tensor;
@@ -31,12 +33,17 @@ class PipelineSchedule {
                                    const std::vector<std::shared_ptr<Tensor>> &target_mbs,
                                    const std::shared_ptr<nn::Module> &loss_fn, DataType dtype);
 
+    using NoSyncFunc = std::function<std::vector<std::unique_ptr<nn::NoSyncGuard>>()>;
+
+    void SetNoSyncFunc(NoSyncFunc func) { no_sync_func_ = std::move(func); }
+
     std::vector<std::shared_ptr<Tensor>> ReceiveFromPrev(int peer_rank);
     std::vector<std::shared_ptr<Tensor>> SendToNext(const std::vector<std::shared_ptr<Tensor>> &tensors, int peer_rank);
 
 protected:
     int num_micro_batches_ = -1;
     std::shared_ptr<PipelineStage> stage_ = nullptr;
+    NoSyncFunc no_sync_func_;
 };
 
 class PipelineParallelScheduler {

infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc

@@ -216,4 +216,13 @@ DistributedDataParallel::Forward(const std::vector<std::shared_ptr<Tensor>> &inp
 }
 
 std::shared_ptr<nn::Module> DistributedDataParallel::module() const { return modules_.at(kModuleName); }
+
+std::unique_ptr<nn::NoSyncGuard> DistributedDataParallel::no_sync() {
+    SetIsLastMicrobatch(false);
+    return std::make_unique<nn::NoSyncGuard>([this] { SetIsLastMicrobatch(true); });
+}
+
+void DistributedDataParallel::SetIsLastMicrobatch(bool is_last_microbatch) {
+    for (auto &group : bucket_groups_) { group->SetIsLastMicrobatch(is_last_microbatch); }
+}
 } // namespace infini_train::nn::parallel

infini_train/src/nn/parallel/ddp/param_and_grad_buffer.cc

@@ -154,19 +154,23 @@ void ParamAndGradBucketGroup::RegisterGradReady(const std::shared_ptr<Tensor> &p
         return;
     }
 
-    // TODO(zbl): Only register grads as ready and trigger grad sync when processing the last microbatch
-    //            For now, is_last_microbatch_ is always true
+    // Only the last microbatch registers ready grads so the reduce can overlap with its backward pass.
     if (is_last_microbatch_) {
         if (!parameter || params_.find(parameter.get()) == params_.end()) {
             return;
         }
 
-        params_with_grad_.insert(parameter.get());
-        // TODO(zbl): check this if sync is only done in last mircobatch
-        // if (!inserted) {
-        //     LOG(FATAL) << "ParamAndGradBucketGroup: RegisterGradReady() was called twice for the same parameter in a
-        //     bucket group."; return;
-        // }
+        if (grad_reduce_dispatched_) {
+            LOG(FATAL) << "ParamAndGradBucketGroup: RegisterGradReady() was called after grad sync was dispatched.";
+            return;
+        }
+
+        auto [_, inserted] = params_with_grad_.insert(parameter.get());
+        if (!inserted) {
+            LOG(FATAL) << "ParamAndGradBucketGroup: RegisterGradReady() was called twice for the same parameter in a "
+                          "bucket group.";
+            return;
+        }
 
         if (params_with_grad_.size() == params_.size()) {
             // All param grads are ready in this group, trigger grad sync
@@ -297,7 +301,6 @@ void ParamAndGradBucketGroup::StartGradSync() {
     }
 
     grad_reduce_dispatched_ = true;
-    // TODO(zbl): no need to clear params_with_grad_ here if grad sync is only done on last microbatch
     params_with_grad_.clear();
 }
 

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

@@ -104,4 +104,8 @@ PipelineParallel::PipelineParallel(const std::shared_ptr<Module> module, int num
 }
 
 std::vector<std::shared_ptr<Module>> *PipelineParallel::mutable_chunks() { return pipeline_stage_->mutable_chunks(); }
+
+void PipelineParallel::SetNoSyncFunc(std::function<std::vector<std::unique_ptr<nn::NoSyncGuard>>()> func) {
+    schedule_->SetNoSyncFunc(std::move(func));
+}
 } // namespace infini_train::nn::parallel

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

@@ -212,6 +212,12 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
     std::vector<std::vector<std::vector<std::shared_ptr<Tensor>>>> activations(
         vpp_size, std::vector<std::vector<std::shared_ptr<Tensor>>>(n));
 
+    std::vector<std::unique_ptr<nn::NoSyncGuard>> no_sync_guards;
+    if (no_sync_func_) {
+        no_sync_guards = no_sync_func_();
+    }
+    std::vector<int> backward_counts(vpp_size, 0);
+
     for (size_t i = 0; i < schedule.size(); ++i) {
         const auto &task = schedule[i];
         if (task.stage_id != stage_idx) {
@@ -244,6 +250,10 @@ float PipelineSchedule::StepMicroBatches(const std::vector<std::shared_ptr<Tenso
                 }
             }
         } else {
+            const bool is_last_microbatch = ++backward_counts[task.local_chunk_idx] == n;
+            if (is_last_microbatch && no_sync_func_) {
+                no_sync_guards[task.local_chunk_idx].reset();
+            }
             if (task.is_last_chunk) {
                 auto target = microbatch_targets[mb];
                 std::shared_ptr<Tensor> loss;