Skip to content

feat: implement ep topology infrastructure #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kilinchange wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: implement ep topology infrastructure #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 37 additions & 19 deletions infini_train/include/nn/parallel/global.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,23 @@ 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 };
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};
int sizes[AXIS_COUNT]{1, 1, 1, 1};
// 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
Axis order[AXIS_COUNT]{TP, EP, DP, PP};
int strides[AXIS_COUNT]{1, 1, 1, 1};

void InitStrides();
int RankOf(int dp, int tp, int pp, int ep) const;
int RankOf(int dp, int tp, int pp) const;
void CoordOf(int rank, int &dp, int &tp, int &pp, int &ep) const;
void CoordOf(int rank, int &dp, int &tp, int &pp) const;
int GroupId(Axis target, int dp, int tp, int pp, int ep) const;
int GroupId(Axis target, int dp, int tp, int pp) const;
std::vector<int> GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp, int fixed_ep) const;
std::vector<int> GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp) const;
};

Expand All @@ -29,7 +33,7 @@ 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);

int nnodes() const;

Expand All @@ -55,6 +59,8 @@ class GlobalEnv {

int virtual_pipeline_parallel_size() const;

int expert_parallel_size() const;

Layout layout() const;

private:
Expand All @@ -80,6 +86,7 @@ class GlobalEnv {

int pipeline_parallel_size_ = 1;
int virtual_pipeline_parallel_size_ = 1;
int expert_parallel_size_ = 1;

mutable std::mutex mutex_;
bool initialized_ = false;
Expand All @@ -88,9 +95,9 @@ class GlobalEnv {
};

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 = 1) {
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);
}
inline int GetNnodes() { return GlobalEnv::Instance().nnodes(); }
inline int GetWorldSize() { return GlobalEnv::Instance().world_size(); }
Expand All @@ -105,6 +112,7 @@ inline bool GetSequenceParallelEnabled() { return GlobalEnv::Instance().sequence
inline int GetDataParallelSize() { return GlobalEnv::Instance().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
Expand All @@ -113,33 +121,43 @@ inline int GetVirtualPipelineParallelSize() { return GlobalEnv::Instance().virtu
/**
* @brief Get the global rank corresponding to the given (dp, tp, pp) coordinate.
*/
inline int GetRankOf(int dp, int tp, int pp, int ep) { return GlobalEnv::Instance().layout().RankOf(dp, tp, pp, ep); }
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.
* @brief Get the (dp, tp, pp, ep) coordinate corresponding to the given global rank.
*/
inline void GetCoordOf(int rank, int &dp, int &tp, int &pp, int &ep) {
return GlobalEnv::Instance().layout().CoordOf(rank, dp, tp, pp, ep);
}
inline void GetCoordOf(int rank, int &dp, int &tp, int &pp) {
return GlobalEnv::Instance().layout().CoordOf(rank, dp, tp, pp);
}

/**
* @brief Get the group ID that the (dp, tp, pp) coordinate belongs to along a given parallel axis.
* @brief Get the group ID that the (dp, tp, pp, ep) coordinate belongs to along a given parallel axis.
*/
inline int GetGroupId(Axis target, int dp, int tp, int pp, int ep) {
return GlobalEnv::Instance().layout().GroupId(target, dp, tp, pp, ep);
}
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);
int dp, tp, pp, ep;
GetCoordOf(rank, dp, tp, pp, ep);
return GlobalEnv::Instance().layout().GroupId(target, dp, tp, pp, ep);
}

/**
* @brief Get all ranks that belong to the same group as the given (dp, tp, pp) coordinate
* @brief Get all ranks that belong to the same group as the given (dp, tp, pp, ep) coordinate
* along a specified parallel axis (e.g., all ranks in the same TP group).
*/
inline std::vector<int> GetGroupRanks(Axis target, int dp, int tp, int pp, int ep) {
return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp, ep);
}
inline std::vector<int> GetGroupRanks(Axis target, int dp, int tp, int pp) {
return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp);
}
Expand All @@ -149,9 +167,9 @@ inline std::vector<int> GetGroupRanks(Axis target, int dp, int tp, int pp) {
* along a specified parallel axis (e.g., all ranks in the same DP group).
*/
inline std::vector<int> GetGroupRanks(Axis target, int rank) {
int dp, tp, pp;
GetCoordOf(rank, dp, tp, pp);
return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp);
int dp, tp, pp, ep;
GetCoordOf(rank, dp, tp, pp, ep);
return GlobalEnv::Instance().layout().GroupRanks(target, dp, tp, pp, ep);
}

/**
Expand All @@ -169,7 +187,7 @@ inline std::vector<int> GetGroupRanks(Axis target, int rank) {
*
* Example:
* === Parallel Communication Groups ===
* world_size = 8, config: {DP=2, TP=4, PP=1}, order: {TP -> DP -> PP}
* world_size = 8, config: {DP=2, TP=4, PP=1, EP=1}, order: {TP -> EP -> 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]
Expand Down
4 changes: 4 additions & 0 deletions infini_train/include/nn/parallel/utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,16 @@ std::string GetTensorParallelProcessGroupName(int global_rank);

std::string GetPipelineParallelProcessGroupName(int global_rank);

std::string GetExpertParallelProcessGroupName(int global_rank);

std::vector<int> GetDataParallelGroupRanks(int global_rank);

std::vector<int> GetTensorParallelGroupRanks(int global_rank);

std::vector<int> GetPipelineParallelGroupRanks(int global_rank);

std::vector<int> GetExpertParallelGroupRanks(int global_rank);

// TP/SP Communication Helper Functions
std::vector<std::shared_ptr<Tensor>> GatherFromTPRegionFunc(const std::shared_ptr<Tensor> &input);
std::vector<std::shared_ptr<Tensor>> ReduceScatterToSPRegionFunc(const std::shared_ptr<Tensor> &input);
Expand Down
81 changes: 58 additions & 23 deletions infini_train/src/nn/parallel/global.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
#include "infini_train/include/nn/parallel/global.h"

#include <algorithm>
#include <cstdlib>
#include <format>
#include <sstream>
#include <string>
#include <tuple>

#include "glog/logging.h"

Expand All @@ -29,9 +32,9 @@ void Layout::InitStrides() {
}
}

int Layout::RankOf(int dp, int tp, int pp) const {
int Layout::RankOf(int dp, int tp, int pp, int ep) const {
// Return the thread rank given layout coords
const int coord[AXIS_COUNT] = {dp, tp, pp};
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<Axis>(i);
Expand All @@ -40,14 +43,22 @@ int Layout::RankOf(int dp, int tp, int pp) const {
return r;
}

void Layout::CoordOf(int rank, int &dp, int &tp, int &pp) const {
int Layout::RankOf(int dp, int tp, int pp) const { return RankOf(dp, tp, pp, 0); }

void Layout::CoordOf(int rank, int &dp, int &tp, int &pp, int &ep) 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];
ep = (rank / strides[EP]) % sizes[EP];
}

int Layout::GroupId(Axis target, int dp, int tp, int pp) const {
void Layout::CoordOf(int rank, int &dp, int &tp, int &pp) const {
int ep;
CoordOf(rank, dp, tp, pp, ep);
}

int Layout::GroupId(Axis target, int dp, int tp, int pp, int ep) const {
// Return the parallel ProcessGroup ID where the rank is in
int id = 0;
int mult = 1;
Expand All @@ -56,38 +67,46 @@ int Layout::GroupId(Axis target, int dp, int tp, int pp) const {
if (ax == target) {
continue;
}
int c = (ax == DP ? dp : (ax == TP ? tp : pp));
int c = (ax == DP ? dp : (ax == TP ? tp : (ax == PP ? pp : ep)));
id += c * mult;
mult *= sizes[ax];
}
return id;
}

std::vector<int> Layout::GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp) const {
int Layout::GroupId(Axis target, int dp, int tp, int pp) const { return GroupId(target, dp, tp, pp, 0); }

std::vector<int> Layout::GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp, int fixed_ep) const {
// Return all the ranks within the same parallel ProcessGroup
std::vector<int> ranks;
ranks.reserve(sizes[target]);
int dp = fixed_dp, tp = fixed_tp, pp = fixed_pp;
int dp = fixed_dp, tp = fixed_tp, pp = fixed_pp, ep = fixed_ep;
for (int v = 0; v < sizes[target]; ++v) {
if (target == DP) {
dp = v;
} else if (target == TP) {
tp = v;
} else {
} else if (target == PP) {
pp = v;
} else if (target == EP) {
ep = v;
}
ranks.push_back(RankOf(dp, tp, pp));
ranks.push_back(RankOf(dp, tp, pp, ep));
}
return ranks;
}

std::vector<int> Layout::GroupRanks(Axis target, int fixed_dp, int fixed_tp, int fixed_pp) const {
return GroupRanks(target, fixed_dp, fixed_tp, fixed_pp, 0);
}

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::lock_guard<std::mutex> lock(mutex_);

CHECK(!initialized_) << "Repeated initialization of GlobalEnv!";
Expand All @@ -100,15 +119,23 @@ 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";
tensor_parallel_size_ = 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_;
expert_parallel_size_ = expert_parallel_size;
const int model_parallel_size = tensor_parallel_size_ * pipeline_parallel_size_ * expert_parallel_size_;
CHECK_EQ(world_size_ % model_parallel_size, 0)
<< "World size must be divisible by tensor_parallel_size * pipeline_parallel_size * expert_parallel_size";
data_parallel_size_ = world_size_ / model_parallel_size;

layout_.sizes[DP] = data_parallel_size_;
layout_.sizes[TP] = tensor_parallel_size_;
layout_.sizes[PP] = pipeline_parallel_size_;
layout_.sizes[EP] = expert_parallel_size_;
layout_.InitStrides();

initialized_ = true;
Expand Down Expand Up @@ -174,13 +201,18 @@ int GlobalEnv::virtual_pipeline_parallel_size() const {
return virtual_pipeline_parallel_size_;
}

int GlobalEnv::expert_parallel_size() const {
CHECK(initialized_) << "GlobalEnv is not initialized!";
return expert_parallel_size_;
}

Layout GlobalEnv::layout() const {
CHECK(initialized_) << "GlobalEnv is not initialized!";
return layout_;
}

namespace {
inline const char *AxisName(Axis a) { return a == DP ? "DP" : (a == TP ? "TP" : "PP"); }
inline const char *AxisName(Axis a) { return a == DP ? "DP" : (a == TP ? "TP" : (a == PP ? "PP" : "EP")); }

inline int NumGroups(const Layout &L, Axis target) {
int n = 1;
Expand All @@ -196,8 +228,8 @@ inline int NumGroups(const Layout &L, Axis target) {
std::string ProcessGroupOverview(const Layout &L, 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={}, EP={}}}, order: {{",
GetWorldSize(), L.sizes[DP], L.sizes[TP], L.sizes[PP], L.sizes[EP]);

for (int i = 0; i < AXIS_COUNT; ++i) { oss << AxisName(L.order[i]) << (i + 1 == AXIS_COUNT ? "" : " -> "); }
oss << "}\n";
Expand All @@ -208,13 +240,15 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) {
oss << std::format("[{}] size={}, unenabled\n", AxisName(ax), L.sizes[ax]);
continue;
}
// Build <Group ID, <DP, TP, PP>> mapping
std::vector<std::pair<int, std::tuple<int, int, int>>> groups;
// Build <Group ID, <DP, TP, PP, EP>> mapping
std::vector<std::pair<int, std::tuple<int, int, int, int>>> 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));
for (int ep = 0; ep < (ax == EP ? 1 : L.sizes[EP]); ++ep) {
int gid = L.GroupId(ax, dp, tp, pp, ep);
groups.emplace_back(gid, std::make_tuple(dp, tp, pp, ep));
}
}
}
}
Expand All @@ -228,14 +262,15 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) {
// 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);
int dp, tp, pp, ep;
std::tie(dp, tp, pp, ep) = pair.second;
auto ranks = L.GroupRanks(ax, dp, tp, pp, ep);
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);
auto ep_size_str = (ax == EP) ? "-" : std::to_string(ep);

std::string ranks_str;
ranks_str.reserve(ranks.size() * 4);
Expand All @@ -245,8 +280,8 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) {
}
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);
oss << std::format(" - {} {} (dp={}, tp={}, pp={}, ep={}): [{}]\n", name, gid, dp_size_str, tp_size_str,
pp_size_str, ep_size_str, ranks_str);
}
if (a + 1 < AXIS_COUNT) {
oss << "\n";
Expand Down
6 changes: 6 additions & 0 deletions infini_train/src/nn/parallel/utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,10 @@ std::string GetPipelineParallelProcessGroupName(int global_rank) {
return "PP" + std::to_string(global::GetGroupId(global::PP, global_rank));
}

std::string GetExpertParallelProcessGroupName(int global_rank) {
return "EP" + std::to_string(global::GetGroupId(global::EP, global_rank));
}

std::vector<int> GetDataParallelGroupRanks(int global_rank) { return global::GetGroupRanks(global::DP, global_rank); }

std::vector<int> GetTensorParallelGroupRanks(int global_rank) { return global::GetGroupRanks(global::TP, global_rank); }
Expand All @@ -24,4 +28,6 @@ std::vector<int> GetPipelineParallelGroupRanks(int global_rank) {
return global::GetGroupRanks(global::PP, global_rank);
}

std::vector<int> GetExpertParallelGroupRanks(int global_rank) { return global::GetGroupRanks(global::EP, global_rank); }

} // namespace infini_train::nn::parallel
Loading
Loading