diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..5551f724 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -563,7 +563,9 @@ int main(int argc, char *argv[]) { auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, - FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); + FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel, + /*expert_parallel_size=*/1, + /*expert_tensor_parallel_size=*/std::nullopt); utils::PrecisionCheckEnv::Instance().Init(precision_config); LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e0808..2c4bc389 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -540,7 +540,9 @@ int main(int argc, char *argv[]) { auto precision_config = utils::PrecisionCheckConfig::Parse(FLAGS_precision_check); nn::parallel::global::InitAllEnv(FLAGS_nthread_per_process, FLAGS_tensor_parallel, FLAGS_sequence_parallel, - FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel); + FLAGS_pipeline_parallel, FLAGS_virtual_pipeline_parallel, + /*expert_parallel_size=*/1, + /*expert_tensor_parallel_size=*/std::nullopt); utils::PrecisionCheckEnv::Instance().Init(precision_config); LOG(INFO) << nn::parallel::global::ProcessGroupOverview(); diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 84853865..ccf8ec4d 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -68,7 +68,9 @@ int main(int argc, char *argv[]) { /*tensor_parallel_size=*/1, /*sequence_parallel_enabled=*/false, /*pipeline_parallel_size=*/1, - /*virtual_pipeline_parallel_size=*/1); + /*virtual_pipeline_parallel_size=*/1, + /*expert_parallel_size=*/1, + /*expert_tensor_parallel_size=*/std::nullopt); infini_train::nn::TransformerConfig model_config = mixtral::TinyMixtralConfig(); mixtral::SanitizeTinyMixtralConfig(model_config); diff --git a/infini_train/include/nn/parallel/global.h b/infini_train/include/nn/parallel/global.h index 9373100f..ed446de3 100644 --- a/infini_train/include/nn/parallel/global.h +++ b/infini_train/include/nn/parallel/global.h @@ -1,6 +1,10 @@ #pragma once +#include +#include +#include #include +#include #include #include @@ -8,20 +12,41 @@ namespace infini_train::nn::parallel::global { extern thread_local int thread_global_rank; -enum Axis : uint8_t { DP = 0, TP = 1, PP = 2, AXIS_COUNT = 3 }; +// TP and DP are view-local axes. In the expert rank generator they represent +// expert tensor parallelism (ETP) and expert data parallelism (EDP), matching +// Megatron-LM's RankGenerator semantics. +enum Axis : uint8_t { DP = 0, TP = 1, PP = 2, EP = 3, AXIS_COUNT = 4 }; -struct Layout { - int sizes[AXIS_COUNT]{1, 1, 1}; - // Default order according to Megatron-LM is TP-DP-PP. Ref: - // https://github.com/NVIDIA/Megatron-LM/blob/e07c4a4450b6faa187a1ef4ec082a35ad7d2f085/megatron/core/parallel_state.py#L618 - Axis order[AXIS_COUNT]{TP, DP, PP}; - int strides[AXIS_COUNT]{1, 1, 1}; +class RankGenerator { +public: + // Default order according to Megatron-LM is tp-cp-ep-dp-pp. Ref: + // https://github.com/NVIDIA/Megatron-LM/blob/cf2f07d7b1315c96c05554c670c43207c6783e5e/megatron/core/parallel_state.py#L561 + explicit RankGenerator(int tensor_parallel_size, int expert_parallel_size, int data_parallel_size, + int pipeline_parallel_size, std::array order = {TP, EP, DP, PP}); + + int AxisSize(Axis axis) const; + int WorldSize() const; + const std::array &order() const; + + // Generate all rank groups by varying the specified axes and fixing all + // other axes, following Megatron-LM's RankGenerator::get_ranks() ordering. + std::vector> GetRanks(std::initializer_list varying_axes) const; + std::vector> GetRanks(Axis varying_axis) const; + int GroupId(std::initializer_list varying_axes, int global_rank) const; + int GroupId(Axis varying_axis, int global_rank) const; + std::vector GroupRanks(std::initializer_list varying_axes, int global_rank) const; + std::vector GroupRanks(Axis varying_axis, int global_rank) const; + +private: void InitStrides(); - int RankOf(int dp, int tp, int pp) const; - void CoordOf(int rank, int &dp, int &tp, int &pp) const; - int GroupId(Axis target, int dp, int tp, int pp) const; - std::vector GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp) const; + + int RankOf(int dp, int tp, int pp, int ep) const; + void CoordOf(int rank, int &dp, int &tp, int &pp, int &ep) const; + + std::array sizes_{1, 1, 1, 1}; + std::array order_{TP, EP, DP, PP}; + std::array strides_{1, 1, 1, 1}; }; class GlobalEnv { @@ -29,7 +54,8 @@ class GlobalEnv { static GlobalEnv &Instance(); void Init(int threads_per_process, int tensor_parallel_size, bool sequence_parallel_enabled, - int pipeline_parallel_size, int virtual_pipeline_parallel_size); + int pipeline_parallel_size, int virtual_pipeline_parallel_size, int expert_parallel_size, + std::optional expert_tensor_parallel_size); int nnodes() const; @@ -45,17 +71,26 @@ class GlobalEnv { int tensor_parallel_size() const; + int expert_tensor_parallel_size() const; + int sequence_parallel_size() const; bool sequence_parallel_enabled() const; int data_parallel_size() const; + int expert_data_parallel_size() const; + int pipeline_parallel_size() const; int virtual_pipeline_parallel_size() const; - Layout layout() const; + int expert_parallel_size() const; + + // Two logical rank views over the same physical world. The dense view uses + // TP/DP, while the expert view exposes the same axes as ETP/EDP. + const RankGenerator &dense_rank_generator() const; + const RankGenerator &expert_rank_generator() const; private: GlobalEnv() = default; @@ -74,23 +109,29 @@ class GlobalEnv { int local_proc_rank_ = 0; int tensor_parallel_size_ = 1; + int expert_tensor_parallel_size_ = 1; bool sequence_parallel_enabled_ = false; int data_parallel_size_ = 1; + int expert_data_parallel_size_ = 1; int pipeline_parallel_size_ = 1; int virtual_pipeline_parallel_size_ = 1; + int expert_parallel_size_ = 1; mutable std::mutex mutex_; bool initialized_ = false; - Layout layout_; + RankGenerator dense_rank_generator_{1, 1, 1, 1}; + RankGenerator expert_rank_generator_{1, 1, 1, 1}; }; inline void InitAllEnv(int nthread_per_process, int tensor_parallel_size, bool sequence_parallel_enabled, - int pipeline_parallel_size, int virtual_pipeline_parallel) { + int pipeline_parallel_size, int virtual_pipeline_parallel, int expert_parallel_size, + std::optional expert_tensor_parallel_size) { GlobalEnv::Instance().Init(nthread_per_process, tensor_parallel_size, sequence_parallel_enabled, - pipeline_parallel_size, virtual_pipeline_parallel); + pipeline_parallel_size, virtual_pipeline_parallel, expert_parallel_size, + expert_tensor_parallel_size); } inline int GetNnodes() { return GlobalEnv::Instance().nnodes(); } inline int GetWorldSize() { return GlobalEnv::Instance().world_size(); } @@ -100,88 +141,90 @@ inline int GetGlobalProcRank() { return GlobalEnv::Instance().global_proc_rank() inline int GetLocalProcRank() { return GlobalEnv::Instance().local_proc_rank(); } inline int GetTensorParallelSize() { return GlobalEnv::Instance().tensor_parallel_size(); } +inline int GetExpertTensorParallelSize() { return GlobalEnv::Instance().expert_tensor_parallel_size(); } inline int GetSequenceParallelSize() { return GlobalEnv::Instance().sequence_parallel_size(); } inline bool GetSequenceParallelEnabled() { return GlobalEnv::Instance().sequence_parallel_enabled(); } inline int GetDataParallelSize() { return GlobalEnv::Instance().data_parallel_size(); } +inline int GetExpertDataParallelSize() { return GlobalEnv::Instance().expert_data_parallel_size(); } inline int GetPipelineParallelSize() { return GlobalEnv::Instance().pipeline_parallel_size(); } inline int GetVirtualPipelineParallelSize() { return GlobalEnv::Instance().virtual_pipeline_parallel_size(); } +inline int GetExpertParallelSize() { return GlobalEnv::Instance().expert_parallel_size(); } -// ========================= -// Layout Helper Functions -// ========================= - -/** - * @brief Get the global rank corresponding to the given (dp, tp, pp) coordinate. - */ -inline int GetRankOf(int dp, int tp, int pp) { return GlobalEnv::Instance().layout().RankOf(dp, tp, pp); } -/** - * @brief Get the (dp, tp, pp) coordinate corresponding to the given global rank. - */ -inline void GetCoordOf(int rank, int &dp, int &tp, int &pp) { - return GlobalEnv::Instance().layout().CoordOf(rank, dp, tp, pp); -} +inline const RankGenerator &GetDenseRankGenerator() { return GlobalEnv::Instance().dense_rank_generator(); } +inline const RankGenerator &GetExpertRankGenerator() { return GlobalEnv::Instance().expert_rank_generator(); } /** - * @brief Get the group ID that the (dp, tp, pp) coordinate belongs to along a given parallel axis. - */ -inline int GetGroupId(Axis target, int dp, int tp, int pp) { - return GlobalEnv::Instance().layout().GroupId(target, dp, tp, pp); -} -/** - * @brief Get the group ID that a given rank belongs to along a specific parallel axis. - */ -inline int GetGroupId(Axis target, int rank) { - int dp, tp, pp; - GetCoordOf(rank, dp, tp, pp); - return GlobalEnv::Instance().layout().GroupId(target, dp, tp, pp); -} - -/** - * @brief Get all ranks that belong to the same group as the given (dp, tp, pp) coordinate - * along a specified parallel axis (e.g., all ranks in the same TP group). - */ -inline std::vector GetGroupRanks(Axis target, int dp, int tp, int pp) { - return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp); -} - -/** - * @brief Get all ranks that belong to the same group as the given rank - * along a specified parallel axis (e.g., all ranks in the same DP group). + * @brief Generates a human-readable overview of dense parallel groups. + * + * The dense view reports only TP, DP, and PP. Its size-one EP axis is an + * implementation detail of RankGenerator and is intentionally omitted. If + * dense_rank_generator is omitted, the global dense rank generator is used. + * + * @param dense_rank_generator Rank generator for the dense TP/DP/PP view. + * @param skip_trivial_axes If true, groups whose size is one are marked as + * "unenabled" and their rank lists are omitted. + * + * Example output for dense {TP=2, DP=2, PP=1}: + * @code + * === Parallel Communication Groups === + * world_size = 4, config: {DP=2, TP=2, PP=1} + * [Dense Rank View] shape={TP=2, DP=2, PP=1}, order={ TP -> DP -> PP } + * [DP] size=2, num_groups=2 + * - DP 0: [0, 2] + * - DP 1: [1, 3] + * [TP] size=2, num_groups=2 + * - TP 0: [0, 1] + * - TP 1: [2, 3] + * [PP] size=1, unenabled + * @endcode + * + * @return A formatted overview suitable for logging and topology validation. */ -inline std::vector GetGroupRanks(Axis target, int rank) { - int dp, tp, pp; - GetCoordOf(rank, dp, tp, pp); - return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp); -} +std::string ProcessGroupOverview(const RankGenerator &dense_rank_generator = GetDenseRankGenerator(), + bool skip_trivial_axes = true); /** - * @brief Generate a human-readable overview of all parallel communication groups. - * - * The output is intended for debugging, logging, and runtime verification of - * distributed parallelism configuration. + * @brief Generates a human-readable overview of dense and expert parallel groups. * - * @param L The Layout describing DP / TP / PP sizes and axis ordering. - * @param skip_trivial_axes - * If true, axes whose size <= 1(i.e. parallel strategy that is not enabled) - * will be marked as "unenabled" and their detailed group listing will be skipped. + * Both rank generators describe logical views over the same physical ranks. + * In the expert view, the TP and DP axes represent ETP and EDP. This overload + * additionally reports EP and combined ETP+EP groups. * - * @return A formatted string containing the full overview of process groups. + * @param dense_rank_generator Rank generator for the dense TP/DP/PP view. + * @param expert_rank_generator Rank generator for the ETP/EP/EDP/PP view. + * @param skip_trivial_axes If true, groups whose size is one are marked as + * "unenabled" and their rank lists are omitted. * - * Example: - * === Parallel Communication Groups === - * world_size = 8, config: {DP=2, TP=4, PP=1}, order: {TP -> DP -> PP} - * [DP] size=2, num_groups=4 - * - DP 0 (dp=-, tp=0, pp=0): [0, 4] - * - DP 1 (dp=-, tp=1, pp=0): [1, 5] - * - DP 2 (dp=-, tp=2, pp=0): [2, 6] - * - DP 3 (dp=-, tp=3, pp=0): [3, 7] + * Example output for dense {TP=2, DP=2, PP=1} and expert + * {ETP=1, EP=2, EDP=2, PP=1} views: + * @code + * === Parallel Communication Groups === + * world_size = 4, config: {DP=2, EDP=2, TP=2, ETP=1, PP=1, EP=2} + * [Dense Rank View] shape={TP=2, DP=2, PP=1}, order={ TP -> DP -> PP } + * [DP] size=2, num_groups=2 + * - DP 0: [0, 2] + * - DP 1: [1, 3] + * [TP] size=2, num_groups=2 + * - TP 0: [0, 1] + * - TP 1: [2, 3] + * [PP] size=1, unenabled * - * [TP] size=4, num_groups=2 - * - TP 0 (dp=0, tp=-, pp=0): [0, 1, 2, 3] - * - TP 1 (dp=1, tp=-, pp=0): [4, 5, 6, 7] + * [Expert Rank View] shape={ETP=1, EP=2, EDP=2, PP=1}, order={ ETP -> EP -> EDP -> PP } + * [EDP] size=2, num_groups=2 + * - EDP 0: [0, 2] + * - EDP 1: [1, 3] + * [ETP] size=1, unenabled + * [EP] size=2, num_groups=2 + * - EP 0: [0, 1] + * - EP 1: [2, 3] + * [ETP_EP] size=2, num_groups=2 + * - ETP_EP 0: [0, 1] + * - ETP_EP 1: [2, 3] + * @endcode * - * [PP] size=1, unenabled + * @return A formatted overview suitable for logging and topology validation. */ -std::string ProcessGroupOverview(const Layout &L = GlobalEnv::Instance().layout(), bool skip_trivial_axes = true); +std::string ProcessGroupOverview(const RankGenerator &dense_rank_generator, const RankGenerator &expert_rank_generator, + bool skip_trivial_axes = true); } // namespace infini_train::nn::parallel::global diff --git a/infini_train/include/nn/parallel/utils.h b/infini_train/include/nn/parallel/utils.h index 4dc737e7..3592f82d 100644 --- a/infini_train/include/nn/parallel/utils.h +++ b/infini_train/include/nn/parallel/utils.h @@ -9,18 +9,38 @@ class Tensor; } // namespace infini_train namespace infini_train::nn::parallel { +// DP group generated from the dense rank view. std::string GetDataParallelProcessGroupName(int global_rank); +// DP group generated from the expert rank view, exposed as EDP. +std::string GetExpertDataParallelProcessGroupName(int global_rank); + std::string GetTensorParallelProcessGroupName(int global_rank); +// TP group generated from the expert rank view, exposed as ETP. +std::string GetExpertTensorParallelProcessGroupName(int global_rank); + std::string GetPipelineParallelProcessGroupName(int global_rank); +std::string GetExpertParallelProcessGroupName(int global_rank); + +// The expert rank view's ETP + EP group. +std::string GetExpertTensorAndExpertParallelProcessGroupName(int global_rank); + std::vector GetDataParallelGroupRanks(int global_rank); +std::vector GetExpertDataParallelGroupRanks(int global_rank); + std::vector GetTensorParallelGroupRanks(int global_rank); +std::vector GetExpertTensorParallelGroupRanks(int global_rank); + std::vector GetPipelineParallelGroupRanks(int global_rank); +std::vector GetExpertParallelGroupRanks(int global_rank); + +std::vector GetExpertTensorAndExpertParallelGroupRanks(int global_rank); + // TP/SP Communication Helper Functions std::vector> GatherFromTPRegionFunc(const std::shared_ptr &input); std::vector> ReduceScatterToSPRegionFunc(const std::shared_ptr &input); diff --git a/infini_train/src/nn/parallel/global.cc b/infini_train/src/nn/parallel/global.cc index 65a3208e..6c087a26 100644 --- a/infini_train/src/nn/parallel/global.cc +++ b/infini_train/src/nn/parallel/global.cc @@ -1,7 +1,9 @@ #include "infini_train/include/nn/parallel/global.h" +#include #include #include +#include #include #include "glog/logging.h" @@ -19,75 +21,245 @@ namespace infini_train::nn::parallel::global { thread_local int thread_global_rank = 0; -void Layout::InitStrides() { - // Calculate strides +namespace { + +std::array MakeAxisMask(std::initializer_list axes) { + CHECK_GT(axes.size(), 0); + std::array mask{}; + for (const Axis axis : axes) { + CHECK_GE(static_cast(axis), 0); + CHECK_LT(static_cast(axis), AXIS_COUNT); + mask[axis] = true; + } + return mask; +} + +const char *AxisName(Axis axis) { + if (axis == DP) { + return "DP"; + } else if (axis == TP) { + return "TP"; + } else if (axis == PP) { + return "PP"; + } else if (axis == EP) { + return "EP"; + } + CHECK(false) << "Invalid Axis value: " << static_cast(axis); +} + +std::string OrderString(const RankGenerator &rank_generator, bool expert_view) { + std::string result; + bool is_first_axis = true; + for (int index = 0; index < AXIS_COUNT; ++index) { + const Axis axis = rank_generator.order()[index]; + if (!expert_view && axis == EP) { + continue; + } + if (!is_first_axis) { + result += " -> "; + } + is_first_axis = false; + if (expert_view && axis == TP) { + result += "ETP"; + } else if (expert_view && axis == DP) { + result += "EDP"; + } else { + result += AxisName(axis); + } + } + return result; +} + +void AppendGroups(std::ostringstream &oss, const char *name, const RankGenerator &rank_generator, + std::initializer_list varying_axes, bool skip_trivial_axes) { + const auto varying_axis_mask = MakeAxisMask(varying_axes); + int group_size = 1; + for (int axis = 0; axis < AXIS_COUNT; ++axis) { + if (varying_axis_mask[axis]) { + group_size *= rank_generator.AxisSize(static_cast(axis)); + } + } + + if (skip_trivial_axes && group_size <= 1) { + oss << std::format("[{}] size={}, unenabled\n", name, group_size); + return; + } + + const auto groups = rank_generator.GetRanks(varying_axes); + oss << std::format("[{}] size={}, num_groups={}\n", name, group_size, groups.size()); + for (size_t group_id = 0; group_id < groups.size(); ++group_id) { + std::string ranks_string; + for (size_t rank_index = 0; rank_index < groups[group_id].size(); ++rank_index) { + if (rank_index > 0) { + ranks_string += ", "; + } + ranks_string += std::to_string(groups[group_id][rank_index]); + } + oss << std::format(" - {} {}: [{}]\n", name, group_id, ranks_string); + } +} + +void AppendDenseRankView(std::ostringstream &oss, const RankGenerator &rank_generator, bool skip_trivial_axes) { + CHECK_EQ(rank_generator.AxisSize(EP), 1) << "The dense rank view must not contain an expert-parallel axis"; + oss << std::format("[Dense Rank View] shape={{TP={}, DP={}, PP={}}}, order={{ {} }}\n", rank_generator.AxisSize(TP), + rank_generator.AxisSize(DP), rank_generator.AxisSize(PP), + OrderString(rank_generator, /*expert_view=*/false)); + AppendGroups(oss, "DP", rank_generator, {DP}, skip_trivial_axes); + AppendGroups(oss, "TP", rank_generator, {TP}, skip_trivial_axes); + AppendGroups(oss, "PP", rank_generator, {PP}, skip_trivial_axes); +} + +void AppendExpertRankView(std::ostringstream &oss, const RankGenerator &rank_generator, bool skip_trivial_axes) { + oss << std::format("[Expert Rank View] shape={{ETP={}, EP={}, EDP={}, PP={}}}, order={{ {} }}\n", + rank_generator.AxisSize(TP), rank_generator.AxisSize(EP), rank_generator.AxisSize(DP), + rank_generator.AxisSize(PP), OrderString(rank_generator, /*expert_view=*/true)); + AppendGroups(oss, "EDP", rank_generator, {DP}, skip_trivial_axes); + AppendGroups(oss, "ETP", rank_generator, {TP}, skip_trivial_axes); + AppendGroups(oss, "EP", rank_generator, {EP}, skip_trivial_axes); + AppendGroups(oss, "ETP_EP", rank_generator, {TP, EP}, skip_trivial_axes); +} + +} // namespace + +RankGenerator::RankGenerator(int tensor_parallel_size, int expert_parallel_size, int data_parallel_size, + int pipeline_parallel_size, std::array order) + : order_(order) { + sizes_[DP] = data_parallel_size; + sizes_[TP] = tensor_parallel_size; + sizes_[PP] = pipeline_parallel_size; + sizes_[EP] = expert_parallel_size; + + for (const int size : sizes_) { CHECK_GE(size, 1) << "Parallel axis size must be >= 1"; } + InitStrides(); +} + +void RankGenerator::InitStrides() { int stride = 1; for (int i = 0; i < AXIS_COUNT; ++i) { - const Axis ax = order[i]; - strides[ax] = stride; - stride *= sizes[ax]; + const Axis ax = order_[i]; + strides_[ax] = stride; + stride *= sizes_[ax]; } } -int Layout::RankOf(int dp, int tp, int pp) const { - // Return the thread rank given layout coords - const int coord[AXIS_COUNT] = {dp, tp, pp}; +int RankGenerator::AxisSize(Axis axis) const { + CHECK_GE(static_cast(axis), 0); + CHECK_LT(static_cast(axis), AXIS_COUNT); + return sizes_[axis]; +} + +int RankGenerator::WorldSize() const { + int world_size = 1; + for (const int size : sizes_) { world_size *= size; } + return world_size; +} + +const std::array &RankGenerator::order() const { return order_; } + +int RankGenerator::RankOf(int dp, int tp, int pp, int ep) const { + const int coord[AXIS_COUNT] = {dp, tp, pp, ep}; int r = 0; for (int i = 0; i < AXIS_COUNT; ++i) { const Axis ax = static_cast(i); - r += coord[ax] * strides[ax]; + r += coord[ax] * strides_[ax]; } return r; } -void Layout::CoordOf(int rank, int &dp, int &tp, int &pp) const { - // Return the layout coords given thread rank - dp = (rank / strides[DP]) % sizes[DP]; - tp = (rank / strides[TP]) % sizes[TP]; - pp = (rank / strides[PP]) % sizes[PP]; +void RankGenerator::CoordOf(int rank, int &dp, int &tp, int &pp, int &ep) const { + dp = (rank / strides_[DP]) % sizes_[DP]; + tp = (rank / strides_[TP]) % sizes_[TP]; + pp = (rank / strides_[PP]) % sizes_[PP]; + ep = (rank / strides_[EP]) % sizes_[EP]; } -int Layout::GroupId(Axis target, int dp, int tp, int pp) const { - // Return the parallel ProcessGroup ID where the rank is in - int id = 0; - int mult = 1; - for (int i = AXIS_COUNT - 1; i >= 0; --i) { - Axis ax = order[i]; - if (ax == target) { +int RankGenerator::GroupId(std::initializer_list varying_axes, int global_rank) const { + const auto varying_axis_mask = MakeAxisMask(varying_axes); + int dp, tp, pp, ep; + CoordOf(global_rank, dp, tp, pp, ep); + const std::array coordinates{dp, tp, pp, ep}; + + // The first unmasked axis in rank order varies fastest in the group + // ordinal, matching generate_masked_orthogonal_rank_groups in Megatron-LM. + int group_id = 0; + int group_stride = 1; + for (const Axis axis : order_) { + if (varying_axis_mask[axis]) { continue; } - int c = (ax == DP ? dp : (ax == TP ? tp : pp)); - id += c * mult; - mult *= sizes[ax]; + group_id += coordinates[axis] * group_stride; + group_stride *= sizes_[axis]; } - return id; + return group_id; } -std::vector Layout::GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp) const { - // Return all the ranks within the same parallel ProcessGroup +int RankGenerator::GroupId(Axis varying_axis, int global_rank) const { return GroupId({varying_axis}, global_rank); } + +std::vector> RankGenerator::GetRanks(std::initializer_list varying_axes) const { + const auto varying_axis_mask = MakeAxisMask(varying_axes); + int group_size = 1; + for (int axis = 0; axis < AXIS_COUNT; ++axis) { + if (varying_axis_mask[axis]) { + group_size *= sizes_[axis]; + } + } + + const int num_groups = WorldSize() / group_size; + std::vector> groups(num_groups); + for (int rank = 0; rank < WorldSize(); ++rank) { + const int group_id = GroupId(varying_axes, rank); + CHECK_GE(group_id, 0); + CHECK_LT(group_id, num_groups); + groups[group_id].push_back(rank); + } + for (const auto &group : groups) { CHECK_EQ(group.size(), group_size); } + return groups; +} + +std::vector> RankGenerator::GetRanks(Axis varying_axis) const { return GetRanks({varying_axis}); } + +std::vector RankGenerator::GroupRanks(std::initializer_list varying_axes, int global_rank) const { + const auto varying_axis_mask = MakeAxisMask(varying_axes); + int dp, tp, pp, ep; + CoordOf(global_rank, dp, tp, pp, ep); + std::array coordinates{dp, tp, pp, ep}; + + int group_size = 1; + for (int axis = 0; axis < AXIS_COUNT; ++axis) { + if (varying_axis_mask[axis]) { + group_size *= sizes_[axis]; + } + } + std::vector ranks; - ranks.reserve(sizes[target]); - int dp = fixed_dp, tp = fixed_tp, pp = fixed_pp; - for (int v = 0; v < sizes[target]; ++v) { - if (target == DP) { - dp = v; - } else if (target == TP) { - tp = v; - } else { - pp = v; + ranks.reserve(group_size); + for (int rank_in_group = 0; rank_in_group < group_size; ++rank_in_group) { + int remaining_index = rank_in_group; + for (const Axis axis : order_) { + if (!varying_axis_mask[axis]) { + continue; + } + coordinates[axis] = remaining_index % sizes_[axis]; + remaining_index /= sizes_[axis]; } - ranks.push_back(RankOf(dp, tp, pp)); + CHECK_EQ(remaining_index, 0); + ranks.push_back(RankOf(coordinates[DP], coordinates[TP], coordinates[PP], coordinates[EP])); } return ranks; } +std::vector RankGenerator::GroupRanks(Axis varying_axis, int global_rank) const { + return GroupRanks({varying_axis}, global_rank); +} + GlobalEnv &GlobalEnv::Instance() { static GlobalEnv instance; return instance; } void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool sequence_parallel_enabled, - int pipeline_parallel_size, int virtual_pipeline_parallel_size) { + int pipeline_parallel_size, int virtual_pipeline_parallel_size, int expert_parallel_size, + std::optional expert_tensor_parallel_size) { std::lock_guard lock(mutex_); CHECK(!initialized_) << "Repeated initialization of GlobalEnv!"; @@ -100,16 +272,40 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq nthread_per_process_ = nthread_per_process; CHECK_GE(tensor_parallel_size, 1) << "Tensor Parallel size must be >= 1"; + CHECK_GE(pipeline_parallel_size, 1) << "Pipeline Parallel size must be >= 1"; + CHECK_GE(virtual_pipeline_parallel_size, 1) << "Virtual Pipeline Parallel size must be >= 1"; + CHECK_GE(expert_parallel_size, 1) << "Expert Parallel size must be >= 1"; + const int resolved_expert_tensor_parallel_size = expert_tensor_parallel_size.value_or(tensor_parallel_size); + CHECK_GE(resolved_expert_tensor_parallel_size, 1) << "Expert Tensor Parallel size must be >= 1"; + tensor_parallel_size_ = tensor_parallel_size; + expert_tensor_parallel_size_ = resolved_expert_tensor_parallel_size; sequence_parallel_enabled_ = sequence_parallel_enabled; pipeline_parallel_size_ = pipeline_parallel_size; virtual_pipeline_parallel_size_ = virtual_pipeline_parallel_size; - data_parallel_size_ = world_size_ / tensor_parallel_size_ / pipeline_parallel_size_; - - layout_.sizes[DP] = data_parallel_size_; - layout_.sizes[TP] = tensor_parallel_size_; - layout_.sizes[PP] = pipeline_parallel_size_; - layout_.InitStrides(); + expert_parallel_size_ = expert_parallel_size; + + const int dense_model_parallel_size = tensor_parallel_size_ * pipeline_parallel_size_; + CHECK_EQ(world_size_ % dense_model_parallel_size, 0) + << "World size must be divisible by tensor_parallel_size * pipeline_parallel_size"; + data_parallel_size_ = world_size_ / dense_model_parallel_size; + + const int expert_model_parallel_size + = expert_tensor_parallel_size_ * expert_parallel_size_ * pipeline_parallel_size_; + CHECK_EQ(world_size_ % expert_model_parallel_size, 0) + << "World size must be divisible by expert_tensor_parallel_size * expert_parallel_size" + " * pipeline_parallel_size"; + expert_data_parallel_size_ = world_size_ / expert_model_parallel_size; + + // These are two logical views over the same physical ranks. TP and DP in + // the expert generator are exposed as ETP and EDP, respectively. + dense_rank_generator_ = RankGenerator(tensor_parallel_size_, /*expert_parallel_size=*/1, data_parallel_size_, + pipeline_parallel_size_); + expert_rank_generator_ = RankGenerator(expert_tensor_parallel_size_, expert_parallel_size_, + expert_data_parallel_size_, pipeline_parallel_size_); + + CHECK(dense_rank_generator_.GetRanks(PP) == expert_rank_generator_.GetRanks(PP)) + << "Dense and expert rank views must generate identical pipeline-parallel groups"; initialized_ = true; } @@ -149,6 +345,11 @@ int GlobalEnv::tensor_parallel_size() const { return tensor_parallel_size_; } +int GlobalEnv::expert_tensor_parallel_size() const { + CHECK(initialized_) << "GlobalEnv is not initialized!"; + return expert_tensor_parallel_size_; +} + int GlobalEnv::sequence_parallel_size() const { CHECK(initialized_) << "GlobalEnv is not initialized!"; return sequence_parallel_enabled_ ? tensor_parallel_size_ : 1; @@ -164,6 +365,11 @@ int GlobalEnv::data_parallel_size() const { return data_parallel_size_; } +int GlobalEnv::expert_data_parallel_size() const { + CHECK(initialized_) << "GlobalEnv is not initialized!"; + return expert_data_parallel_size_; +} + int GlobalEnv::pipeline_parallel_size() const { CHECK(initialized_) << "GlobalEnv is not initialized!"; return pipeline_parallel_size_; @@ -174,84 +380,48 @@ int GlobalEnv::virtual_pipeline_parallel_size() const { return virtual_pipeline_parallel_size_; } -Layout GlobalEnv::layout() const { +int GlobalEnv::expert_parallel_size() const { CHECK(initialized_) << "GlobalEnv is not initialized!"; - return layout_; + return expert_parallel_size_; } -namespace { -inline const char *AxisName(Axis a) { return a == DP ? "DP" : (a == TP ? "TP" : "PP"); } +const RankGenerator &GlobalEnv::dense_rank_generator() const { + CHECK(initialized_) << "GlobalEnv is not initialized!"; + return dense_rank_generator_; +} -inline int NumGroups(const Layout &L, Axis target) { - int n = 1; - for (int i = 0; i < AXIS_COUNT; ++i) { - if (i != target) { - n *= L.sizes[i]; - } - } - return n; +const RankGenerator &GlobalEnv::expert_rank_generator() const { + CHECK(initialized_) << "GlobalEnv is not initialized!"; + return expert_rank_generator_; } -} // namespace -std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) { +std::string ProcessGroupOverview(const RankGenerator &dense_rank_generator, bool skip_trivial_axes) { std::ostringstream oss; oss << std::format("\n=== Parallel Communication Groups ===\n" - "world_size = {}, config: {{DP={}, TP={}, PP={}}}, order: {{", - GetWorldSize(), L.sizes[DP], L.sizes[TP], L.sizes[PP]); + "world_size = {}, config: {{DP={}, TP={}, PP={}}}\n", + dense_rank_generator.WorldSize(), dense_rank_generator.AxisSize(DP), + dense_rank_generator.AxisSize(TP), dense_rank_generator.AxisSize(PP)); + AppendDenseRankView(oss, dense_rank_generator, skip_trivial_axes); + oss << "\n"; + return oss.str(); +} - for (int i = 0; i < AXIS_COUNT; ++i) { oss << AxisName(L.order[i]) << (i + 1 == AXIS_COUNT ? "" : " -> "); } - oss << "}\n"; +std::string ProcessGroupOverview(const RankGenerator &dense_rank_generator, const RankGenerator &expert_rank_generator, + bool skip_trivial_axes) { + CHECK_EQ(dense_rank_generator.WorldSize(), expert_rank_generator.WorldSize()) + << "Dense and expert rank views must cover the same physical world"; - for (int a = 0; a < AXIS_COUNT; ++a) { - Axis ax = static_cast(a); - if (skip_trivial_axes && L.sizes[ax] <= 1) { - oss << std::format("[{}] size={}, unenabled\n", AxisName(ax), L.sizes[ax]); - continue; - } - // Build > mapping - std::vector>> groups; - for (int dp = 0; dp < (ax == DP ? 1 : L.sizes[DP]); ++dp) { - for (int tp = 0; tp < (ax == TP ? 1 : L.sizes[TP]); ++tp) { - for (int pp = 0; pp < (ax == PP ? 1 : L.sizes[PP]); ++pp) { - int gid = L.GroupId(ax, dp, tp, pp); - groups.emplace_back(gid, std::make_tuple(dp, tp, pp)); - } - } - } - // Sort by the order of Group ID - std::sort(groups.begin(), groups.end(), [](const auto &a, const auto &b) { return a.first < b.first; }); - - const int num_groups = NumGroups(L, ax); - const auto name = AxisName(ax); - oss << std::format("[{}] size={}, num_groups={}\n", name, L.sizes[ax], num_groups); - - // Iterate and print in the order of Group ID - for (const auto &pair : groups) { - int gid = pair.first; - int dp, tp, pp; - std::tie(dp, tp, pp) = pair.second; - auto ranks = L.GroupRanks(ax, dp, tp, pp); - std::sort(ranks.begin(), ranks.end()); - - auto dp_size_str = (ax == DP) ? "-" : std::to_string(dp); - auto tp_size_str = (ax == TP) ? "-" : std::to_string(tp); - auto pp_size_str = (ax == PP) ? "-" : std::to_string(pp); - - std::string ranks_str; - ranks_str.reserve(ranks.size() * 4); - for (size_t i = 0; i < ranks.size(); ++i) { - if (i > 0) { - ranks_str += ", "; - } - ranks_str += std::to_string(ranks[i]); - } - oss << std::format(" - {} {} (dp={}, tp={}, pp={}): [{}]\n", name, gid, dp_size_str, tp_size_str, - pp_size_str, ranks_str); - } - if (a + 1 < AXIS_COUNT) { - oss << "\n"; - } - } + std::ostringstream oss; + oss << std::format("\n=== Parallel Communication Groups ===\n" + "world_size = {}, config: {{DP={}, EDP={}, TP={}, ETP={}, PP={}, EP={}}}\n", + dense_rank_generator.WorldSize(), dense_rank_generator.AxisSize(DP), + expert_rank_generator.AxisSize(DP), dense_rank_generator.AxisSize(TP), + expert_rank_generator.AxisSize(TP), dense_rank_generator.AxisSize(PP), + expert_rank_generator.AxisSize(EP)); + + AppendDenseRankView(oss, dense_rank_generator, skip_trivial_axes); + oss << "\n"; + AppendExpertRankView(oss, expert_rank_generator, skip_trivial_axes); oss << "\n"; return oss.str(); } diff --git a/infini_train/src/nn/parallel/utils.cc b/infini_train/src/nn/parallel/utils.cc index 93c6ae31..27335c24 100644 --- a/infini_train/src/nn/parallel/utils.cc +++ b/infini_train/src/nn/parallel/utils.cc @@ -5,23 +5,59 @@ namespace infini_train::nn::parallel { std::string GetDataParallelProcessGroupName(int global_rank) { - return "DP" + std::to_string(global::GetGroupId(global::DP, global_rank)); + return "DP" + std::to_string(global::GetDenseRankGenerator().GroupId(global::DP, global_rank)); +} + +std::string GetExpertDataParallelProcessGroupName(int global_rank) { + return "EDP" + std::to_string(global::GetExpertRankGenerator().GroupId(global::DP, global_rank)); } std::string GetTensorParallelProcessGroupName(int global_rank) { - return "TP" + std::to_string(global::GetGroupId(global::TP, global_rank)); + return "TP" + std::to_string(global::GetDenseRankGenerator().GroupId(global::TP, global_rank)); +} + +std::string GetExpertTensorParallelProcessGroupName(int global_rank) { + return "ETP" + std::to_string(global::GetExpertRankGenerator().GroupId(global::TP, global_rank)); } std::string GetPipelineParallelProcessGroupName(int global_rank) { - return "PP" + std::to_string(global::GetGroupId(global::PP, global_rank)); + return "PP" + std::to_string(global::GetDenseRankGenerator().GroupId(global::PP, global_rank)); +} + +std::string GetExpertParallelProcessGroupName(int global_rank) { + return "EP" + std::to_string(global::GetExpertRankGenerator().GroupId(global::EP, global_rank)); +} + +std::string GetExpertTensorAndExpertParallelProcessGroupName(int global_rank) { + return "ETP_EP" + std::to_string(global::GetExpertRankGenerator().GroupId({global::TP, global::EP}, global_rank)); +} + +std::vector GetDataParallelGroupRanks(int global_rank) { + return global::GetDenseRankGenerator().GroupRanks(global::DP, global_rank); } -std::vector GetDataParallelGroupRanks(int global_rank) { return global::GetGroupRanks(global::DP, global_rank); } +std::vector GetExpertDataParallelGroupRanks(int global_rank) { + return global::GetExpertRankGenerator().GroupRanks(global::DP, global_rank); +} -std::vector GetTensorParallelGroupRanks(int global_rank) { return global::GetGroupRanks(global::TP, global_rank); } +std::vector GetTensorParallelGroupRanks(int global_rank) { + return global::GetDenseRankGenerator().GroupRanks(global::TP, global_rank); +} + +std::vector GetExpertTensorParallelGroupRanks(int global_rank) { + return global::GetExpertRankGenerator().GroupRanks(global::TP, global_rank); +} std::vector GetPipelineParallelGroupRanks(int global_rank) { - return global::GetGroupRanks(global::PP, global_rank); + return global::GetDenseRankGenerator().GroupRanks(global::PP, global_rank); +} + +std::vector GetExpertParallelGroupRanks(int global_rank) { + return global::GetExpertRankGenerator().GroupRanks(global::EP, global_rank); +} + +std::vector GetExpertTensorAndExpertParallelGroupRanks(int global_rank) { + return global::GetExpertRankGenerator().GroupRanks({global::TP, global::EP}, global_rank); } } // namespace infini_train::nn::parallel diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96776585..fb9580d8 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -16,6 +16,9 @@ add_subdirectory(optimizer) # Autograd operator tests add_subdirectory(autograd) +# Parallel layout tests +add_subdirectory(parallel) + # LoRA tests add_subdirectory(lora) diff --git a/tests/common/test_main.cc b/tests/common/test_main.cc index 54ef48dd..01c5aa87 100644 --- a/tests/common/test_main.cc +++ b/tests/common/test_main.cc @@ -4,6 +4,6 @@ int main(int argc, char **argv) { ::testing::InitGoogleTest(&argc, argv); - infini_train::nn::parallel::global::GlobalEnv::Instance().Init(1, 1, false, 1, 1); + infini_train::nn::parallel::global::GlobalEnv::Instance().Init(1, 1, false, 1, 1, 1, std::nullopt); return RUN_ALL_TESTS(); } diff --git a/tests/parallel/CMakeLists.txt b/tests/parallel/CMakeLists.txt new file mode 100644 index 00000000..fd2dd77c --- /dev/null +++ b/tests/parallel/CMakeLists.txt @@ -0,0 +1,28 @@ +# ========================================================================== +# Parallel layout tests +# ========================================================================== + +file(GLOB PARALLEL_SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc) +list(FILTER PARALLEL_SOURCES EXCLUDE REGEX "test_global_env(_etp)?\\.cc$") + +infini_train_add_test(test_parallel + SOURCES ${PARALLEL_SOURCES} + LABELS cpu +) + +# This target owns its main function so it can initialize GlobalEnv with a +# non-trivial TP/EP topology in a fresh process. +add_executable(test_parallel_env test_global_env.cc) +target_link_libraries(test_parallel_env PRIVATE GTest::gtest) +link_infini_train_exe(test_parallel_env) +gtest_discover_tests(test_parallel_env + PROPERTIES LABELS cpu TIMEOUT 10 +) + +# Validate a topology where dense TP and expert TP (ETP) differ. +add_executable(test_parallel_etp_env test_global_env_etp.cc) +target_link_libraries(test_parallel_etp_env PRIVATE GTest::gtest) +link_infini_train_exe(test_parallel_etp_env) +gtest_discover_tests(test_parallel_etp_env + PROPERTIES LABELS cpu TIMEOUT 10 +) diff --git a/tests/parallel/test_global_env.cc b/tests/parallel/test_global_env.cc new file mode 100644 index 00000000..2036d413 --- /dev/null +++ b/tests/parallel/test_global_env.cc @@ -0,0 +1,136 @@ +#include +#include +#include +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/utils.h" + +namespace { + +using namespace infini_train::nn::parallel; + +using GroupNameFn = std::string (*)(int); +using GroupRanksFn = std::vector (*)(int); + +void ExpectNamesIdentifyExactRankSets(GroupNameFn get_name, GroupRanksFn get_ranks) { + std::map> name_to_ranks; + std::map, std::string> ranks_to_name; + + for (int rank = 0; rank < global::GetWorldSize(); ++rank) { + const auto name = get_name(rank); + const auto ranks = get_ranks(rank); + EXPECT_NE(std::find(ranks.begin(), ranks.end(), rank), ranks.end()); + + const auto [name_it, inserted_name] = name_to_ranks.emplace(name, ranks); + if (!inserted_name) { + EXPECT_EQ(name_it->second, ranks); + } + + const auto [ranks_it, inserted_ranks] = ranks_to_name.emplace(ranks, name); + if (!inserted_ranks) { + EXPECT_EQ(ranks_it->second, name); + } + } +} + +TEST(GlobalEnvTest, DistinguishesDenseAndExpertDataParallelSizes) { + EXPECT_EQ(global::GetWorldSize(), 16); + EXPECT_EQ(global::GetTensorParallelSize(), 2); + EXPECT_EQ(global::GetExpertTensorParallelSize(), 2); + EXPECT_EQ(global::GetPipelineParallelSize(), 2); + EXPECT_EQ(global::GetExpertParallelSize(), 2); + EXPECT_EQ(global::GetDataParallelSize(), 4); + EXPECT_EQ(global::GetExpertDataParallelSize(), 2); +} + +TEST(GlobalEnvTest, OwnsDenseAndExpertRankViewsOverTheSameWorld) { + const auto &dense = global::GetDenseRankGenerator(); + const auto &expert = global::GetExpertRankGenerator(); + + EXPECT_EQ(dense.WorldSize(), global::GetWorldSize()); + EXPECT_EQ(expert.WorldSize(), global::GetWorldSize()); + EXPECT_EQ(dense.AxisSize(global::DP), 4); + EXPECT_EQ(dense.AxisSize(global::TP), 2); + EXPECT_EQ(dense.AxisSize(global::EP), 1); + EXPECT_EQ(expert.AxisSize(global::DP), 2); + EXPECT_EQ(expert.AxisSize(global::TP), 2); + EXPECT_EQ(expert.AxisSize(global::EP), 2); + + EXPECT_EQ(dense.GetRanks(global::PP), expert.GetRanks(global::PP)); +} + +TEST(GlobalEnvTest, PublicGroupHelpersUseMainstreamMoeDomains) { + constexpr int kRank = 7; + + EXPECT_EQ(GetDataParallelGroupRanks(kRank), (std::vector{1, 3, 5, 7})); + EXPECT_EQ(GetDataParallelProcessGroupName(kRank), "DP1"); + + EXPECT_EQ(GetTensorParallelGroupRanks(kRank), (std::vector{6, 7})); + EXPECT_EQ(GetTensorParallelProcessGroupName(kRank), "TP3"); + + EXPECT_EQ(GetExpertTensorParallelGroupRanks(kRank), (std::vector{6, 7})); + EXPECT_EQ(GetExpertTensorParallelProcessGroupName(kRank), "ETP3"); + + EXPECT_EQ(GetPipelineParallelGroupRanks(kRank), (std::vector{7, 15})); + EXPECT_EQ(GetPipelineParallelProcessGroupName(kRank), "PP7"); + + EXPECT_EQ(GetExpertDataParallelGroupRanks(kRank), (std::vector{3, 7})); + EXPECT_EQ(GetExpertDataParallelProcessGroupName(kRank), "EDP3"); + + EXPECT_EQ(GetExpertTensorAndExpertParallelGroupRanks(kRank), (std::vector{4, 5, 6, 7})); + EXPECT_EQ(GetExpertTensorAndExpertParallelProcessGroupName(kRank), "ETP_EP1"); + + EXPECT_EQ(GetExpertParallelGroupRanks(kRank), (std::vector{5, 7})); + EXPECT_EQ(GetExpertParallelProcessGroupName(kRank), "EP3"); + + constexpr int kSecondPipelineStageRank = 15; + EXPECT_EQ(GetDataParallelGroupRanks(kSecondPipelineStageRank), (std::vector{9, 11, 13, 15})); + EXPECT_EQ(GetDataParallelProcessGroupName(kSecondPipelineStageRank), "DP3"); + EXPECT_EQ(GetExpertDataParallelGroupRanks(kSecondPipelineStageRank), (std::vector{11, 15})); + EXPECT_EQ(GetExpertDataParallelProcessGroupName(kSecondPipelineStageRank), "EDP7"); + EXPECT_EQ(GetExpertTensorAndExpertParallelGroupRanks(kSecondPipelineStageRank), (std::vector{12, 13, 14, 15})); + EXPECT_EQ(GetExpertTensorAndExpertParallelProcessGroupName(kSecondPipelineStageRank), "ETP_EP3"); +} + +TEST(GlobalEnvTest, GroupNamesAreStableForEveryMember) { + ExpectNamesIdentifyExactRankSets(GetDataParallelProcessGroupName, GetDataParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetExpertDataParallelProcessGroupName, GetExpertDataParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetTensorParallelProcessGroupName, GetTensorParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetExpertTensorParallelProcessGroupName, GetExpertTensorParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetPipelineParallelProcessGroupName, GetPipelineParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetExpertParallelProcessGroupName, GetExpertParallelGroupRanks); + ExpectNamesIdentifyExactRankSets(GetExpertTensorAndExpertParallelProcessGroupName, + GetExpertTensorAndExpertParallelGroupRanks); +} + +TEST(GlobalEnvTest, DefaultOverviewReportsOnlyTheDenseView) { + const std::string overview = global::ProcessGroupOverview(); + + EXPECT_NE(overview.find("config: {DP=4, TP=2, PP=2}"), std::string::npos); + EXPECT_NE(overview.find("[Dense Rank View] shape={TP=2, DP=4, PP=2}"), std::string::npos); + EXPECT_EQ(overview.find("[Expert Rank View]"), std::string::npos); + EXPECT_EQ(overview.find("[EDP]"), std::string::npos); + EXPECT_EQ(overview.find("[ETP]"), std::string::npos); + EXPECT_EQ(overview.find("[EP]"), std::string::npos); + EXPECT_EQ(overview.find("[ETP_EP]"), std::string::npos); +} + +} // namespace + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + setenv("PROC_WORLD_SIZE", "16", 1); + infini_train::nn::parallel::global::GlobalEnv::Instance().Init( + /*nthread_per_process=*/1, + /*tensor_parallel_size=*/2, + /*sequence_parallel_enabled=*/false, + /*pipeline_parallel_size=*/2, + /*virtual_pipeline_parallel_size=*/1, + /*expert_parallel_size=*/2, + /*expert_tensor_parallel_size=*/std::nullopt); + return RUN_ALL_TESTS(); +} diff --git a/tests/parallel/test_global_env_etp.cc b/tests/parallel/test_global_env_etp.cc new file mode 100644 index 00000000..19097083 --- /dev/null +++ b/tests/parallel/test_global_env_etp.cc @@ -0,0 +1,76 @@ +#include +#include + +#include "gtest/gtest.h" + +#include "infini_train/include/nn/parallel/global.h" +#include "infini_train/include/nn/parallel/utils.h" + +namespace { + +using namespace infini_train::nn::parallel; + +TEST(GlobalEnvEtpTest, UsesIndependentDenseAndExpertTensorParallelSizes) { + EXPECT_EQ(global::GetWorldSize(), 48); + EXPECT_EQ(global::GetTensorParallelSize(), 3); + EXPECT_EQ(global::GetExpertTensorParallelSize(), 2); + EXPECT_EQ(global::GetExpertParallelSize(), 2); + EXPECT_EQ(global::GetPipelineParallelSize(), 2); + EXPECT_EQ(global::GetDataParallelSize(), 8); + EXPECT_EQ(global::GetExpertDataParallelSize(), 6); + + const auto &dense = global::GetDenseRankGenerator(); + const auto &expert = global::GetExpertRankGenerator(); + EXPECT_EQ(dense.AxisSize(global::TP), 3); + EXPECT_EQ(expert.AxisSize(global::TP), 2); + EXPECT_EQ(dense.AxisSize(global::DP), 8); + EXPECT_EQ(expert.AxisSize(global::DP), 6); + EXPECT_EQ(dense.WorldSize(), global::GetWorldSize()); + EXPECT_EQ(expert.WorldSize(), global::GetWorldSize()); +} + +TEST(GlobalEnvEtpTest, ExpertGroupsUseTheExpertTensorParallelAxis) { + constexpr int kRank = 11; + + EXPECT_EQ(GetTensorParallelGroupRanks(kRank), (std::vector{9, 10, 11})); + EXPECT_EQ(GetTensorParallelProcessGroupName(kRank), "TP3"); + + EXPECT_EQ(GetExpertTensorParallelGroupRanks(kRank), (std::vector{10, 11})); + EXPECT_EQ(GetExpertTensorParallelProcessGroupName(kRank), "ETP5"); + + EXPECT_EQ(GetExpertParallelGroupRanks(kRank), (std::vector{9, 11})); + EXPECT_EQ(GetExpertParallelProcessGroupName(kRank), "EP5"); + + EXPECT_EQ(GetExpertTensorAndExpertParallelGroupRanks(kRank), (std::vector{8, 9, 10, 11})); + EXPECT_EQ(GetExpertTensorAndExpertParallelProcessGroupName(kRank), "ETP_EP2"); + + EXPECT_EQ(GetExpertDataParallelGroupRanks(kRank), (std::vector{3, 7, 11, 15, 19, 23})); + EXPECT_EQ(GetExpertDataParallelProcessGroupName(kRank), "EDP3"); +} + +TEST(GlobalEnvEtpTest, DenseAndExpertViewsPreservePipelineGroups) { + const auto &dense = global::GetDenseRankGenerator(); + const auto &expert = global::GetExpertRankGenerator(); + EXPECT_EQ(dense.GetRanks(global::PP), expert.GetRanks(global::PP)); + + constexpr int kRank = 35; + EXPECT_EQ(GetPipelineParallelGroupRanks(kRank), (std::vector{11, 35})); + EXPECT_EQ(GetExpertTensorParallelGroupRanks(kRank), (std::vector{34, 35})); + EXPECT_EQ(GetExpertDataParallelGroupRanks(kRank), (std::vector{27, 31, 35, 39, 43, 47})); +} + +} // namespace + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + setenv("PROC_WORLD_SIZE", "48", 1); + infini_train::nn::parallel::global::GlobalEnv::Instance().Init( + /*nthread_per_process=*/1, + /*tensor_parallel_size=*/3, + /*sequence_parallel_enabled=*/false, + /*pipeline_parallel_size=*/2, + /*virtual_pipeline_parallel_size=*/1, + /*expert_parallel_size=*/2, + /*expert_tensor_parallel_size=*/2); + return RUN_ALL_TESTS(); +} diff --git a/tests/parallel/test_global_layout.cc b/tests/parallel/test_global_layout.cc new file mode 100644 index 00000000..0342c90a --- /dev/null +++ b/tests/parallel/test_global_layout.cc @@ -0,0 +1,238 @@ +#include "infini_train/include/nn/parallel/global.h" + +#include +#include +#include + +#include "gtest/gtest.h" + +namespace { +using namespace infini_train::nn::parallel::global; + +TEST(RankGeneratorTest, DefaultGeneratorUsesViewLocalDpAxis) { + const RankGenerator rank_generator(/*tp=*/1, /*ep=*/1, /*dp=*/1, /*pp=*/1); + + EXPECT_EQ(rank_generator.AxisSize(DP), 1); + EXPECT_EQ(rank_generator.AxisSize(TP), 1); + EXPECT_EQ(rank_generator.AxisSize(PP), 1); + EXPECT_EQ(rank_generator.AxisSize(EP), 1); + EXPECT_EQ(rank_generator.order(), (std::array{TP, EP, DP, PP})); + EXPECT_EQ(rank_generator.WorldSize(), 1); +} + +TEST(RankGeneratorTest, ExpertGroupsMatchMegatronOrthogonalOrdering) { + const RankGenerator expert(/*tp=*/2, /*ep=*/2, /*dp=*/2, /*pp=*/2); + + EXPECT_EQ(expert.GetRanks(DP), + (std::vector>{{0, 4}, {1, 5}, {2, 6}, {3, 7}, {8, 12}, {9, 13}, {10, 14}, {11, 15}})); + EXPECT_EQ(expert.GetRanks(EP), + (std::vector>{{0, 2}, {1, 3}, {4, 6}, {5, 7}, {8, 10}, {9, 11}, {12, 14}, {13, 15}})); + EXPECT_EQ(expert.GetRanks({TP, EP}), + (std::vector>{{0, 1, 2, 3}, {4, 5, 6, 7}, {8, 9, 10, 11}, {12, 13, 14, 15}})); +} + +TEST(RankGeneratorTest, AsymmetricGroupsMatchMegatronOrthogonalOrdering) { + const RankGenerator rank_generator(/*tp=*/2, /*ep=*/3, /*dp=*/2, /*pp=*/1); + + EXPECT_EQ(rank_generator.GetRanks(DP), + (std::vector>{{0, 6}, {1, 7}, {2, 8}, {3, 9}, {4, 10}, {5, 11}})); + EXPECT_EQ(rank_generator.GetRanks(EP), + (std::vector>{{0, 2, 4}, {1, 3, 5}, {6, 8, 10}, {7, 9, 11}})); + EXPECT_EQ(rank_generator.GetRanks({TP, EP}), + (std::vector>{{0, 1, 2, 3, 4, 5}, {6, 7, 8, 9, 10, 11}})); + EXPECT_EQ(rank_generator.GroupId(DP, 11), 5); + EXPECT_EQ(rank_generator.GroupId(EP, 11), 3); + EXPECT_EQ(rank_generator.GroupId({TP, EP}, 11), 1); +} + +TEST(RankGeneratorTest, DenseAndExpertViewsShareTheSamePhysicalRanks) { + const RankGenerator dense(/*tp=*/2, /*ep=*/1, /*dp=*/4, /*pp=*/2); + const RankGenerator expert(/*tp=*/2, /*ep=*/2, /*dp=*/2, /*pp=*/2); + + ASSERT_EQ(dense.WorldSize(), 16); + ASSERT_EQ(expert.WorldSize(), dense.WorldSize()); + + constexpr int kRank = 7; + EXPECT_EQ(dense.GroupRanks(DP, kRank), (std::vector{1, 3, 5, 7})); + EXPECT_EQ(expert.GroupRanks(DP, kRank), (std::vector{3, 7})); + EXPECT_EQ(expert.GroupRanks(EP, kRank), (std::vector{5, 7})); + EXPECT_EQ(expert.GroupRanks({TP, EP}, kRank), (std::vector{4, 5, 6, 7})); + + // Dense DP and expert DPxEP cover the same physical ranks, but are + // intentionally generated from different logical views. + EXPECT_EQ(dense.GetRanks(DP), expert.GetRanks({DP, EP})); + EXPECT_EQ(dense.GetRanks(TP), expert.GetRanks(TP)); + EXPECT_EQ(dense.GetRanks(PP), expert.GetRanks(PP)); +} + +TEST(RankGeneratorTest, DenseAndExpertViewsUseIndependentTensorParallelSizes) { + const RankGenerator dense(/*tp=*/3, /*ep=*/1, /*dp=*/8, /*pp=*/2); + const RankGenerator expert(/*etp=*/2, /*ep=*/2, /*edp=*/6, /*pp=*/2); + + ASSERT_EQ(dense.WorldSize(), 48); + ASSERT_EQ(expert.WorldSize(), dense.WorldSize()); + EXPECT_EQ(dense.AxisSize(TP), 3); + EXPECT_EQ(expert.AxisSize(TP), 2); + EXPECT_EQ(dense.AxisSize(DP), 8); + EXPECT_EQ(expert.AxisSize(DP), 6); + EXPECT_EQ(dense.GetRanks(PP), expert.GetRanks(PP)); + + constexpr int kRank = 11; + EXPECT_EQ(dense.GroupRanks(TP, kRank), (std::vector{9, 10, 11})); + EXPECT_EQ(expert.GroupRanks(TP, kRank), (std::vector{10, 11})); + EXPECT_EQ(expert.GroupRanks(EP, kRank), (std::vector{9, 11})); + EXPECT_EQ(expert.GroupRanks({TP, EP}, kRank), (std::vector{8, 9, 10, 11})); + EXPECT_EQ(expert.GroupRanks(DP, kRank), (std::vector{3, 7, 11, 15, 19, 23})); +} + +TEST(RankGeneratorTest, DenseAndExpertViewsAreIdenticalWhenEpIsOne) { + const RankGenerator dense(/*tp=*/2, /*ep=*/1, /*dp=*/4, /*pp=*/2); + const RankGenerator expert(/*tp=*/2, /*ep=*/1, /*dp=*/4, /*pp=*/2); + + EXPECT_EQ(dense.GetRanks(DP), expert.GetRanks(DP)); + EXPECT_EQ(dense.GetRanks(TP), expert.GetRanks(TP)); + EXPECT_EQ(dense.GetRanks(PP), expert.GetRanks(PP)); + EXPECT_EQ(dense.GetRanks(EP), expert.GetRanks(EP)); +} + +TEST(RankGeneratorTest, EdpOneDoesNotCollapseTheDenseDpView) { + const RankGenerator dense(/*tp=*/2, /*ep=*/1, /*dp=*/2, /*pp=*/2); + const RankGenerator expert(/*tp=*/2, /*ep=*/2, /*dp=*/1, /*pp=*/2); + + constexpr int kRank = 7; + EXPECT_EQ(dense.GroupRanks(DP, kRank), (std::vector{5, 7})); + EXPECT_EQ(expert.GroupRanks(DP, kRank), (std::vector{7})); + EXPECT_EQ(expert.GroupRanks(EP, kRank), (std::vector{5, 7})); + EXPECT_EQ(dense.GroupRanks(DP, kRank), expert.GroupRanks({DP, EP}, kRank)); +} + +TEST(RankGeneratorTest, GroupIdsAreIndicesIntoGeneratedRankGroups) { + const RankGenerator expert(/*tp=*/2, /*ep=*/2, /*dp=*/2, /*pp=*/2); + + const auto edp_groups = expert.GetRanks(DP); + for (size_t group_id = 0; group_id < edp_groups.size(); ++group_id) { + for (const int rank : edp_groups[group_id]) { + EXPECT_EQ(expert.GroupId(DP, rank), group_id); + EXPECT_EQ(expert.GroupRanks(DP, rank), edp_groups[group_id]); + } + } + + const auto tensor_expert_groups = expert.GetRanks({TP, EP}); + for (size_t group_id = 0; group_id < tensor_expert_groups.size(); ++group_id) { + for (const int rank : tensor_expert_groups[group_id]) { + EXPECT_EQ(expert.GroupId({TP, EP}, rank), group_id); + EXPECT_EQ(expert.GroupRanks({TP, EP}, rank), tensor_expert_groups[group_id]); + } + } +} + +TEST(RankGeneratorTest, AsymmetricShapeCoversEveryRankForAllAxisSubsets) { + const RankGenerator rank_generator(/*tp=*/2, /*ep=*/3, /*dp=*/4, /*pp=*/5); + + const auto verify_groups = [&rank_generator](std::initializer_list varying_axes) { + std::array varying_axis_mask{}; + for (const Axis axis : varying_axes) { varying_axis_mask[axis] = true; } + + int expected_group_size = 1; + for (int axis = 0; axis < AXIS_COUNT; ++axis) { + if (varying_axis_mask[axis]) { + expected_group_size *= rank_generator.AxisSize(static_cast(axis)); + } + } + + const auto groups = rank_generator.GetRanks(varying_axes); + ASSERT_EQ(groups.size(), rank_generator.WorldSize() / expected_group_size); + std::vector rank_occurrences(rank_generator.WorldSize(), 0); + for (size_t group_id = 0; group_id < groups.size(); ++group_id) { + ASSERT_EQ(groups[group_id].size(), expected_group_size); + for (const int rank : groups[group_id]) { + ASSERT_GE(rank, 0); + ASSERT_LT(rank, rank_generator.WorldSize()); + ++rank_occurrences[rank]; + EXPECT_EQ(rank_generator.GroupId(varying_axes, rank), static_cast(group_id)); + EXPECT_EQ(rank_generator.GroupRanks(varying_axes, rank), groups[group_id]); + } + } + for (const int occurrence_count : rank_occurrences) { EXPECT_EQ(occurrence_count, 1); } + }; + + verify_groups({DP}); + verify_groups({TP}); + verify_groups({PP}); + verify_groups({EP}); + verify_groups({DP, TP}); + verify_groups({DP, PP}); + verify_groups({DP, EP}); + verify_groups({TP, PP}); + verify_groups({TP, EP}); + verify_groups({PP, EP}); + verify_groups({DP, TP, PP}); + verify_groups({DP, TP, EP}); + verify_groups({DP, PP, EP}); + verify_groups({TP, PP, EP}); + verify_groups({DP, TP, PP, EP}); +} + +TEST(RankGeneratorTest, CompositeAxisOrderDoesNotAffectGroups) { + const RankGenerator expert(/*tp=*/2, /*ep=*/2, /*dp=*/2, /*pp=*/1); + + EXPECT_EQ(expert.GetRanks({DP, EP}), expert.GetRanks({EP, DP})); + EXPECT_EQ(expert.GetRanks({DP, EP, DP}), expert.GetRanks({DP, EP})); + EXPECT_EQ(expert.GroupId({DP, EP}, 7), expert.GroupId({EP, DP}, 7)); +} + +TEST(RankGeneratorTest, CompositeGroupsSupportNonDefaultRankOrder) { + const RankGenerator rank_generator( + /*tp=*/2, /*ep=*/2, /*dp=*/2, /*pp=*/1, std::array{EP, TP, DP, PP}); + + EXPECT_EQ(rank_generator.GroupRanks({DP, EP}, 7), (std::vector{2, 3, 6, 7})); + EXPECT_EQ(rank_generator.GroupRanks({TP, EP}, 7), (std::vector{4, 5, 6, 7})); +} + +TEST(RankGeneratorTest, ProcessGroupOverviewReportsOnlyTheDenseViewByDefault) { + const RankGenerator dense(/*tp=*/2, /*ep=*/1, /*dp=*/2, /*pp=*/1); + + const std::string overview = ProcessGroupOverview(dense); + EXPECT_NE(overview.find("config: {DP=2, TP=2, PP=1}"), std::string::npos); + EXPECT_NE(overview.find("[Dense Rank View] shape={TP=2, DP=2, PP=1}"), std::string::npos); + EXPECT_NE(overview.find("order={ TP -> DP -> PP }"), std::string::npos); + EXPECT_NE(overview.find("[DP] size=2, num_groups=2"), std::string::npos); + EXPECT_NE(overview.find("[TP] size=2, num_groups=2"), std::string::npos); + EXPECT_EQ(overview.find("[Expert Rank View]"), std::string::npos); + EXPECT_EQ(overview.find("[EDP]"), std::string::npos); + EXPECT_EQ(overview.find("[ETP]"), std::string::npos); + EXPECT_EQ(overview.find("[EP]"), std::string::npos); + EXPECT_EQ(overview.find("[ETP_EP]"), std::string::npos); +} + +TEST(RankGeneratorTest, ProcessGroupOverviewSeparatesDenseAndExpertViews) { + const RankGenerator dense(/*tp=*/2, /*ep=*/1, /*dp=*/4, /*pp=*/2); + const RankGenerator expert(/*etp=*/1, /*ep=*/2, /*edp=*/4, /*pp=*/2); + + const std::string overview = ProcessGroupOverview(dense, expert); + EXPECT_NE(overview.find("config: {DP=4, EDP=4, TP=2, ETP=1, PP=2, EP=2}"), std::string::npos); + EXPECT_NE(overview.find("[Dense Rank View] shape={TP=2, DP=4, PP=2}"), std::string::npos); + EXPECT_NE(overview.find("order={ TP -> DP -> PP }"), std::string::npos); + EXPECT_NE(overview.find("[Expert Rank View] shape={ETP=1, EP=2, EDP=4, PP=2}"), std::string::npos); + EXPECT_NE(overview.find("order={ ETP -> EP -> EDP -> PP }"), std::string::npos); + EXPECT_NE(overview.find("[DP] size=4, num_groups=4"), std::string::npos); + EXPECT_NE(overview.find("[EDP] size=4, num_groups=4"), std::string::npos); + EXPECT_NE(overview.find("[ETP] size=1, unenabled"), std::string::npos); + EXPECT_NE(overview.find("[EP] size=2, num_groups=8"), std::string::npos); + EXPECT_NE(overview.find("[ETP_EP] size=2, num_groups=8"), std::string::npos); +} + +TEST(RankGeneratorTest, ProcessGroupOverviewReportsTrivialGroupsAsUnenabled) { + const RankGenerator dense(/*tp=*/1, /*ep=*/1, /*dp=*/1, /*pp=*/1); + const RankGenerator expert(/*tp=*/1, /*ep=*/1, /*dp=*/1, /*pp=*/1); + + const std::string overview = ProcessGroupOverview(dense, expert); + EXPECT_NE(overview.find("config: {DP=1, EDP=1, TP=1, ETP=1, PP=1, EP=1}"), std::string::npos); + EXPECT_NE(overview.find("[DP] size=1, unenabled"), std::string::npos); + EXPECT_NE(overview.find("[EDP] size=1, unenabled"), std::string::npos); + EXPECT_NE(overview.find("[ETP] size=1, unenabled"), std::string::npos); + EXPECT_NE(overview.find("[EP] size=1, unenabled"), std::string::npos); + EXPECT_NE(overview.find("[ETP_EP] size=1, unenabled"), std::string::npos); +} + +} // namespace