品牌宣传型网站构成,wordpress数据库更改账号密码,电子商务网站建设与原理,网页制作师就业班摘要#xff1a;本文重新审视了在训练混合专家#xff08;Mixture-of-Experts, MoEs#xff09;模型时负载均衡损失#xff08;Load-Balancing Loss, LBL#xff09;的实现。具体来说#xff0c;MoEs的LBL定义为N_E乘以从1到N_E的所有专家i的频率f_i与门控得分平均值p_i的…
摘要本文重新审视了在训练混合专家Mixture-of-Experts, MoEs模型时负载均衡损失Load-Balancing Loss, LBL的实现。具体来说MoEs的LBL定义为N_E乘以从1到N_E的所有专家i的频率f_i与门控得分平均值p_i的乘积之和其中N_E是专家的总数f_i表示专家i被选择的频率p_i表示专家i的平均门控得分。现有的MoE训练框架通常采用并行训练策略以便在微批次micro-batch内计算f_i和LBL然后在并行组之间进行平均。本质上用于训练百亿级大规模语言模型LLMs的微批次通常只包含非常少的序列。因此微批次的LBL几乎是在序列级别上进行的这促使路由器在每个序列内均匀分配token。在这种严格约束下即使是来自特定领域序列例如代码的token也会被均匀路由到所有专家从而抑制了专家的专业化。在本文中我们提出使用全局批次global-batch来计算LBL以放宽这一约束。因为全局批次包含的序列比微批次更加多样这将鼓励在语料库级别上进行负载均衡。具体来说我们引入了一个额外的通信步骤来跨微批次同步f_i然后使用它来计算LBL。通过在基于MoEs的大规模语言模型总参数高达428亿训练token数达4000亿上的实验我们惊讶地发现全局批次LBL策略在预训练困惑度和下游任务方面都带来了显著的性能提升。我们的分析还表明全局批次LBL也极大地提高了MoE专家的领域专业化能力。Huggingface链接Paper page论文链接2501.11873
1. 引言
背景与动机
混合专家Mixture-of-Experts, MoEs框架已成为扩展模型参数规模的一种流行技术。它通过路由器网络将输入分配给一组并行专家模块实现了条件稀疏激活从而提高了训练和推理效率。训练MoE模型的一个关键因素是鼓励路由器以均衡的方式分配输入到专家这既是为了保证训练的有效性避免参数冗余也是为了提高推理的效率避免专家利用不均衡导致的推理速度减慢。负载均衡损失Load-Balancing Loss, LBL作为一种辅助损失函数被广泛应用于MoE训练中以鼓励均衡的路由决策。然而现有MoE训练框架在计算LBL时多采用微批次级别这在训练大规模语言模型LLMs时存在局限性。
研究问题
现有微批次级别的LBL计算方法在训练大规模MoE模型时由于每个微批次包含的序列数量有限导致路由器被推向在每个序列内均匀分配token这抑制了专家的专业化并可能损害模型性能。
研究目标
本文旨在提出一种基于全局批次计算LBL的方法以放宽微批次级别的严格约束鼓励在语料库级别上进行负载均衡从而提高模型性能和专家专业化能力。
2. 预备知识
混合专家MoEs
MoEs由一组并行专家模块和一个路由器组成。路由器根据输入为每个专家分配权重通常只激活得分最高的前K个专家。输出是所有激活专家的加权和权重由路由器根据输入为每个专家分配的得分决定。
负载均衡损失LBL
LBL是一种正则化技术用于鼓励专家利用的均衡性防止专家崩溃。它通过对路由器过度将token路由到少数特定专家进行惩罚来实现。LBL的计算公式为N_E乘以所有专家i的频率f_i与门控得分平均值p_i的乘积之和然后除以专家总数N_E。
3. 方法
全局批次LBL计算方法
现有MoE训练框架在计算LBL时通常在每个微批次内计算f_i和LBL然后在并行组之间进行平均。这种方法在训练大规模LLMs时存在局限性因为每个微批次包含的序列数量有限。本文提出使用全局批次来计算LBL以放宽微批次级别的严格约束。具体做法是引入一个额外的通信步骤来跨微批次同步f_i然后使用同步后的f_i来计算LBL。由于全局批次包含的序列比微批次更加多样这将鼓励在语料库级别上进行负载均衡从而促进专家的专业化。
针对计算节点数量有限的情况
当训练LLMs时全局批次大小可能非常大而每个微批次大小可能相对较小。由于计算节点数量有限所有微批次的总和可能小于全局批次大小。在这种情况下本文提出使用一个缓冲区来存储每个梯度累积GA步骤中同步的专家选择计数c_i。然后在GA步骤中使用缓冲区中的信息来计算当前的f_i。完成GA后重置缓冲区。
4. 实验
实验设置
本文在三种不同大小的MoE模型上进行了实验包括3.4B总参数0.6B激活、15B总参数2.54B激活和43B总参数6.6B激活的模型。所有模型均采用细粒度专家和共享专家方法并使用softmax门控、微批次LBL和z-loss。训练数据包含1200亿和4000亿高质量token涵盖多语言、数学和一般知识内容。
实验结果
全局负载均衡提升模型性能实验结果显示随着计算LBL时考虑的token数量Balance BSZ的增加所有指标均得到一致提升。特别是在4000亿token训练设置下使用全局批次LBL的模型性能显著优于使用微批次LBL的模型。全局负载均衡鼓励专家专业化通过分析不同领域下每层专家的选择频率发现使用全局批次LBL的模型在特定领域如数学和代码下出现了更多高频专家表明全局批次LBL更有利于专家的专业化。缓冲区机制的有效性在计算节点数量有限的情况下使用缓冲区机制可以近似全局批次LBL从而在性能上接近使用全局批次LBL的模型并显著优于使用微批次LBL的模型。
5. 分析
消融研究
本文进行了消融研究以验证全局批次LBL相对于微批次LBL的优势。结果显示即使在微批次内随机选择一批与全局批次具有相同token分布和数量的token来计算LBLShuffle LBL micro其性能仍然接近全局批次LBL并显著优于微批次LBL。这表明全局批次LBL的优势主要在于其更宽松的约束条件而不是因为考虑了更多的token。
负载均衡损失和语言建模损失的变化
通过分析训练过程中负载均衡损失和语言建模损失的变化发现从微批次平衡切换到全局批次平衡后负载均衡损失迅速下降并接近从头开始使用全局批次平衡的情况而语言建模损失也保持较低水平。这表明全局批次平衡是一个更宽松的约束条件。
训练过程中改变Balance BSZ的影响
实验结果显示在训练过程中改变Balance BSZ会对最终结果产生影响。较早地切换到全局批次平衡可以获得更好的性能提升而较晚地切换则提升有限。此外从全局批次平衡切换到微批次平衡会导致性能下降表明训练过程中专家选择的变化会显著影响模型性能。
全局批次平衡的计算成本和效率
由于采用了无丢失策略不同方法在计算浮点运算数FLOPs时是相同的。然而由于局部平衡条件的不同使用全局批次平衡的方法可能会经历局部计算不平衡。实验结果显示使用全局批次平衡的方法在计算速度上比使用微批次平衡的方法慢约5.8%但通过引入少量微批次平衡损失可以将其提高到接近使用微批次平衡的速度同时仅对性能产生微小影响。
全局批次平衡带来的可解释性专业化
通过分析使用全局批次平衡的模型在不同领域下的专家选择频率和路由得分发现全局批次平衡促进了专家在不同领域下的专业化并且路由得分与语言建模任务更加一致。
6. 相关工作
负载均衡
现有的工作主要关注如何通过不同的方法来实现负载均衡包括引入辅助损失函数、允许专家根据负载能力选择token等。然而这些方法在计算专家选择频率时多采用微批次级别存在局限性。
专家专业化
初始的MoE设计旨在通过允许不同专家专注于特定任务来实现高效的参数利用。然而由于微批次级别的负载均衡约束大多数MoE模型并未展现出领域级别的专业化。本文提出的全局批次LBL方法有助于促进专家的专业化。
7. 结论
主要贡献
本文提出了使用全局批次来计算LBL的方法以放宽微批次级别的严格约束从而促进专家的专业化和提高模型性能。通过在多种规模的MoE模型上的实验验证了全局批次LBL方法的有效性。提供了对全局批次LBL方法如何影响模型性能和专家专业化的详细分析。
局限性
本文主要关注分析预训练阶段微批次LBL对LLMs的影响未进一步探讨其在微调阶段或在视觉和多模态领域的影响。对专业化的分析主要集中在不同领域下的专家选择频率未进行更严格的验证。放宽微批次LBL可能引入一些延迟未来工作可以考虑在每个微批次内包含更多样化的序列以减轻局部不平衡问题。
8. 未来工作
进一步研究全局批次LBL在微调阶段和在不同领域如视觉和多模态中的应用。对专家专业化进行更严格的验证例如通过人工评估专家在不同任务上的表现。探索其他方法来减轻全局批次LBL可能引入的延迟问题。
