跳到主要内容

数据截至 (上游 commit e79cb4c1bae1)

05 · MoE 支持

这一章讲什么： 混合专家（MoE）层在 Megatron 里的完整实现：TopKRouter 怎么选专家、token 怎么靠 All-to-All 在专家并行（EP）组里搬家、专家计算为什么用分组 GEMM、以及负载均衡 Loss 挂在哪里。读完你会理解「route → dispatch → compute → combine」这四拍为什么是 MoE 系统的标准骨架。

1. 它要解决的小问题​

MoE 把每层的 FFN 换成 N 个「专家」小 FFN，每个 token 只激活其中 top-k 个。参数量翻 N 倍，计算量几乎不变——这是它的全部吸引力。

但工程上引入两个新问题：

  • token 要搬家。 专家摊在 EP 组的不同卡上（01 章），每张卡的 token 可能要去任何一张卡找它的专家——这是一个 All-to-All 通信模式，且每次前向的通信量由路由结果动态决定。
  • 计算是变长的。 每个专家收到的 token 数不同（路由不均匀），一次把「多个小矩阵乘」高效算完，需要专门的分组 GEMM，而不是循环调用普通 Linear。

2. 思路：四拍流水线​

Megatron 把 MoE 层的前向稳定成四拍，接口固定在 MoELayer.forward（megatron/core/transformer/moe/moe_layer.py:637）的 docstring 里（:641-646）：

hidden_states [s, b, h]
     │
     ▼
① route      TopKRouter: 每个 token 打分、选 top-k 专家
             产出 probs + routing_map [num_tokens, num_experts]
     │
     ▼
② dispatch   本地 permute(按目标专家排序 token)
             → All-to-All 搬到专家所在的 EP 卡
     │
     ▼
③ compute    GroupedMLP: 本地 num_local_experts 个专家一次 grouped GEMM 算完
     │
     ▼
④ combine    All-to-All 搬回 → unpermute 还原 token 顺序
             → 按 probs 加权求和
     │
     ▼
output [s, b, h]  (形状与输入完全相同)

怎么读这张图： ②④ 互为镜像（permute 配 unpermute，两次 All-to-All 方向相反），①③ 是纯本地计算。通信只在 ②④ 的边界上。

四拍各自的代码落点：route（moe_layer.py:466）、dispatch（:513）、routed_experts_compute（:553）、combine（:590），在 custom_forward 里顺序调用（:681-704）。

3. 原理演示：permute / All-to-All / unpermute​

为什么先 permute 再 All-to-All？因为 All-to-All 按「目标 rank」切分发送，而路由结果按「目标专家」给出——先按专家把 token 排序聚拢，发送边界就变成连续切片。示意：

# 示意，非源码
# routing_map: [num_tokens, num_experts] 的布尔矩阵
perm = argsort_by_expert(routing_map)      # ① 按目标专家排序的置换
sorted_tokens = tokens[perm]               #    聚拢后, 同专家的 token 连续

received = all_to_all(sorted_tokens)       # ② 每卡收到"来找本地专家"的 token
expert_out = grouped_gemm(received)        # ③ 分组 GEMM: 各专家各算各的段

returned = all_to_all(expert_out)          # ④ 搬回原卡
output = empty_like(tokens)
output[perm] = returned                    #    unpermute: 逆置换还原顺序
output = (output * probs).sum_over_k()     #    top-k 加权求和

重点看：两次 permute 是同一置换的正反两个方向，All-to-All 只看见「已排序的连续段」。这个不变量让 dispatcher 实现可以整体替换（AllGather 式、DeepEP 式）而不动 router 和 experts。

4. 真实实现​

4.1 Router：打分、选 k、挂 Loss​

TopKRouter（megatron/core/transformer/moe/router.py:148）的核心是 routing（:750）。顺序值得记住：

  1. z-loss：apply_z_loss（router.py:646）先压 logits 量级（数值稳定性）。
  2. 选路：默认走 topk_routing_with_score_function（:785-797），支持 softmax/sigmoid 打分、group-limited top-k（DeepSeek 式 moe_router_group_topk）、可学的 expert_bias（无辅助 Loss 的负载均衡）；也有 sinkhorn 等备选（:779-784）。
  3. token dropping：moe_expert_capacity_factor 非 None 时调 apply_router_token_dropping（:800-808）——专家容量有限时丢弃溢出 token。
  4. 辅助 Loss：训练态下挂 _apply_aux_loss / _apply_seq_aux_loss / _apply_global_aux_loss（:818-841），通过 autograd function 附在 probs 上，不污染主流程。

产出是 probs 和 routing_map 两件套——后者就是 dispatch 的全部输入。

4.2 Dispatcher：permute + All-to-All 的真实形态​

默认实现 MoEAlltoAllTokenDispatcher（megatron/core/transformer/moe/token_dispatcher.py:375），dispatch 拆成三步：

  • preprocess（dispatch_preprocess，:624）：展平 token、算通信元数据，然后 permute（megatron/core/transformer/moe/moe_utils.py:344）按 routing_map 把本地 token 排序聚拢，同时保存 reversed_local_input_permutation_mapping 供还原（:667-678）。
  • All-to-All（token_dispatch，:682）：对 token 和 probs 各做一次 All-to-All（:705、:718），底层是 all_to_all（megatron/core/tensor_parallel/mappings.py:558），按 output_splits/input_splits 变长切分——split 尺寸正来自路由结果。
  • postprocess（dispatch_postprocess，:728）：TP 维的 All-Gather、多本地专家时再按专家排序。

combine 方向是镜像：combine_preprocess（:796）先 All-to-All 收回，combine_postprocess（:876）用保存的逆映射 unpermute（moe_utils.py:497）还原 [s, b, h] 形状（:891-903）。

替代实现：MoEAllGatherTokenDispatcher（:233，EP 组内先 all-gather 全量 token 再本地选，通信换显存）、以及 DeepEP/Flex 变体。四拍接口不变。

4.3 Experts：分组 GEMM​

默认 TEGroupedMLP（megatron/core/transformer/moe/experts.py:189）：用 Transformer Engine 的 GroupedLinear，把「num_local_experts 个小 MLP」表达成一次分组 GEMM——输入是 permute 后按专家连续排布的 token 段，每个专家权重是 GEMM 的一段。备选 SequentialMLP（experts.py:1480）逐专家循环算（省显存、慢）。两者与 megablocks 的 block-sparse GEMM 是同一数学的三种工程表达。

4.4 共享专家的重叠​

dispatch 期间有个精巧的算子排序：token 的 All-to-All 发起后，共享专家的 fc1 被插在两次 All-to-All 之间启动（token_dispatch 里 :709-717 的注释写明 launch 顺序「tokens A2A → shared experts fc1 → probs A2A」），让共享专家计算与通信重叠。combine 末尾再把共享专家输出加回去（:906-909）。

5. 关键细节与坑​

  • MoE + TP 必须配 SP。 MoELayer.forward 开头直接报错：训练时 TP>1 而不开 sequence parallelism，性能会退化（moe_layer.py:660-664）。原因：EP 卡间搬运的 token 数随 TP 复制倍数膨胀。
  • routing_map 必须是 bool。 dispatcher preprocess 有硬断言（token_dispatcher.py:645-647）——用概率矩阵冒充 mask 会在 permute 处静默出错。
  • All-to-All 的 split 需要 host 同步。 tokens_per_expert 要从 GPU 搬到 CPU 才能算 split 尺寸，代码里用 _maybe_dtoh_and_synchronize（:913）把同步点收敛到固定位置、避免每步多次隐式同步——这是 MoE 性能调优里最真实的坑之一。
  • 负载均衡靠三件套，不是一个 Loss。 z-loss（压 logits）、aux/seq-aux/global-aux Loss（罚不均）、expert_bias（无 Loss 偏置）可同时存在；router 里按固定顺序施加（§4.1）。搞混它们的生效条件（is_aux_loss_enabled、训练态）是常见 bug 源。
  • capacity 与 drop_and_pad 改变张量形状语义。 开了容量限制后 token 数被 pad/drop 到固定容量，后面的 grouped GEMM 才能用定长段——这也让 All-to-All 变成定长，性能可预测，代价是溢出 token 被丢弃（:800-808）。
  • 梯度缩放要补偿 EP。 MoE 参数在 DP 维的「有效份数」与稠密参数不同，buffer 层的 gradient_scaling_factor 注释里专门提到 MoE 修正（megatron/core/distributed/param_and_grad_buffer.py:1065-1067）——跨章节联动点，见 04 章。

6. 代码地图​

主题文件路径符号名
MoE 层四拍megatron/core/transformer/moe/moe_layer.pyMoELayer（route/dispatch/routed_experts_compute/combine）
路由megatron/core/transformer/moe/router.pyTopKRouter、routing、apply_z_loss
top-k 打分实现megatron/core/transformer/moe/router.pytopk_routing_with_score_function
token 丢弃megatron/core/transformer/moe/router.pyapply_router_token_dropping
All-to-All 调度器megatron/core/transformer/moe/token_dispatcher.pyMoEAlltoAllTokenDispatcher
All-Gather 调度器megatron/core/transformer/moe/token_dispatcher.pyMoEAllGatherTokenDispatcher
分组 GEMM 专家megatron/core/transformer/moe/experts.pyTEGroupedMLP、GroupedMLPSubmodules
逐专家循环专家megatron/core/transformer/moe/experts.pySequentialMLP
permute/unpermutemegatron/core/transformer/moe/moe_utils.pypermute、unpermute
All-to-All 原语megatron/core/tensor_parallel/mappings.pyall_to_all

下一章：06 · GPT 模型组装与端到端主线——把前五章串成一个训练 step。