24年3月来自Meta FAIR的论文“Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM”。
有效地训练大语言模型 (LLM) ,可以使其具备多个专业领域的能力,例如编码、数学推理和世界知识。该方法名为 Branch-Train-MiX (BTX),从种子模型开始,以高吞吐量和低通信成本高度并行地进行专家训练。在对各个专家进行异步训练后,BTX 将其前馈参数汇集在混合专家 (MoE) 层中作为专家,并对剩余参数取平均值,然后进行 MoE 微调阶段学习 token 级的路由。BTX 泛化两种特殊情况,一是 Branch-Train-Merge 方法,没有经过 MoE 微调阶段学习路由,二是稀疏升级改造(sparse upcycling),省略这个异步训练专家的阶段。与其他方法相比,BTX 实现最佳的准确率-效率权衡。
近年来,大语言模型 (LLM) 在各种任务中表现出色 (Brown,2020;Touvron,2023;Achiam,2023),包括代码生成 (Li,2022b;Rozière,2023)、解决数学问题 (Azerbayev,2023)和多语言性 (Zhao,2024) 等。训练此类 LLM 需要大量计算和数据,超过数千个 GPU 和数万亿个 token。训练并行化通常是在不同的 GPU 上维护模型的多个副本并在每次权重更新后保持它们同步。这种频繁通信成本是将训练扩展到更多 GPU 的主要瓶颈。除了这个问题之外,同步训练更容易受到硬件故障的影响,因为单个 GPU 故障可能会导致整个训练停止(Zhang,2022;Gemini,2023)。
Li(2022a)最近的研究,提出Branch-Train-Merge (BTM) 方法,用于在没有任何同步的情况下令人尴尬的LLM并行训练,提高预训练的吞吐量。它首先创建种子 LLM 的多个副本,然后在不同的数据子集上分别训练每个副本。这会产生多个不共享任何参数的独立 LLM,每个 LLM 都是专门研究其自身数据分布的专家,例如知识域、语言甚至某种模态。在测试时,输入提示被分类到一个或多个域中,然后从相应的专家模型中形成最终输出,这些模型组合起来以预测下一个token。虽然这种方法可以提高训练效率,但其主要缺点是缺乏统一的单一模型,因此无法进行进一步的监督微调 (SFT) 或从人类反馈中进行强化学习 (RLHF) 微调 (Ouyang,2022),这两者都可以进一步提高性能,并且是构建对齐 LLM 的关键步骤。
减少 LLM 计算占用空间的另一项工作是混合专家 (MoE) 方法 (Jacobs,1991;Shazeer,2017),其中只有一部分参数在任何给定时间处于活动状态。特别是,MoE 应用于 Transformers 的前馈子层 (Fedus,2022;Roller,2021;Lewis,2021),允许参数总数增长而无需额外的计算。以这种方式扩展的 LLM 在下游任务上表现出色(Jiang,2024;Xue,2024)。与 Branch-Train-Merge 不同,混合专家方法通常以完全同步的方式进行训练,并且由于全对全(all-to-all)通信,通信成本会随着专家数量的增加而增加。
如图所示,提出的Branch-Train-MiX (BTX) 方法分为三个步骤:1) 通过复制多个预训练种子 LLM 进行分支;2) 在不同的数据子集上分别训练这些副本以获得专家 LLM;3) 使用混合专家的前馈 (FF) 层将这些专家 LLM 组合成单个 LLM,并对整体统一模型进行微调。
但是,路由决策可能会因token而异,因此,一个输入序列可以根据需要使用所有域专家 FF 层,即使在任何给定token处仅访问少数层。在实验中,用 Top-k (k=2) 路由。
对于自注意子层,通过简单地平均权重来组合不同域专家。这样做的动机是假设自注意层在领域上的专业性不如前馈层。对其余参数(嵌入等)也进行相同的平均。
注意,引入的唯一新参数是路由器的转换参数 Wl,与网络的其余部分相比,其大小可以忽略不计。尽管如此,这些新参数需要进行微调,以便路由器能够在选择使用哪个域 FFi 层时做出最佳决策。此外,微调很有用,因为自注意权重是通过平均构建的,而且可能不是最优的。总体而言,整个系统在高度并行的训练框架中根本没有针对协同工作进行优化,但假设是,即使是少量的组合微调也可能带来很大的改进。
MoE 的一个常见问题是出现死专家,这些专家根本不会被路由器激活。像 Top-k 这样的常见路由方法不太可能摆脱这种情况,因为死专家永远不会出现在 top-k 选择中,因此永远不会收到训练信号。
负载平衡
提供了一个简单的解决方案,即添加一个额外的损失项,鼓励平等地利用专家。