Support torchrun-style InfiniTrain multi-process launch

README.md

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

example/gpt2/main.cc

@@ -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();

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();

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(); }

infini_train/include/nn/parallel/rank.h

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

infini_train/src/core/ccl/ccl_utils.cc

@@ -2,6 +2,7 @@
 
 #include <chrono>
 #include <cstdio>
+#include <cstdlib>
 #include <filesystem>
 #include <fstream>
 #include <iterator>
@@ -11,21 +12,31 @@
 
 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;
     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 = 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

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()};
 }
 

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);

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,15 +33,16 @@ 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_);
             param->RegisterPostAccumulateGradHook(std::move(hook));
         }
     }
     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);

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;

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;

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; }
 

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"