MoE 训练工程：专家并行与负载均衡实战
实验目标
MoE（Mixture of Experts）模型把参数量做大的同时保持激活量可控，但代价是把训练工程复杂度推到新高度：
- 专家分散在多张卡上，怎么分才能让通信与显存都可控？
- 路由层把 token 分给专家，负载天然不均衡，怎么平衡而不伤模型质量？
- All-to-All 通信在专家并行中占多大开销，怎么评估与优化？
本篇聚焦 MoE 训练工程侧（模型架构解析见 10 章 DeepSeek 专篇），以配置与调度为主线拆解专家并行、负载均衡与通信开销。说明：本篇为综述 + 架构解析路线，配置片段为通用实践示例，无第一手实验数据；具体实现以所用框架（Megatron-LM/DeepSpeed 等）版本为准。
第一段：专家并行（EP）的放置逻辑
MoE 的并行策略核心是专家并行（Expert Parallelism, EP）：把 FFN 专家按 EP 维度切分到不同 GPU，路由层把 token 发送到对应专家的卡上计算。与稠密模型对比：
| 维度 | 稠密模型 TP/PP | MoE 增加 EP |
|---|---|---|
| 切分对象 | 权重/层 | 专家集合 |
| 通信模式 | AllReduce（TP）、点对点（PP） | All-to-All（路由交换） |
| 显存收益 | 权重分片 | 专家权重分片（大 MoE 显存主要吃在这里） |
| 负载风险 | 各卡计算量近似 | 路由不均衡 → 卡间负载倾斜 |
典型配置思路（以 Megatron-LM 风格为例）：
yaml
# 假设 64 卡、8 专家、每专家 FFN 很大
# 方式 A：EP=8，专家平均放 8 个 rank，TP=4，DP=2
model:
num_experts: 8
moe_expert_parallel_size: 8 # EP=8：8 个专家各放一块卡
tensor_model_parallel_size: 4
data_parallel_size: 2
# 方式 B：EP=2（专家 2 卡一组），适合专家数少、卡多场景
model:
num_experts: 8
moe_expert_parallel_size: 2
tensor_model_parallel_size: 4工程要点：
- EP 与 TP 的乘积是"单组计算域"：token 路由只在 EP 组内交换，EP×TP 决定 All-to-All 的通信规模，组内越小通信越少；
- 专家数 vs EP 的关系：EP 通常取"专家数 / 每卡专家数"；专家数少于 EP 组数时要做专家复制（每个 rank 多份副本）——配置前先算清除法关系；
- 显存与通信的权衡：EP 越大单卡显存越省（专家分散），但 All-to-All 规模越大；EP 越小通信越少但显存压力越大。
第二段：负载不均衡——MoE 训练的"天然敌人"
路由层按 token 与专家的亲和度分发，热门专家被大量 token 命中是常态，导致：
- 热门专家所在卡算力打满，冷门专家卡空闲——整组吞吐被最热卡决定；
- 显存/带宽使用倾斜，极端时卡 OOM。
缓解手段一：负载均衡损失（Auxiliary Load Balancing Loss）
在训练损失上叠加一项，惩罚"路由分布与均匀分布的距离"：
python
# 概念示意（真实实现见框架源码）：aux_loss 项
# f_i: 路由分配给专家 i 的 token 比例；P_i: 专家 i 被选择的平均路由概率
aux_loss = num_experts * sum(f_i * P_i for i in range(num_experts))
total_loss = task_loss + aux_coef * aux_loss # aux_coef 通常很小（如 0.01）- 系数要小：负载均衡是"辅助"不是"主目标"，系数过大伤模型质量；常见在 0.001~0.01 量级，需实验调；
- 两种经典变体：开关专家负载（Switch Transformer 的 router z-loss 变体）与分组均衡（Grouped, 保证每个数据组内也近似均衡）——按框架支持选择。
缓解手段二：Token 丢弃（Drop Token）
当某专家收到的 token 超过上限时，丢弃多余 token（其损失不计入）：
yaml
model:
moe_token_dropping: true
moe_token_drop_policy: "probs" # 按路由概率丢低概率 token
moe_token_drop_coef: 0.0 # 训练早期建议关闭，收敛稳定后按需开启- 丢弃率是敏感旋钮：丢得太多模型质量受损，训练后期通常关闭或控制在很低水平；
- 丢弃策略：优先丢低置信 token，比随机丢对质量影响小。
工程观察：负载均衡调整是"质量 × 吞吐"的平衡——追求极致均衡会让路由偏离数据分布，得不偿失。实践上先开小系数 aux loss 跑通，观察各专家利用率分布（框架日志/Profiler），再逐步微调，而不是一开始就上大系数。
第三段：All-to-All 通信——MoE 的隐形开销
token 路由需要把 token 数据从"源卡"搬到"专家所在卡"，这是 All-to-All 集合通信：
- 通信量与 token 批量 × 隐藏维度成正比（搬的是激活，不是权重）；
- 与 EP 组大小成正比（组内卡越多，交换面越大）；
- 高频执行（每个 transformer 层的 MoE FFN 都做一次），对网络带宽与拓扑敏感——这正是 03 章多轨网络（Rail-Optimized）设计要服务的负载。
评估方法：
bash
# 用框架自带的通信基准/Profiler 观察 All-to-All 耗时占比
# Megatron/DeepSpeed 训练日志中看 MoE 层 vs 稠密层耗时分布
# 结合网络监控（03 章 NCCL 测试）确认交换机/链路带宽利用率优化路径：
- 减少通信量：token 本地化路由（如 Switch 的 local dispatch）、分组路由减少跨组交换；
- 提升通信效率：多轨网络拓扑、RDMA（RoCE/IB）就绪、All-to-All 的实现优化（框架版本差异明显）；
- 隐藏通信：与计算重叠（pipeline 内 MoE 层通信与相邻层计算并行）。
第四段：MoE 训练的稳定性与可观测
- Checkpoint 与容错：MoE 权重分片复杂，checkpoint 需保证 EP 分片的一致性——复用 10 章 checkpoint 策略，但注意专家分片 + TP 分片的双重恢复，失败恢复后核对各 rank 专家分布一致；
- 观测指标：训练日志中增加专家利用率分布（min/max/方差）、路由分布熵、token 丢弃率三个指标，作为负载均衡调参与健康检查的输入；
- 与 11 章联动：MoE 的 TP/PP/DP/EP 四维并行组合、NCCL 调优（03 章）、loss spike 排查（11 章）同样适用——MoE 是叠加在通用分布式训练之上的工程层。
MoE 训练配置检查清单
- [ ] EP/TP/DP 乘积 = 总卡数，专家数与 EP 整除关系确认
- [ ] aux loss 系数初始化小（0.001~0.01 量级），实验微调
- [ ] token 丢弃训练早期关闭，收敛稳定后按需开启
- [ ] All-to-All 通信耗时已观测（Profiler/日志），与网络带宽对照
- [ ] checkpoint 支持 EP+TP 双重分片恢复，恢复后专家分布核对
- [ ] 专家利用率分布/路由熵/丢弃率三个指标纳入训练观测
工程启示
- MoE 训练工程 = 通信与负载的平衡艺术：EP 放法决定通信与显存，路由策略决定负载均衡，两者共同决定实际吞吐。
- 先跑通再调优：MoE 配置旋钮多（EP 大小、aux 系数、丢弃率），先小规模跑通观察指标，再逐步放大与微调，避免一步到位踩坑。
- 可观测是调优前提：没有专家利用率分布数据，负载均衡调参就是盲调；把观测指标提前建好，调优才有依据。
下一篇预告：MoE 工程就绪后，预训练进入长跑——下一批内容将补 13 应用章的工程闭环（科学软件从论文到可复现工具），敬请期待。