feat: Move the logic for cleaning up the ncclUniqueId file into the ProcessGroup destructor.

Author: kilinchange

infini_train/include/nn/parallel/process_group.h

@@ -29,6 +29,7 @@ namespace infini_train::nn::parallel {
 class ProcessGroup {
 public:
     explicit ProcessGroup(const std::string &process_group_name, const std::vector<int> &device_indices);
+    ~ProcessGroup();
 
     int GetGroupRank(int thread_rank) const;
 
@@ -69,6 +70,8 @@ class ProcessGroup {
     int world_size_ = 0;
 
     const std::string name_ = "";
+
+    bool is_main_process_ = false;
 };
 #endif
 

infini_train/src/nn/parallel/process_group.cc

@@ -68,6 +68,15 @@ void ReadNcclUniqueId(ncclUniqueId &nccl_id, const std::string &pg_name) {
     ifs.read(reinterpret_cast<char *>(&nccl_id), sizeof(nccl_id));
     ifs.close();
 }
+
+void CleanupNcclIdFile(const std::string &pg_name) {
+    const std::filesystem::path cwd = std::filesystem::current_path();
+    std::string file_path = NcclFileName(pg_name);
+
+    if (std::filesystem::exists(file_path)) {
+        std::filesystem::remove(file_path);
+    }
+}
 #endif
 
 } // namespace
@@ -86,6 +95,12 @@ ProcessGroup::ProcessGroup(const std::string &process_group_name, const std::vec
     }
 }
 
+ProcessGroup::~ProcessGroup() {
+    if (is_main_process_) {
+        CleanupNcclIdFile(name_);
+    }
+}
+
 void ProcessGroup::InitSingleProcess(const std::vector<int> &ranks) {
     comms_.resize(world_size_);
     NCCL_CHECK(ncclCommInitAll(comms_.data(), world_size_, ranks.data()));
@@ -100,13 +115,17 @@ void ProcessGroup::InitSingleProcess(const std::vector<int> &ranks) {
 
 void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
     int n_threads = global::GetNthreadPerProc();
+    int global_proc_rank = global::GetGlobalProcRank();
+    int lower_rank = global_proc_rank * n_threads;
+    int upper_rank = (global_proc_rank + 1) * n_threads;
 
     ncclUniqueId nccl_id;
 
-    if (std::ranges::min(ranks) < (global::GetGlobalProcRank() + 1) * global::GetNthreadPerProc()
-        && std::ranges::min(ranks) >= global::GetGlobalProcRank() * global::GetNthreadPerProc()) {
-        ncclGetUniqueId(&nccl_id);
+    int min_rank = std::ranges::min(ranks);
+    if (min_rank < upper_rank && min_rank >= lower_rank) {
+        is_main_process_ = true;
 
+        ncclGetUniqueId(&nccl_id);
         WriteNcclUniqueId(nccl_id, name_);
     } else {
         ReadNcclUniqueId(nccl_id, name_);
@@ -115,8 +134,8 @@ void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
     std::vector<int> device_indices;
     NCCL_CHECK(ncclGroupStart());
     for (int i = 0; i < n_threads; ++i) {
-        int global_rank = global::GetGlobalProcRank() * global::GetNthreadPerProc() + i;
-        auto it = std::ranges::find(ranks, global_rank);
+        int global_thread_rank = lower_rank + i;
+        auto it = std::ranges::find(ranks, global_thread_rank);
         if (it != ranks.end()) {
             cudaSetDevice(i);
 

tools/infini_run/infini_run.cc

@@ -15,20 +15,6 @@ DEFINE_int32(nproc_per_node, 1, "Number of processes per node");
 DEFINE_int32(node_rank, 0, "Rank of this node");
 DEFINE_string(rdzv_endpoint, "127.0.0.1:29500", "Rendezvous endpoint (host:port)");
 
-void CleanupNcclIdFiles() {
-    const std::filesystem::path cwd = std::filesystem::current_path();
-    std::regex pattern(R"(ncclUniqueId_.*\.bin)");
-
-    for (const auto &entry : std::filesystem::directory_iterator(cwd)) {
-        if (entry.is_regular_file()) {
-            const std::string filename = entry.path().filename().string();
-            if (std::regex_match(filename, pattern)) {
-                std::filesystem::remove(entry.path());
-            }
-        }
-    }
-}
-
 int main(int argc, char **argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, true);
     google::InitGoogleLogging(argv[0]);
@@ -70,9 +56,5 @@ int main(int argc, char **argv) {
         wait(&status);
     }
 
-    if (FLAGS_node_rank == 0) {
-        CleanupNcclIdFiles();
-    }
-
     return 0;
 }