feat: implement ep topology infrastructure

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();

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();

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);

infini_train/include/nn/parallel/global.h

@@ -1,35 +1,61 @@
 #pragma once
 
+#include <array>
+#include <cstdint>
+#include <initializer_list>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <vector>
 
 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<Axis, AXIS_COUNT> order = {TP, EP, DP, PP});
+
+    int AxisSize(Axis axis) const;
+    int WorldSize() const;
+    const std::array<Axis, AXIS_COUNT> &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<std::vector<int>> GetRanks(std::initializer_list<Axis> varying_axes) const;
+    std::vector<std::vector<int>> GetRanks(Axis varying_axis) const;
 
+    int GroupId(std::initializer_list<Axis> varying_axes, int global_rank) const;
+    int GroupId(Axis varying_axis, int global_rank) const;
+    std::vector<int> GroupRanks(std::initializer_list<Axis> varying_axes, int global_rank) const;
+    std::vector<int> 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<int> 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<int, AXIS_COUNT> sizes_{1, 1, 1, 1};
+    std::array<Axis, AXIS_COUNT> order_{TP, EP, DP, PP};
+    std::array<int, AXIS_COUNT> strides_{1, 1, 1, 1};
 };
 
 class GlobalEnv {
 public:
     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<int> 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<int> 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<int> 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<int> 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

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<int> GetDataParallelGroupRanks(int global_rank);
 
+std::vector<int> GetExpertDataParallelGroupRanks(int global_rank);
+
 std::vector<int> GetTensorParallelGroupRanks(int global_rank);
 
+std::vector<int> GetExpertTensorParallelGroupRanks(int global_rank);
+
 std::vector<int> GetPipelineParallelGroupRanks(int global_rank);
 
+std::vector<int> GetExpertParallelGroupRanks(int global_rank);
+
+std::vector<int> GetExpertTensorAndExpertParallelGroupRanks(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);