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

DeepSeek:LLM在MoE训练中的无损平衡_第1张图片

标题: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训练中负载平衡与模型性能之间的困境。
无损平衡与专家并行性兼容,随着计算批量大小的增加,负载平衡持续改善,而辅助损失控制方法的负载平衡水平在大计算批量时没有提升。

个人观点

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

附录

DeepSeek:LLM在MoE训练中的无损平衡_第2张图片
DeepSeek:LLM在MoE训练中的无损平衡_第3张图片
DeepSeek:LLM在MoE训练中的无损平衡_第4张图片
DeepSeek:LLM在MoE训练中的无损平衡_第5张图片

你可能感兴趣的:(大模型-模型训练,人工智能,自然语言处理,语言模型,论文笔记)