Skip to content

Commit b176ede

Browse files
kilinchangekilinchange
committed
feat: add test bash
1 parent d786e95 commit b176ede

5 files changed

Lines changed: 549 additions & 71 deletions

File tree

‎example/gpt2/main.cc‎

Lines changed: 10 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,6 @@
1-
#include <algorithm>
21
#include <chrono>
32
#include <cstdlib>
43
#include <format>
5-
#include <iterator>
64
#include <memory>
75
#include <optional>
86
#include <thread>
@@ -151,33 +149,28 @@ void Train(const nn::parallel::Rank &rank) {
151149
const ProcessGroup *tp_pg = nullptr;
152150
const ProcessGroup *pp_pg = nullptr;
153151

154-
auto rank_in_group = [&](const std::vector<int> &group_ranks) {
155-
auto it = std::find(group_ranks.begin(), group_ranks.end(), rank.GlobalRank());
156-
CHECK(it != group_ranks.end());
157-
return static_cast<int>(std::distance(group_ranks.begin(), it));
158-
};
159-
160152
if (rank.IsParallel()) {
161153
auto parallel_device_type
162154
= (FLAGS_device == kDeviceMACA) ? Device::DeviceType::kMACA : Device::DeviceType::kCUDA;
163155
device = Device(parallel_device_type, rank.thread_rank());
164156

165-
// NOTE(dcj): On MACA, defer ProcessGroup creation until AFTER the model
166-
// has been uploaded to the device. MCCL init registers internal P2P
167-
// buffers that leave stale read-only mappings in the address ranges
168-
// mcMalloc later hands out; allocating the model first keeps it in a
169-
// P2P-clean region of the VA space and avoids the init-time race on
170-
// multi-thread DDP+TP. Mirrors the llama3 fix combo.
157+
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
171158
if (ddp_world_size > 1) {
172-
ddp_rank = rank_in_group(GetDataParallelGroupRanks(rank.GlobalRank()));
159+
ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()),
160+
GetDataParallelGroupRanks(rank.GlobalRank()));
161+
ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank());
173162
}
174163
if (tp_world_size > 1) {
175-
tp_rank = rank_in_group(GetTensorParallelGroupRanks(rank.GlobalRank()));
164+
tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()),
165+
GetTensorParallelGroupRanks(rank.GlobalRank()));
166+
tp_rank = tp_pg->GetGroupRank(rank.GlobalRank());
176167
// NOTE(zbl): Reserved for VocabParallelEmbedding
177168
nn::parallel::tp_rank = tp_rank;
178169
}
179170
if (pp_world_size > 1) {
180-
pp_rank = rank_in_group(GetPipelineParallelGroupRanks(rank.GlobalRank()));
171+
pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()),
172+
GetPipelineParallelGroupRanks(rank.GlobalRank()));
173+
pp_rank = pp_pg->GetGroupRank(rank.GlobalRank());
181174
nn::parallel::pp_rank = pp_rank;
182175
}
183176
} else {
@@ -213,22 +206,6 @@ void Train(const nn::parallel::Rank &rank) {
213206

214207
model->To(device);
215208

216-
if (rank.IsParallel()) {
217-
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
218-
if (ddp_world_size > 1) {
219-
ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()),
220-
GetDataParallelGroupRanks(rank.GlobalRank()));
221-
}
222-
if (tp_world_size > 1) {
223-
tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()),
224-
GetTensorParallelGroupRanks(rank.GlobalRank()));
225-
}
226-
if (pp_world_size > 1) {
227-
pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()),
228-
GetPipelineParallelGroupRanks(rank.GlobalRank()));
229-
}
230-
}
231-
232209
utils::PrecisionChecker::BuildNameMap(model.get());
233210

234211
// Get chunk size before wrapping with LoRA (needed for PipelineParallel)

‎example/llama3/main.cc‎

Lines changed: 10 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
#include <algorithm>
21
#include <cstdlib>
32
#include <format>
4-
#include <iterator>
53
#include <memory>
64
#include <optional>
75
#include <thread>
@@ -133,36 +131,28 @@ void Train(const nn::parallel::Rank &rank) {
133131
const ProcessGroup *tp_pg = nullptr;
134132
const ProcessGroup *pp_pg = nullptr;
135133

136-
auto rank_in_group = [&](const std::vector<int> &group_ranks) {
137-
auto it = std::find(group_ranks.begin(), group_ranks.end(), rank.GlobalRank());
138-
CHECK(it != group_ranks.end());
139-
return static_cast<int>(std::distance(group_ranks.begin(), it));
140-
};
141-
142134
if (rank.IsParallel()) {
143135
auto parallel_device_type
144136
= (FLAGS_device == kDeviceMACA) ? Device::DeviceType::kMACA : Device::DeviceType::kCUDA;
145137
device = Device(parallel_device_type, rank.thread_rank());
146138

147-
// NOTE(dcj): On MACA, defer ProcessGroup creation until AFTER the model
148-
// has been uploaded to the device. MCCL init registers internal P2P
149-
// buffers that leave stale read-only mappings in the address ranges
150-
// mcMalloc later hands out; allocating the model first keeps it in a
151-
// P2P-clean region of the VA space and avoids the "Writing to readonly
152-
// page" race on multi-thread DDP.
153-
//
154-
// Compute the in-group ranks now so model loading (which reads
155-
// nn::parallel::tp_rank) gets the correct shard.
139+
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
156140
if (ddp_world_size > 1) {
157-
ddp_rank = rank_in_group(GetDataParallelGroupRanks(rank.GlobalRank()));
141+
ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()),
142+
GetDataParallelGroupRanks(rank.GlobalRank()));
143+
ddp_rank = ddp_pg->GetGroupRank(rank.GlobalRank());
158144
}
159145
if (tp_world_size > 1) {
160-
tp_rank = rank_in_group(GetTensorParallelGroupRanks(rank.GlobalRank()));
146+
tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()),
147+
GetTensorParallelGroupRanks(rank.GlobalRank()));
148+
tp_rank = tp_pg->GetGroupRank(rank.GlobalRank());
161149
// NOTE(zbl): Reserved for VocabParallelEmbedding
162150
nn::parallel::tp_rank = tp_rank;
163151
}
164152
if (pp_world_size > 1) {
165-
pp_rank = rank_in_group(GetPipelineParallelGroupRanks(rank.GlobalRank()));
153+
pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()),
154+
GetPipelineParallelGroupRanks(rank.GlobalRank()));
155+
pp_rank = pp_pg->GetGroupRank(rank.GlobalRank());
166156
nn::parallel::pp_rank = pp_rank;
167157
}
168158
} else {
@@ -197,22 +187,6 @@ void Train(const nn::parallel::Rank &rank) {
197187

198188
model->To(device);
199189

200-
if (rank.IsParallel()) {
201-
auto *pg_factory = ProcessGroupFactory::Instance(device.type());
202-
if (ddp_world_size > 1) {
203-
ddp_pg = pg_factory->GetOrCreate(GetDataParallelProcessGroupName(rank.GlobalRank()),
204-
GetDataParallelGroupRanks(rank.GlobalRank()));
205-
}
206-
if (tp_world_size > 1) {
207-
tp_pg = pg_factory->GetOrCreate(GetTensorParallelProcessGroupName(rank.GlobalRank()),
208-
GetTensorParallelGroupRanks(rank.GlobalRank()));
209-
}
210-
if (pp_world_size > 1) {
211-
pp_pg = pg_factory->GetOrCreate(GetPipelineParallelProcessGroupName(rank.GlobalRank()),
212-
GetPipelineParallelGroupRanks(rank.GlobalRank()));
213-
}
214-
}
215-
216190
utils::PrecisionChecker::BuildNameMap(model.get());
217191

218192
// Apply LoRA using GetLoRAModel (in-place injection)

‎infini_train/src/core/runtime/maca/maca_guard_impl.cc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ MacaGuardImpl::MacaGuardImpl() {
129129
// static initialization (before main), so we setenv here in the ctor
130130
// just prior to mcInit(0). Users can override by setting the env var
131131
// themselves before launch.
132-
// setenv("MACA_LAUNCH_BLOCKING", "1", 0);
132+
setenv("MACA_LAUNCH_BLOCKING", "1", 0);
133133

134134
// When TP > 1 on MACA, disable the MCCL-level P2P path to prevent multi-PG
135135
// init deadlocks (threads concurrently creating both DP and TP comms hang

‎scripts/run_models_and_profile.bash‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ Options:
1313
EOF
1414
}
1515

16-
CONFIG_FILE="test_config.json"
16+
CONFIG_FILE="test_config_maca.json"
1717
ONLY_RUN_TAGS=""
1818

1919
while [[ $# -gt 0 ]]; do

0 commit comments

Comments
 (0)