Skip to content

Commit 50c8039

Browse files
chen2021673chen2021673
committed
fix: refine multi-process launcher integration
- support infini_run with or without the optional -- separator - validate node rank bounds - use infini_run only for the new 8_proc test group - standardize torchrun environment variables and device index mapping - clarify NCCL unique ID filename helpers
1 parent 0431ee7 commit 50c8039

12 files changed

Lines changed: 49 additions & 47 deletions

File tree

‎example/gpt2/main.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ void Train(const nn::parallel::Rank &rank) {
180180
const ProcessGroup *pp_pg = nullptr;
181181

182182
if (rank.IsParallel()) {
183-
device = Device(Device::DeviceType::kCUDA, global::GetLocalDeviceIndex(rank.thread_rank()));
183+
device = Device(Device::DeviceType::kCUDA, global::GetDeviceIndex(rank.thread_rank()));
184184
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
185185

186186
if (ddp_world_size > 1) {

‎example/llama3/main.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -168,7 +168,7 @@ void Train(const nn::parallel::Rank &rank) {
168168
const ProcessGroup *pp_pg = nullptr;
169169

170170
if (rank.IsParallel()) {
171-
device = Device(Device::DeviceType::kCUDA, global::GetLocalDeviceIndex(rank.thread_rank()));
171+
device = Device(Device::DeviceType::kCUDA, global::GetDeviceIndex(rank.thread_rank()));
172172
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
173173

174174
if (ddp_world_size > 1) {

‎infini_train/include/nn/parallel/global.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ inline int GetNprocPerNode() { return GlobalEnv::Instance().nproc_per_node(); }
9898
inline int GetNthreadPerProc() { return GlobalEnv::Instance().nthread_per_process(); }
9999
inline int GetGlobalProcRank() { return GlobalEnv::Instance().global_proc_rank(); }
100100
inline int GetLocalProcRank() { return GlobalEnv::Instance().local_proc_rank(); }
101-
inline int GetLocalDeviceIndex(int thread_rank = 0) { return GetLocalProcRank() * GetNthreadPerProc() + thread_rank; }
101+
inline int GetDeviceIndex(int thread_rank) { return GetLocalProcRank() * GetNthreadPerProc() + thread_rank; }
102102

103103
inline int GetTensorParallelSize() { return GlobalEnv::Instance().tensor_parallel_size(); }
104104
inline int GetSequenceParallelSize() { return GlobalEnv::Instance().sequence_parallel_size(); }

‎infini_train/src/core/ccl/ccl_utils.cc‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,36 +12,36 @@
1212

1313
namespace infini_train::core {
1414
namespace {
15-
std::string UniqueIdPath(const std::string &pg_name) {
15+
std::string UniqueIdFileName(const std::string &pg_name) {
1616
const char *run_id = std::getenv("INFINI_RUN_ID");
1717
const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
1818
return "cclUniqueId_" + prefix + pg_name + ".bin";
1919
}
2020

21-
std::string UniqueIdTmpPath(const std::string &pg_name) {
21+
std::string UniqueIdTmpFileName(const std::string &pg_name) {
2222
const char *run_id = std::getenv("INFINI_RUN_ID");
2323
const std::string prefix = run_id == nullptr ? "" : std::string(run_id) + "_";
2424
return "cclUniqueId_" + prefix + pg_name + ".tmp";
2525
}
2626
} // namespace
2727

2828
void WriteUniqueIdFile(const CclUniqueId &unique_id, const std::string &pg_name) {
29-
const std::string tmp_path = UniqueIdTmpPath(pg_name);
29+
const std::string tmp_path = UniqueIdTmpFileName(pg_name);
3030

3131
std::ofstream ofs(tmp_path, std::ios::binary);
3232
CHECK(ofs.good()) << "Failed to open unique_id tmp file for write: " << tmp_path;
3333
const size_t size = unique_id.Size();
3434
ofs.write(reinterpret_cast<const char *>(unique_id.Data()), static_cast<std::streamsize>(size));
3535
ofs.close();
3636

37-
const std::string file_path = UniqueIdPath(pg_name);
37+
const std::string file_path = UniqueIdFileName(pg_name);
3838
CHECK_EQ(std::rename(tmp_path.c_str(), file_path.c_str()), 0)
3939
<< "Failed to rename unique_id file from " << tmp_path << " to " << file_path;
4040
}
4141

4242
void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) {
4343
CHECK_NOTNULL(unique_id);
44-
const std::string file_path = UniqueIdPath(pg_name);
44+
const std::string file_path = UniqueIdFileName(pg_name);
4545

4646
while (!std::filesystem::exists(file_path)) { std::this_thread::sleep_for(std::chrono::microseconds(1000)); }
4747

@@ -57,12 +57,12 @@ void ReadUniqueIdFile(CclUniqueId *unique_id, const std::string &pg_name) {
5757
}
5858

5959
void CleanupUniqueIdFile(const std::string &pg_name) {
60-
const std::string file_path = UniqueIdPath(pg_name);
60+
const std::string file_path = UniqueIdFileName(pg_name);
6161
if (std::filesystem::exists(file_path)) {
6262
std::filesystem::remove(file_path);
6363
}
6464

65-
const std::string tmp_path = UniqueIdTmpPath(pg_name);
65+
const std::string tmp_path = UniqueIdTmpFileName(pg_name);
6666
if (std::filesystem::exists(tmp_path)) {
6767
std::filesystem::remove(tmp_path);
6868
}

‎infini_train/src/device.cc‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,15 +33,10 @@ std::string Device::ToString() const {
3333
}
3434

3535
nn::parallel::Rank Device::Rank() const {
36-
if (IsCPU()) {
37-
return {nn::parallel::global::GetGlobalProcRank(), 0, nn::parallel::global::GetNprocPerNode(),
38-
nn::parallel::global::GetNthreadPerProc()};
39-
}
40-
41-
const int thread_rank = index_ - nn::parallel::global::GetLocalDeviceIndex();
42-
CHECK_GE(thread_rank, 0) << "CUDA device index is outside the current process rank range";
36+
const int thread_rank = index_ - nn::parallel::global::GetDeviceIndex(0);
37+
CHECK_GE(thread_rank, 0) << "Device index is outside the current process rank range";
4338
CHECK_LT(thread_rank, nn::parallel::global::GetNthreadPerProc())
44-
<< "CUDA device index is outside the current process rank range";
39+
<< "Device index is outside the current process rank range";
4540
return {nn::parallel::global::GetGlobalProcRank(), thread_rank, nn::parallel::global::GetNprocPerNode(),
4641
nn::parallel::global::GetNthreadPerProc()};
4742
}

‎infini_train/src/nn/parallel/data_parallel.cc‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,7 @@ ParallelApply(const std::vector<std::shared_ptr<Module>> &modules,
6161
DataParallel::DataParallel(const std::shared_ptr<Module> &module, int dim, Device::DeviceType device_type) : dim_(dim) {
6262
devices_.reserve(global::GetNthreadPerProc());
6363
for (int thread_rank = 0; thread_rank < global::GetNthreadPerProc(); ++thread_rank) {
64-
const int device_index
65-
= device_type == Device::DeviceType::kCUDA ? global::GetLocalDeviceIndex(thread_rank) : thread_rank;
64+
const int device_index = global::GetDeviceIndex(thread_rank);
6665
devices_.emplace_back(device_type, device_index);
6766
}
6867

‎infini_train/src/nn/parallel/ddp/distributed_data_parallel.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
3333
continue;
3434
}
3535
auto device = param->GetDevice();
36-
CHECK_EQ(device.index(), global::GetLocalDeviceIndex(rank.thread_rank()))
36+
CHECK_EQ(device.index(), global::GetDeviceIndex(rank.thread_rank()))
3737
<< "All parameters must be on the same device as the module";
3838
if (!ddp_config.gradient_bucketing_enabled && ddp_config.zero_stage < 1) {
3939
auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>(
@@ -42,7 +42,7 @@ DistributedDataParallel::DistributedDataParallel(std::shared_ptr<nn::Module> mod
4242
}
4343
}
4444
for (auto &buffer : module->Buffers()) {
45-
CHECK_EQ(buffer->GetDevice().index(), global::GetLocalDeviceIndex(rank.thread_rank()))
45+
CHECK_EQ(buffer->GetDevice().index(), global::GetDeviceIndex(rank.thread_rank()))
4646
<< "All buffers must be on the same device as the module";
4747
}
4848
modules_[kModuleName] = std::move(module);

‎infini_train/src/nn/parallel/global.cc‎

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -92,21 +92,18 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq
9292

9393
CHECK(!initialized_) << "Repeated initialization of GlobalEnv!";
9494

95-
const int proc_world_size = GetEnvAsInt("WORLD_SIZE", GetEnvAsInt("PROC_WORLD_SIZE", 1));
96-
nproc_per_node_ = GetEnvAsInt("LOCAL_WORLD_SIZE", GetEnvAsInt("NPROC_PER_NODE", 1));
97-
CHECK_GT(nproc_per_node_, 0) << "NPROC_PER_NODE/LOCAL_WORLD_SIZE must be positive";
98-
CHECK_GT(proc_world_size, 0) << "PROC_WORLD_SIZE/WORLD_SIZE must be positive";
99-
CHECK_EQ(proc_world_size % nproc_per_node_, 0)
100-
<< "PROC_WORLD_SIZE/WORLD_SIZE must be divisible by NPROC_PER_NODE/LOCAL_WORLD_SIZE";
95+
const int proc_world_size = GetEnvAsInt("WORLD_SIZE", 1);
96+
nproc_per_node_ = GetEnvAsInt("LOCAL_WORLD_SIZE", 1);
97+
CHECK_GT(nproc_per_node_, 0) << "LOCAL_WORLD_SIZE must be positive";
98+
CHECK_GT(proc_world_size, 0) << "WORLD_SIZE must be positive";
99+
CHECK_EQ(proc_world_size % nproc_per_node_, 0) << "WORLD_SIZE must be divisible by LOCAL_WORLD_SIZE";
101100
nnodes_ = proc_world_size / nproc_per_node_;
102-
global_proc_rank_ = GetEnvAsInt("RANK", GetEnvAsInt("GLOBAL_PROC_RANK", 0));
103-
local_proc_rank_ = GetEnvAsInt("LOCAL_RANK", GetEnvAsInt("LOCAL_PROC_RANK", 0));
104-
CHECK_GE(global_proc_rank_, 0) << "GLOBAL_PROC_RANK/RANK must be non-negative";
105-
CHECK_LT(global_proc_rank_, proc_world_size)
106-
<< "GLOBAL_PROC_RANK/RANK must be less than PROC_WORLD_SIZE/WORLD_SIZE";
107-
CHECK_GE(local_proc_rank_, 0) << "LOCAL_PROC_RANK/LOCAL_RANK must be non-negative";
108-
CHECK_LT(local_proc_rank_, nproc_per_node_)
109-
<< "LOCAL_PROC_RANK/LOCAL_RANK must be less than NPROC_PER_NODE/LOCAL_WORLD_SIZE";
101+
global_proc_rank_ = GetEnvAsInt("RANK", 0);
102+
local_proc_rank_ = GetEnvAsInt("LOCAL_RANK", 0);
103+
CHECK_GE(global_proc_rank_, 0) << "RANK must be non-negative";
104+
CHECK_LT(global_proc_rank_, proc_world_size) << "RANK must be less than WORLD_SIZE";
105+
CHECK_GE(local_proc_rank_, 0) << "LOCAL_RANK must be non-negative";
106+
CHECK_LT(local_proc_rank_, nproc_per_node_) << "LOCAL_RANK must be less than LOCAL_WORLD_SIZE";
110107

111108
nthread_per_process_ = nthread_per_process;
112109
world_size_ = proc_world_size * nthread_per_process;

‎infini_train/src/nn/parallel/process_group.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
9999
int global_thread_rank = lower_rank + i;
100100
auto it = std::ranges::find(ranks, global_thread_rank);
101101
if (it != ranks.end()) {
102-
auto device = Device(backend_, global::GetLocalDeviceIndex(i));
102+
auto device = Device(backend_, global::GetDeviceIndex(i));
103103
core::DeviceGuard guard(device);
104104

105105
core::CclComm *comm_raw = nullptr;

‎scripts/run_models_and_profile.bash‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -439,7 +439,7 @@ check_model_inputs() {
439439
fi
440440
}
441441

442-
model_cmd_for_test() {
442+
infini_run_cmd_for_test() {
443443
local model_bin="$1"
444444
local input_bin="$2"
445445
local llmc_filepath="$3"
@@ -522,25 +522,36 @@ for ((id=0; id<num_basic_compile_commands; ++id)); do
522522

523523
for ((ti=0; ti<num_tests; ++ti)); do
524524
test_id=$(jq -r ".test_groups[$gi].tests[$ti].id" "$CONFIG_FILE")
525-
nproc_per_node="$(jq -r ".test_groups[$gi].launcher_args.nproc_per_node // empty" "$CONFIG_FILE")"
526-
: "${nproc_per_node:=${INFINI_NPROC_PER_NODE:-1}}"
525+
nproc_per_node="$(jq -r ".test_groups[$gi].infini_run_args.nproc_per_node // empty" "$CONFIG_FILE")"
527526
if tag_enabled_for_model "$group_tag" "$GPT2_TEST_GROUPS"; then
528527
LORA_WEIGHTS_DIR="$GPT2_LORA_WEIGHTS_DIR"
529528
gpt2_arg_str="$(args_string_for_test "$gi" "$ti" "gpt2" "$test_id")"
530-
gpt2_cmd="$(model_cmd_for_test "./gpt2" "$GPT2_INPUT_BIN" "$GPT2_LLMC_FILEPATH" "$gpt2_arg_str" "$nproc_per_node")"
529+
if [[ -n "$nproc_per_node" ]]; then
530+
gpt2_cmd="$(infini_run_cmd_for_test "./gpt2" "$GPT2_INPUT_BIN" "$GPT2_LLMC_FILEPATH" "$gpt2_arg_str" "$nproc_per_node")"
531+
else
532+
gpt2_cmd="${prefix}./gpt2 --input_bin ${GPT2_INPUT_BIN} --llmc_filepath ${GPT2_LLMC_FILEPATH} --device cuda ${gpt2_arg_str}"
533+
fi
531534
run_and_log "$gpt2_cmd" "gpt2_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
532535
fi
533536

534537
if tag_enabled_for_model "$group_tag" "$LLAMA3_TEST_GROUPS"; then
535538
LORA_WEIGHTS_DIR="$LLAMA3_LORA_WEIGHTS_DIR"
536539
llama3_arg_str="$(args_string_for_test "$gi" "$ti" "llama3" "$test_id")"
537-
llama3_cmd="$(model_cmd_for_test "./llama3" "$LLAMA3_INPUT_BIN" "$LLAMA3_LLMC_FILEPATH" "$llama3_arg_str" "$nproc_per_node")"
540+
if [[ -n "$nproc_per_node" ]]; then
541+
llama3_cmd="$(infini_run_cmd_for_test "./llama3" "$LLAMA3_INPUT_BIN" "$LLAMA3_LLMC_FILEPATH" "$llama3_arg_str" "$nproc_per_node")"
542+
else
543+
llama3_cmd="${prefix}./llama3 --input_bin ${LLAMA3_INPUT_BIN} --llmc_filepath ${LLAMA3_LLMC_FILEPATH} --device cuda ${llama3_arg_str}"
544+
fi
538545
run_and_log "$llama3_cmd" "llama3_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
539546
fi
540547

541548
if tag_enabled_for_model "$group_tag" "$MIXTRAL_TEST_GROUPS"; then
542549
mixtral_arg_str="$(args_string_for_test "$gi" "$ti" "mixtral" "$test_id")"
543-
mixtral_cmd="$(model_cmd_for_test "./mixtral" "$MIXTRAL_INPUT_BIN" "$MIXTRAL_LLMC_FILEPATH" "$mixtral_arg_str" "$nproc_per_node")"
550+
if [[ -n "$nproc_per_node" ]]; then
551+
mixtral_cmd="$(infini_run_cmd_for_test "./mixtral" "$MIXTRAL_INPUT_BIN" "$MIXTRAL_LLMC_FILEPATH" "$mixtral_arg_str" "$nproc_per_node")"
552+
else
553+
mixtral_cmd="${prefix}./mixtral --input_bin ${MIXTRAL_INPUT_BIN} --llmc_filepath ${MIXTRAL_LLMC_FILEPATH} --device cuda ${mixtral_arg_str}"
554+
fi
544555
run_and_log "$mixtral_cmd" "mixtral_${test_id}${log_suffix}" "$profile_flag" "$group_tag"
545556
fi
546557
done

0 commit comments

Comments
 (0)