From 53d91cb0c8a634fe1fb90d84f788f9e6f69f2378 Mon Sep 17 00:00:00 2001 From: cx Date: Fri, 3 Jul 2026 09:11:44 +0000 Subject: [PATCH 1/3] Support torchrun-style InfiniTrain multi-process launch --- example/gpt2/main.cc | 7 +- example/llama3/main.cc | 7 +- infini_train/include/nn/parallel/global.h | 1 + infini_train/src/device.cc | 11 +- infini_train/src/nn/parallel/data_parallel.cc | 6 +- .../parallel/ddp/distributed_data_parallel.cc | 6 +- infini_train/src/nn/parallel/global.cc | 29 ++++- infini_train/src/nn/parallel/process_group.cc | 2 +- scripts/run_models_and_profile.bash | 19 ++- scripts/test_config.json | 120 ++++++++++++------ tools/infini_run/infini_run.cc | 100 +++++++++++++-- 11 files changed, 238 insertions(+), 70 deletions(-) diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..e7c7a1c0 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) { @@ -568,7 +568,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 threads; for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { @@ -579,7 +578,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..216ece24 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) { @@ -545,7 +545,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 threads; for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) { @@ -556,7 +555,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..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/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> &modules, DataParallel::DataParallel(const std::shared_ptr &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 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( function::ReduceOpType::kAvg, ddp_pg_); @@ -44,7 +46,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr 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..e4399128 100644 --- a/infini_train/src/nn/parallel/global.cc +++ b/infini_train/src/nn/parallel/global.cc @@ -13,6 +13,8 @@ int GetEnvAsInt(const std::string &name, int default_value) { return value ? std::atoi(value) : default_value; } +bool HasEnv(const std::string &name) { return std::getenv(name.c_str()) != nullptr; } + } // namespace namespace infini_train::nn::parallel::global { @@ -92,13 +94,30 @@ 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("PROC_WORLD_SIZE", GetEnvAsInt("WORLD_SIZE", 1)); + nproc_per_node_ = GetEnvAsInt("NPROC_PER_NODE", GetEnvAsInt("LOCAL_WORLD_SIZE", 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"; + const bool nnodes_env_set = HasEnv("NNODES"); + nnodes_ = GetEnvAsInt("NNODES", proc_world_size / nproc_per_node_); + CHECK_GT(nnodes_, 0) << "NNODES must be positive"; + if (nnodes_env_set) { + CHECK_EQ(nnodes_ * nproc_per_node_, proc_world_size) + << "NNODES * NPROC_PER_NODE/LOCAL_WORLD_SIZE must equal PROC_WORLD_SIZE/WORLD_SIZE"; + } + global_proc_rank_ = GetEnvAsInt("GLOBAL_PROC_RANK", GetEnvAsInt("RANK", 0)); + local_proc_rank_ = GetEnvAsInt("LOCAL_PROC_RANK", GetEnvAsInt("LOCAL_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 &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/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash index 4b31546d..987feba9 100755 --- a/scripts/run_models_and_profile.bash +++ b/scripts/run_models_and_profile.bash @@ -379,12 +379,25 @@ args_string_for_test() { | $args | (if has("save") then .save = namespaced_path(.save; $model; $run_mode) else . end) | (if has("load") then .load = namespaced_path(.load; $model; $resume_src_mode) else . end) + | del(.nproc_per_node) | to_entries[] | "--\(.key) \(.value|tostring)" ' "$CONFIG_FILE" | paste -sd' ' - | \ sed "s|@CKPT_ROOT_DIR@|${CKPT_ROOT_DIR}|g" } +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") @@ -454,15 +467,17 @@ for ((id=0; id BuildLauncherArgv(int train_program_index, char **argv) { + std::vector 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); +} + +} // namespace + int main(int argc, char **argv) { - gflags::ParseCommandLineFlags(&argc, &argv, true); + const int train_program_index = FindTrainProgramIndex(argc, argv); + std::vector launcher_argv = BuildLauncherArgv(train_program_index, argv); + int launcher_argc = static_cast(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"; - std::string train_program = argv[1]; + CHECK_LT(train_program_index, argc) << "No training program specified!"; + + 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 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); @@ -32,16 +82,23 @@ int main(int argc, char **argv) { pid_t pid = fork(); 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("NNODES", FLAGS_nnodes); + SetEnvInt("NPROC_PER_NODE", FLAGS_nproc_per_node); + 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("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("GLOBAL_PROC_RANK", global_proc_rank); + SetEnvInt("LOCAL_PROC_RANK", local_proc_rank); + SetEnvInt("RANK", global_proc_rank); + SetEnvInt("LOCAL_RANK", local_proc_rank); - setenv("PROC_WORLD_SIZE", std::to_string(world_size).c_str(), 1); + SetEnvInt("PROC_WORLD_SIZE", proc_world_size); + SetEnvInt("WORLD_SIZE", proc_world_size); + SetEnvInt("GROUP_RANK", FLAGS_node_rank); + SetEnvInt("ROLE_RANK", global_proc_rank); + SetEnvInt("ROLE_WORLD_SIZE", proc_world_size); execvp(train_program.c_str(), train_argv.data()); perror("exec failed"); @@ -49,10 +106,29 @@ int main(int argc, char **argv) { } } + int exit_code = 0; for (int i = 0; i < FLAGS_nproc_per_node; ++i) { int status; - wait(&status); + pid_t child = wait(&status); + if (child < 0) { + perror("wait failed"); + return 1; + } + + if (WIFEXITED(status)) { + int child_exit_code = WEXITSTATUS(status); + if (child_exit_code != 0 && exit_code == 0) { + exit_code = child_exit_code; + } + } else if (WIFSIGNALED(status)) { + int signal = WTERMSIG(status); + if (exit_code == 0) { + exit_code = 128 + signal; + } + } else if (exit_code == 0) { + exit_code = 1; + } } - return 0; + return exit_code; } From ee4695c2b2807bfb9b40d2487e3de90145c0c551 Mon Sep 17 00:00:00 2001 From: cx Date: Thu, 9 Jul 2026 02:32:10 +0000 Subject: [PATCH 2/3] test: add basic 8-process config group Add a dedicated 8_proc test group containing the 8-process variants of the original basic multi-GPU cases. --- scripts/test_config.json | 251 ++++++++++++++++++++++++++------------- 1 file changed, 171 insertions(+), 80 deletions(-) diff --git a/scripts/test_config.json b/scripts/test_config.json index 0ecc30fd..70b686fc 100644 --- a/scripts/test_config.json +++ b/scripts/test_config.json @@ -56,8 +56,7 @@ "id": "3", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120 @@ -67,8 +66,7 @@ "id": "3_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120 @@ -78,8 +76,7 @@ "id": "4", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -90,8 +87,7 @@ "id": "4_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -102,8 +98,7 @@ "id": "5", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -115,8 +110,7 @@ "id": "5_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -128,8 +122,7 @@ "id": "6", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -140,8 +133,7 @@ "id": "6_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -152,8 +144,7 @@ "id": "7", "args": { "dtype": "float32", - "nproc_per_node": 4, - "nthread_per_process": 1, + "nthread_per_process": 4, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -165,8 +156,7 @@ "id": "7_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 4, - "nthread_per_process": 1, + "nthread_per_process": 4, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -178,8 +168,7 @@ "id": "8", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -193,8 +182,7 @@ "id": "8_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -213,8 +201,7 @@ "id": "3_distopt", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -225,8 +212,7 @@ "id": "3_zero2", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -237,8 +223,7 @@ "id": "3_bfloat16_distopt", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -249,8 +234,7 @@ "id": "3_bfloat16_zero2", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -261,8 +245,7 @@ "id": "4_distopt", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -274,8 +257,7 @@ "id": "4_zero2", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -287,8 +269,7 @@ "id": "4_bfloat16_distopt", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -300,8 +281,7 @@ "id": "4_bfloat16_zero2", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -313,8 +293,7 @@ "id": "5_distopt", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -327,8 +306,7 @@ "id": "5_zero2", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -341,8 +319,7 @@ "id": "5_bfloat16_distopt", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -355,8 +332,7 @@ "id": "5_bfloat16_zero2", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -369,8 +345,7 @@ "id": "8_distopt", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -385,8 +360,7 @@ "id": "8_bfloat16_distopt", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -624,8 +598,7 @@ "id": "3_lora", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -638,8 +611,7 @@ "id": "3_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -652,8 +624,7 @@ "id": "4_lora", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -667,8 +638,7 @@ "id": "4_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -682,8 +652,7 @@ "id": "5_lora", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -698,8 +667,7 @@ "id": "5_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -714,8 +682,7 @@ "id": "6_lora", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -729,8 +696,7 @@ "id": "6_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -744,8 +710,7 @@ "id": "7_lora", "args": { "dtype": "float32", - "nproc_per_node": 4, - "nthread_per_process": 1, + "nthread_per_process": 4, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -760,8 +725,7 @@ "id": "7_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 4, - "nthread_per_process": 1, + "nthread_per_process": 4, "num_iteration": 10, "batch_size": 10, "total_batch_size": 5120, @@ -776,8 +740,7 @@ "id": "8_lora", "args": { "dtype": "float32", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -794,8 +757,7 @@ "id": "8_lora_bfloat16", "args": { "dtype": "bfloat16", - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 10, "batch_size": 40, "total_batch_size": 5120, @@ -852,8 +814,7 @@ { "id": "ckpt_3d_ddp_tp_pp", "args": { - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 50, "save_interval": 10, "batch_size": 40, @@ -870,8 +831,7 @@ { "id": "ckpt_3d_ddp_tp_pp_resume", "args": { - "nproc_per_node": 8, - "nthread_per_process": 1, + "nthread_per_process": 8, "num_iteration": 50, "save_interval": 10, "batch_size": 40, @@ -887,6 +847,137 @@ } } ] + }, + { + "tag": "8_proc", + "tests": [ + { + "id": "3", + "args": { + "dtype": "float32", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "3_bfloat16", + "args": { + "dtype": "bfloat16", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "4", + "args": { + "dtype": "float32", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "4_bfloat16", + "args": { + "dtype": "bfloat16", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "5", + "args": { + "dtype": "float32", + "nproc_per_node": 8, + "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", + "nproc_per_node": 8, + "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", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "6_bfloat16", + "args": { + "dtype": "bfloat16", + "nproc_per_node": 8, + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "8", + "args": { + "dtype": "float32", + "nproc_per_node": 8, + "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", + "nproc_per_node": 8, + "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 + } + } + ] } ] } From 5314e45f159eec3b8643c4a29dea67dd67e2fc36 Mon Sep 17 00:00:00 2001 From: cx Date: Fri, 10 Jul 2026 11:08:50 +0000 Subject: [PATCH 3/3] fix: stabilize distributed training resume flow Track DataLoader progress by global batches so distributed ranks slice data consistently and can resume/cycle from saved consumption counts. Also scope CCL unique ID files per run, generate NCCL IDs only on the main rank, clean up run-local rendezvous files, and add DataLoader coverage. --- README.md | 5 +- example/gpt2/main.cc | 32 +++--- example/llama3/main.cc | 32 +++--- infini_train/include/core/ccl/ccl.h | 2 +- infini_train/include/dataloader.h | 13 ++- infini_train/src/core/ccl/ccl.cc | 4 +- infini_train/src/core/ccl/ccl_utils.cc | 28 +++-- infini_train/src/core/ccl/cuda/nccl_impl.cc | 6 +- infini_train/src/core/ccl/cuda/nccl_impl.h | 2 +- infini_train/src/dataloader.cc | 75 ++++++++++--- infini_train/src/nn/parallel/process_group.cc | 8 +- tests/CMakeLists.txt | 3 + tests/dataloader/CMakeLists.txt | 8 ++ tests/dataloader/test_dataloader.cc | 101 ++++++++++++++++++ tools/infini_run/infini_run.cc | 35 +++++- 15 files changed, 277 insertions(+), 77 deletions(-) create mode 100644 tests/dataloader/CMakeLists.txt create mode 100644 tests/dataloader/test_dataloader.cc diff --git a/README.md b/README.md index 1a51d1a8..befab7a0 100644 --- a/README.md +++ b/README.md @@ -121,8 +121,9 @@ The following examples demonstrate **LLaMA 3 supervised fine-tuning (SFT)** usin ./infini_run \ --nnodes=2 \ --nproc_per_node=1 \ - --node_rank=[rank_id] \ - -- ./llama3 \ + --node_rank=[rank_id] \ + --rdzv_endpoint=[master_addr]:29500 \ + ./llama3 \ --device cuda \ --input_bin [training_data_path] \ --llmc_filepath [model_path] \ diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index e7c7a1c0..9ff2cc70 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -374,15 +374,17 @@ void Train(const nn::parallel::Rank &rank) { start_step = resume_result.global_step; size_t consumed_batches = resume_result.consumed_batches; - // TODO(jym): Replace with Sampler abstraction when available. - // Skip dataloader to resume from the correct batch position. - if (consumed_batches > 0) { - size_t start = train_iter.BatchIndex(); - // Each rank processes every ddp_world_size-th batch starting from its own rank. - // num_skips calculates how many ++ iterations to reach the saved batch position. - size_t num_skips = (consumed_batches - start) / ddp_world_size; - for (size_t i = 0; i < num_skips; ++i) { ++train_iter; } - } + // consumed_batches is the number of global dataloader batches consumed across cyclic epochs. + train_iter.SeekGlobalBatch(consumed_batches % train_loader.NumGlobalBatches()); + auto next_train_batch = [&]() { + auto batch = *train_iter; + ++train_iter; + if (train_iter == train_loader.end()) { + train_iter = train_loader.begin(); + } + ++consumed_batches; + return batch; + }; auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) { SaveCheckpoint({ @@ -457,11 +459,7 @@ void Train(const nn::parallel::Rank &rank) { infini_train::AutocastGuard autocast_guard(device.type(), dtype); // (bs, seq_len), (bs, seq_len) - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; - consumed_batches = train_iter.BatchIndex(); + auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); @@ -491,11 +489,7 @@ void Train(const nn::parallel::Rank &rank) { scheduler->Step(); } } else { - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; - consumed_batches = train_iter.BatchIndex(); + auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 216ece24..e782a6f6 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -354,15 +354,17 @@ void Train(const nn::parallel::Rank &rank) { start_step = resume_result.global_step; size_t consumed_batches = resume_result.consumed_batches; - // TODO(jym): Replace with Sampler abstraction when available. - // Skip dataloader to resume from the correct batch position. - if (consumed_batches > 0) { - size_t start = train_iter.BatchIndex(); - // Each rank processes every ddp_world_size-th batch starting from its own rank. - // num_skips calculates how many ++ iterations to reach the saved batch position. - size_t num_skips = (consumed_batches - start) / ddp_world_size; - for (size_t i = 0; i < num_skips; ++i) { ++train_iter; } - } + // consumed_batches is the number of global dataloader batches consumed across cyclic epochs. + train_iter.SeekGlobalBatch(consumed_batches % train_loader.NumGlobalBatches()); + auto next_train_batch = [&]() { + auto batch = *train_iter; + ++train_iter; + if (train_iter == train_loader.end()) { + train_iter = train_loader.begin(); + } + ++consumed_batches; + return batch; + }; auto save_checkpoint = [&](const std::filesystem::path &save_dir, int64_t global_step) { SaveCheckpoint({ @@ -435,11 +437,7 @@ void Train(const nn::parallel::Rank &rank) { infini_train::AutocastGuard autocast_guard(device.type(), dtype); // (bs, seq_len), (bs, seq_len) - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; - consumed_batches = train_iter.BatchIndex(); + auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); @@ -468,11 +466,7 @@ void Train(const nn::parallel::Rank &rank) { scheduler->Step(); } } else { - auto [x, y] = *train_iter; - // if we are trying to overfit a single batch, we reset the loader here by commenting out the line below - // TODO(dcj): support dataloader.reset() later - ++train_iter; - consumed_batches = train_iter.BatchIndex(); + auto [x, y] = next_train_batch(); x = std::make_shared(x->To(device)); y = std::make_shared(y->To(device)); diff --git a/infini_train/include/core/ccl/ccl.h b/infini_train/include/core/ccl/ccl.h index 626cb078..5cce524d 100644 --- a/infini_train/include/core/ccl/ccl.h +++ b/infini_train/include/core/ccl/ccl.h @@ -28,7 +28,7 @@ class CclImpl { virtual void GetAsyncError(const CclComm *comm, CclStatus *async_error) const; - virtual void GetUniqueId(CclUniqueId **unique_id) const; + virtual void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const; virtual void CommInitAll(CclComm **comms, int ndev, const int *devlist) const; 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> 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_; 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/src/core/ccl/ccl.cc b/infini_train/src/core/ccl/ccl.cc index 92c14cc6..1c0f2b6e 100644 --- a/infini_train/src/core/ccl/ccl.cc +++ b/infini_train/src/core/ccl/ccl.cc @@ -16,7 +16,9 @@ void CclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const { LOG(FATAL) << "CclImpl::GetAsyncError is not implemented."; } -void CclImpl::GetUniqueId(CclUniqueId **unique_id) const { LOG(FATAL) << "CclImpl::GetUniqueId is not implemented."; } +void CclImpl::CreateUniqueId(CclUniqueId **, bool) const { + LOG(FATAL) << "CclImpl::CreateUniqueId is not implemented."; +} void CclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const { LOG(FATAL) << "CclImpl::CommInitAll is not implemented."; 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 #include +#include #include #include #include @@ -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(unique_id.Data()), static_cast(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/core/ccl/cuda/nccl_impl.cc b/infini_train/src/core/ccl/cuda/nccl_impl.cc index 9e4b1a0d..d17cee46 100644 --- a/infini_train/src/core/ccl/cuda/nccl_impl.cc +++ b/infini_train/src/core/ccl/cuda/nccl_impl.cc @@ -68,14 +68,16 @@ void NcclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const } } -void NcclImpl::GetUniqueId(CclUniqueId **unique_id) const { +void NcclImpl::CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const { CHECK_NOTNULL(unique_id); if (*unique_id == nullptr) { *unique_id = new NcclUniqueId(); } auto *nccl_unique_id = dynamic_cast(*unique_id); CHECK_NOTNULL(nccl_unique_id); - NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id())); + if (generate_id) { + NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id())); + } } void NcclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const { diff --git a/infini_train/src/core/ccl/cuda/nccl_impl.h b/infini_train/src/core/ccl/cuda/nccl_impl.h index fca177fd..6fda67c4 100644 --- a/infini_train/src/core/ccl/cuda/nccl_impl.h +++ b/infini_train/src/core/ccl/cuda/nccl_impl.h @@ -17,7 +17,7 @@ class NcclImpl final : public CclImpl { void GetAsyncError(const CclComm *comm, CclStatus *async_error) const override; - void GetUniqueId(CclUniqueId **unique_id) const override; + void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const override; void CommInitAll(CclComm **comms, int ndev, const int *devlist) const override; 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 #include #include +#include #include #include @@ -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 Stack(const std::vector> &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()); @@ -35,21 +42,31 @@ std::shared_ptr Stack(const std::vector> &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> 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> data_vec; std::vector> 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> 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, 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, 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/nn/parallel/process_group.cc b/infini_train/src/nn/parallel/process_group.cc index 9dab7848..14c46241 100644 --- a/infini_train/src/nn/parallel/process_group.cc +++ b/infini_train/src/nn/parallel/process_group.cc @@ -83,11 +83,13 @@ void ProcessGroup::InitMultiProcess(const std::vector &ranks) { int upper_rank = (global_proc_rank + 1) * n_threads; core::CclUniqueId *unique_id_raw = nullptr; - ccl_impl_->GetUniqueId(&unique_id_raw); - std::unique_ptr unique_id(unique_id_raw); int min_rank = std::ranges::min(ranks); - if (min_rank < upper_rank && min_rank >= lower_rank) { + bool is_main_rank = min_rank < upper_rank && min_rank >= lower_rank; + ccl_impl_->CreateUniqueId(&unique_id_raw, is_main_rank); + std::unique_ptr unique_id(unique_id_raw); + + if (is_main_rank) { is_main_process_ = true; core::WriteUniqueIdFile(*unique_id, name_); } else { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585..a26d27d8 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) +# 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 +#include +#include +#include + +#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> operator[](size_t idx) const override { + auto data = std::make_shared(std::vector{1}, DataType::kINT64); + auto label = std::make_shared(std::vector{1}, DataType::kINT64); + *static_cast(data->DataPtr()) = static_cast(idx); + *static_cast(label->DataPtr()) = static_cast(idx + 1000); + return {data, label}; + } + + size_t Size() const override { return size_; } + +private: + size_t size_ = 0; +}; + +std::vector TensorValues(const std::shared_ptr &tensor) { + const auto *data = static_cast(tensor->DataPtr()); + return std::vector(data, data + tensor->NumElements()); +} +} // namespace + +TEST(DataLoaderTest, RegularDataLoaderKeepsPartialLastBatch) { + DataLoader loader(std::make_shared(5), 2); + + std::vector> batches; + for (const auto &[x, y] : loader) { batches.push_back(TensorValues(x)); } + + ASSERT_EQ(batches.size(), 3); + EXPECT_EQ(batches[0], (std::vector{0, 1})); + EXPECT_EQ(batches[1], (std::vector{2, 3})); + EXPECT_EQ(batches[2], (std::vector{4})); +} + +TEST(DataLoaderTest, DistributedDataLoaderSlicesFullGlobalBatchesByRank) { + const auto dataset = std::make_shared(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{0, 1})); + EXPECT_EQ(TensorValues((*r1).first), (std::vector{2, 3})); + EXPECT_EQ(TensorValues((*r2).first), (std::vector{4, 5})); + + ++r0; + ++r1; + ++r2; + + EXPECT_EQ(TensorValues((*r0).first), (std::vector{6, 7})); + EXPECT_EQ(TensorValues((*r1).first), (std::vector{8, 9})); + EXPECT_EQ(TensorValues((*r2).first), (std::vector{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(13), 2, 2, 3); + auto iter = loader.begin(); + + EXPECT_EQ(loader.NumGlobalBatches(), 2); + + EXPECT_EQ(TensorValues((*iter).first), (std::vector{4, 5})); + ++iter; + EXPECT_EQ(iter.GlobalBatchIndex(), 1); + + EXPECT_EQ(TensorValues((*iter).first), (std::vector{10, 11})); + ++iter; + EXPECT_EQ(iter, loader.end()); + + iter = loader.begin(); + EXPECT_EQ(TensorValues((*iter).first), (std::vector{4, 5})); +} diff --git a/tools/infini_run/infini_run.cc b/tools/infini_run/infini_run.cc index b632f64d..189d2752 100644 --- a/tools/infini_run/infini_run.cc +++ b/tools/infini_run/infini_run.cc @@ -1,7 +1,10 @@ +#include #include #include +#include #include #include +#include #include #include @@ -51,6 +54,28 @@ void SetEnvInt(const char *name, int value) { setenv(name, value_str.c_str(), 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) { @@ -77,6 +102,7 @@ int main(int argc, char **argv) { 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_nnodes == 1 ? GenerateLocalRunId() : ""; for (int local_proc_rank = 0; local_proc_rank < FLAGS_nproc_per_node; ++local_proc_rank) { pid_t pid = fork(); @@ -88,6 +114,9 @@ int main(int argc, char **argv) { setenv("MASTER_ADDR", master_addr.c_str(), 1); setenv("MASTER_PORT", master_port.c_str(), 1); + if (!run_id.empty()) { + setenv("INFINI_RUN_ID", run_id.c_str(), 1); + } SetEnvInt("GLOBAL_PROC_RANK", global_proc_rank); SetEnvInt("LOCAL_PROC_RANK", local_proc_rank); @@ -112,7 +141,8 @@ int main(int argc, char **argv) { pid_t child = wait(&status); if (child < 0) { perror("wait failed"); - return 1; + exit_code = 1; + break; } if (WIFEXITED(status)) { @@ -130,5 +160,8 @@ int main(int argc, char **argv) { } } + if (!run_id.empty()) { + CleanupRunUniqueIdFiles(run_id); + } return exit_code; }