通过全局负载均衡提升混合专家模型的性能和特异化程度
GITHUB HUGGING FACE MODELSCOPE DISCORD 引言 混合专家模型(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 都有明显提升。...