ABSTRACT
对于混合专家（MoE）模型，不平衡的专家负载将导致路由崩溃或增加计算开销。现有方法通常采用辅助损失来促进负载平衡，但较大的辅助损耗会在训练中引入不可忽略的干扰梯度，从而影响模型的性能 。为了控制负载平衡，同时在训练过程中不产生不希望的梯度 ，我们提出了一种 Loss-Free Balancing 平衡策略。具体来说，在top-K路由决策之前，Loss-Free Balancing将首先对每个专家的路由分数应用专家bias 。通过根据每个专家最近的负载动态更新其偏差，Loss-Free Balancing 可以始终保持专家负载的平衡分布 。此外，由于Loss-Free Balancing 不产生任何干扰梯度，它也提高了MoE训练获得的模型性能的上界。我们在MoE模型上验证了** Loss-Free Balancing **的性能，该模型在多达200B个令牌上训练了多达3B个参数。实验结果表明，与传统的辅助丢控负载均衡策略相比，该策略具有更好的性能和负载均衡性。
INTRODUCTION
混合专家（MoE）架构已经成为在大型语言模型（llm）中扩展参数时管理计算成本的有前途的解决方案。最近MoE在基于transformer的模型中的应用（Vaswani等人，2017）已经成功地尝试将语言模型缩放到相当大的尺寸(Shao等人，2024；DeepSeek-AI等，2024；Dai et al., 2024；Fedus等人，2021；Lepikhin等人，2020)，导致显著的性能改进。然而，训练MoE模型总是面临负载不平衡的情况，这可能导致路由崩溃（Shazeer等人，2017）或计算开销增加(Fedus等人，2021；Lepikhin et al., 2020；Shazeer et al., 2017)。为了避免路由不均衡，现有方法(Fedus et al., 2021；Lepikhin等人，2020)通常使用辅助损失来鼓励平衡专家负载。虽然辅助损失可以缓解训练过程中的负载不平衡，但它也引入了与语言建模目标相冲突的不希望的梯度 。这些干扰梯度会影响模型的性能，因此现有的MoE方法总是需要考虑负载平衡和模型性能之间的权衡。
在这里插入图片描述
在本文中，我们提出了一种辅助的无损耗负载均衡策略，旨在保持对专家负载均衡的控制，同时不引入干扰梯度。无损耗平衡具有令牌路由和偏差更新的迭代过程。如图1所示，在MoE的top-K路由决策之前，Loss-Free Balancing将首先对原始路由分数应用专家偏差以产生有偏差的门测分数，这决定了训练期间每个令牌的实际路由目标 。这些专家偏见将根据在最近的训练令牌上观察到的专家负载不断更新，其中重载专家的偏见将被抑制，而轻负荷专家的偏见将被提升 。通过这种动态更新策略，Loss-Free Balancing保证了有偏的门控分数能够一致地导致均衡的路由结果。与辅助损失控制负载平衡策略相比，无损失平衡不会引入干扰主要语言建模目标的不希望的梯度，因此其训练过程更加无噪声和友好 。
为了验证Loss-Free Balancing的性能，我们从零开始在100B个token上训练具有1B参数的MoE语言模型，在200B个token上训练具有3B参数的MoE语言模型。实验结果表明，与传统的辅助损耗控制模型相比，无损耗平衡产生的MoE模型具有更好的验证损耗。同时，在保持性能优势的同时，Loss-Free Balancing在全局和批处理级别上也实现了明显更好的负载平衡，并且与专家并行自然兼容，通常用于训练超大型MoE模型。
BACKGROUND
MIXTURE-OF-EXPERTS
当前主流的MoE架构用MoE层代替标准transformer中的MLP层。在MoE层，采用Top-K路由为每个令牌选择专家 。设
表示第
个令牌到
个专家MoE层的输入，输出
计算如下：
其中 G 是非线性门函数，
是第
位专家的 centroid.
AUXILIARY LOSS FOR LOAD BALANC
Auxiliary Loss 不受控制的路由策略可能会遇到负载不平衡，这有两个明显的缺点。首先，存在路由崩溃 的风险，其中模型始终只选择少数专家，阻碍了其他专家的充分训练。其次，当专家分布在多个设备上时，负载不平衡会加剧计算瓶颈 。为了解决这些问题，辅助损失通常用于控制负载平衡。对于长度为T 的序列，将辅助损耗定义为
在这里插入图片描述
The Dilemma Between Load Balance and Model Performance 上面提到的辅助损失可以促进负载平衡，但作为一个额外的正则化项，它也会干扰语言建模训练。缺乏辅助损失或辅助损失系数α小会导致平衡性差，而较大的α会损害训练，导致次优表现 。为了说明这种困境，我们在图2中展示了负载平衡和模型性能之间的关系。我们在1e-2, 1e-3, 1e-4和0中改变α，并给出相应的
，它测量负载平衡的程度，其计算细节在§4.1中描述。如图所示，较小的α会导致路由崩溃，影响模型效率，并可能导致一些专家学习或利用不足；较大的α值可以控制负载平衡，但会显著降低模型的性能。为了打破这一困境，我们提出了一种解决方案，即直接控制专家负载平衡，但除了语言建模损失的梯度外，不引入意外梯度 。
在这里插入图片描述
AUXILIARY-LOSS-FREE LOAD BALANCING STRATEG
为了更好地实现负载平衡，并且不直接干扰训练目标的主梯度，我们提出了Loss-Free Balancing，即根据每个专家的平衡情况直接调整门控分数。如图1所示，我们在每个专家的门限分数
中加入一个专家偏向项
，并使用偏向分数来确定top-K的选择：
为了得到合适的 biases，我们根据以下原则对每个偏差
进行双迭代调整：当相应专家的负载相对较大时减小偏差，反之亦然 。具体来说，对于每个
，我们持续监控其对应的上一批专家负载。如果专家在前一批中有很大的负载，我们将减少其偏差。否则，我们将增加它。算法1描述了我们针对专家bias的更新算法的细节。值得注意的是，我们基于历史平衡条件来更新偏差，因为利用当前序列的负载信息会打破语言建模的因果约束，导致未来token信息的泄漏 。通过对偏差的动态调整，可以达到较好的专家负载平衡，但不像辅助损失控制方法那样直接将噪声梯度引入模型。
在这里插入图片描述
分配给每个export i 的 token 数
, 这里的
可能是每个sample的。
可能是batch上的平均。 根据
和
的相对关系，来决定
的调整方向。
Comparison with Other Load Balancing Methods. 为了展示无损耗均衡的理论优势，我们将其与其他两种主流负载均衡方法进行了比较，即 auxiliary-loss-controlled 方法 和Expert Choice (EC) 方法。如§2.2所述，辅助损失控制方法面临负载平衡和模型性能之间的两难境地，并且可能不存在完美的权衡。对于EC方法，它将打破语言建模的因果约束，因为每个标记的目标专家都以同一序列或批次中的未来标记为条件。这将导致有关未来令牌的信息泄露，从而破坏模型的泛化 。表1总结了不同负载平衡方法的属性。
在这里插入图片描述
EXPERIMENTS
Model Architecture. 与DeepSeekMoE略有不同的是，在我们的主要实验中，我们选择sigmoid 而不是softmax作为门控函数G，因为我们发现sigmoid基线比softmax基线性能更好。
Metrics. 我们从训练语料库中保留一个验证集来评估模型的性能和负载平衡。对于模型性能，我们以困惑度作为度量。对于负载平衡，我们引入了一个称为maximal violation（MaxVio） 的度量来量化MoE层的负载平衡程度：
在这里插入图片描述
MAIN RESULTS
表2显示了使用辅助损失或我们的辅助-无损失负载平衡策略训练的1B和3B MoE模型的验证困惑度和
。如表所示，与辅助控损方法相比，我们的Loss-Free Balancing在1B和3B模型上都实现了更好的perplexity和更好的全局负载均衡。此外，为了展示训练过程中的负载平衡情况，我们在图3中提供了描述
在训练步骤上的负载平衡曲线，这表明了Loss-Free balancing在负载平衡上的持久优势。综上所述，我们的Loss-Free Balancing方法避免了训练过程中的干扰梯度，有效地控制了负载平衡，打破了MoE训练中负载平衡与模型性能之间的两难困境.
在这里插入图片描述
EMPIRICAL STUDIES ON BIAS UPDATE ALGORITHM
我们对偏差更新算法的更新率和变体进行了实证研究，以验证我们主要实验中使用的最优配置
Update rate. 算法1中的更新率
控制专家偏差
收敛到“合适偏差”的速度。图4表明，过低的更新率u = 0.0001可能导致收敛缓慢，而不必要的高更新率u = 0.01可能导致专家偏差
在训练后期出现不希望的波动，从而恶化该阶段的负载平衡。这两种情况都会损害性能。一个合适的选择是u = 0.001，这显示了良好的训练平衡和验证困惑 。
Update rule. 我们研究了一种不同的专家偏差更新规则。具体来说，我们尝试将
的更新规则更改为
，以鼓励具有高违反错误的专家的偏差更快地改变。尽管这种变体略微改善了负载平衡，但它并没有带来更好的性能，如表3所示。因此，我们维持符号版本。
Multiplicative bias. 除了将专家偏见添加到门控分数之外，使用乘法偏见也是一种潜在的变体：
这些
可以使用与算法1类似的过程进行更新，只是它们应该初始化为1而不是0。表4显示，与使用加性偏差相比，使用乘法偏差导致的模型性能稍微差一些，在负载平衡方面没有显著改善。基于这些发现，我们得出结论，加性偏差是我们的方法更合适的选择。
在这里插入图片描述
DISCUSSION
LOSS-FREE BALANCING IS COMPATIBLE WITH EXPERT PARALLELISM
超大规模的MoE模型通常采用专家并行（Lepikhin et al., 2020）进行训练或推理，将专家分布在不同的设备上，以减少内存需求。在这种情况下，单个计算步骤中的数据负载平衡对于效率至关重要。由于专家并行性，每个计算步骤涉及micro_batch_size * ep_data_parallel_size样本，我们将其称为计算批处理。这里，micro_batch_size表示在单个设备上的一个梯度累积步骤中处理的样本数量。
Loss-Free Balancing可以实现近乎最优的全局负载平衡，并且随着计算批大小的增加，每个计算步骤的负载平衡将越来越接近全局负载平衡。在图5中，我们使用
度量来检查计算批级负载平衡。结果表明，随着计算批量的增加，我们的无损耗均衡方法的负载均衡总是不断提高，而辅助损耗控制方法在计算批量较大时，负载均衡大致保持在一个恒定的水平。由于专家并行会通过ep_data_parallel_size的倍数显著增加计算批大小，因此Loss-Free Balancing与大规模MoE训练自然兼容，并且随着专家并行大小的增加，其在负载平衡上的优势将进一步增强。
在这里插入图片描述
LOAD BALANCING AND FUTURE TOKEN LEAKAGE
对于临时语言模型，负载平衡方法必须遵守语言建模的因果约束，以避免将来的令牌泄漏。虽然传统的辅助控制平衡和我们的无损耗平衡遵守这一约束，但专家选择（EC） （Zhou et al., 2022）违反了这一约束。EC通过为每个专家分配完全相同数量的令牌来确保完美的负载平衡。然而，这种方法固有地导致了未来令牌泄漏的严重问题。
我们设计了实验来证明现实模型训练中未来token泄漏的存在。(1)我们减少了块大小，其中执行top-K选择，从8192个令牌（4个句子）到512个（1/4个句子），期望暴露这种泄漏。我们观察到异常的损耗下降（约10%），证实了泄漏的存在。(2)我们通过在top-K选择步骤中跨块洗牌使泄漏更加困难，并观察到异常损失下降得到缓解。关于EC信息泄漏的详细实验结果见附录D.2。
未来的令牌泄漏是致命的，因为它会破坏模型的泛化并阻止对模型性能的可靠评估。因此，与EC相比，用我们的无损平衡放大MoE模型更安全.
CONCLUSION
在这项工作中，我们介绍了一种新的MoE负载平衡控制方法，即不引入辅助损耗梯度。无损耗平衡解决了传统的辅助损耗负载平衡控制的问题，这在训练过程中引入了额外的梯度，并且在执行负载平衡时可能会损害模型的性能。分别在100B和300B令牌上训练的1B和3B MoE模型上进行的实验表明，与传统的辅助损耗训练相比，Loss-Free Balancing实现了更好的模型性能和负载平衡。