fix: make multi-node DDP precision work

Author: kilinchange

example/llama3/main.cc

@@ -107,7 +107,7 @@ void Train(const nn::parallel::Rank &rank) {
         if (ddp_world_size > 1) {
             ddp_pg = ProcessGroupFactory::Instance()->GetOrCreate(GetDataParallelProcessGroupName(rank.thread_rank()),
                                                                   GetDataParallelGroupRanks(rank.thread_rank()));
-            ddp_rank = ddp_pg->GetGroupRank(rank.thread_rank());
+            ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank());
         }
 
         if (tp_world_size > 1) {

infini_train/include/nn/parallel/rank.h

@@ -10,6 +10,8 @@ class Rank {
     int process_size() const;
     int thread_size() const;
 
+    int GlobalRank() const;
+
     bool IsParallel() const;
 
     bool IsMainRank() const;

infini_train/src/nn/parallel/distributed_data_parallel.cc

@@ -24,8 +24,8 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
 
         auto ddp_pg
             = ProcessGroupFactory::Instance()->Get(GetDataParallelProcessGroupName(device->rank().thread_rank()));
-        auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(function::ReduceOpType::kAvg,
-                                                                                          ddp_pg);
+        // FIXME(dcj): use multi-node ddp_pg here
+        auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(function::ReduceOpType::kAvg);
         param->RegisterPostAccumulateGradHook(std::move(hook));
     }
     for (auto &buffer : module->Buffers()) {

infini_train/src/nn/parallel/process_group.cc

@@ -1,5 +1,6 @@
 #include "infini_train/include/nn/parallel/process_group.h"
 
+#include <algorithm>
 #include <memory>
 #include <numeric>
 #include <vector>
@@ -51,6 +52,8 @@ ProcessGroup::ProcessGroup(const ncclUniqueId &nccl_id) : world_size_(global::Ge
         device_indices[i] = i;
 
         int global_rank = global::GetGlobalProcRank() * global::GetNthreadPerProc() + i;
+
+        cudaSetDevice(i);
         NCCL_CHECK(ncclCommInitRank(&comms_[i], world_size_, nccl_id, global_rank));
     }
     NCCL_CHECK(ncclGroupEnd());
@@ -59,11 +62,13 @@ ProcessGroup::ProcessGroup(const ncclUniqueId &nccl_id) : world_size_(global::Ge
 }
 
 void ProcessGroup::Init(const std::vector<int> &device_indices) {
-    for (int i = 0; i < world_size_; ++i) {
+    // FIXME(dcj): This is a temporary solution to get the device and comm for each thread.
+    int local_comm_size = std::min(static_cast<int>(device_indices.size()), global::GetNthreadPerProc());
+    for (int i = 0; i < local_comm_size; ++i) {
         auto device = DeviceManager::Instance()->GetDevice(DeviceType::kCUDA, device_indices[i]);
         devices_.push_back(device);
         device_comm_map_[device] = comms_[i];
-        thread_group_rank_map_[device->rank().thread_rank()] = i;
+        thread_group_rank_map_[device->rank().thread_rank()] = i + global::GetGlobalProcRank() * local_comm_size;
     }
 }
 

infini_train/src/nn/parallel/rank.cc

@@ -10,6 +10,8 @@ int Rank::thread_rank() const { return thread_rank_; }
 int Rank::process_size() const { return process_size_; }
 int Rank::thread_size() const { return thread_size_; }
 
+int Rank::GlobalRank() const { return process_rank_ * thread_size_ + thread_rank_; }
+
 bool Rank::IsParallel() const { return thread_size_ * process_size_ > 1; }
 bool Rank::IsMainRank() const { return thread_rank_ == 0; }
 } // namespace infini_train::nn::parallel