Mindspeed-LLM 基于 Megatron-Core 的 MoE 框架，在运行前通过一系列 Monkey Patch / Wrapper 注册， 将 Megatron 的核心 MoE 模块（Experts 层、Router 路由层、通信调度层） 替换为华为自研的 高效异步 + 通信重叠 + TP/EP 扩展版本。
I. Experts 层（GroupedMLP / SequentialMLP / TEGroupedMLP）
展示 Experts 模块（GroupedMLP / SequentialMLP / TEGroupedMLP）在 Mindspeed 中的替换逻辑与最终调用流向：
┌──────────────────────────────────────────────┐
│ Megatron 原始 Experts 模块 │
│ (megatron.core.transformer.moe.experts) │
└──────────────────┬────────────────────────────┘
│
▼
┌──────────────────────────────────────────────┐
│ Mindspeed Patch 注册逻辑 │
│ （在 patch_core_transformers 中注册） │
└──────────────────┬────────────────────────────┘
│
┌────────────────────────────────────────────────────────────────┐
│ 被替换函数映射一览 │
└────────────────────────────────────────────────────────────────┘
│
├──► GroupedMLP.__init__ ───────────────► groupedmlp_init_wrapper
│
├──► GroupedMLP.forward ───────────────► groupedmlp_forward
│ │
│ ├─(当 overlap on)──► group_mlp_forward
│ └─(当 overlap off)─► groupedmlp_forward
│
├──► SequentialMLP.forward ─────► sequential_mlp_forward
│ (仅当 dispatcher == alltoall)
│
└──► grouped_gemm_util.* ──────► Ops / grouped_gemm_is_available / assert_grouped_gemm_is_available (自定义替换)
（注：TEGroupedMLP 保持原状，不参与替换）
┌────────────────────────────────────────────────────────────────┐
│ GroupedMLP 前向主流程图 │
└────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────┐
│ (1) permuted_local_hidden_states 输入 │
│ (2) 权重 reshape: │
│ w1: [num_local_experts, hidden, *] │
│ w2: [num_local_experts, *, hidden] │
├────────────────────────────────────────────┤
│ (3) 根据 overlap / moe_bmm_mc2 / recompute │
│ 三种情况选择路径： │
│ │
│ ├ 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 │
│ └ recompute on → CheckpointWithoutOutput 管理激活重计算 │
├────────────────────────────────────────────┤
│ (4) 输出 fc2_output │
└────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────────┐
│ SequentialMLP 前向主流程图 │
└────────────────────────────────────────────────────────────────┘
│
▼
┌────────────────────────────────────────────┐
│ (1) 根据 tokens_per_expert 累积偏移 │
│ (2) 逐专家执行 forward │
│ (3) 若 TP>1 → COMM_STREAM + all_gather_last_dim │
│ (4) 输出拼接 │
└────────────────────────────────────────────┘
┌────────────────────────────────────────────────────────────────┐
│ 实际运行函数路径一览 │
└────────────────────────────────────────────────────────────────┘
GroupedMLP.__init__ → groupedmlp_init_wrapper
│
└─► GroupedMLP.forward
├─ overlap on → group_mlp_forward
│ ├─ grouped_mlp_with_comp_and_comm_overlap_all2all
│ └─ grouped_mlp_with_comp_and_comm_overlap_allgather
└─ overlap off → groupedmlp_forward
├─ gg.ops.gmm
├─ fused_alltoall_gather_bmm
└─ fused_bmm_reducescatter_alltoall
SequentialMLP.forward
├─ dispatcher==alltoall → sequential_mlp_forward
│ └─ 内部 all_gather_last_dim
└─ dispatcher==allgather → 原生 forward
TEGroupedMLP → 保留原始实现 (无改动)
┌────────────────────────────────────────────────────────────────┐
│ 被淘汰 / 不再调用的原函数 │
└────────────────────────────────────────────────────────────────┘
✗ GroupedMLP.__init__ (原版)
✗ GroupedMLP.forward (原版)
✗ SequentialMLP.forward (在 alltoall 模式下)
✗ grouped_gemm_util 默认实现 (替换为 Ops 等)
✓ TEGroupedMLP (保留原版)
🧠 一句话总结：
在 Mindspeed 中，GroupedMLP 的初始化与前向都被彻底接管； SequentialMLP 只在
dispatcher == alltoall时改走新路径； 所有 overlap（AllGather / AllToAll）场景下统一进入group_mlp_forward； 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_alltoall
SequentialMLP.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_forward
overlap off → groupedmlp_forward
dispatcher=alltoall → sequential_mlp_forward
dispatcher=allgather → 原生 forward
🧩 被淘汰函数
| 原始函数 | 状态 |
|---|---|
GroupedMLP.__init__ |
❌ 永久替换 |
GroupedMLP.forward |
❌ 永久替换 |
SequentialMLP.forward |
⚙️ alltoall 模式下替换 |
grouped_gemm_util 默认实现 |
❌ 替换为自定义 Ops |
TEGroupedMLP |
✅ 保留原版 |
II. Router 层（TopKRouter）
下面示意图，展示 MoE Router 部分（TopKRouter）在 Mindspeed 中被替换后的整体调用与函数接管关系
┌──────────────────────────────┐
│ Megatron 原始 TopKRouter │
│ (megatron.core.transformer │
│ .moe.router.TopKRouter)│
└─────────────┬────────────────┘
│
▼
┌──────────────────────────────┐
│ Mindspeed Patch │
│ MOERouter.register_patches │
│ (from moe_router.py) │
└─────────────┬────────────────┘
│
┌─────────────────────────────────────────────────────────────┐
│ 被替换函数映射一览 │
└─────────────────────────────────────────────────────────────┘
│
├──► TopKRouter.__init__ ───────────────► topk_router_init_wrapper
│
├──► TopKRouter.forward ────────────────► topk_router_forward
│
├──► TopKRouter.gating ────────────────► topk_router_gating_func
│
├──► TopKRouter.routing ────────────────► topk_router_routing
│
└──► z_loss_func (moe_utils) ────────────► mindspeed 自定义 z_loss_func
┌─────────────────────────────────────────────────────────────┐
│ topk_router_forward 主调用链 │
└─────────────────────────────────────────────────────────────┘
│
▼
┌──────────────────────────────┐
│ apply_input_jitter (可选) │
├──────────────────────────────┤
│ topk_router_gating_func │
│ → 线性变换 (FP32/FP16 可选)│
├──────────────────────────────┤
│ topk_router_routing │
│ → 根据 routing_type 分派： │
│ │
│ ├ sinkhorn → sinkhorn_load_balancing
│ ├ aux_loss → aux_loss_load_balancing
│ ├ softmax_topk → softmax + topk
│ ├ group_limited_greedy → group_limited_greedy_topKgating
│ ├ pai_megatron_aux_loss → pai_megatron_aux_loss
│ ├ sparsemixer_topk → sparsemixer_top2
│ ├ none / noaux_tc → topk_softmax_with_capacity
│ │
├──────────────────────────────┤
│ optional: apply_z_loss, gather_from_sequence_parallel_region │
└──────────────────────────────┘
┌─────────────────────────────────────────────────────────────┐
│ 实际运行函数 (最终执行路径) │
└─────────────────────────────────────────────────────────────┘
TopKRouter.forward → topk_router_forward
├─→ gating → topk_router_gating_func
├─→ routing → topk_router_routing
│ ├ sinkhorn_load_balancing
│ ├ aux_loss_load_balancing
│ ├ topk_softmax_with_capacity
│ ├ group_limited_greedy_topKgating
│ ├ pai_megatron_aux_loss
│ └ sparsemixer_top2
└─→ apply_z_loss (替换版)
┌─────────────────────────────────────────────────────────────┐
│ 被淘汰 / 不再调用的原函数 │
└─────────────────────────────────────────────────────────────┘
✗ TopKRouter.__init__ (Megatron 原版)
✗ TopKRouter.forward
✗ TopKRouter.gating
✗ TopKRouter.routing
✗ moe_utils.topk_softmax_with_capacity (原版)
可总结为一句话：
原 Megatron TopKRouter 的四个核心方法全部被替换为 Mindspeed 对应 wrapper，最终所有 gating/routing 行为都经由
topk_router_forward → topk_router_routing流程统一调度。
🔧 替换注册点
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 版本仅保留结构壳体。