🧩 Mindspeed-LLM 的 MoE 实现机制与重写逻辑总览
Mindspeed-LLM 基于 Megatron-Core 的 MoE 框架，在运行前通过一系列 Monkey Patch / Wrapper 注册，
将 Megatron 的核心 MoE 模块（Experts 层、Router 路由层、通信调度层）
替换为华为自研的 高效异步 + 通信重叠 + TP/EP 扩展版本。
I. Experts 层（GroupedMLP / SequentialMLP / TEGroupedMLP）
🔧 核心修改入口
patch_core_transformers() 中注册：
MegatronAdaptation.register('megatron.core.transformer.moe.experts.GroupedMLP.__init__',groupedmlp_init_wrapper)
MegatronAdaptation.register('megatron.core.transformer.moe.experts.GroupedMLP.forward',groupedmlp_forward 或 group_mlp_forward)
MegatronAdaptation.register('megatron.core.transformer.moe.experts.SequentialMLP.forward',sequential_mlp_forward)
🧠 逻辑核心
| 模块 | 替换后函数 | 触发条件 | 功能概述 |
|---|---|---|---|
| GroupedMLP.init | groupedmlp_init_wrapper |
总是生效 | 临时改 TP=1 初始化 GMM 权重，注册 gmm_weight 标记，若 GLU 激活→换为 fused_swiglu |
| GroupedMLP.forward | groupedmlp_forward / group_mlp_forward |
overlap off / on | - off：调用 gg.ops.gmm 或 fused_alltoall_gather_bmm - on：走通信重叠核 (grouped_mlp_with_comp_and_comm_overlap_alltoall/allgather) |
| SequentialMLP.forward | sequential_mlp_forward |
dispatcher==alltoall | 多专家串行执行 + TP all-gather 通信流优化 |
| TEGroupedMLP | 不变 | 无 | 保持原 TransformerEngine 路径 |
⚙️ 调用路径（概览）
GroupedMLP.__init__ → groupedmlp_init_wrapper
GroupedMLP.forward ├─ overlap on → group_mlp_forward│ └─ grouped_mlp_with_comp_and_comm_overlap_all2all/allgather└─ overlap off → groupedmlp_forward├─ gg.ops.gmm├─ fused_alltoall_gather_bmm└─ fused_bmm_reducescatter_alltoallSequentialMLP.forward ├─ dispatcher==alltoall → sequential_mlp_forward│ └─ COMM_STREAM + all_gather_last_dim└─ dispatcher==allgather → 原生 forward
🗺️ Experts 模块字符示意图
┌────────────────────────────┐
│ Megatron 原始 Experts │
│ (GroupedMLP / Sequential) │
└──────────┬────────────────┘│▼
┌────────────────────────────┐
│ Mindspeed Patch 注册逻辑 │
│ (patch_core_transformers) │
└──────────┬────────────────┘│┌──────────────────────────────────────────┐│ 被替换函数一览 │├──────────────────────────────────────────┤│ GroupedMLP.__init__ → groupedmlp_init_wrapper ││ GroupedMLP.forward → groupedmlp_forward / group_mlp_forward ││ SequentialMLP.forward → sequential_mlp_forward (仅 alltoall) ││ grouped_gemm_util.* → 自定义 Ops 体系 │└──────────────────────────────────────────┘↓实际执行路径（动态判定）overlap on → group_mlp_forwardoverlap off → groupedmlp_forwarddispatcher=alltoall → sequential_mlp_forwarddispatcher=allgather → 原生 forward
🧩 被淘汰函数
| 原始函数 | 状态 |
|---|---|
GroupedMLP.__init__ |
❌ 永久替换 |
GroupedMLP.forward |
❌ 永久替换 |
SequentialMLP.forward |
⚙️ alltoall 模式下替换 |
grouped_gemm_util 默认实现 |
❌ 替换为自定义 Ops |
TEGroupedMLP |
✅ 保留原版 |
II. Router 层（TopKRouter）
🔧 替换注册点
MOERouter.register_patches() 中：
MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.__init__',topk_router_init_wrapper)
MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.forward',topk_router_forward)
MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.gating',topk_router_gating_func)
MegatronAdaptation.register('megatron.core.transformer.moe.router.TopKRouter.routing',topk_router_routing)
MegatronAdaptation.register('megatron.core.transformer.moe.moe_utils.topk_softmax_with_capacity',topk_softmax_with_capacity)
🧠 TopKRouter 逻辑结构
| 模块 | 替换后函数 | 功能概述 |
|---|---|---|
TopKRouter.__init__ |
topk_router_init_wrapper |
初始化 group / n_group / score_fn 等参数 |
TopKRouter.forward |
topk_router_forward |
主入口：input jitter → gating → routing |
TopKRouter.gating |
topk_router_gating_func |
logits 线性变换，控制激活方式 |
TopKRouter.routing |
topk_router_routing |
调度多种路由算法分支 |
moe_utils.topk_softmax_with_capacity |
Mindspeed 实现 | 增强版 softmax+容量控制 |
（条件）aux_loss_load_balancing |
自定义实现 | 支持异步 overlap / async log loss 等场景 |
🧮 路由算法分支
在 topk_router_routing 中，根据 --moe-router-load-balancing-type 分派到：
| 类型 | 实际函数 | 说明 |
|---|---|---|
sinkhorn |
sinkhorn_load_balancing |
Sinkhorn 约束分配 |
aux_loss |
aux_loss_load_balancing |
含负载均衡损失 |
softmax_topk |
softmax + topk |
普通 softmax 策略 |
group_limited_greedy |
group_limited_greedy_topKgating |
先选组后选专家 |
pai_megatron_aux_loss |
pai_megatron_aux_loss |
PAI 风格 loss |
sparsemixer_topk |
sparsemixer_top2 |
稀疏混合器路由 |
none/noaux_tc |
topk_softmax_with_capacity |
无 loss 简化路由 |
⚙️ Router 调用主链
TopKRouter.forward├─ topk_router_forward│ ├─ apply_input_jitter (可选)│ ├─ topk_router_gating_func│ ├─ topk_router_routing│ │ ├─ sinkhorn_load_balancing│ │ ├─ aux_loss_load_balancing│ │ ├─ group_limited_greedy_topKgating│ │ ├─ sparsemixer_top2│ │ ├─ pai_megatron_aux_loss│ │ └─ topk_softmax_with_capacity│ └─ apply_z_loss / gather_from_sequence_parallel_region (可选)
🗺️ Router 模块字符示意图
┌────────────────────────────┐
│ Megatron 原始 TopKRouter │
└──────────┬────────────────┘│▼
┌────────────────────────────┐
│ Mindspeed Patch 注册逻辑 │
│ (MOERouter.register_patches) │
└──────────┬────────────────┘│┌──────────────────────────────────────────┐│ 被替换函数一览 │├──────────────────────────────────────────┤│ TopKRouter.__init__ → topk_router_init_wrapper ││ TopKRouter.forward → topk_router_forward ││ TopKRouter.gating → topk_router_gating_func ││ TopKRouter.routing → topk_router_routing ││ moe_utils.topk_softmax_with_capacity → 自定义版 ││ (条件) TopKRouter.aux_loss_load_balancing → Mindspeed 版 │└──────────────────────────────────────────┘↓实际执行链topk_router_forward├─ gating → topk_router_gating_func├─ routing → topk_router_routing│ ├ sinkhorn_load_balancing│ ├ aux_loss_load_balancing│ ├ group_limited_greedy_topKgating│ ├ sparsemixer_top2│ ├ pai_megatron_aux_loss│ └ topk_softmax_with_capacity└─ apply_z_loss / gather_sequence_parallel (可选)被淘汰：✗ TopKRouter.__init__ (原版)✗ TopKRouter.forward✗ TopKRouter.gating✗ TopKRouter.routing✗ 原版 topk_softmax_with_capacity
III. 🔩 总结：Mindspeed-LLM MoE 的总体演化逻辑
| 模块 | 原 Megatron 功能 | Mindspeed 替换方案 | 增强点 |
|---|---|---|---|
| Experts (GroupedMLP) | 多专家独立 GEMM | 通信-计算重叠 / TP-EP 扩展 / fused_swiglu | 提升并行效率、减通信瓶颈 |
| SequentialMLP | 串行专家推理 | TP all-gather + 异步流重叠 | 支持多专家 alltoall 模式 |
| Router (TopKRouter) | softmax+topk 简单 gating | 多种策略（Sinkhorn, SparseMixer, Greedy）+ 异步均衡 loss | 路由更鲁棒、更灵活 |
| 通信核 | 普通 alltoall / allgather | grouped_mlp_with_comp_and_comm_overlap_* |
计算与通信异步融合 |
| Checkpoint / 重计算 | 默认 PyTorch 机制 | CheckpointWithoutOutput 自定义 |
精确控制激活保存与释放 |
✅ 一句话总结：
Mindspeed-LLM 的 MoE 实现在保持 Megatron 接口兼容的同时，
通过 Patch 替换实现了“专家层异步通信 + 重计算优化 + 多策略路由”的闭环体系。
核心思路：GroupedMLP/SequentialMLP 接管专家执行；TopKRouter 接管路由调度。
所有前向路径最终都经由 Mindspeed 自定义函数运行，原始 Megatron 版本仅保留结构壳体