[WIP] feat: implement ep topology infrastructure

example/gpt2/main.cc

@@ -20,10 +20,10 @@
 #include "infini_train/include/nn/modules/loss.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
+#include "infini_train/include/nn/parallel/comm.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
 #include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/nn/parallel/parallel_functional.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 #include "infini_train/include/nn/parallel/rank.h"
 #include "infini_train/include/nn/parallel/reduce_op_type.h"
@@ -507,7 +507,7 @@ void Train(const nn::parallel::Rank &rank) {
 
         if (ddp_world_size > 1) {
             auto lossf_tensor = std::make_shared<Tensor>(&lossf, std::vector<int64_t>{}, DataType::kFLOAT32, device);
-            function::AllReduce(lossf_tensor, function::ReduceOpType::kAvg, ddp_pg);
+            comm::AllReduce(lossf_tensor, comm::ReduceOpType::kAvg, ddp_pg);
             lossf = static_cast<const float *>(lossf_tensor->To(Device()).DataPtr())[0];
         }
 
@@ -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

@@ -19,10 +19,10 @@
 #include "infini_train/include/nn/modules/loss.h"
 #include "infini_train/include/nn/modules/module.h"
 #include "infini_train/include/nn/modules/transformer/transformer.h"
+#include "infini_train/include/nn/parallel/comm.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_optimizer.h"
 #include "infini_train/include/nn/parallel/global.h"
-#include "infini_train/include/nn/parallel/parallel_functional.h"
 #include "infini_train/include/nn/parallel/pp/pipeline_parallel.h"
 #include "infini_train/include/nn/parallel/process_group.h"
 #include "infini_train/include/nn/parallel/rank.h"
@@ -484,7 +484,7 @@ void Train(const nn::parallel::Rank &rank) {
 
         if (ddp_world_size > 1) {
             auto lossf_tensor = std::make_shared<Tensor>(&lossf, std::vector<int64_t>{}, DataType::kFLOAT32, device);
-            function::AllReduce(lossf_tensor, function::ReduceOpType::kAvg, ddp_pg);
+            comm::AllReduce(lossf_tensor, comm::ReduceOpType::kAvg, ddp_pg);
             lossf = static_cast<const float *>(lossf_tensor->To(Device()).DataPtr())[0];
         }
 
@@ -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/checkpoint_loader.cc

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

example/mixtral/config.h

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

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/autograd/comm.h

@@ -1,10 +1,12 @@
 #pragma once
 
+#include <cstdint>
 #include <memory>
 #include <vector>
 
 #include "infini_train/include/autograd/function.h"
 #include "infini_train/include/device.h"
+#include "infini_train/include/nn/parallel/reduce_op_type.h"
 
 namespace infini_train {
 class Tensor;
@@ -16,6 +18,35 @@ class ProcessGroup;
 } // namespace infini_train
 
 namespace infini_train::autograd::comm {
+class AllGather : public autograd::Function {
+public:
+    static constexpr char kType[] = "AllGatherFunction";
+
+    explicit AllGather(const nn::parallel::ProcessGroup *pg);
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    const nn::parallel::ProcessGroup *pg_ = nullptr;
+};
+
+class ReduceScatter : public autograd::Function {
+public:
+    static constexpr char kType[] = "ReduceScatterFunction";
+
+    ReduceScatter(nn::parallel::comm::ReduceOpType reduce_op, const nn::parallel::ProcessGroup *pg);
+
+    std::vector<std::shared_ptr<Tensor>> Forward(const std::vector<std::shared_ptr<Tensor>> &input_tensors) override;
+
+    std::vector<std::shared_ptr<Tensor>> Backward(const std::vector<std::shared_ptr<Tensor>> &grad_outputs) override;
+
+private:
+    nn::parallel::comm::ReduceOpType reduce_op_;
+    const nn::parallel::ProcessGroup *pg_ = nullptr;
+};
+
 class Scatter : public autograd::Function {
 public:
     static constexpr char kType[] = "ScatterFunction";

infini_train/include/autograd/function_hook.h

@@ -35,13 +35,13 @@ class PostAccumulateGradHook {
 
 class AllReducePostAccumulateHook : public PostAccumulateGradHook {
 public:
-    AllReducePostAccumulateHook(infini_train::nn::parallel::function::ReduceOpType reduce_op,
+    AllReducePostAccumulateHook(infini_train::nn::parallel::comm::ReduceOpType reduce_op,
                                 const infini_train::nn::parallel::ProcessGroup *pg = nullptr);
 
     void operator()(const std::shared_ptr<Tensor> &tensor) override;
 
 private:
-    infini_train::nn::parallel::function::ReduceOpType reduce_op_;
+    infini_train::nn::parallel::comm::ReduceOpType reduce_op_;
     const infini_train::nn::parallel::ProcessGroup *pg_ = nullptr;
 };
 } // namespace infini_train::autograd

infini_train/include/core/ccl/ccl.h

@@ -37,21 +37,20 @@ class CclImpl {
     virtual void CommDestroy(CclComm *comm) const;
 
     virtual void AllReduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype,
-                           nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const;
+                           nn::parallel::comm::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const;
 
     virtual void Broadcast(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, int root,
                            const CclComm *comm, Stream *stream) const;
 
     virtual void Reduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype,
-                        nn::parallel::function::ReduceOpType reduce_op, int root, const CclComm *comm,
+                        nn::parallel::comm::ReduceOpType reduce_op, int root, const CclComm *comm,
                         Stream *stream) const;
 
     virtual void AllGather(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, const CclComm *comm,
                            Stream *stream) const;
 
     virtual void ReduceScatter(const void *sendbuff, void *recvbuff, size_t recv_count, DataType dtype,
-                               nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm,
-                               Stream *stream) const;
+                               nn::parallel::comm::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const;
 
     virtual void Send(const void *buff, size_t count, DataType dtype, int peer, const CclComm *comm,
                       Stream *stream) const;

infini_train/include/nn/modules/transformer/moe/experts.h

@@ -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;
 };
 

infini_train/include/nn/modules/transformer/moe/moe_layer.h

@@ -20,6 +20,7 @@ class MoELayer : public CloneableModule<MoELayer> {
 
 private:
     TransformerConfig config_;
+    int64_t num_local_experts_ = 0;
 };
 
 } // namespace infini_train::nn::moe

infini_train/include/nn/modules/transformer/moe/moe_utils.h

@@ -32,7 +32,7 @@ struct PermutationMetadata {
     // Number of dispatched tokens assigned to each expert. Kept on host because the
     // current local SequentialMLP uses it only for CPU-side slicing decisions.
     // Shape: [num_experts].
-    std::vector<int64_t> tokens_per_expert;
+    std::shared_ptr<Tensor> tokens_per_expert;
 };
 
 struct PermutationResult {
@@ -82,13 +82,10 @@ PermutationResult Permute(const std::shared_ptr<Tensor> &tokens, const std::shar
                           const std::shared_ptr<Tensor> &routing_map);
 
 // Restore permuted expert outputs to token order, matching Megatron moe_utils.unpermute.
-// Each expert output row is weighted by its aligned permuted_probs value before scatter-add.
 // Args:
 //   permuted_tokens: [num_dispatched_tokens, hidden_size] expert outputs in permuted order.
-//   permuted_probs:  [num_dispatched_tokens] routing probability for each expert output row.
 //   restore_shape:   [num_tokens, hidden_size] output shape before permutation.
-std::shared_ptr<Tensor> Unpermute(const std::shared_ptr<Tensor> &permuted_tokens,
-                                  const std::shared_ptr<Tensor> &permuted_probs, const PermutationMetadata &metadata,
+std::shared_ptr<Tensor> Unpermute(const std::shared_ptr<Tensor> &permuted_tokens, const PermutationMetadata &metadata,
                                   const std::vector<int64_t> &restore_shape);
 
 } // namespace infini_train::nn::moe

infini_train/include/nn/modules/transformer/moe/token_dispatcher.h

@@ -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

infini_train/include/nn/modules/transformer/transformer_config.h

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

infini_train/include/nn/parallel/comm.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <memory>
+
+#include "infini_train/include/nn/parallel/reduce_op_type.h"
+
+namespace infini_train {
+class Tensor;
+}
+
+namespace infini_train::nn::parallel {
+class ProcessGroup;
+class Work;
+} // namespace infini_train::nn::parallel
+
+namespace infini_train::nn::parallel::comm {
+
+std::shared_ptr<Work> AllReduce(const std::shared_ptr<Tensor> &tensor, ReduceOpType reduce_op,
+                                const ProcessGroup *pg = nullptr, bool async_op = false);
+
+std::shared_ptr<Work> AllGather(const std::shared_ptr<Tensor> &output, const std::shared_ptr<Tensor> &input,
+                                const ProcessGroup *pg = nullptr, bool async_op = false);
+
+std::shared_ptr<Work> ReduceScatter(const std::shared_ptr<Tensor> &output, const std::shared_ptr<Tensor> &input,
+                                    ReduceOpType reduce_op, const ProcessGroup *pg = nullptr, bool async_op = false);
+
+} // namespace infini_train::nn::parallel::comm

infini_train/include/nn/parallel/ddp/reducer.h

@@ -7,7 +7,6 @@
 
 #include "infini_train/include/datatype.h"
 #include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h"
-#include "infini_train/include/nn/parallel/parallel_functional.h"
 
 namespace infini_train {
 class Tensor;