华为昇腾笔记之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 版本仅保留结构壳体