fix: fix grad accumulation/overlap under ZeRO-2

example/gpt2/main.cc

@@ -290,11 +290,18 @@ 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);
             }
+            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 +493,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 && 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,18 @@ 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);
             }
+            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 +472,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 && 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/autograd/function_hook.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <atomic>
 #include <memory>
 
 #include "infini_train/include/nn/parallel/reduce_op_type.h"
@@ -36,12 +37,14 @@ class PostAccumulateGradHook {
 class AllReducePostAccumulateHook : public PostAccumulateGradHook {
 public:
     AllReducePostAccumulateHook(infini_train::nn::parallel::function::ReduceOpType reduce_op,
-                                const infini_train::nn::parallel::ProcessGroup *pg = nullptr);
+                                const infini_train::nn::parallel::ProcessGroup *pg = nullptr,
+                                std::shared_ptr<const std::atomic_bool> enabled = nullptr);
 
     void operator()(const std::shared_ptr<Tensor> &tensor) override;
 
 private:
     infini_train::nn::parallel::function::ReduceOpType reduce_op_;
     const infini_train::nn::parallel::ProcessGroup *pg_ = nullptr;
+    std::shared_ptr<const std::atomic_bool> enabled_;
 };
 } // namespace infini_train::autograd

infini_train/include/nn/modules/module.h

@@ -20,6 +20,18 @@ template <typename HookType> class HookHandleImpl;
 namespace infini_train::nn {
 class Module;
 
+class NoSyncGuard {
+public:
+    explicit NoSyncGuard(std::function<void()> exit_func);
+    ~NoSyncGuard();
+
+    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 +94,8 @@ class Module : public std::enable_shared_from_this<Module> {
         return 0.0f;
     };
 
+    virtual std::unique_ptr<NoSyncGuard> no_sync();
+
     virtual void To(Device device);
 
     virtual void To(DataType dtype);

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

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <atomic>
 #include <memory>
 
 #include "infini_train/include/nn/modules/module.h"
@@ -31,6 +32,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,9 +44,12 @@ 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;
+    // Whether to enable grad sync on last microbatch (DDP naive path)
+    std::shared_ptr<std::atomic_bool> is_last_microbatch_ = std::make_shared<std::atomic_bool>(true);
 
     DistributedDataParallelConfig ddp_config_;
     const ProcessGroup *ddp_pg_ = 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);
+
     // Start grad reduce
     void StartGradSync();
 
@@ -150,6 +152,7 @@ class ParamAndGradBucketGroup {
     std::vector<std::vector<std::shared_ptr<Tensor>>> param_buffer_shard_list_;
     std::vector<std::vector<std::shared_ptr<Tensor>>> grad_buffer_shard_list_;
 
+    // Whether to enable grad sync on last microbatch (DDP + ZeRO path)
     bool is_last_microbatch_ = true;
 
     bool grad_reduce_dispatched_ = false;

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

@@ -63,6 +63,8 @@ class Reducer : public std::enable_shared_from_this<Reducer> {
     // Prepare bucket info for next step
     void PrepareForBackward();
 
+    void SetIsLastMicrobatch(bool is_last_microbatch);
+
     // For custom DDP hook to overwrite the default AllReduce.
     // This can be used for algorithms like Gradient Compression/GossipGrad.
     // Hook is registered using `Reducer::RegisterCommHook()`.
@@ -149,10 +151,12 @@ class Reducer : public std::enable_shared_from_this<Reducer> {
     std::vector<uint8_t> ready_seen_this_iter_;
     // Whether to rebuild buckets on next train step
     bool need_rebuild_ = false;
-    // Whether to buckets have already been rebuilt on the second step
+    // Whether buckets have already been rebuilt on the second step
     bool has_rebuilt_bucket_ = false;
     // Whether all buckets are ready and backward can be finalized
     bool all_buckets_ready_this_iter_ = false;
+    // Whether to enable grad sync on last microbatch (DDP gradient bucketing path)
+    bool is_last_microbatch_ = true;
 };
 
 } // namespace infini_train::nn::parallel

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,5 +1,6 @@
 #pragma once
 
+#include <functional>
 #include <memory>
 #include <vector>
 
@@ -10,7 +11,8 @@ class Tensor;
 class Optimizer;
 namespace nn {
 class Module;
-}
+class NoSyncGuard;
+} // namespace nn
 } // namespace infini_train
 
 namespace infini_train::nn::parallel {
@@ -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);
+
     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/autograd/function_hook.cc

@@ -1,16 +1,23 @@
 #include "infini_train/include/autograd/function_hook.h"
 
+#include <utility>
+
 #include "infini_train/include/nn/parallel/parallel_functional.h"
 #include "infini_train/include/nn/parallel/process_group.h"
 #include "infini_train/include/tensor.h"
 
 namespace infini_train::autograd {
 AllReducePostAccumulateHook::AllReducePostAccumulateHook(infini_train::nn::parallel::function::ReduceOpType reduce_op,
-                                                         const infini_train::nn::parallel::ProcessGroup *pg)
+                                                         const infini_train::nn::parallel::ProcessGroup *pg,
+                                                         std::shared_ptr<const std::atomic_bool> enabled)
     : reduce_op_(reduce_op),
-      pg_(pg ? pg : infini_train::nn::parallel::ProcessGroupFactory::Instance()->GetDefaultProcessGroup()) {}
+      pg_(pg ? pg : infini_train::nn::parallel::ProcessGroupFactory::Instance()->GetDefaultProcessGroup()),
+      enabled_(std::move(enabled)) {}
 
 void AllReducePostAccumulateHook::operator()(const std::shared_ptr<Tensor> &tensor) {
+    if (enabled_ && !enabled_->load(std::memory_order_relaxed)) {
+        return;
+    }
     infini_train::nn::parallel::function::AllReduce(tensor, reduce_op_, pg_);
 }
 } // namespace infini_train::autograd

infini_train/src/nn/modules/module.cc

@@ -21,10 +21,22 @@
 
 namespace infini_train::nn {
 
+NoSyncGuard::NoSyncGuard(std::function<void()> exit_func) : exit_func_(std::move(exit_func)) {}
+
+NoSyncGuard::~NoSyncGuard() {
+    if (exit_func_) {
+        exit_func_();
+    }
+}
+
 Module::Module() : Module(kUndefinedType) {}
 
 Module::Module(const std::string &type) : type_(type), device_(Device()) {}
 
+std::unique_ptr<NoSyncGuard> Module::no_sync() {
+    return std::make_unique<NoSyncGuard>([] {});
+}
+
 const std::string &Module::type() const { return type_; }
 
 std::vector<std::shared_ptr<Tensor>> Module::Parameters() const {

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

@@ -37,7 +37,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
             << "All parameters must be on the same device as the module";
         if (!ddp_config.gradient_bucketing_enabled && ddp_config.zero_stage < 1) {
             auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(
-                function::ReduceOpType::kAvg, ddp_pg_);
+                function::ReduceOpType::kAvg, ddp_pg_, is_last_microbatch_);
             param->RegisterPostAccumulateGradHook(std::move(hook));
         }
     }
@@ -216,4 +216,17 @@ 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) {
+    is_last_microbatch_->store(is_last_microbatch, std::memory_order_relaxed);
+    if (reducer_) {
+        reducer_->SetIsLastMicrobatch(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

@@ -146,6 +146,8 @@ void ParamAndGradBucketGroup::Reset() {
     }
 }
 
+void ParamAndGradBucketGroup::SetIsLastMicrobatch(bool is_last_microbatch) { is_last_microbatch_ = is_last_microbatch; }
+
 void ParamAndGradBucketGroup::RegisterGradReady(const std::shared_ptr<Tensor> &parameter) {
     if (!ddp_config_.overlap_grad_reduce) {
         LOG(WARNING)
@@ -154,19 +156,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,8 +303,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();
 }
 
 void ParamAndGradBucketGroup::FinishGradSync() {

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

@@ -304,6 +304,11 @@ void Reducer::PrepareForBackward() {
     }
 }
 
+void Reducer::SetIsLastMicrobatch(bool is_last_microbatch) {
+    std::lock_guard<std::mutex> lock(mutex_);
+    is_last_microbatch_ = is_last_microbatch;
+}
+
 void Reducer::AttachHooksToParameters() {
     for (size_t param_idx = 0; param_idx < params_.size(); ++param_idx) {
         class BucketHook final : public autograd::PostAccumulateGradHook {
@@ -332,6 +337,10 @@ void Reducer::AttachHooksToParameters() {
 
 void Reducer::MarkVariableReadyDense(size_t variable_index) {
     std::unique_lock<std::mutex> lock(mutex_);
+    if (!is_last_microbatch_) {
+        return;
+    }
+
     const auto loc = locators_.at(variable_index);
     auto &bucket = buckets_.at(loc.bucket_index);
 

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