细节中的恶魔：关于训练专用混合专家模型时负载平衡损失的实现

本文重访训练混合专家（MoE）模型时，如何实现负载平衡损失（LBL）的实现方式。具体而言，MoE的LBL被定义为 $N_E \sum_{i=1}^{N_E} f_i p_i$，其中 $N_E$ 为专家总数，$f_i$ 表示第 $i$ 个专家被选中频率，$p_i$ 表示第 $i$ 个专家的平均门控得分。现有的MoE训练框架通常采用并行训练策略，因此在 $\textbf{微批次}$ 内计算 $f_i$ 和 LBL，然后在并行组之间进行平均。在本质上，训练大规模LLM的微批次通常包含极少的序列。因此，微批次LBL几乎处于序列级别，路由器被迫在每个序列内部均匀分配 token。在这种严格约束下，即便是来自特定领域序列（例如代码）的 token 也被均匀路由到所有专家，从而抑制了专家的专业化。在本工作中，我们提出使用 $\textbf{全局批次}$ 来计算LBL，以放宽这一约束。因为全局批次包含远多于微批次的多样化序列，这将在语料库层面鼓励负载平衡。具体而言，我们引入额外的通信步骤，在微批次之间同步 $f_i$，并据此计算LBL。通过在训练包含最高 $42.8B$ 总参数和 $400B$ token 的MoE基LLM时进行实验，我们惊讶地发现，全局批次LBL策略在预训练困惑度和下游任务方面均取得优异的性能提升。我们的分析表明，全局批次LBL也大大提高了MoE专家的领域专业化程度。