华为昇腾笔记之Mindspeed-LLM 中 MoE 实现机制与重写逻辑总览

🧩 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 版本仅保留结构壳体