Branch-Train-MiX：将专家LLM混合为混合专家LLM

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，并对整体统一模型进行微调。

其采用混合专家方法来组合域专家模型 Mi。但是，并没有使用混合 Mi 最终输出的传统程序，而是通过在 Transformer 的每一层内执行 MoE 来进行更细粒度的混合。具体来说，将来自域专家的不同前馈子层组合成单个 MoE 子层。

但是，路由决策可能会因token而异，因此，一个输入序列可以根据需要使用所有域专家 FF 层，即使在任何给定token处仅访问少数层。在实验中，用 Top-k (k=2) 路由。

对于自注意子层，通过简单地平均权重来组合不同域专家。这样做的动机是假设自注意层在领域上的专业性不如前馈层。对其余参数（嵌入等）也进行相同的平均。

注意，引入的唯一新参数是路由器的转换参数 Wl，与网络的其余部分相比，其大小可以忽略不计。尽管如此，这些新参数需要进行微调，以便路由器能够在选择使用哪个域 FFi 层时做出最佳决策。此外，微调很有用，因为自注意权重是通过平均构建的，而且可能不是最优的。总体而言，整个系统在高度并行的训练框架中根本没有针对协同工作进行优化，但假设是，即使是少量的组合微调也可能带来很大的改进。
MoE 的一个常见问题是出现死专家，这些专家根本不会被路由器激活。像 Top-k 这样的常见路由方法不太可能摆脱这种情况，因为死专家永远不会出现在 top-k 选择中，因此永远不会收到训练信号。负载平衡提供了一个简单的解决方案，即添加一个额外的损失项，鼓励平等地利用专家。

除了 Top-k 路由，还尝试了其他路由方法：

• 开关：这是 Fedus (2022) 提出的 Top-1 路由方法。
• 软路由：用 softmax 作为路由函数 g，因此所有专家在训练和推理期间都被激活。虽然它可能提供最佳性能，但代价是增加计算量。
• 样本 Top-1：使用 gumbel softmax (Jang et al., 2016)。在训练时，从 gumbel softmax 生成一个软样本，但除了最大值以外其他值置零。然后只计算与这个最大值相对应的一个专家，省略其他专家计算。在推理时，只是进行硬采样。在训练结束时将温度退火为一个尖锐的分布，以逐渐减少训练和推理之间的差异。

MoE 层中的模块数量与训练的域数量相匹配，因为每个模块对应一个域。但是，可以将每个域 FF 子层拆分为多个块，以简单的方式增加模块数量。给定 N 个域和 FF 激活大小，将每个 FF 层拆分为 C 个块。因此，最终的 MoE 层将具有 MC 模块。

没有直接以一对一的方式从域专家中初始化 MoE 专家，而是尝试在每个 MoE 专家中包含所有域。这样做的动机是，以标准方式训练的 MoE 专家不会表现出域专业化，而是在不同域中均匀激活（Jiang，2024）。相比之下，域专家通过其训练数据专门针对特定域。为了打破这种域专业化，将每个域专家的 FF 层拆分为 N 个块，然后合并来自所有域的第 n 个块构建第 n 个 MoE 专家。这样，每个 MoE 专家都包含来自所有域相同数量的参数。