Support torchrun-style InfiniTrain multi-process launch

README.md

@@ -121,8 +121,10 @@ The following examples demonstrate **LLaMA 3 supervised fine-tuning (SFT)** usin
 ./infini_run \
   --nnodes=2 \
   --nproc_per_node=1 \
-  --node_rank=[rank_id] \
-  -- ./llama3 \
+  --node_rank=[rank_id] \
+  --rdzv_endpoint=[master_addr]:29500 \
+  --rdzv_id=[job_id] \
+  ./llama3 \
      --device cuda \
      --input_bin [training_data_path] \
      --llmc_filepath [model_path] \

example/gpt2/main.cc

@@ -181,7 +181,7 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        device = Device(Device::DeviceType::kCUDA, global::GetLocalDeviceIndex(rank.thread_rank()));
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -374,15 +374,19 @@ void Train(const nn::parallel::Rank &rank) {
     start_step = resume_result.global_step;
     size_t consumed_batches = resume_result.consumed_batches;
 
-    // TODO(jym): Replace with Sampler abstraction when available.
-    // Skip dataloader to resume from the correct batch position.
-    if (consumed_batches > 0) {
-        size_t start = train_iter.BatchIndex();
-        // Each rank processes every ddp_world_size-th batch starting from its own rank.
-        // num_skips calculates how many ++ iterations to reach the saved batch position.
-        size_t num_skips = (consumed_batches - start) / ddp_world_size;
-        for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
-    }
+    // consumed_batches is the number of global dataloader batches consumed across cyclic epochs.
+    train_iter.SeekGlobalBatch(consumed_batches % train_loader.NumGlobalBatches());
+    auto next_train_batch = [&]() {
+        auto batch = *train_iter;
+        // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
+        // TODO(dcj): support dataloader.reset() later
+        ++train_iter;
+        if (train_iter == train_loader.end()) {
+            train_iter = train_loader.begin();
+        }
+        ++consumed_batches;
+        return batch;
+    };
 
     auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) {
         SaveCheckpoint({
@@ -457,11 +461,7 @@ void Train(const nn::parallel::Rank &rank) {
                 infini_train::AutocastGuard autocast_guard(device.type(), dtype);
 
                 // (bs, seq_len), (bs, seq_len)
-                auto [x, y] = *train_iter;
-                // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
-                // TODO(dcj): support dataloader.reset() later
-                ++train_iter;
-                consumed_batches = train_iter.BatchIndex();
+                auto [x, y] = next_train_batch();
                 x = std::make_shared<Tensor>(x->To(device));
                 y = std::make_shared<Tensor>(y->To(device));
 
@@ -491,11 +491,7 @@ void Train(const nn::parallel::Rank &rank) {
                 scheduler->Step();
             }
         } else {
-            auto [x, y] = *train_iter;
-            // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
-            // TODO(dcj): support dataloader.reset() later
-            ++train_iter;
-            consumed_batches = train_iter.BatchIndex();
+            auto [x, y] = next_train_batch();
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 
@@ -568,7 +564,6 @@ int main(int argc, char *argv[]) {
 
     LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
 
-    // NOTE(dcj): currently we only support single process
     if (FLAGS_nthread_per_process > 1) {
         std::vector<std::thread> threads;
         for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) {
@@ -579,7 +574,9 @@ int main(int argc, char *argv[]) {
 
         for (auto &thread : threads) { thread.join(); }
     } else {
-        Train({0, 0, 1, 1});
+        nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
+                                FLAGS_nthread_per_process);
+        Train(rank);
     }
 
     gflags::ShutDownCommandLineFlags();

example/llama3/main.cc

@@ -167,7 +167,7 @@ void Train(const nn::parallel::Rank &rank) {
     const ProcessGroup *pp_pg = nullptr;
 
     if (rank.IsParallel()) {
-        device = Device(Device::DeviceType::kCUDA, rank.thread_rank());
+        device = Device(Device::DeviceType::kCUDA, global::GetLocalDeviceIndex(rank.thread_rank()));
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -354,15 +354,19 @@ void Train(const nn::parallel::Rank &rank) {
     start_step = resume_result.global_step;
     size_t consumed_batches = resume_result.consumed_batches;
 
-    // TODO(jym): Replace with Sampler abstraction when available.
-    // Skip dataloader to resume from the correct batch position.
-    if (consumed_batches > 0) {
-        size_t start = train_iter.BatchIndex();
-        // Each rank processes every ddp_world_size-th batch starting from its own rank.
-        // num_skips calculates how many ++ iterations to reach the saved batch position.
-        size_t num_skips = (consumed_batches - start) / ddp_world_size;
-        for (size_t i = 0; i < num_skips; ++i) { ++train_iter; }
-    }
+    // consumed_batches is the number of global dataloader batches consumed across cyclic epochs.
+    train_iter.SeekGlobalBatch(consumed_batches % train_loader.NumGlobalBatches());
+    auto next_train_batch = [&]() {
+        auto batch = *train_iter;
+        // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
+        // TODO(dcj): support dataloader.reset() later
+        ++train_iter;
+        if (train_iter == train_loader.end()) {
+            train_iter = train_loader.begin();
+        }
+        ++consumed_batches;
+        return batch;
+    };
 
     auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) {
         SaveCheckpoint({
@@ -435,11 +439,7 @@ void Train(const nn::parallel::Rank &rank) {
                 infini_train::AutocastGuard autocast_guard(device.type(), dtype);
 
                 // (bs, seq_len), (bs, seq_len)
-                auto [x, y] = *train_iter;
-                // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
-                // TODO(dcj): support dataloader.reset() later
-                ++train_iter;
-                consumed_batches = train_iter.BatchIndex();
+                auto [x, y] = next_train_batch();
                 x = std::make_shared<Tensor>(x->To(device));
                 y = std::make_shared<Tensor>(y->To(device));
 
@@ -468,11 +468,7 @@ void Train(const nn::parallel::Rank &rank) {
                 scheduler->Step();
             }
         } else {
-            auto [x, y] = *train_iter;
-            // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below
-            // TODO(dcj): support dataloader.reset() later
-            ++train_iter;
-            consumed_batches = train_iter.BatchIndex();
+            auto [x, y] = next_train_batch();
             x = std::make_shared<Tensor>(x->To(device));
             y = std::make_shared<Tensor>(y->To(device));
 
@@ -545,7 +541,6 @@ int main(int argc, char *argv[]) {
 
     LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
 
-    // NOTE(dcj): currently we only support single process
     if (FLAGS_nthread_per_process > 1) {
         std::vector<std::thread> threads;
         for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) {
@@ -556,7 +551,9 @@ int main(int argc, char *argv[]) {
 
         for (auto &thread : threads) { thread.join(); }
     } else {
-        Train({0, 0, 1, 1});
+        nn::parallel::Rank rank(nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
+                                FLAGS_nthread_per_process);
+        Train(rank);
     }
 
     gflags::ShutDownCommandLineFlags();

infini_train/include/dataloader.h

@@ -12,7 +12,7 @@ class Tensor;
 namespace infini_train {
 class DataLoaderIterator {
 public:
-    DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx, size_t max_batch_idx,
+    DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t global_batch_idx, size_t num_global_batches,
                        size_t ddp_rank = 0, size_t ddp_world_size = 1);
 
     std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> operator*() const;
@@ -24,13 +24,14 @@ class DataLoaderIterator {
     friend bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs);
     friend bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs);
 
-    size_t BatchIndex() const;
+    size_t GlobalBatchIndex() const;
+    DataLoaderIterator &SeekGlobalBatch(size_t global_batch_idx);
 
 private:
     const Dataset *dataset_ = nullptr; // not owned
     size_t batch_size_ = 0;
-    size_t batch_idx_ = 0;
-    size_t max_batch_idx_ = 0;
+    size_t global_batch_idx_ = 0;
+    size_t num_global_batches_ = 0;
     size_t ddp_rank_ = 0;
     size_t ddp_world_size_ = 1;
 };
@@ -42,10 +43,12 @@ class DataLoader {
     virtual DataLoaderIterator begin() const;
     virtual DataLoaderIterator end() const;
 
+    size_t NumGlobalBatches() const;
+
 protected:
     std::shared_ptr<Dataset> dataset_;
     size_t batch_size_ = 0;
-    size_t max_batch_idx_ = 0;
+    size_t num_global_batches_ = 0;
 };
 
 class DistributedDataLoader : public DataLoader {

infini_train/include/nn/parallel/global.h

@@ -98,6 +98,7 @@ inline int GetNprocPerNode() { return GlobalEnv::Instance().nproc_per_node(); }
 inline int GetNthreadPerProc() { return GlobalEnv::Instance().nthread_per_process(); }
 inline int GetGlobalProcRank() { return GlobalEnv::Instance().global_proc_rank(); }
 inline int GetLocalProcRank() { return GlobalEnv::Instance().local_proc_rank(); }
+inline int GetLocalDeviceIndex(int thread_rank = 0) { return GetLocalProcRank() * GetNthreadPerProc() + thread_rank; }
 
 inline int GetTensorParallelSize() { return GlobalEnv::Instance().tensor_parallel_size(); }
 inline int GetSequenceParallelSize() { return GlobalEnv::Instance().sequence_parallel_size(); }

infini_train/include/nn/parallel/rank.h

@@ -3,7 +3,7 @@
 namespace infini_train::nn::parallel {
 class Rank {
 public:
-    Rank(int process_rank, int thread_rank, int process_size, int thread_size);
+    Rank(int global_process_rank, int thread_rank, int processes_per_node, int threads_per_process);
 
     int process_rank() const;
     int thread_rank() const;
@@ -19,9 +19,9 @@ class Rank {
     bool IsLastRank() const;
 
 private:
-    const int process_rank_ = 0; // Rank of the current process within the node
-    const int thread_rank_ = 0;  // Rank of the current thread within the process
-    const int process_size_ = 1; // Total number of processes on this node
-    const int thread_size_ = 1;  // Total number of threads in the current process
+    const int global_process_rank_ = 0; // Global rank of the current process
+    const int thread_rank_ = 0;         // Rank of the current thread within the process
+    const int processes_per_node_ = 1;  // Number of processes on each node
+    const int threads_per_process_ = 1; // Number of threads in each process
 };
 } // namespace infini_train::nn::parallel

infini_train/src/core/ccl/ccl_utils.cc

@@ -2,6 +2,7 @@
 
 #include <chrono>
 #include <cstdio>
+#include <cstdlib>
 #include <filesystem>
 #include <fstream>
 #include <iterator>
@@ -11,26 +12,36 @@
 
 namespace infini_train::core {
 namespace {
-std::string UniqueIdFileName(const std::string &name, bool tmp = false) {
-    return "cclUniqueId_" + name + (tmp ? ".tmp" : ".bin");
+std::string UniqueIdPath(const std::string &pg_name) {
+    const char *run_id = std::getenv("INFINI_RUN_ID");
+    const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
+    return "cclUniqueId_" + prefix + pg_name + ".bin";
+}
+
+std::string UniqueIdTmpPath(const std::string &pg_name) {
+    const char *run_id = std::getenv("INFINI_RUN_ID");
+    const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
+    return "cclUniqueId_" + prefix + pg_name + ".tmp";
 }
 } // namespace
 
 void WriteUniqueIdFile(const CclUniqueId &unique_id, const std::string &pg_name) {
-    const std::string tmp_path = UniqueIdFileName(pg_name, true);
+    const std::string tmp_path = UniqueIdTmpPath(pg_name);
 
     std::ofstream ofs(tmp_path, std::ios::binary);
     CHECK(ofs.good()) << "Failed to open unique_id tmp file for write: " << tmp_path;
     const size_t size = unique_id.Size();
     ofs.write(reinterpret_cast<const char *>(unique_id.Data()), static_cast<std::streamsize>(size));
     ofs.close();
 
-    std::rename(tmp_path.c_str(), UniqueIdFileName(pg_name).c_str());
+    const std::string file_path = UniqueIdPath(pg_name);
+    CHECK_EQ(std::rename(tmp_path.c_str(), file_path.c_str()), 0)
+        << "Failed to rename unique_id file from " << tmp_path << " to " << file_path;
 }
 
 void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) {
     CHECK_NOTNULL(unique_id);
-    const std::string file_path = UniqueIdFileName(pg_name);
+    const std::string file_path = UniqueIdPath(pg_name);
 
     while (!std::filesystem::exists(file_path)) { std::this_thread::sleep_for(std::chrono::microseconds(1000)); }
 
@@ -46,10 +57,15 @@ void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) {
 }
 
 void CleanupUniqueIdFile(const std::string &pg_name) {
-    const std::string file_path = UniqueIdFileName(pg_name);
+    const std::string file_path = UniqueIdPath(pg_name);
     if (std::filesystem::exists(file_path)) {
         std::filesystem::remove(file_path);
     }
+
+    const std::string tmp_path = UniqueIdTmpPath(pg_name);
+    if (std::filesystem::exists(tmp_path)) {
+        std::filesystem::remove(tmp_path);
+    }
 }
 
 } // namespace infini_train::core