Skip to content

[WIP] 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 2 commits into
base: master
Choose a base branch
from
Open

[WIP] feat: implement ep topology infrastructure #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d79f114
feat(parallel): add dense and expert rank generators
kilinchange Jul 20, 2026
96d04a3
feat: support ep for moe_layer, experts and token_dispatcher
kilinchange Jul 23, 2026
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
4 changes: 3 additions & 1 deletion example/gpt2/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
4 changes: 3 additions & 1 deletion example/llama3/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
1 change: 0 additions & 1 deletion example/mixtral/checkpoint_loader.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,6 @@ nn::TransformerConfig ConfigFromLLMC(const std::string &filepath) {

nn::MoEConfig moe_config;
moe_config.num_experts = infini_train::BytesToType<int32_t>(header, 8 * sizeof(int32_t));
moe_config.expert_parallel_size = 1;
moe_config.router_topk = infini_train::BytesToType<int32_t>(header, 15 * sizeof(int32_t));
moe_config.moe_ffn_hidden_size = infini_train::BytesToType<int32_t>(header, 16 * sizeof(int32_t));
moe_config.token_dispatcher_type = nn::MoEConfig::TokenDispatcherType::kAllGather;
Expand Down
2 changes: 0 additions & 2 deletions example/mixtral/config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ inline nn::TransformerConfig TinyMixtralConfig() {

nn::MoEConfig moe_config;
moe_config.num_experts = 8;
moe_config.expert_parallel_size = 1; // Single-rank validation scale.
moe_config.router_topk = 2;
moe_config.moe_ffn_hidden_size = 1792; // Scaled down as 512 * 3.5; real Mixtral uses 14336.
moe_config.token_dispatcher_type = nn::MoEConfig::TokenDispatcherType::kAllGather; // Single-rank validation path.
Expand Down Expand Up @@ -63,7 +62,6 @@ inline void SanitizeTinyMixtralConfig(const nn::TransformerConfig &c) {

const auto &moe = c.moe_config.value();
CHECK_GT(moe.num_experts, 0);
CHECK_EQ(moe.expert_parallel_size, 1) << "tiny Mixtral single-rank validation expects EP=1";
CHECK_GT(moe.router_topk, 0);
CHECK_LE(moe.router_topk, moe.num_experts);
CHECK_GT(moe.moe_ffn_hidden_size, 0);
Expand Down
4 changes: 3 additions & 1 deletion example/mixtral/main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
3 changes: 1 addition & 2 deletions infini_train/include/nn/modules/transformer/moe/experts.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,11 @@ class SequentialMLP : public CloneableModule<SequentialMLP> {
static constexpr char kType[] = "SequentialMLP";
static constexpr char kExpertNamePrefix[] = "expert_";

explicit SequentialMLP(const TransformerConfig &config);
SequentialMLP(int64_t num_local_experts, const TransformerConfig &config);

std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;

private:
TransformerConfig config_;
int64_t num_local_experts_ = 0;
};

Expand Down
1 change: 1 addition & 0 deletions infini_train/include/nn/modules/transformer/moe/moe_layer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ class MoELayer : public CloneableModule<MoELayer> {

private:
TransformerConfig config_;
int64_t num_local_experts_ = 0;
};

} // namespace infini_train::nn::moe
50 changes: 27 additions & 23 deletions infini_train/include/nn/modules/transformer/moe/token_dispatcher.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,67 +4,71 @@
#include <memory>
#include <vector>

#include "infini_train/include/datatype.h"
#include "infini_train/include/nn/modules/transformer/moe/moe_utils.h"
#include "infini_train/include/nn/modules/transformer/transformer_config.h"

namespace infini_train {
class Tensor;
} // namespace infini_train

namespace infini_train::nn::parallel {
class ProcessGroup;
}

namespace infini_train::nn::moe {

class MoETokenDispatcher {
public:
virtual ~MoETokenDispatcher() = default;

const PermutationResult &Dispatch(const std::shared_ptr<Tensor> &tokens, const std::shared_ptr<Tensor> &routing_map,
const std::shared_ptr<Tensor> &probs);
std::shared_ptr<Tensor> Combine(const std::shared_ptr<Tensor> &hidden_states) const;

protected:
explicit MoETokenDispatcher(const TransformerConfig &config);

virtual std::vector<std::shared_ptr<Tensor>> DispatchPreprocess(const std::shared_ptr<Tensor> &tokens,
const std::shared_ptr<Tensor> &routing_map,
const std::shared_ptr<Tensor> &probs)
= 0;
virtual std::vector<std::shared_ptr<Tensor>> TokenDispatch(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs) const
const std::shared_ptr<Tensor> &probs)
= 0;
virtual const PermutationResult &DispatchPostprocess(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs)
virtual PermutationResult DispatchPostprocess(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs)
= 0;
virtual std::shared_ptr<Tensor> CombinePreprocess(const std::shared_ptr<Tensor> &hidden_states) const = 0;
virtual std::shared_ptr<Tensor> TokenCombine(const std::shared_ptr<Tensor> &hidden_states) const = 0;
virtual std::shared_ptr<Tensor> CombinePostprocess(const std::shared_ptr<Tensor> &hidden_states) const = 0;

protected:
explicit MoETokenDispatcher(const TransformerConfig &config);

TransformerConfig config_;
PermutationResult dispatch_;
std::vector<int64_t> hidden_dims_;
std::shared_ptr<Tensor> routing_map_;
std::shared_ptr<Tensor> local_map_;
std::shared_ptr<Tensor> local_probs_;
int64_t num_tokens_ = 0;
int64_t hidden_size_ = 0;
};

class MoEAllGatherTokenDispatcher : public MoETokenDispatcher {
public:
MoEAllGatherTokenDispatcher(int64_t num_local_experts, const TransformerConfig &config);
MoEAllGatherTokenDispatcher(int64_t num_local_experts, const std::vector<int64_t> &local_expert_ids,
const TransformerConfig &config, const parallel::ProcessGroup *process_group);

private:
std::vector<std::shared_ptr<Tensor>> DispatchPreprocess(const std::shared_ptr<Tensor> &tokens,
const std::shared_ptr<Tensor> &routing_map,
const std::shared_ptr<Tensor> &probs) override;
std::vector<std::shared_ptr<Tensor>> TokenDispatch(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs) const override;
const PermutationResult &DispatchPostprocess(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs) override;
const std::shared_ptr<Tensor> &probs) override;
PermutationResult DispatchPostprocess(const std::shared_ptr<Tensor> &hidden_states,
const std::shared_ptr<Tensor> &probs) override;
std::shared_ptr<Tensor> CombinePreprocess(const std::shared_ptr<Tensor> &hidden_states) const override;
std::shared_ptr<Tensor> TokenCombine(const std::shared_ptr<Tensor> &hidden_states) const override;
std::shared_ptr<Tensor> CombinePostprocess(const std::shared_ptr<Tensor> &hidden_states) const override;

int64_t num_local_experts_ = 0;
private:
// Global expert IDs owned by this rank, in local expert order.
std::vector<int64_t> local_expert_ids_;
const parallel::ProcessGroup *process_group_ = nullptr;
PermutationMetadata permutation_metadata_;
std::vector<int64_t> hidden_shape_;
std::vector<int64_t> hidden_shape_before_permute_;
std::shared_ptr<Tensor> routing_map_;
std::shared_ptr<Tensor> empty_route_hidden_states_;
std::shared_ptr<Tensor> empty_route_probs_;
DataType local_probs_dtype_;
};

} // namespace infini_train::nn::moe
1 change: 0 additions & 1 deletion infini_train/include/nn/modules/transformer/transformer_config.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@ struct MoEConfig {
};

int64_t num_experts = 0;
int64_t expert_parallel_size = 1;
int64_t router_topk = 1;
bool router_pre_softmax = false;
std::optional<float> router_topk_scaling_factor = std::nullopt;
Expand Down
Loading
Loading