通过全局负载均衡提升混合专家模型的性能和特异化程度

通过全局负载均衡提升混合专家模型的性能和特异化程度

混合专家模型（MoEs）通过路由机制动态并稀疏地激活模型参数，使得能高效地增大模型参数规模。然而，基于TopK机制的稀疏激活会在训练中遇到专家激活不均衡的问题。本文提出将局部负载均衡放松到全局负载均衡，能显著增强专家的特异化并提高模型性能。

引言

混合专家模型（MoEs）通过路由机制动态并稀疏地激活模型参数，使得能高效地增大模型参数规模。基于TopK机制的稀疏激活会在训练中会遇到专家激活不均衡的问题：少数被频繁选择的专家会被优化得更多，进一步使得这些专家被更频繁地选择，最终导致只选择少数专家，造成剩余专家的冗余。因此，MoE在训练中需要引入额外的辅助损失（load balance loss，LBL）来鼓励专家的选择趋于均衡。

目前主流MoE训练框架中LBL实现其实是局部的负载均衡，这使得模型需要将局部的输入都均匀分配给不同的专家。然而，局部的输入往往只来自个别领域，局部负载均衡会让模型将所有领域的输入都均匀分配。这种均匀分配会阻碍某些专家更多处理特定领域的数据，也即阻碍专家出现领域层次的分化特征。我们发现，将局部的负载均衡放松到全局的负载均衡，能显著增强专家的特异化并提高模型性能。

从局部均衡到全局均衡

LBL的计算公式为$N_E \sum_{i=1}^{N_E} f_ip_i$，其中$N_E$为专家数，$f_i$为专家i被选择的频率，$p_i$为路由赋予专家i的平均分数。目前Megatron-mcore等主流框架中的LBL都是在每一张卡上统计计算后再全局平均，这使得$f_i$也是在局部统计，而优化LBL也鼓励模型将每个局部的输入都均匀分配给所有的专家。这也解释了为什么目前大部分MoE工作没有观察到领域层面的专家分化。

得益于LBL计算的格式，我们可以通过通信不同节点的$f_i$来将局部的LBL转化为全局的LBL。因为$f_i$只是一个专家数大小的向量，即使是在全局通信的情况下也不会带来明显的开销。此外由于LBL的计算与模型其它部分的计算相对独立，还可以用计算掩盖等策略进一步消除同步$f_i$的通信开销。

扩大均衡的范围带来稳定的提升

我们在三种参数规模（3.4B激活0.6B, 15B激活2.54B，43B激活6.6B）下分别训练了120B和400B tokens，对比了不同的均衡范围（Balance BSZ）对模型性能的影响。所有模型都使用了细粒度专家、共享专家及dropless策略（专家不会抛弃超过容量的tokens）。可以看到，将均衡范围从一般框架实现的4，8或者16增大到128以上后模型在Benchmark指标和PPL都有明显提升。

我们在3.4B激活0.6B的模型训练400B tokens到设置上进一步对比了模型效果随着均衡范围的变化，可以看到balance BSZ从2到128模型的PPL在快速降低，在128后逐渐饱和。目前主流MoE框架中即使是进行了机内通信，对于较大的模型balance BSZ也一般在8到16的，这进一步体现了我们通信方法的意义。

添加少量局部均衡损失能提高模型效率

只使用全局均衡会导致局部均衡状况有所降低，这会一定程度影响MoE的计算效率。我们进一步实验了在主要使用全局均衡的情况下，在训练过程中添加局部均衡（默认实现的LBL，损失权重为全局LBL的1%）限制对于模型性能和效率的影响。可以看到，添加局部均衡能提升模型的速度（每个更新步耗时从1.64秒提升到1.59秒），同时模型的效果也几乎不受影响。

结论

我们回顾了目前MoE训练框架中均衡损失，发现目前的实现方式会将所有来自相同领域的局部输入都均匀分配，限制了专家的分化。通过轻量的通信将局部均衡放松为全局均衡，MoE模型的性能和专家特异性都得到了显著的提升。我们认为这一进展解决了现有MoE训练中的一个关键问题，为MoE模型的优化提供了新的视角，并有助于构建更加可解释的模型。尽管我们的实验主要集中在基于语言的任务上，我们希望我们的工作能够为在不同领域训练更大规模、更有效的MoE模型提供帮助。