fix: stabilize distributed training resume flow

Track DataLoader progress by global batches so distributed ranks slice data
consistently and can resume/cycle from saved consumption counts.

Also scope CCL unique ID files per run, generate NCCL IDs only on the main
rank, clean up run-local rendezvous files, and add DataLoader coverage.

Author: chen2021673

README.md

@@ -121,8 +121,9 @@ 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 \
+  ./llama3 \
      --device cuda \
      --input_bin [training_data_path] \
      --llmc_filepath [model_path] \

example/gpt2/main.cc

@@ -374,15 +374,17 @@ 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;
+        ++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 +459,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 +489,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));
 

example/llama3/main.cc

@@ -354,15 +354,17 @@ 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;
+        ++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 +437,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 +466,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));
 

infini_train/include/core/ccl/ccl.h

@@ -28,7 +28,7 @@ class CclImpl {
 
     virtual void GetAsyncError(const CclComm *comm, CclStatus *async_error) const;
 
-    virtual void GetUniqueId(CclUniqueId **unique_id) const;
+    virtual void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const;
 
     virtual void CommInitAll(CclComm **comms, int ndev, const int *devlist) const;
 

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/src/core/ccl/ccl.cc

@@ -16,7 +16,9 @@ void CclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const {
     LOG(FATAL) << "CclImpl::GetAsyncError is not implemented.";
 }
 
-void CclImpl::GetUniqueId(CclUniqueId **unique_id) const { LOG(FATAL) << "CclImpl::GetUniqueId is not implemented."; }
+void CclImpl::CreateUniqueId(CclUniqueId **, bool) const {
+    LOG(FATAL) << "CclImpl::CreateUniqueId is not implemented.";
+}
 
 void CclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const {
     LOG(FATAL) << "CclImpl::CommInitAll is not implemented.";

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

infini_train/src/core/ccl/cuda/nccl_impl.cc

@@ -68,14 +68,16 @@ void NcclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const
     }
 }
 
-void NcclImpl::GetUniqueId(CclUniqueId **unique_id) const {
+void NcclImpl::CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const {
     CHECK_NOTNULL(unique_id);
     if (*unique_id == nullptr) {
         *unique_id = new NcclUniqueId();
     }
     auto *nccl_unique_id = dynamic_cast<NcclUniqueId *>(*unique_id);
     CHECK_NOTNULL(nccl_unique_id);
-    NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id()));
+    if (generate_id) {
+        NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id()));
+    }
 }
 
 void NcclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const {

infini_train/src/core/ccl/cuda/nccl_impl.h

@@ -17,7 +17,7 @@ class NcclImpl final : public CclImpl {
 
     void GetAsyncError(const CclComm *comm, CclStatus *async_error) const override;
 
-    void GetUniqueId(CclUniqueId **unique_id) const override;
+    void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const override;
 
     void CommInitAll(CclComm **comms, int ndev, const int *devlist) const override;