diff --git a/README.md b/README.md
index 1a51d1a8..ed6f0dd9 100644
--- a/README.md
+++ b/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] \
diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc
index 60c0c908..7d183efa 100644
--- a/example/gpt2/main.cc
+++ b/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();
diff --git a/example/llama3/main.cc b/example/llama3/main.cc
index 302e0808..a7d8ee0a 100644
--- a/example/llama3/main.cc
+++ b/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();
diff --git a/infini_train/include/dataloader.h b/infini_train/include/dataloader.h
index 38fa02a4..34dd04b7 100644
--- a/infini_train/include/dataloader.h
+++ b/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 {
diff --git a/infini_train/include/nn/parallel/global.h b/infini_train/include/nn/parallel/global.h
index 9373100f..c7c74a99 100644
--- a/infini_train/include/nn/parallel/global.h
+++ b/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(); }
diff --git a/infini_train/include/nn/parallel/rank.h b/infini_train/include/nn/parallel/rank.h
index eb61417d..baa623db 100644
--- a/infini_train/include/nn/parallel/rank.h
+++ b/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
diff --git a/infini_train/src/core/ccl/ccl_utils.cc b/infini_train/src/core/ccl/ccl_utils.cc
index f8968c6b..ebf1ed19 100644
--- a/infini_train/src/core/ccl/ccl_utils.cc
+++ b/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,13 +12,21 @@
 
 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;
@@ -25,12 +34,14 @@ void WriteUniqueIdFile(const CclUniqueId &unique_id, const std::string &pg_name)
     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
diff --git a/infini_train/src/dataloader.cc b/infini_train/src/dataloader.cc
index b7cc94f2..c7f64b41 100644
--- a/infini_train/src/dataloader.cc
+++ b/infini_train/src/dataloader.cc
@@ -3,6 +3,7 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstring>
+#include <functional>
 #include <numeric>
 #include <utility>
 
@@ -13,8 +14,14 @@
 
 namespace infini_train {
 namespace {
+size_t CheckedCeilDiv(size_t numerator, size_t denominator) {
+    CHECK_GT(denominator, 0);
+    return (numerator + denominator - 1) / denominator;
+}
+
 // TODO(dcj): Use official stack implementation later.
 std::shared_ptr<Tensor> Stack(const std::vector<std::shared_ptr<Tensor>> &tensors) {
+    CHECK(!tensors.empty()) << "Cannot stack an empty batch. Check DataLoader iterator end handling.";
     const int batch_size = tensors.size();
     const auto &dims = tensors[0]->Dims();
     const int stacked_dim = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies<int64_t>());
@@ -35,21 +42,31 @@ std::shared_ptr<Tensor> Stack(const std::vector<std::shared_ptr<Tensor>> &tensor
 }
 } // namespace
 
-DataLoaderIterator::DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t batch_idx,
-                                       size_t max_batch_idx, size_t ddp_rank, size_t ddp_world_size)
-    : dataset_(&dataset), batch_size_(batch_size), batch_idx_(batch_idx), max_batch_idx_(max_batch_idx),
-      ddp_rank_(ddp_rank), ddp_world_size_(ddp_world_size){};
+DataLoaderIterator::DataLoaderIterator(const Dataset &dataset, size_t batch_size, size_t global_batch_idx,
+                                       size_t num_global_batches, size_t ddp_rank, size_t ddp_world_size)
+    : dataset_(&dataset), batch_size_(batch_size), global_batch_idx_(global_batch_idx),
+      num_global_batches_(num_global_batches), ddp_rank_(ddp_rank), ddp_world_size_(ddp_world_size){};
 
 std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> DataLoaderIterator::operator*() const {
     /*
       0,         1,            ..., x,                  ...
       [0, bs-1], [bs, 2*bs-1], ..., [x*bs, (x+1)*bs-1], ...
                                     ^
-                                    batch_idx
+                                    global_batch_idx
     */
     std::vector<std::shared_ptr<Tensor>> data_vec;
     std::vector<std::shared_ptr<Tensor>> label_vec;
-    for (int idx = batch_idx_ * batch_size_; idx < (batch_idx_ + 1) * batch_size_ && idx < dataset_->Size(); ++idx) {
+    CHECK_LT(global_batch_idx_, num_global_batches_)
+        << "Cannot dereference DataLoader end iterator. global_batch_idx=" << global_batch_idx_
+        << ", num_global_batches=" << num_global_batches_ << ", ddp_rank=" << ddp_rank_
+        << ", ddp_world_size=" << ddp_world_size_;
+    const size_t start_idx = (global_batch_idx_ * ddp_world_size_ + ddp_rank_) * batch_size_;
+    CHECK_LT(start_idx, dataset_->Size())
+        << "DataLoader batch starts past dataset end. global_batch_idx=" << global_batch_idx_
+        << ", start_idx=" << start_idx << ", dataset_size=" << dataset_->Size() << ", batch_size=" << batch_size_
+        << ", ddp_rank=" << ddp_rank_ << ", ddp_world_size=" << ddp_world_size_;
+    const size_t end_idx = std::min(start_idx + batch_size_, dataset_->Size());
+    for (size_t idx = start_idx; idx < end_idx; ++idx) {
         auto &&[data, label] = dataset_->operator[](idx);
         data_vec.push_back(std::move(data));
         label_vec.push_back(std::move(label));
@@ -58,7 +75,7 @@ std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> DataLoaderIterator::
 }
 
 DataLoaderIterator &DataLoaderIterator::operator++() {
-    batch_idx_ = std::min(batch_idx_ + ddp_world_size_, max_batch_idx_);
+    global_batch_idx_ = std::min(global_batch_idx_ + 1, num_global_batches_);
     return *this;
 }
 
@@ -68,36 +85,60 @@ DataLoaderIterator DataLoaderIterator::operator++(int) {
     return tmp;
 }
 
-bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) { return lhs.batch_idx_ < rhs.batch_idx_; }
+bool operator<(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) {
+    return lhs.global_batch_idx_ < rhs.global_batch_idx_;
+}
 
 bool operator!=(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) {
-    return lhs.batch_idx_ != rhs.batch_idx_;
+    return lhs.global_batch_idx_ != rhs.global_batch_idx_;
 }
 
 bool operator==(const DataLoaderIterator &lhs, const DataLoaderIterator &rhs) {
-    return lhs.batch_idx_ == rhs.batch_idx_;
+    return lhs.global_batch_idx_ == rhs.global_batch_idx_;
 }
 
-size_t DataLoaderIterator::BatchIndex() const { return batch_idx_; }
+size_t DataLoaderIterator::GlobalBatchIndex() const { return global_batch_idx_; }
+
+DataLoaderIterator &DataLoaderIterator::SeekGlobalBatch(size_t global_batch_idx) {
+    CHECK_LE(global_batch_idx, num_global_batches_)
+        << "Cannot seek past DataLoader end. global_batch_idx=" << global_batch_idx
+        << ", num_global_batches=" << num_global_batches_;
+    global_batch_idx_ = global_batch_idx;
+    return *this;
+}
 
 DataLoader::DataLoader(const std::shared_ptr<Dataset> &dataset, size_t batch_size)
-    : dataset_(dataset), batch_size_(batch_size), max_batch_idx_((dataset_->Size() + batch_size_ - 1) / batch_size_) {}
+    : dataset_(dataset), batch_size_(batch_size), num_global_batches_(CheckedCeilDiv(dataset_->Size(), batch_size_)) {}
 
-DataLoaderIterator DataLoader::begin() const { return DataLoaderIterator(*dataset_, batch_size_, 0, max_batch_idx_); }
+DataLoaderIterator DataLoader::begin() const {
+    return DataLoaderIterator(*dataset_, batch_size_, 0, num_global_batches_);
+}
 
 DataLoaderIterator DataLoader::end() const {
-    return DataLoaderIterator(*dataset_, batch_size_, max_batch_idx_, max_batch_idx_);
+    return DataLoaderIterator(*dataset_, batch_size_, num_global_batches_, num_global_batches_);
 }
 
+size_t DataLoader::NumGlobalBatches() const { return num_global_batches_; }
+
 DistributedDataLoader::DistributedDataLoader(const std::shared_ptr<Dataset> &dataset, size_t batch_size,
                                              size_t ddp_rank, size_t ddp_world_size)
-    : DataLoader(dataset, batch_size), ddp_rank_(ddp_rank), ddp_world_size_(ddp_world_size) {}
+    : DataLoader(dataset, batch_size), ddp_rank_(ddp_rank), ddp_world_size_(ddp_world_size) {
+    CHECK_GT(ddp_world_size_, 0);
+    CHECK_LT(ddp_rank_, ddp_world_size_);
+    const size_t global_batch_size = ddp_world_size_ * batch_size_;
+    CHECK_GE(dataset_->Size(), global_batch_size)
+        << "DistributedDataLoader needs at least one full global batch. dataset_size=" << dataset_->Size()
+        << ", global_batch_size=" << global_batch_size << " (" << batch_size_ << " per rank * " << ddp_world_size_
+        << " ranks). Reduce batch size/world size or use a larger dataset.";
+    num_global_batches_ = dataset_->Size() / global_batch_size;
+}
 
 DataLoaderIterator DistributedDataLoader::begin() const {
-    return DataLoaderIterator(*dataset_, batch_size_, ddp_rank_, max_batch_idx_, ddp_rank_, ddp_world_size_);
+    return DataLoaderIterator(*dataset_, batch_size_, 0, num_global_batches_, ddp_rank_, ddp_world_size_);
 }
 
 DataLoaderIterator DistributedDataLoader::end() const {
-    return DataLoaderIterator(*dataset_, batch_size_, max_batch_idx_, max_batch_idx_, ddp_rank_, ddp_world_size_);
+    return DataLoaderIterator(*dataset_, batch_size_, num_global_batches_, num_global_batches_, ddp_rank_,
+                              ddp_world_size_);
 }
 } // namespace infini_train
diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc
index 1bb3aaad..db10cd54 100644
--- a/infini_train/src/device.cc
+++ b/infini_train/src/device.cc
@@ -33,7 +33,16 @@ std::string Device::ToString() const {
 }
 
 nn::parallel::Rank Device::Rank() const {
-    return {nn::parallel::global::GetGlobalProcRank(), index_, nn::parallel::global::GetNprocPerNode(),
+    if (IsCPU()) {
+        return {nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
+                nn::parallel::global::GetNthreadPerProc()};
+    }
+
+    const int thread_rank = index_ - nn::parallel::global::GetLocalDeviceIndex();
+    CHECK_GE(thread_rank, 0) << "CUDA device index is outside the current process rank range";
+    CHECK_LT(thread_rank, nn::parallel::global::GetNthreadPerProc())
+        << "CUDA device index is outside the current process rank range";
+    return {nn::parallel::global::GetGlobalProcRank(), thread_rank, nn::parallel::global::GetNprocPerNode(),
             nn::parallel::global::GetNthreadPerProc()};
 }
 
diff --git a/infini_train/src/nn/parallel/data_parallel.cc b/infini_train/src/nn/parallel/data_parallel.cc
index c48761b7..6199d39f 100644
--- a/infini_train/src/nn/parallel/data_parallel.cc
+++ b/infini_train/src/nn/parallel/data_parallel.cc
@@ -60,7 +60,11 @@ ParallelApply(const std::vector<std::shared_ptr<Module>> &modules,
 
 DataParallel::DataParallel(const std::shared_ptr<Module> &module, int dim, Device::DeviceType device_type) : dim_(dim) {
     devices_.reserve(global::GetNthreadPerProc());
-    for (int index = 0; index < global::GetNthreadPerProc(); ++index) { devices_.emplace_back(device_type, index); }
+    for (int thread_rank = 0; thread_rank < global::GetNthreadPerProc(); ++thread_rank) {
+        const int device_index
+            = device_type == Device::DeviceType::kCUDA ? global::GetLocalDeviceIndex(thread_rank) : thread_rank;
+        devices_.emplace_back(device_type, device_index);
+    }
 
     CHECK_GT(devices_.size(), 0) << "No available devices found";
     output_device_ = devices_.at(0);
diff --git a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc
index a3bfe008..1d4aa467 100644
--- a/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc
+++ b/infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc
@@ -10,6 +10,7 @@
 
 #include "infini_train/include/autograd/function_hook.h"
 #include "infini_train/include/nn/modules/module.h"
+#include "infini_train/include/nn/parallel/global.h"
 #include "infini_train/include/nn/parallel/parallel_functional.h"
 #include "infini_train/include/nn/parallel/process_group.h"
 #include "infini_train/include/nn/parallel/rank.h"
@@ -36,7 +37,8 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
             continue;
         }
         auto device = param->GetDevice();
-        CHECK_EQ(device.index(), rank.thread_rank()) << "All parameters must be on the same device as the module";
+        CHECK_EQ(device.index(), global::GetLocalDeviceIndex(rank.thread_rank()))
+            << "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_);
@@ -44,7 +46,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
         }
     }
     for (auto &buffer : module->Buffers()) {
-        CHECK_EQ(buffer->GetDevice().index(), rank.thread_rank())
+        CHECK_EQ(buffer->GetDevice().index(), global::GetLocalDeviceIndex(rank.thread_rank()))
             << "All buffers must be on the same device as the module";
     }
     modules_[kModuleName] = std::move(module);
diff --git a/infini_train/src/nn/parallel/global.cc b/infini_train/src/nn/parallel/global.cc
index 65a3208e..e56cac72 100644
--- a/infini_train/src/nn/parallel/global.cc
+++ b/infini_train/src/nn/parallel/global.cc
@@ -92,13 +92,24 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq
 
     CHECK(!initialized_) << "Repeated initialization of GlobalEnv!";
 
-    nnodes_ = GetEnvAsInt("NNODES", 1);
-    nproc_per_node_ = GetEnvAsInt("NPROC_PER_NODE", 1);
-    world_size_ = GetEnvAsInt("PROC_WORLD_SIZE", 1) * nthread_per_process;
-    global_proc_rank_ = GetEnvAsInt("GLOBAL_PROC_RANK", 0);
-    local_proc_rank_ = GetEnvAsInt("LOCAL_PROC_RANK", 0);
+    const int proc_world_size = GetEnvAsInt("WORLD_SIZE", GetEnvAsInt("PROC_WORLD_SIZE", 1));
+    nproc_per_node_ = GetEnvAsInt("LOCAL_WORLD_SIZE", GetEnvAsInt("NPROC_PER_NODE", 1));
+    CHECK_GT(nproc_per_node_, 0) << "NPROC_PER_NODE/LOCAL_WORLD_SIZE must be positive";
+    CHECK_GT(proc_world_size, 0) << "PROC_WORLD_SIZE/WORLD_SIZE must be positive";
+    CHECK_EQ(proc_world_size % nproc_per_node_, 0)
+        << "PROC_WORLD_SIZE/WORLD_SIZE must be divisible by NPROC_PER_NODE/LOCAL_WORLD_SIZE";
+    nnodes_ = proc_world_size / nproc_per_node_;
+    global_proc_rank_ = GetEnvAsInt("RANK", GetEnvAsInt("GLOBAL_PROC_RANK", 0));
+    local_proc_rank_ = GetEnvAsInt("LOCAL_RANK", GetEnvAsInt("LOCAL_PROC_RANK", 0));
+    CHECK_GE(global_proc_rank_, 0) << "GLOBAL_PROC_RANK/RANK must be non-negative";
+    CHECK_LT(global_proc_rank_, proc_world_size)
+        << "GLOBAL_PROC_RANK/RANK must be less than PROC_WORLD_SIZE/WORLD_SIZE";
+    CHECK_GE(local_proc_rank_, 0) << "LOCAL_PROC_RANK/LOCAL_RANK must be non-negative";
+    CHECK_LT(local_proc_rank_, nproc_per_node_)
+        << "LOCAL_PROC_RANK/LOCAL_RANK must be less than NPROC_PER_NODE/LOCAL_WORLD_SIZE";
 
     nthread_per_process_ = nthread_per_process;
+    world_size_ = proc_world_size * nthread_per_process;
     CHECK_GE(tensor_parallel_size, 1) << "Tensor Parallel size must be >= 1";
     tensor_parallel_size_ = tensor_parallel_size;
     sequence_parallel_enabled_ = sequence_parallel_enabled;
diff --git a/infini_train/src/nn/parallel/process_group.cc b/infini_train/src/nn/parallel/process_group.cc
index 3c4c4910..9dab7848 100644
--- a/infini_train/src/nn/parallel/process_group.cc
+++ b/infini_train/src/nn/parallel/process_group.cc
@@ -99,7 +99,7 @@ void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
         int global_thread_rank = lower_rank + i;
         auto it = std::ranges::find(ranks, global_thread_rank);
         if (it != ranks.end()) {
-            auto device = Device(backend_, i);
+            auto device = Device(backend_, global::GetLocalDeviceIndex(i));
             core::DeviceGuard guard(device);
 
             core::CclComm *comm_raw = nullptr;
diff --git a/infini_train/src/nn/parallel/rank.cc b/infini_train/src/nn/parallel/rank.cc
index 617a8e50..d5b4a788 100644
--- a/infini_train/src/nn/parallel/rank.cc
+++ b/infini_train/src/nn/parallel/rank.cc
@@ -2,18 +2,19 @@
 #include "infini_train/include/nn/parallel/global.h"
 
 namespace infini_train::nn::parallel {
-Rank::Rank(int process_rank, int thread_rank, int process_size, int thread_size)
-    : process_rank_(process_rank), thread_rank_(thread_rank), process_size_(process_size), thread_size_(thread_size) {}
+Rank::Rank(int global_process_rank, int thread_rank, int processes_per_node, int threads_per_process)
+    : global_process_rank_(global_process_rank), thread_rank_(thread_rank), processes_per_node_(processes_per_node),
+      threads_per_process_(threads_per_process) {}
 
-int Rank::process_rank() const { return process_rank_; }
+int Rank::process_rank() const { return global_process_rank_; }
 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::process_size() const { return processes_per_node_; }
+int Rank::thread_size() const { return threads_per_process_; }
 
-int Rank::GlobalRank() const { return process_rank_ * thread_size_ + thread_rank_; }
+int Rank::GlobalRank() const { return global_process_rank_ * threads_per_process_ + thread_rank_; }
 
-bool Rank::IsParallel() const { return thread_size_ * process_size_ > 1; }
+bool Rank::IsParallel() const { return global::GetWorldSize() > 1; }
 
 bool Rank::IsMainRank() const { return GlobalRank() == 0; }
 
diff --git a/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash
index ac2b4931..62cf1c84 100755
--- a/scripts/run_models_and_profile.bash
+++ b/scripts/run_models_and_profile.bash
@@ -438,6 +438,18 @@ check_model_inputs() {
     fi
 }
 
+model_cmd_for_test() {
+    local model_bin="$1"
+    local input_bin="$2"
+    local llmc_filepath="$3"
+    local arg_str="$4"
+    local nproc_per_node="$5"
+
+    printf './infini_run --nproc_per_node=%s %s --input_bin %q --llmc_filepath %q --device cuda %s' \
+        "$nproc_per_node" "$model_bin" \
+        "$input_bin" "$llmc_filepath" "$arg_str"
+}
+
 # Run tests
 num_basic_compile_commands=$(jq '.basic_compile_commands | length' "$CONFIG_FILE")
 num_groups=$(jq '.test_groups | length' "$CONFIG_FILE")
@@ -509,21 +521,23 @@ for ((id=0; id<num_basic_compile_commands; ++id)); do
 
             for ((ti=0; ti<num_tests; ++ti)); do
                 test_id=$(jq -r ".test_groups[$gi].tests[$ti].id" "$CONFIG_FILE")
+                nproc_per_node="$(jq -r ".test_groups[$gi].launcher_args.nproc_per_node // empty" "$CONFIG_FILE")"
+                : "${nproc_per_node:=${INFINI_NPROC_PER_NODE:-1}}"
                 if tag_enabled_for_model "$group_tag" "$GPT2_TEST_GROUPS"; then
                     gpt2_arg_str="$(args_string_for_test "$gi" "$ti" "gpt2" "$test_id")"
-                    gpt2_cmd="${prefix}./gpt2 --input_bin ${GPT2_INPUT_BIN} --llmc_filepath ${GPT2_LLMC_FILEPATH} --device cuda ${gpt2_arg_str}"
+                    gpt2_cmd="$(model_cmd_for_test "./gpt2" "$GPT2_INPUT_BIN" "$GPT2_LLMC_FILEPATH" "$gpt2_arg_str" "$nproc_per_node")"
                     run_and_log "$gpt2_cmd" "gpt2_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
                 fi
 
                 if tag_enabled_for_model "$group_tag" "$LLAMA3_TEST_GROUPS"; then
                     llama3_arg_str="$(args_string_for_test "$gi" "$ti" "llama3" "$test_id")"
-                    llama3_cmd="${prefix}./llama3 --input_bin ${LLAMA3_INPUT_BIN} --llmc_filepath ${LLAMA3_LLMC_FILEPATH} --device cuda ${llama3_arg_str}"
+                    llama3_cmd="$(model_cmd_for_test "./llama3" "$LLAMA3_INPUT_BIN" "$LLAMA3_LLMC_FILEPATH" "$llama3_arg_str" "$nproc_per_node")"
                     run_and_log "$llama3_cmd" "llama3_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
                 fi
 
                 if tag_enabled_for_model "$group_tag" "$MIXTRAL_TEST_GROUPS"; then
                     mixtral_arg_str="$(args_string_for_test "$gi" "$ti" "mixtral" "$test_id")"
-                    mixtral_cmd="${prefix}./mixtral --input_bin ${MIXTRAL_INPUT_BIN} --llmc_filepath ${MIXTRAL_LLMC_FILEPATH} --device cuda ${mixtral_arg_str}"
+                    mixtral_cmd="$(model_cmd_for_test "./mixtral" "$MIXTRAL_INPUT_BIN" "$MIXTRAL_LLMC_FILEPATH" "$mixtral_arg_str" "$nproc_per_node")"
                     run_and_log "$mixtral_cmd" "mixtral_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
                 fi
             done
diff --git a/scripts/test_config.json b/scripts/test_config.json
index 13b41095..7dd75192 100644
--- a/scripts/test_config.json
+++ b/scripts/test_config.json
@@ -13,8 +13,8 @@
         "RUN_PROFILE_TEST": "true",
         "MIXTRAL_INPUT_BIN": "/data1/shared/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin",
         "MIXTRAL_LLMC_FILEPATH": "/data1/shared/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin",
-        "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler",
-        "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler",
+        "GPT2_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc",
+        "LLAMA3_TEST_GROUPS": "basic,zero,lora,checkpoint,lr_scheduler,8_proc",
         "MIXTRAL_TEST_GROUPS": "moe"
     },
     "basic_compile_commands": [
@@ -887,6 +887,130 @@
                     }
                 }
             ]
+        },
+        {
+            "tag": "8_proc",
+            "launcher_args": {
+                "nproc_per_node": 8
+            },
+            "tests": [
+                {
+                    "id": "3",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120
+                    }
+                },
+                {
+                    "id": "3_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120
+                    }
+                },
+                {
+                    "id": "4",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4
+                    }
+                },
+                {
+                    "id": "4_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4
+                    }
+                },
+                {
+                    "id": "5",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4,
+                        "sequence_parallel": true
+                    }
+                },
+                {
+                    "id": "5_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4,
+                        "sequence_parallel": true
+                    }
+                },
+                {
+                    "id": "6",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8
+                    }
+                },
+                {
+                    "id": "6_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8
+                    }
+                },
+                {
+                    "id": "8",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2
+                    }
+                },
+                {
+                    "id": "8_bfloat16",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 2,
+                        "sequence_parallel": true,
+                        "pipeline_parallel": 2,
+                        "virtual_pipeline_parallel": 2
+                    }
+                }
+            ]
         }
     ]
 }
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 96776585..eed77360 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -7,6 +7,12 @@ include(${CMAKE_SOURCE_DIR}/cmake/test_macros.cmake)
 # Common test utilities
 add_subdirectory(common)
 
+# Distributed tests
+add_subdirectory(distributed)
+
+# DataLoader tests
+add_subdirectory(dataloader)
+
 # Tensor tests
 add_subdirectory(tensor)
 
diff --git a/tests/dataloader/CMakeLists.txt b/tests/dataloader/CMakeLists.txt
new file mode 100644
index 00000000..7dcf5c88
--- /dev/null
+++ b/tests/dataloader/CMakeLists.txt
@@ -0,0 +1,8 @@
+# ==========================================================================
+# DataLoader tests
+# ==========================================================================
+
+infini_train_add_test(test_dataloader
+  SOURCES test_dataloader.cc
+  LABELS cpu
+)
diff --git a/tests/dataloader/test_dataloader.cc b/tests/dataloader/test_dataloader.cc
new file mode 100644
index 00000000..837232d1
--- /dev/null
+++ b/tests/dataloader/test_dataloader.cc
@@ -0,0 +1,101 @@
+#include <cstdint>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "gtest/gtest.h"
+
+#include "infini_train/include/dataloader.h"
+#include "infini_train/include/dataset.h"
+#include "infini_train/include/tensor.h"
+
+using namespace infini_train;
+
+namespace {
+class IndexDataset : public Dataset {
+public:
+    explicit IndexDataset(size_t size) : size_(size) {}
+
+    std::pair<std::shared_ptr<Tensor>, std::shared_ptr<Tensor>> operator[](size_t idx) const override {
+        auto data = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kINT64);
+        auto label = std::make_shared<Tensor>(std::vector<int64_t>{1}, DataType::kINT64);
+        *static_cast<int64_t *>(data->DataPtr()) = static_cast<int64_t>(idx);
+        *static_cast<int64_t *>(label->DataPtr()) = static_cast<int64_t>(idx + 1000);
+        return {data, label};
+    }
+
+    size_t Size() const override { return size_; }
+
+private:
+    size_t size_ = 0;
+};
+
+std::vector<int64_t> TensorValues(const std::shared_ptr<Tensor> &tensor) {
+    const auto *data = static_cast<const int64_t *>(tensor->DataPtr());
+    return std::vector<int64_t>(data, data + tensor->NumElements());
+}
+} // namespace
+
+TEST(DataLoaderTest, RegularDataLoaderKeepsPartialLastBatch) {
+    DataLoader loader(std::make_shared<IndexDataset>(5), 2);
+
+    std::vector<std::vector<int64_t>> batches;
+    for (const auto &[x, y] : loader) { batches.push_back(TensorValues(x)); }
+
+    ASSERT_EQ(batches.size(), 3);
+    EXPECT_EQ(batches[0], (std::vector<int64_t>{0, 1}));
+    EXPECT_EQ(batches[1], (std::vector<int64_t>{2, 3}));
+    EXPECT_EQ(batches[2], (std::vector<int64_t>{4}));
+}
+
+TEST(DataLoaderTest, DistributedDataLoaderSlicesFullGlobalBatchesByRank) {
+    const auto dataset = std::make_shared<IndexDataset>(13);
+    const size_t batch_size = 2;
+    const size_t world_size = 3;
+
+    DistributedDataLoader rank0(dataset, batch_size, 0, world_size);
+    DistributedDataLoader rank1(dataset, batch_size, 1, world_size);
+    DistributedDataLoader rank2(dataset, batch_size, 2, world_size);
+
+    auto r0 = rank0.begin();
+    auto r1 = rank1.begin();
+    auto r2 = rank2.begin();
+
+    EXPECT_EQ(TensorValues((*r0).first), (std::vector<int64_t>{0, 1}));
+    EXPECT_EQ(TensorValues((*r1).first), (std::vector<int64_t>{2, 3}));
+    EXPECT_EQ(TensorValues((*r2).first), (std::vector<int64_t>{4, 5}));
+
+    ++r0;
+    ++r1;
+    ++r2;
+
+    EXPECT_EQ(TensorValues((*r0).first), (std::vector<int64_t>{6, 7}));
+    EXPECT_EQ(TensorValues((*r1).first), (std::vector<int64_t>{8, 9}));
+    EXPECT_EQ(TensorValues((*r2).first), (std::vector<int64_t>{10, 11}));
+
+    ++r0;
+    ++r1;
+    ++r2;
+
+    EXPECT_EQ(r0, rank0.end());
+    EXPECT_EQ(r1, rank1.end());
+    EXPECT_EQ(r2, rank2.end());
+}
+
+TEST(DataLoaderTest, DistributedDataLoaderCyclesOnGlobalBatchBoundary) {
+    DistributedDataLoader loader(std::make_shared<IndexDataset>(13), 2, 2, 3);
+    auto iter = loader.begin();
+
+    EXPECT_EQ(loader.NumGlobalBatches(), 2);
+
+    EXPECT_EQ(TensorValues((*iter).first), (std::vector<int64_t>{4, 5}));
+    ++iter;
+    EXPECT_EQ(iter.GlobalBatchIndex(), 1);
+
+    EXPECT_EQ(TensorValues((*iter).first), (std::vector<int64_t>{10, 11}));
+    ++iter;
+    EXPECT_EQ(iter, loader.end());
+
+    iter = loader.begin();
+    EXPECT_EQ(TensorValues((*iter).first), (std::vector<int64_t>{4, 5}));
+}
diff --git a/tests/distributed/CMakeLists.txt b/tests/distributed/CMakeLists.txt
new file mode 100644
index 00000000..b8ed4970
--- /dev/null
+++ b/tests/distributed/CMakeLists.txt
@@ -0,0 +1,25 @@
+# ==========================================================================
+# Distributed tests
+# ==========================================================================
+
+infini_train_add_test(test_rank
+  SOURCES test_rank.cc
+  LABELS cpu
+)
+
+add_test(
+  NAME RankTest.MultiNodeSingleProcessIsParallel
+  COMMAND ${CMAKE_COMMAND} -E env
+    "WORLD_SIZE=2"
+    "LOCAL_WORLD_SIZE=1"
+    "RANK=0"
+    "LOCAL_RANK=0"
+    $<TARGET_FILE:test_rank>
+    --gtest_filter=RankTest.DetectsParallelismFromGlobalWorldSize
+)
+
+set_tests_properties(RankTest.MultiNodeSingleProcessIsParallel
+  PROPERTIES
+    LABELS cpu
+    TIMEOUT 10
+)
diff --git a/tests/distributed/test_rank.cc b/tests/distributed/test_rank.cc
new file mode 100644
index 00000000..e0b37614
--- /dev/null
+++ b/tests/distributed/test_rank.cc
@@ -0,0 +1,16 @@
+#include "gtest/gtest.h"
+
+#include "infini_train/include/nn/parallel/global.h"
+#include "infini_train/include/nn/parallel/rank.h"
+
+namespace infini_train::nn::parallel {
+namespace {
+
+TEST(RankTest, DetectsParallelismFromGlobalWorldSize) {
+    const Rank rank(global::GetGlobalProcRank(), 0, global::GetNprocPerNode(), global::GetNthreadPerProc());
+
+    EXPECT_EQ(rank.IsParallel(), global::GetWorldSize() > 1);
+}
+
+} // namespace
+} // namespace infini_train::nn::parallel
diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc
index f6fe2f5e..69f8f0dc 100644
--- a/tools/infini_run/infini_run.cc
+++ b/tools/infini_run/infini_run.cc
@@ -1,8 +1,14 @@
+#include <cerrno>
+#include <chrono>
+#include <csignal>
 #include <cstdio>
 #include <cstdlib>
+#include <filesystem>
 #include <string>
 #include <sys/wait.h>
+#include <system_error>
 #include <unistd.h>
+#include <unordered_set>
 #include <vector>
 
 #include "gflags/gflags.h"
@@ -13,46 +19,169 @@ 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)");
 
+DEFINE_string(rdzv_id, "", "Unique job ID shared by all nodes");
+
+namespace {
+
+bool IsLauncherFlag(const std::string &flag) {
+    return flag == "--nnodes" || flag == "--nproc_per_node" || flag == "--node_rank" || flag == "--rdzv_endpoint"
+        || flag == "--rdzv_id";
+}
+
+int FindTrainProgramIndex(int argc, char **argv) {
+    for (int i = 1; i < argc; ++i) {
+        const std::string arg = argv[i];
+        if (arg == "--") {
+            return i;
+        }
+        if (arg.rfind("--", 0) != 0) {
+            return i;
+        }
+
+        const size_t eq_pos = arg.find('=');
+        const std::string flag = arg.substr(0, eq_pos);
+        if (IsLauncherFlag(flag) && eq_pos == std::string::npos) {
+            ++i;
+        }
+    }
+    return argc;
+}
+
+std::vector<char *> BuildLauncherArgv(int train_program_index, char **argv) {
+    std::vector<char *> launcher_argv;
+    launcher_argv.push_back(argv[0]);
+    for (int i = 1; i < train_program_index; ++i) { launcher_argv.push_back(argv[i]); }
+    launcher_argv.push_back(nullptr);
+    return launcher_argv;
+}
+
+void SetEnvInt(const char *name, int value) {
+    const auto value_str = std::to_string(value);
+    setenv(name, value_str.c_str(), 1);
+}
+
+void TerminateChildren(const std::unordered_set<pid_t> &child_pids) {
+    for (pid_t child_pid : child_pids) { kill(child_pid, SIGTERM); }
+}
+
+int ExitCodeFromStatus(int status) {
+    if (WIFEXITED(status)) {
+        return WEXITSTATUS(status);
+    }
+    if (WIFSIGNALED(status)) {
+        return 128 + WTERMSIG(status);
+    }
+    return 1;
+}
+
+void CleanupRunUniqueIdFiles(const std::string &run_id) {
+    const std::string prefix = "cclUniqueId_" + run_id + "_";
+    for (const auto &entry : std::filesystem::directory_iterator(std::filesystem::current_path())) {
+        if (!entry.is_regular_file()) {
+            continue;
+        }
+        const std::string filename = entry.path().filename().string();
+        if (filename.rfind(prefix, 0) == 0) {
+            std::error_code ec;
+            std::filesystem::remove(entry.path(), ec);
+            if (ec) {
+                LOG(WARNING) << "Failed to remove unique-id file " << entry.path() << ": " << ec.message();
+            }
+        }
+    }
+}
+
+std::string GenerateLocalRunId() {
+    const auto now = std::chrono::steady_clock::now().time_since_epoch().count();
+    return std::to_string(getpid()) + "_" + std::to_string(now);
+}
+
+} // namespace
+
 int main(int argc, char **argv) {
-    gflags::ParseCommandLineFlags(&argc, &argv, true);
+    const int train_program_index = FindTrainProgramIndex(argc, argv);
+    std::vector<char *> launcher_argv = BuildLauncherArgv(train_program_index, argv);
+    int launcher_argc = static_cast<int>(launcher_argv.size()) - 1;
+    char **launcher_argv_ptr = launcher_argv.data();
+    gflags::ParseCommandLineFlags(&launcher_argc, &launcher_argv_ptr, true);
     google::InitGoogleLogging(argv[0]);
 
-    CHECK_GE(argc, 2) << "No training prgram specified!";
+    CHECK_GT(FLAGS_nnodes, 0) << "nnodes must be positive";
+    CHECK_GT(FLAGS_nproc_per_node, 0) << "nproc_per_node must be positive";
+    CHECK_NE(FLAGS_rdzv_endpoint.find(':'), std::string::npos) << "rdzv_endpoint must be host:port";
+    CHECK(FLAGS_nnodes == 1 || !FLAGS_rdzv_id.empty())
+        << "rdzv_id must be set to the same unique job ID on every node for multi-node training";
+
+    CHECK_LT(train_program_index, argc) << "No training program specified!";
 
-    std::string train_program = argv[1];
+    std::string train_program = argv[train_program_index];
+    CHECK_NE(train_program, "--") << "Explicit '--' separator is not supported; pass the training program directly "
+                                     "after infini_run launcher flags";
     std::vector<char *> train_argv;
-    for (int i = 1; i < argc; ++i) { train_argv.push_back(argv[i]); }
+    for (int i = train_program_index; i < argc; ++i) { train_argv.push_back(argv[i]); }
     train_argv.push_back(nullptr);
 
-    int world_size = FLAGS_nnodes * FLAGS_nproc_per_node;
+    int proc_world_size = FLAGS_nnodes * FLAGS_nproc_per_node;
     std::string master_addr = FLAGS_rdzv_endpoint.substr(0, FLAGS_rdzv_endpoint.find(':'));
     std::string master_port = FLAGS_rdzv_endpoint.substr(FLAGS_rdzv_endpoint.find(':') + 1);
+    const std::string run_id = FLAGS_rdzv_id.empty() ? GenerateLocalRunId() : FLAGS_rdzv_id;
+
+    std::unordered_set<pid_t> running_children;
+    int exit_code = 0;
 
     for (int local_proc_rank = 0; local_proc_rank < FLAGS_nproc_per_node; ++local_proc_rank) {
         pid_t pid = fork();
+        if (pid < 0) {
+            perror("fork failed");
+            exit_code = 1;
+            TerminateChildren(running_children);
+            break;
+        }
         if (pid == 0) {
             int global_proc_rank = FLAGS_node_rank * FLAGS_nproc_per_node + local_proc_rank;
-            setenv("NNODES", std::to_string(FLAGS_nnodes).c_str(), 1);
-            setenv("NPROC_PER_NODE", std::to_string(FLAGS_nproc_per_node).c_str(), 1);
+            SetEnvInt("LOCAL_WORLD_SIZE", FLAGS_nproc_per_node);
 
             setenv("MASTER_ADDR", master_addr.c_str(), 1);
             setenv("MASTER_PORT", master_port.c_str(), 1);
+            setenv("INFINI_RUN_ID", run_id.c_str(), 1);
 
-            setenv("GLOBAL_PROC_RANK", std::to_string(global_proc_rank).c_str(), 1);
-            setenv("LOCAL_PROC_RANK", std::to_string(local_proc_rank).c_str(), 1);
+            SetEnvInt("RANK", global_proc_rank);
+            SetEnvInt("LOCAL_RANK", local_proc_rank);
 
-            setenv("PROC_WORLD_SIZE", std::to_string(world_size).c_str(), 1);
+            SetEnvInt("WORLD_SIZE", proc_world_size);
 
             execvp(train_program.c_str(), train_argv.data());
             perror("exec failed");
             exit(1);
         }
+        running_children.insert(pid);
     }
 
-    for (int i = 0; i < FLAGS_nproc_per_node; ++i) {
+    while (!running_children.empty()) {
         int status;
-        wait(&status);
+        pid_t child = wait(&status);
+        if (child < 0) {
+            if (errno == EINTR) {
+                continue;
+            }
+            perror("wait failed");
+            if (exit_code == 0) {
+                exit_code = 1;
+            }
+            TerminateChildren(running_children);
+            break;
+        }
+
+        running_children.erase(child);
+        const int child_exit_code = ExitCodeFromStatus(status);
+        if (child_exit_code != 0 && exit_code == 0) {
+            exit_code = child_exit_code;
+            TerminateChildren(running_children);
+        }
     }
 
-    return 0;
+    if (FLAGS_nnodes == 1) {
+        CleanupRunUniqueIdFiles(run_id);
+    }
+    return exit_code;
 }