Skip to content

Commit c716a34

Browse files
kilinchangekilinchange
committed
refactor(parallel): simplify topology process group setup
- remove centralized model-parallel process group initialization - preserve caller-owned GetOrCreate group setup - require explicit expert parallel dimensions in GlobalEnv - preserve legacy process group IDs through RankGenerator - restore DDP-specific buffer naming and simplify topology tests
1 parent 09ce22d commit c716a34

13 files changed

Lines changed: 66 additions & 206 deletions

File tree

‎example/gpt2/main.cc‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@
3131
#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
3232
#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
3333
#include "infini_train/include/nn/parallel/global.h"
34-
#include "infini_train/include/nn/parallel/parallel_state.h"
3534
#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
3635
#include "infini_train/include/nn/parallel/process_group.h"
3736
#include "infini_train/include/nn/parallel/rank.h"
@@ -612,17 +611,12 @@ int main(int argc, char *argv[]) {
612611

613612
auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check);
614613
nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,
615-
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel);
614+
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel,
615+
/*expert_parallel_size=*/1,
616+
/*expert_tensor_parallel_size=*/1);
616617
utils::PrecisionCheckEnv::Instance().Init(precision_config);
617618

618619
LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
619-
if (nn::parallel::global::GetWorldSize() > 1) {
620-
nn::parallel::InitializeModelParallel({
621-
.backend = Device::ParseType(FLAGS_device).value(),
622-
.initialize_expert_groups = false,
623-
});
624-
}
625-
626620
if (FLAGS_nthread_per_process > 1) {
627621
std::vector<std::thread> threads;
628622
for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) {

‎example/llama3/main.cc‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
3131
#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
3232
#include "infini_train/include/nn/parallel/global.h"
33-
#include "infini_train/include/nn/parallel/parallel_state.h"
3433
#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
3534
#include "infini_train/include/nn/parallel/process_group.h"
3635
#include "infini_train/include/nn/parallel/rank.h"
@@ -601,17 +600,12 @@ int main(int argc, char *argv[]) {
601600

602601
auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check);
603602
nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,
604-
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel);
603+
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel,
604+
/*expert_parallel_size=*/1,
605+
/*expert_tensor_parallel_size=*/1);
605606
utils::PrecisionCheckEnv::Instance().Init(precision_config);
606607

607608
LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
608-
if (nn::parallel::global::GetWorldSize() > 1) {
609-
nn::parallel::InitializeModelParallel({
610-
.backend = Device::ParseType(FLAGS_device).value(),
611-
.initialize_expert_groups = false,
612-
});
613-
}
614-
615609
if (FLAGS_nthread_per_process > 1) {
616610
std::vector<std::thread> threads;
617611
for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) {

‎example/mixtral/main.cc‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,9 @@ int main(int argc, char *argv[]) {
8080
/*tensor_parallel_size=*/1,
8181
/*sequence_parallel_enabled=*/false,
8282
/*pipeline_parallel_size=*/1,
83-
/*virtual_pipeline_parallel_size=*/1);
83+
/*virtual_pipeline_parallel_size=*/1,
84+
/*expert_parallel_size=*/1,
85+
/*expert_tensor_parallel_size=*/1);
8486

8587
infini_train::nn::TransformerConfig model_config = mixtral::TinyMixtralConfig();
8688
mixtral::SanitizeTinyMixtralConfig(model_config);

‎example/qwen3/main.cc‎

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
3131
#include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
3232
#include "infini_train/include/nn/parallel/global.h"
33-
#include "infini_train/include/nn/parallel/parallel_state.h"
3433
#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
3534
#include "infini_train/include/nn/parallel/process_group.h"
3635
#include "infini_train/include/nn/parallel/rank.h"
@@ -599,17 +598,12 @@ int main(int argc, char *argv[]) {
599598

600599
auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check);
601600
nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel,
602-
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel);
601+
FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel,
602+
/*expert_parallel_size=*/1,
603+
/*expert_tensor_parallel_size=*/1);
603604
utils::PrecisionCheckEnv::Instance().Init(precision_config);
604605

605606
LOG(INFO) << nn::parallel::global::ProcessGroupOverview();
606-
if (nn::parallel::global::GetWorldSize() > 1) {
607-
nn::parallel::InitializeModelParallel({
608-
.backend = Device::ParseType(FLAGS_device).value(),
609-
.initialize_expert_groups = false,
610-
});
611-
}
612-
613607
if (FLAGS_nthread_per_process > 1) {
614608
std::vector<std::thread> threads;
615609
for (int idx = 0; idx < FLAGS_nthread_per_process; ++idx) {

‎infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -183,21 +183,21 @@ class ParamAndGradBuffer {
183183
* @param params Parameters with the same parameter/gradient dtype pair.
184184
* @param param_dtype Flat parameter-buffer dtype.
185185
* @param grad_dtype Gradient storage dtype.
186-
* @param process_group Parameter synchronization group used by derived bucket groups.
187-
* @param ddp_config DDP/DistributedOptimizer bucketing and padding config.
186+
* @param ddp_pg Data-parallel process group used by derived bucket groups.
187+
* @param ddp_config DDP/DistributedOptimizer bucketing and padding config.
188188
*/
189189
ParamAndGradBuffer(const std::vector<std::shared_ptr<Tensor>> &params, DataType &param_dtype, DataType &grad_dtype,
190-
const ProcessGroup *process_group, DistributedDataParallelConfig ddp_config);
190+
const ProcessGroup *ddp_pg, DistributedDataParallelConfig ddp_config);
191191

192192
DistributedDataParallelConfig ddp_config() const { return ddp_config_; }
193193

194194
std::shared_ptr<Tensor> param_buffer() const { return param_buffer_; }
195195

196196
std::shared_ptr<Tensor> grad_buffer() const { return grad_buffer_; }
197197

198-
const ProcessGroup *process_group() const { return process_group_; }
198+
const ProcessGroup *ddp_pg() const { return ddp_pg_; }
199199

200-
size_t process_group_size() const { return process_group_size_; }
200+
size_t ddp_world_size() const { return ddp_world_size_; }
201201

202202
std::vector<std::shared_ptr<ParamAndGradBucket>> buckets() const { return buckets_; }
203203

@@ -220,8 +220,8 @@ class ParamAndGradBuffer {
220220
size_t numel_ = 0;
221221
size_t numel_unpadded_ = 0;
222222

223-
const ProcessGroup *process_group_ = nullptr;
224-
size_t process_group_size_ = 1;
223+
const ProcessGroup *ddp_pg_ = nullptr;
224+
size_t ddp_world_size_ = 1;
225225
std::vector<std::shared_ptr<ParamAndGradBucket>> buckets_;
226226

227227
bool need_rebind_grad_views_ = true;

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

Lines changed: 5 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ class GlobalEnv {
5454
static GlobalEnv &Instance();
5555

5656
void Init(int threads_per_process, int tensor_parallel_size, bool sequence_parallel_enabled,
57-
int pipeline_parallel_size, int virtual_pipeline_parallel_size, int expert_parallel_size = 1,
58-
std::optional<int> expert_tensor_parallel_size = std::nullopt);
57+
int pipeline_parallel_size, int virtual_pipeline_parallel_size, int expert_parallel_size,
58+
std::optional<int> expert_tensor_parallel_size);
5959

6060
int nnodes() const;
6161

@@ -127,8 +127,8 @@ class GlobalEnv {
127127
};
128128

129129
inline void InitAllEnv(int nthread_per_process, int tensor_parallel_size, bool sequence_parallel_enabled,
130-
int pipeline_parallel_size, int virtual_pipeline_parallel, int expert_parallel_size = 1,
131-
std::optional<int> expert_tensor_parallel_size = std::nullopt) {
130+
int pipeline_parallel_size, int virtual_pipeline_parallel, int expert_parallel_size,
131+
std::optional<int> expert_tensor_parallel_size) {
132132
GlobalEnv::Instance().Init(nthread_per_process, tensor_parallel_size, sequence_parallel_enabled,
133133
pipeline_parallel_size, virtual_pipeline_parallel, expert_parallel_size,
134134
expert_tensor_parallel_size);
@@ -162,28 +162,7 @@ inline void GetCoordOf(int rank, int &dp, int &tp, int &pp) {
162162
GetDenseRankGenerator().CoordOf(rank, dp, tp, pp, ep);
163163
}
164164

165-
inline int GetGroupId(Axis target, int rank) {
166-
const auto &rank_generator = GetDenseRankGenerator();
167-
int dp = 0;
168-
int tp = 0;
169-
int pp = 0;
170-
int ep = 0;
171-
rank_generator.CoordOf(rank, dp, tp, pp, ep);
172-
const std::array<int, AXIS_COUNT> coordinates{dp, tp, pp, ep};
173-
174-
// Preserve the legacy dense group numbering used in process-group names.
175-
int group_id = 0;
176-
int group_stride = 1;
177-
for (int index = AXIS_COUNT - 1; index >= 0; --index) {
178-
const Axis axis = rank_generator.order()[index];
179-
if (axis == target || axis == EP) {
180-
continue;
181-
}
182-
group_id += coordinates[axis] * group_stride;
183-
group_stride *= rank_generator.AxisSize(axis);
184-
}
185-
return group_id;
186-
}
165+
inline int GetGroupId(Axis target, int rank) { return GetDenseRankGenerator().GroupId(target, rank); }
187166

188167
inline int GetGroupId(Axis target, int dp, int tp, int pp) { return GetGroupId(target, GetRankOf(dp, tp, pp)); }
189168

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

Lines changed: 0 additions & 18 deletions
This file was deleted.

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

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -418,11 +418,11 @@ void ParamAndGradBucketGroup::SetNextParamGatherBucketGroup(std::shared_ptr<Para
418418
}
419419

420420
ParamAndGradBuffer::ParamAndGradBuffer(const std::vector<std::shared_ptr<Tensor>> &params, DataType &param_dtype,
421-
DataType &grad_dtype, const ProcessGroup *process_group,
421+
DataType &grad_dtype, const ProcessGroup *ddp_pg,
422422
DistributedDataParallelConfig ddp_config)
423-
: params_(std::move(params)), process_group_(process_group), ddp_config_(ddp_config) {
424-
CHECK_NOTNULL(process_group_);
425-
process_group_size_ = static_cast<size_t>(process_group_->GetGroupSize());
423+
: params_(std::move(params)), ddp_pg_(ddp_pg), ddp_config_(ddp_config) {
424+
CHECK_NOTNULL(ddp_pg_);
425+
ddp_world_size_ = static_cast<size_t>(ddp_pg_->GetGroupSize());
426426

427427
grads_.clear();
428428
grads_.resize(params_.size());
@@ -448,7 +448,7 @@ void ParamAndGradBuffer::BuildBuckets(DataType param_dtype, DataType grad_dtype)
448448
if (ddp_config_.zero_stage >= 1) {
449449
// According to Megatron-LM, ensure that all buckets start at a memory address that is 256B
450450
// aligned(128 values since params and grads use >= 16-bit precision)
451-
size_t lcm_val = std::lcm(process_group_size_, kBucketEndAlignElements);
451+
size_t lcm_val = std::lcm(ddp_world_size_, kBucketEndAlignElements);
452452
if (ddp_config_.pad_buckets_for_high_nccl_busbw) {
453453
// According to Megatron-LM, when the bucket size is divisible by a large power of 2 (2^16),
454454
// NCCL collectives can have high bus bandwidth at large DP counts
@@ -519,7 +519,7 @@ void ParamAndGradBuffer::BuildBuckets(DataType param_dtype, DataType grad_dtype)
519519
CHECK(numel_unpadded_ <= numel_);
520520
if (ddp_config_.zero_stage >= 1) {
521521
// numel must be multiple of ddp size (so that reduce-scatter could easily shard the buffer among ranks)
522-
CHECK_EQ(numel_ % process_group_size_, 0);
522+
CHECK_EQ(numel_ % ddp_world_size_, 0);
523523
} else {
524524
CHECK_EQ(numel_, numel_unpadded_);
525525
}
@@ -550,8 +550,8 @@ void ParamAndGradBuffer::BuildBuckets(DataType param_dtype, DataType grad_dtype)
550550
= [&](const std::vector<std::shared_ptr<Tensor>> &bucket_params, size_t start_index, size_t end_index,
551551
size_t num_elements_unpadded, size_t bucket_id) -> std::shared_ptr<ParamAndGradBucket> {
552552
if (ddp_config_.zero_stage >= 1) {
553-
CHECK_EQ(start_index % process_group_size_, 0);
554-
CHECK_EQ(end_index % process_group_size_, 0);
553+
CHECK_EQ(start_index % ddp_world_size_, 0);
554+
CHECK_EQ(end_index % ddp_world_size_, 0);
555555
CHECK_EQ(bucket_indices_.at(bucket_id).first, start_index);
556556
CHECK_EQ(bucket_indices_.at(bucket_id).second, end_index);
557557
}
@@ -691,20 +691,19 @@ PartitionBuckets(const std::vector<std::shared_ptr<ParamAndGradBuffer>> &buffers
691691
if (force_single_bucket_group) {
692692
std::vector<std::shared_ptr<ParamAndGradBucket>> all_buckets;
693693
auto ddp_config = buffers.front()->ddp_config();
694-
auto process_group = buffers.front()->process_group();
695-
auto process_group_size = buffers.front()->process_group_size();
694+
auto ddp_pg = buffers.front()->ddp_pg();
695+
auto ddp_world_size = buffers.front()->ddp_world_size();
696696

697697
for (const auto &buffer : buffers) {
698-
CHECK(buffer->process_group() == process_group)
699-
<< "PartitionBuckets: buffers have different process groups.";
700-
CHECK(buffer->process_group_size() == process_group_size)
701-
<< "PartitionBuckets: buffers have different process group sizes.";
698+
CHECK(buffer->ddp_pg() == ddp_pg) << "PartitionBuckets: buffers have different ddp_pg.";
699+
CHECK(buffer->ddp_world_size() == ddp_world_size)
700+
<< "PartitionBuckets: buffers have different ddp_world_size.";
702701

703702
all_buckets.insert(all_buckets.end(), buffer->buckets().begin(), buffer->buckets().end());
704703
}
705704

706705
bucket_groups.push_back(
707-
std::make_shared<ParamAndGradBucketGroup>(all_buckets, process_group, process_group_size, ddp_config));
706+
std::make_shared<ParamAndGradBucketGroup>(all_buckets, ddp_pg, ddp_world_size, ddp_config));
708707
return bucket_groups;
709708
}
710709

@@ -715,7 +714,7 @@ PartitionBuckets(const std::vector<std::shared_ptr<ParamAndGradBuffer>> &buffers
715714
std::vector<std::shared_ptr<ParamAndGradBucket>> single_bucket_list;
716715
single_bucket_list.push_back(bucket);
717716
bucket_groups.push_back(std::make_shared<ParamAndGradBucketGroup>(
718-
single_bucket_list, buffer->process_group(), buffer->process_group_size(), buffer->ddp_config()));
717+
single_bucket_list, buffer->ddp_pg(), buffer->ddp_world_size(), buffer->ddp_config()));
719718
}
720719
}
721720

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -179,11 +179,12 @@ int RankGenerator::GroupId(std::initializer_list<Axis> varying_axes, int global_
179179
CoordOf(global_rank, dp, tp, pp, ep);
180180
const std::array<int, AXIS_COUNT> coordinates{dp, tp, pp, ep};
181181

182-
// The first unmasked axis in rank order varies fastest in the group
183-
// ordinal, matching generate_masked_orthogonal_rank_groups in Megatron-LM.
182+
// Preserve the legacy Layout::GroupId numbering used in process-group
183+
// names: fixed axes later in rank order vary faster in the group ordinal.
184184
int group_id = 0;
185185
int group_stride = 1;
186-
for (const Axis axis : order_) {
186+
for (int index = AXIS_COUNT - 1; index >= 0; --index) {
187+
const Axis axis = order_[index];
187188
if (varying_axis_mask[axis]) {
188189
continue;
189190
}

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

Lines changed: 0 additions & 61 deletions
This file was deleted.

0 commit comments

Comments
 (0)