07 MoE Load Balancing Loss

LLM 算子深度解析07 MoE Load Balancing Loss — 让 8 个专家都有活干1. 路由崩塌MoE 训练的最大敌人1.1 问题从哪来06 节我们实现了一个 Top-K Router每个 token 打分 → 选 Top-K 专家 → 加权输出。一切看起来很完美。但 Router 的W_gate是被梯度优化的。这意味着它会偷懒——一旦某个专家初始权重稍微有利Router 就把更多 token 分给它 → 那个专家接收更多梯度 → 学得更快 → Router 觉得它更靠谱 → 分更多 token 给它……马太效应强者越强弱者饿死。训练前理想 Expert_0: ████████ Expert_1: ████████ Expert_2: ████████ Expert_3: ████████ 训练 N 步后崩塌 Expert_0: ████████████████████████████████ Expert_1: ████████████████████████████ Expert_2: ▌ Expert_3: Expert_4: Expert_5: Expert_6: Expert_7: 6 个专家被饿死——几乎收不到 token参数不再更新后果算力不均衡 → 部分 GPU 空转 部分 OOM、MoE 退化成了只用 2 个专家的 Dense 模型。1.2 解法在 Loss 中加交警在 CrossEntropy Loss 之外额外加一个惩罚项Total LossCE Lossα⋅Laux\text{Total Loss} \text{CE Loss} \alpha \cdot L_{\text{aux}}Total LossCE Lossα⋅Laux​α\alphaα很小如 0.01确保辅助损失引导但不主导训练。2. 数学原理负载均衡损失公式的均值不等式魔法2.1 核心公式Mixtral / Switch Transformer 使用的经典公式Lauxα⋅E⋅∑i1Efi⋅PiL_{\text{aux}} \alpha \cdot E \cdot \sum_{i1}^E f_i \cdot P_iLaux​α⋅E⋅i1∑E​fi​⋅Pi​逐符号含义EEE专家总数如 8fif_ifi​专家iii被路由到的token 比例实际分配比例∑fi1\sum f_i 1∑fi​1PiP_iPi​专家iii在所有 token 上的平均路由概率得分∑Pi1\sum P_i 1∑Pi​1α\alphaα超参数通常 0.012.2 为什么这个公式能防止崩塌——均值不等式的力量给定两个概率分布fff和PPP∑fi1,∑Pi1\sum f_i 1, \sum P_i 1∑fi​1,∑Pi​1它们的内积∑fi⋅Pi\sum f_i \cdot P_i∑fi​⋅Pi​在完全均匀分布时取最小值。直觉验证如果f[1,0,0,0]f [1, 0, 0, 0]f[1,0,0,0]全去 Expert_0且P[0.9,0.02,0.02,...]P [0.9, 0.02, 0.02, ...]P[0.9,0.02,0.02,...]→ 内积 ≈ 0.9大如果f[0.25,0.25,0.25,0.25]f [0.25, 0.25, 0.25, 0.25]f[0.25,0.25,0.25,0.25]且P[0.25,0.25,0.25,0.25]P [0.25, 0.25, 0.25, 0.25]P[0.25,0.25,0.25,0.25]→ 内积 4 × 0.0625 0.25小 3.6 倍**优化器为了降低这个 Loss会被迫把 token 往不同的专家赶**这就是整个机制的底层逻辑。2.3 理论最小值对于 Top-K 路由每个 token 选 K 个专家当负载完全均匀时LauxminαKL_{\text{aux}}^{\text{min}} \frac{\alpha}{K}Lauxmin​Kα​验证E8, K2, α0.01 → 理论最小值 0.005。这在测试代码中被精确验证。3. 代码实现scatter_add_ 是灵魂3.1 完整实现defcompute_load_balancing_loss(routing_weights:torch.Tensor,# [N, top_k] 权重已重归一化每行和1selected_experts:torch.Tensor,# [N, top_k] 专家索引 (LongTensor)num_experts:int,# 专家总数 Etop_k:int,# 每个 token 的专家数alpha:float0.01# 损失系数):N,_selected_experts.shape# ---- Step 1: 计算 P_i — 每个专家的平均路由概率得分 ----P_itorch.zeros(num_experts,dtyperouting_weights.dtype,devicerouting_weights.device)P_i.scatter_add_(0,selected_experts.flatten(),routing_weights.flatten())# scatter_add_ 把每个 (专家索引, 权重) 累加到 P_i 的对应位置P_iP_i/(N*top_k)# 归一化 → 概率分布# ---- Step 2: 计算 f_i — 每个专家的实际 token 比例 ----expert_maskF.one_hot(selected_experts,num_classesnum_experts)# [N, top_k, E]tokens_per_expertexpert_mask.sum(dim(0,1)).float()# [E]f_itokens_per_expert/(N*top_k)# 归一化 → 概率分布# ---- Step 3: 公式直译 ----aux_lossalpha*num_experts*(f_i*P_i).sum()returnaux_loss3.2 scatter_add_ 深入解析——面试最爱问P_i.scatter_add_(0,selected_experts.flatten(),routing_weights.flatten())scatter_add_(dim, index, src)沿dim维把src的每个值累加到P_i[index[j]]。# 具体例子3 tokens, 2 experts, top_k1selected_experts[[0],[1],[0]]# token 0→expert0, token1→expert1, token2→expert0routing_weights[[0.8],[0.9],[0.7]]# flatten 后: index [0, 1, 0], src [0.8, 0.9, 0.7]# scatter_add_(dim0):# P_i[0] 0.8 (来自 token 0)# P_i[1] 0.9 (来自 token 1)# P_i[0] 0.7 (来自 token 2, 累加)# 结果: P_i [1.5, 0.9]为什么是scatter_add_而非scatter_scatter_是覆盖后来的值覆盖先前的多个 token 选同一专家时只有最后一个生效。scatter_add_是累加——这是正确行为。3.3 F.one_hot 的路由统计妙用expert_maskF.one_hot(selected_experts,num_classesnum_experts)# [N, top_k] → [N, top_k, E]# [[3, 7], [1, 3], ...] → 三维 one-hot 张量tokens_per_expertexpert_mask.sum(dim(0,1)).float()# 沿 token 维和 top_k 维求和 → [E] → 每个专家被选中的总次数维度追踪输入: routing_weights: [N, top_k] 如 [1000, 2] selected_experts: [N, top_k] 如 [1000, 2] P_i 计算: flatten(): [N*top_k] 如 [2000] scatter_add_: [E] 如 [8] / (N*top_k): [E] 如 [8] f_i 计算: one_hot: [N, top_k, E] 如 [1000, 2, 8] sum(dim(0,1)): [E] 如 [8] / (N*top_k): [E] 如 [8] f_i * P_i: [E] 逐元素乘 .sum(): 标量 * alpha * E: 标量 ← L_aux4. 工业对照4.1 Mixtral 的做法完全一致HuggingFace 的 Mixtral 实现modeling_mixtral.py与我们的代码逻辑完全一致——load_balancing_loss_func同样使用 scatter_add 和 one_hot 统计。4.2 DeepSeek 的改进去辅助损失化的负载均衡DeepSeek-V2/V3 做了一个关键创新——Auxiliary-Loss-Free Load Balancing传统Mixtral加辅助损失 → 但 α 太大影响主任务太小无效 → 精细调参 DeepSeek给每个专家维护一个 bias → 太忙就减 bias太闲就加 bias → 零额外损失# DeepSeek 的动态 Bias 方法概念上expert_biastorch.zeros(num_experts)foreach training step:ifexpert_load[i]mean_load:expert_bias[i]-bias_update_step# 忙 → 降分else:expert_bias[i]bias_update_step# 闲 → 加分# router_logits expert_bias (bias 不参与梯度传播)好处不需要额外损失项不用调 α负载均衡直接在 Router 的输出层面解决。这是 MoE 负载均衡的下一代方案。4.3 α 超参数选择指南α 值效果适用场景0.001极弱几乎等于没加不推荐0.01标准值Mixtral 默认通用推荐0.1较强可能影响主任务E 64 时考虑1.0太强主任务性能明显受损不推荐5. 踩坑实录坑现象根因解决scatter_代替scatter_add_Loss 值不稳定结果随机scatter_后写入覆盖先写入必须用scatter_add_做累加忘记除以 (N × top_k)Loss 异常大没有归一化值域是 O(N)P_i / (N * top_k),f_i / (N * top_k)dtype 不一致RuntimeError: dtype mismatchrouting_weights 是 FP16P_i 默认 FP32torch.zeros(..., dtyperouting_weights.dtype)推理时还在算 aux_loss显存多占一块辅助损失只在训练时需要if self.training:包裹 aux_loss 计算同时加了多个辅助损失各损失互相打架α 之间未协调梯度方向冲突总辅助损失不应超过主损失的 5%6. 延伸思考6.1 辅助损失的副作用自由与平等的权衡本质上看负载均衡损失是一种对 Router 自由的限制。有的 token 真的更适合 Expert_0但辅助损失强行把它赶到 Expert_6 → 模型效果略降。这就是 MoE 训练的永恒矛盾“选最好的专家”效果最优vs “让所有专家都有活干”算力均衡。α 就是这个天平的砝码。6.2 Router Z-Loss另一个常用辅助损失除了负载均衡损失还有一个叫Z-Loss的辅助损失Lz1N∑i1N(log⁡∑j1Eehij)2L_z \frac{1}{N} \sum_{i1}^N \left( \log \sum_{j1}^E e^{h_{ij}} \right)^2Lz​N1​i1∑N​(logj1∑E​ehij​)2它惩罚 Router logits 的 log-sum-exp——防止 Router 输出极端大的 logit 值从而让 Softmax 数值更稳定。经常和负载均衡损失组合使用。6.3 值得深挖的方向Expert Choice Routing反过来让专家挑 token每个专家固定处理 top-C 个——天然负载均衡不需要辅助损失Adaptive Aux Loss Coefficient根据当前负载不均衡程度动态调整 α——均衡时不加惩罚不均衡时加大Load Balancing via Expert Capacity硬限制每个专家每批次最多处理 C 个 token超出的溢出或跳过DeepSeek 的 Auxiliary-Loss-Free 方案用 expert bias 替代辅助损失——当前 SOTA下一篇[[08 Architecture Tricks]] — 两行代码的架构变体Qwen 的权重绑定与 Gemma 的 1 RMSNorm。