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 fd2522e3..3fef5bdf 100644
--- a/example/gpt2/main.cc
+++ b/example/gpt2/main.cc
@@ -180,7 +180,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::GetDeviceIndex(rank.thread_rank()));
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -565,7 +565,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) {
@@ -576,7 +575,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 05fad4a8..ca627741 100644
--- a/example/llama3/main.cc
+++ b/example/llama3/main.cc
@@ -168,7 +168,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::GetDeviceIndex(rank.thread_rank()));
         auto *pg_factory = ProcessGroupFactory::Instance(device.type());
 
         if (ddp_world_size > 1) {
@@ -544,7 +544,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) {
@@ -555,7 +554,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/nn/parallel/global.h b/infini_train/include/nn/parallel/global.h
index 9373100f..38694a91 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 GetDeviceIndex(int thread_rank) { 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..1368993a 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 UniqueIdFileName(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 UniqueIdTmpFileName(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 = UniqueIdTmpFileName(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,7 +34,9 @@ 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 = UniqueIdFileName(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) {
@@ -50,6 +61,11 @@ void CleanupUniqueIdFile(const std::string &pg_name) {
     if (std::filesystem::exists(file_path)) {
         std::filesystem::remove(file_path);
     }
+
+    const std::string tmp_path = UniqueIdTmpFileName(pg_name);
+    if (std::filesystem::exists(tmp_path)) {
+        std::filesystem::remove(tmp_path);
+    }
 }
 
 } // namespace infini_train::core
diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc
index 1bb3aaad..98021db5 100644
--- a/infini_train/src/device.cc
+++ b/infini_train/src/device.cc
@@ -33,7 +33,11 @@ std::string Device::ToString() const {
 }
 
 nn::parallel::Rank Device::Rank() const {
-    return {nn::parallel::global::GetGlobalProcRank(), index_, nn::parallel::global::GetNprocPerNode(),
+    const int thread_rank = index_ - nn::parallel::global::GetDeviceIndex(0);
+    CHECK_GE(thread_rank, 0) << "Device index is outside the current process rank range";
+    CHECK_LT(thread_rank, nn::parallel::global::GetNthreadPerProc())
+        << "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..cf16f6c3 100644
--- a/infini_train/src/nn/parallel/data_parallel.cc
+++ b/infini_train/src/nn/parallel/data_parallel.cc
@@ -60,7 +60,10 @@ 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 = global::GetDeviceIndex(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 b08b64fb..57460ba0 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"
@@ -32,7 +33,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::GetDeviceIndex(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_);
@@ -40,7 +42,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::GetDeviceIndex(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..655b4bce 100644
--- a/infini_train/src/nn/parallel/global.cc
+++ b/infini_train/src/nn/parallel/global.cc
@@ -92,13 +92,21 @@ 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", 1);
+    nproc_per_node_ = GetEnvAsInt("LOCAL_WORLD_SIZE", 1);
+    CHECK_GT(nproc_per_node_, 0) << "LOCAL_WORLD_SIZE must be positive";
+    CHECK_GT(proc_world_size, 0) << "WORLD_SIZE must be positive";
+    CHECK_EQ(proc_world_size % nproc_per_node_, 0) << "WORLD_SIZE must be divisible by LOCAL_WORLD_SIZE";
+    nnodes_ = proc_world_size / nproc_per_node_;
+    global_proc_rank_ = GetEnvAsInt("RANK", 0);
+    local_proc_rank_ = GetEnvAsInt("LOCAL_RANK", 0);
+    CHECK_GE(global_proc_rank_, 0) << "RANK must be non-negative";
+    CHECK_LT(global_proc_rank_, proc_world_size) << "RANK must be less than WORLD_SIZE";
+    CHECK_GE(local_proc_rank_, 0) << "LOCAL_RANK must be non-negative";
+    CHECK_LT(local_proc_rank_, nproc_per_node_) << "LOCAL_RANK must be less than 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 174aa645..caa08981 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::GetDeviceIndex(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 b7a7b5cc..67bc7acd 100755
--- a/scripts/run_models_and_profile.bash
+++ b/scripts/run_models_and_profile.bash
@@ -439,6 +439,18 @@ check_model_inputs() {
     fi
 }
 
+infini_run_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")
@@ -510,23 +522,36 @@ 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].infini_run_args.nproc_per_node // empty" "$CONFIG_FILE")"
                 if tag_enabled_for_model "$group_tag" "$GPT2_TEST_GROUPS"; then
                     LORA_WEIGHTS_DIR="$GPT2_LORA_WEIGHTS_DIR"
                     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}"
+                    if [[ -n "$nproc_per_node" ]]; then
+                        gpt2_cmd="$(infini_run_cmd_for_test "./gpt2" "$GPT2_INPUT_BIN" "$GPT2_LLMC_FILEPATH" "$gpt2_arg_str" "$nproc_per_node")"
+                    else
+                        gpt2_cmd="${prefix}./gpt2 --input_bin ${GPT2_INPUT_BIN} --llmc_filepath ${GPT2_LLMC_FILEPATH} --device cuda ${gpt2_arg_str}"
+                    fi
                     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
                     LORA_WEIGHTS_DIR="$LLAMA3_LORA_WEIGHTS_DIR"
                     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}"
+                    if [[ -n "$nproc_per_node" ]]; then
+                        llama3_cmd="$(infini_run_cmd_for_test "./llama3" "$LLAMA3_INPUT_BIN" "$LLAMA3_LLMC_FILEPATH" "$llama3_arg_str" "$nproc_per_node")"
+                    else
+                        llama3_cmd="${prefix}./llama3 --input_bin ${LLAMA3_INPUT_BIN} --llmc_filepath ${LLAMA3_LLMC_FILEPATH} --device cuda ${llama3_arg_str}"
+                    fi
                     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}"
+                    if [[ -n "$nproc_per_node" ]]; then
+                        mixtral_cmd="$(infini_run_cmd_for_test "./mixtral" "$MIXTRAL_INPUT_BIN" "$MIXTRAL_LLMC_FILEPATH" "$mixtral_arg_str" "$nproc_per_node")"
+                    else
+                        mixtral_cmd="${prefix}./mixtral --input_bin ${MIXTRAL_INPUT_BIN} --llmc_filepath ${MIXTRAL_LLMC_FILEPATH} --device cuda ${mixtral_arg_str}"
+                    fi
                     run_and_log "$mixtral_cmd" "mixtral_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
                 fi
             done
@@ -562,6 +587,10 @@ else
     echo -e "\033[1;33m         To enable comparison, set 'variables.COMPARE_LOG_DIR' in ${CONFIG_FILE}\033[0m"
     echo -e "\033[1;33m         or export COMPARE_LOG_DIR=/path/to/baseline_logs before running.\033[0m"
     echo -e "\033[1;33m============================================================\033[0m"
+
+    echo -e "\n\033[1;33mRun comparison manually:\033[0m"
+    echo "python3 compare_loss.py \"/path/to/baseline_logs\" \"$(realpath "$LOG_DIR")\""
+    echo "python3 compare_tps.py \"/path/to/baseline_logs\" \"$(realpath "$LOG_DIR")\""
 fi
 
 echo -e "\n\033[1;36m[END OF TEST] Cleaning build directory after all tests\033[0m"
diff --git a/scripts/test_config.json b/scripts/test_config.json
index 9bd7780e..2b8202d2 100644
--- a/scripts/test_config.json
+++ b/scripts/test_config.json
@@ -15,8 +15,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": [
@@ -895,6 +895,130 @@
                     }
                 }
             ]
+        },
+        {
+            "tag": "8_proc",
+            "infini_run_args": {
+                "nproc_per_node": 8
+            },
+            "tests": [
+                {
+                    "id": "3_8proc",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120
+                    }
+                },
+                {
+                    "id": "3_bfloat16_8proc",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120
+                    }
+                },
+                {
+                    "id": "4_8proc",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4
+                    }
+                },
+                {
+                    "id": "4_bfloat16_8proc",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 40,
+                        "total_batch_size": 5120,
+                        "tensor_parallel": 4
+                    }
+                },
+                {
+                    "id": "5_8proc",
+                    "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_8proc",
+                    "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_8proc",
+                    "args": {
+                        "dtype": "float32",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8
+                    }
+                },
+                {
+                    "id": "6_bfloat16_8proc",
+                    "args": {
+                        "dtype": "bfloat16",
+                        "nthread_per_process": 1,
+                        "num_iteration": 10,
+                        "batch_size": 10,
+                        "total_batch_size": 5120,
+                        "pipeline_parallel": 8
+                    }
+                },
+                {
+                    "id": "8_8proc",
+                    "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_8proc",
+                    "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..f7ced03d 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -7,6 +7,9 @@ include(${CMAKE_SOURCE_DIR}/cmake/test_macros.cmake)
 # Common test utilities
 add_subdirectory(common)
 
+# Distributed tests
+add_subdirectory(distributed)
+
 # Tensor tests
 add_subdirectory(tensor)
 
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..c2d6c8fa 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 + 1;
+        }
+        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_GE(FLAGS_node_rank, 0) << "node_rank must be non-negative";
+    CHECK_LT(FLAGS_node_rank, FLAGS_nnodes) << "node_rank must be less than nnodes";
+    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];
     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;
 }