DeepSeek：LLM在MoE训练中的无损平衡

📖标题：AUXILIARY-LOSS-FREE LOAD BALANCING STRAT-EGY FOR MIXTURE-OF-EXPERTS
🌐来源：arXiv, 2408.15664

摘要

🔸对于混合专家（MoE）模型，不平衡的专家负载将导致路由崩溃或计算开销增加。现有方法通常采用辅助损耗来促进负载平衡，但较大的辅助损耗会在训练中引入不可忽略的干扰梯度，从而损害模型性能。
🔸为了在训练过程中控制负载平衡，同时不产生不希望的梯度，我们提出了无损耗平衡，其特征是辅助无损耗负载平衡策略。具体来说，在做出前K个路由决策之前，无损平衡将首先对每个专家的路由得分应用专家偏见。通过根据每个专家最近的负载动态更新其偏差，无损耗平衡可以始终如一地保持专家负载的平衡分布。此外，由于无损耗平衡不会产生任何干扰梯度，它还提高了从MoE训练中获得的模型性能的上限。我们在MoE模型上验证了无损耗平衡的性能，该模型具有在多达200B个令牌上训练的多达3B个参数。实验结果表明，与传统的辅助损耗控制负载平衡策略相比，无损耗平衡实现了更好的性能和更好的负载平衡。

🛎️文章简介

🔸研究问题：在混合专家系统（MoE）中，不平衡的专家负载将导致路由崩溃或计算开销增加。
🔸主要贡献：论文提出了一种无损平衡的训练策略（Loss-Free Balancing），通过偏差更新来调整每个专家的门控分数，实现了负载均衡，同时不引入干扰梯度。

📝重点思路

🔺相关工作

🔸MoE架构已成为在大语言模型 (LLM) 中扩展参数时，管理计算成本的一种有前景的解决方案，包括DeepSeek在内的模型都将模型扩展到相当大的规模。
🔸MoE架构用MoE层替换了标准Transformer中的MLP层，基于非线性门函数，采用Top-K路由来为每个令牌选择专家。
🔸训练MoE模型总是面临负载不平衡的情况，现有方法通常使用辅助损失来鼓励平衡专家负载，但也会引入与语言建模目标相冲突的不需要的干扰梯度，需要权负载平衡和模型性能。

🔺论文方案

🔸无损平衡在MoE的top-K路由决策之前，需要将专家偏差引入到原始路由分数上，生成偏置门控分数，这些分数决定了每个令牌在训练期间的实际路由目标。
🔸为了得到适当的偏差，根据最近训练令牌上的专家负载进行更新，负载重的专家偏差会降低，负载轻的专家偏差会提高。
🔸值得注意的是，需要根据历史平衡条件更新偏差，因为利用当前序列的负载信息会打破语言建模的因果约束，导致未来令牌的信息泄漏。
🔸通过这种动态更新策略，确保偏置门控分数能够持续导致平衡的路由结果。

🔎分析总结

🔸与使用辅助损失控制的方法相比无损平衡在MoE模型上实现了更好的验证困惑度和更优的全局负载平衡。
🔸通过绘制负载平衡曲线，展示了无损平衡训练步骤中的持续负载平衡优势。
🔸无损平衡避免了训练过程中的干扰梯度，有效控制了负载平衡，打破了MoE训练中负载平衡与模型性能之间的困境。
🔸无损平衡与专家并行性兼容，随着计算批量大小的增加，负载平衡持续改善，而辅助损失控制方法的负载平衡水平在大计算批量时没有提升。

💡个人观点

论文的核心在于动态调整专家偏差，实现了无辅助损失的负载平衡，避免了干扰梯度的引入，从而在保持模型性能的同时优化了负载平衡。

附录