系列定位与本文目标

本系列《Ai大模型训练教程：从入门到实战落地的系统课程》会从“能跑通”到“能规模化稳定训练/上线”逐步深入。本文聚焦 MoE（Mixture of Experts）大模型训练入门中最容易踩坑、但也最决定工程成败的三件事：路由器（Router）、专家并行（Expert Parallel / EP） 与 负载均衡（Load Balancing）。

如果你已经能训练一个普通的 Dense Transformer，那么进入 MoE 的关键问题不再是“能不能训练”，而是：

• 路由是否稳定（是否抖动、是否塌陷到少数专家）
• 专家并行是否高效（通信是否拖垮吞吐、all-to-all 是否爆显存）
• 负载是否均衡（是否出现严重溢出、丢 token、训练不收敛）

下面按“概念—实现步骤—排障建议—可落地的配置”来讲。

MoE 基本结构：Dense + 稀疏专家层

MoE 通常是在 Transformer 的 FFN 位置，把一个 FFN 替换成多个“专家 FFN”。每个 token 只会被送到少数几个专家（常见 Top-1 或 Top-2），因此理论上参数量可以大幅增加，但计算量相对可控。

一个典型 MoE 层包含：

1. Router（路由器）：对每个 token 计算其该去哪些专家（产生门控分数）。
2. Dispatch（分发）：把 token 按专家分桶（bucket），重排为每个专家的 batch。
3. Experts（专家网络）：各自处理各自 token。
4. Combine（聚合）：把专家输出按原 token 顺序合并回去，并按门控权重加权。

工程上最难的通常不在专家 FFN 本身（就是两层 MLP），而在 Router、Dispatch/Combine 的通信与内存、以及负载均衡策略。

路由器（Router）：你训练不稳的“第一嫌疑人”

Router 在做什么

Router 一般是一个线性层（或小 MLP）：

• 输入：token hidden state h（维度 d_model）
• 输出：对 E 个专家的 logits z = W_r h（维度 E）
• 选择：取 Top-k 专家，得到专家 id 与对应 gate 权重

常见的 gate 权重计算：

• Softmax over experts：p = softmax(z)，取 top-k 的 p
• 带噪声路由（Noisy Top-k）：给 logits 加噪声，防止早期塌陷

工程难点 1：路由塌陷（Expert Collapse）

现象：训练一段时间后，大多数 token 都被分到少数专家，其他专家几乎不更新。

常见原因：

• Router 过强/过早收敛，导致探索不足
• 数据分布偏、序列长度/任务类型导致某些模式主导
• 没有或弱负载均衡损失

可落地的解决动作：

1. 启用负载均衡损失（见后文），并调大其权重（但不要过大）。
2. Noisy Top-k：训练早期加噪声，让 router 不那么“自信”。
3. Router z-loss / logit regularization：抑制 logits 过大导致的过度确定。
4. Router 学习率单独设小：例如主干 LR=2e-4，router LR=1e-4 或更低。

工程难点 2：路由抖动（Routing Jitter）

现象：同一类 token 在不同 step/不同卡上被分配到不同专家，导致梯度噪声大、loss 波动。

排查与建议：

• 检查 Top-k 前的数值稳定性：混合精度下 logits 是否溢出/下溢。
• Router 的 dropout/噪声策略：噪声太大或一直开启会导致后期不稳定；可采用“前 N steps 开启噪声，后续关闭或衰减”。
• 温度/缩放：对 logits 做 scale 或使用温度 softmax，让分布不至于过尖。

工程难点 3：Top-2 的 combine 与反向传播细节

Top-2 会把一个 token 发给两个专家，输出加权合并。工程上要注意：

• combine 权重的归一化：Top-2 的两条路径权重是否归一，否则输出尺度会漂。
• 反向传播路径：Router 参数更新来自 gate 权重对损失的梯度，若实现里对 top-k 做了不恰当的 stop-gradient，会导致 router 学不到。

建议：优先使用成熟实现（Megatron-LM MoE、DeepSpeed-MoE、Fairseq MoE、Tutel 等）的 routing kernel；自己写 dispatch/combine 极易出错。

专家并行（Expert Parallel, EP）：all-to-all 的吞吐与显存战争

EP 的基本思想

在数据并行（DP）之外，把专家分布到不同 GPU：

• 每张卡只存一部分专家（或若干专家）
• token 根据 router 结果，通过 all-to-all 发送到持有对应专家的 GPU
• 专家计算后再 all-to-all 把结果发回

因此 EP 的核心成本是：
1) token 的跨卡搬运（通信）
2) token 重排/打包（dispatch）
3) 专家侧的 padding（容量不足/溢出处理）

工程难点 1：all-to-all 通信成为瓶颈

当 MoE 层数量多、序列长、batch 大时，通信量会非常可观。

实操优化建议（按优先级）：

1. 把 EP group 限制在同一节点/同一 NVLink 域：跨节点 all-to-all 极痛。
2. 减少 MoE 层频率：不是每层都放 MoE；例如 1/2 或 1/4 的 FFN 层替换为 MoE。
3. 使用高性能 MoE 库/内核：例如 Tutel / TransformerEngine 的 MoE 实现通常比手写 PyTorch 版本快得多。
4. 通信与计算重叠：检查实现是否支持 overlap（很多框架提供开关）。
5. 减少 token dtype/压缩：部分实现可在通信时用 FP16/BF16；但要确认数值稳定。

工程难点 2：容量（capacity）与 padding 造成的浪费

MoE dispatch 通常会给每个专家预留容量：

• 设本 step 总 token 数为 T
• 专家数 E
• Top-1 时，期望每个专家收到 T/E
• 实际会有波动，因此设置 capacity factor：

capacity = ceil((T/E) * capacity_factor)

当某个专家收到 token 超过 capacity，会发生：

• 丢 token（drop）：被丢弃或走旁路（影响收敛）
• 溢出到其他专家：实现复杂
• 动态扩容：可能导致额外开销与不确定性

可落地的配置建议：

• 初次跑通可用 capacity_factor=1.25~2.0（Top-1 常取 1.25-1.5，Top-2 可能更高）
• 观察 dropped_tokens_rate，目标一般 < 0.1%（任务敏感时希望更低）
• 如果 dropped 高：优先增强负载均衡损失，其次才是加 capacity（加 capacity 会加显存/算力浪费）

工程难点 3：专家参数与优化器状态的显存压力

虽然 MoE 计算稀疏，但参数量巨大；优化器状态（Adam 的 m/v）会把显存放大 2-4 倍。

实操建议：

• 专家参数用 ZeRO/FSDP 分片：把优化器状态分摊到更多卡。
• 专家用 FP8/INT8 存储（若框架支持）：训练端通常更难，但部分系统支持混合精度专家。
• 把 MoE 层做 activation checkpoint：节省激活显存，但会增加计算。
• 专家与 Dense 部分分开策略：Dense 用较大 batch，MoE 受容量约束时可能要调小 micro-batch。

负载均衡（Load Balancing）：让每个专家都“有活干”

为什么必须做负载均衡

没有负载均衡，router 会倾向于把 token 分给少数“更擅长”的专家，导致：

• 大多数专家训练不足（浪费参数）
• 少数专家过载，capacity 溢出导致 dropped token
• all-to-all 通信不均衡，产生 straggler，整体 step 时间被最慢的专家拖住

常见负载均衡损失：importance + load

一个经典做法（在多个实现中都能见到变体）：

• Importance：对每个专家，统计 gate 概率之和（软分配强度）
• Load：对每个专家，统计实际被分配 token 数（硬分配计数）

希望二者都接近均匀分布。损失通常形如鼓励均匀、惩罚偏斜的项。

实操要点：

1. 先把指标打出来：

   • tokens_per_expert（每 step 或滑动平均）
   • gate_prob_sum_per_expert
   • dropped_tokens、overflow_rate

2. 损失权重从小到大调：

   • 太小：不起作用，仍塌陷
   • 太大：router 被迫“平均分配”，但不管语义，可能损害困惑度/下游效果

3. 分阶段策略：

   • 训练前期：较强的均衡约束 + 路由噪声
   • 训练中后期：逐步减弱均衡约束，让专家形成分工

工程难点：均衡与效果的拉扯

MoE 的目标不是让专家永远平均，而是让专家形成可解释/不可解释的“功能分化”，但又不能过载。

经验法则：

• 如果你看到 dropped_tokens_rate 上升、step time 抖动，优先加均衡约束或调 capacity。
• 如果困惑度明显变差、下游指标掉，且 dropped 很低，可能是均衡太强导致路由被“人为打散”。

一套可执行的训练落地流程（从跑通到稳定）

第 1 步：先定最小可行配置（MVP）

建议从“保守但稳定”的配置起步：

• Top-k：Top-1（先少一半复杂度）
• 专家数 E：每 EP group 8 或 16（结合单机 GPU 数）
• capacity_factor：1.25~1.5
• MoE 层比例：每 2~4 个 Transformer block 放 1 个 MoE
• 启用 load balancing loss（中等权重）

目的：先保证不炸显存、吞吐还能接受、loss 能下降。

第 2 步：把“必备监控”接到日志里

至少记录这些（每 step 或每 N step）：

• 每层 tokens_per_expert 的均值/方差、max/min
• overflow/dropped token 数量
• router logits 的均值/方差（检查是否数值爆炸）
• all-to-all 通信耗时占比（如框架提供 profiler）
• 每个专家的计算耗时（排查 straggler）

没有这些，你只能“感觉”在调参，效率会极低。

第 3 步：处理三类常见故障

故障 A：训练发散或 loss 大幅抖动

排查顺序：

1. 混合精度：router logits 是否溢出（尤其 FP16）；优先用 BF16
2. 关闭/减弱路由噪声，看是否稳定
3. 调小 router 学习率，或给 router 加梯度裁剪
4. 检查 combine 权重归一化是否正确（Top-2 尤其重要）

故障 B：吞吐很低，profiling 显示 all-to-all 很慢

优化顺序：

1. 确保 EP group 在单机/同域互联内
2. 减少 MoE 层数或把 MoE 位置后移（减少早期大激活的通信）
3. 降低 token 数（短序列预训练阶段）或减少 micro-batch
4. 使用更高性能 MoE kernel/库

故障 C：dropped token 很高

处理顺序：

1. 增强负载均衡损失（先解决分配偏斜）
2. 提高 capacity_factor（兜底，但会浪费）
3. 从 Top-2 改回 Top-1（Top-2 更容易超容量）
4. 检查是否某些 token 被强制路由到少数专家（例如实现 bug 或 mask 处理错误）

一个“可参考”的 MoE 配置模板（概念层面）

你可以把下面当作调参起点（不同框架字段名会不同）：

• Router：Linear(d_model -> E)，Noisy Top-1（前 10% steps 开启噪声，之后关闭）
• Top-k：1（稳定后再尝试 2）
• capacity_factor：1.25（dropped>0.1% 则 1.5）
• load_balance_loss_weight：从 0.01 起步，观察塌陷再加到 0.05/0.1
• router_lr：主干 lr 的 0.5 倍或更低
• precision：BF16 优先
• EP group：单节点内做 all-to-all，跨节点尽量只做 DP/TP

落地时关键不是“照抄参数”，而是围绕监控指标闭环：均衡指标 → dropped → 吞吐 → 任务效果。

本文小结：MoE 工程的三根主线

• 路由器决定 token 去向，稳定性与塌陷防控是第一优先级。
• 专家并行带来 all-to-all 通信与容量管理，吞吐常被通信与重排拖住。
• 负载均衡是 MoE 能否“用起来”的底层保障：既要避免专家闲置与过载，也要避免过度平均损害效果。

在《Ai大模型训练教程》后续文章里，如果你要进一步实战落地，可以沿着本文的工程主线继续深入：用 profiler 把 all-to-all 占比压下去；用更成熟的 MoE kernel 解决 dispatch/combine 性能；用更系统的指标与自动调参策略降低 dropped 与抖动，最终得到可复现、可扩展的 MoE 训练流程。