密集训练稀疏推理：重新思考混合专家语言模型的训练

24年4月来自MIT、MIT-IBM实验室、Toronto大学和Vector Institute的论文“Dense Training, Sparse Inference: Rethinking Training of Mixture-of-Experts Language Models”。

与密集模型相比，混合专家 (MoE) 语言模型可以将计算成本降低 2-4 倍，而不会牺牲性能，从而使其在计算受限场景中更加高效。但是，MoE 模型通常需要 2-4 倍以上的参数才能实现与密集模型相当的性能，这会导致更大的 GPU 内存需求，并使 MoE 模型在自回归生成等 I/O 受限场景中效率较低。提出的一种用于 MoE 模型的混合密集训练和稀疏推理框架 (DS-MoE)，通过在训练期间对所有专家进行密集计算并在推理期间进行稀疏计算来实现计算和参数方面强大的效率。DS-MoE 模型比标准稀疏 MoE 具有更高的参数效率，并且在总参数大小和性能方面与密集模型相当，同时计算成本更低（激活 30-40% 的模型参数）。使用 vLLM 的性能测试表明，DS-MoE-6B 模型的运行速度比 Mistral-7B 等类似的密集模型快 1.86 倍，比 DeepSeekMoE-16B 和 Qwen1.5-MoE-A2.7B 等同类 MoE 快 1.50 倍到 1.71 倍。

传统的 MoE 语言模型虽然能够以大约 40-50% 的计算量达到密集模型的性能，但需要的参数却要多得多，通常是密集模型的 2-3 倍。参数需求的增加，源于 MoE 模型中的反向传播过程，该过程经过稀疏优化，因此效率不如密集模型。

DS-MoE 框架有两个关键元素：互信息 (MI) 损失和混合注意头 (MoA) 块。DS-MOE方法以密集方式训练 MoE 模型，并附加 MI 损失，然后以稀疏方式执行推理。此策略旨在保留传统 MoE 模型的运行效率，同时实现密集模型的参数效率。

给定输入token X，MoE计算输出如下：

其中 S 是专家的得分向量，Ai 是得分最高的 K 位专家中第 i 个索引。在训练期间，梯度仅通过选定的专家 ei 和相应的得分 Si 进行反向传播。传统MoE中 S 的梯度如下：

在此上下文中，M ∈ {0, 1}^N 用作识别已激活专家的二进制掩码。具体而言，Mi = 1 表示专家 ei 在前向传递中处于活动状态，否则 Mi = 0。
密集训练的基本概念是围绕全梯度优化路由器。为了保留 S 的所有梯度，确保在前向传递期间计算每个专家的输出，并保留此输出在反向传播期间使用。其中 S 和 ei(X) 的梯度表达如下：

其中 j 表示专家索引。

如图所示，对 MoE 模型进行密集训练的方法涉及激活所有专家。

在推理过程中，仅使用根据得分确定前 K 名专家。K 的选择基于预定值或自适应，取决于有多少专家的得分高于指定阈值 ε。用 SimpleMoE （Tan 2024) 实现稀疏推理。

为了在所有专家之间实现负载平衡并防止模型容量利用不足，在路由器中集成互信息 (MI) 损失。该方法遵循 Shen(2023d) 的方法，旨在最大化专家分布熵，以确保工作负载均匀分布：

从而促进专家之间的负载平衡。此外，为了避免路由器采用过于简单的解决方案并确保专家集中度，最小化专家分布的条件熵 H(e|X)。因此，整体损失函数定义为：

其中 X 表示小批次（minibatch）的tokens。这种方法不仅确保专家之间的公平负载平衡，而且还保持对适当解决方案的高度集中，从而优化路由器的性能。计算总损失为

其中 Llm是标准自回归语言建模损失，α 是相互信息损失的权重。

与大多数当前保留一个密集层进行自注意的 MoE 语言模型不同，用混合注意 (MoA) 头层的方法（Zhang et al. 2022) 取代自注意层。MoA 头是按照组查询注意 (GQA) 方式构建的（Ainslie et al. 2023) ，其中K和V对在一组查询向量之间共享。在实现中，MoA 头中的每个专家负责计算 Nhead 个查询向量 Q。对于给定的输入token X，专家 ei 的输出如下：

其中 Wq 是专家 ei 的查询投影权重。专家的输出投影由 Wo 表示。该层的最终输出计算如下：

其中A是激活专家的索引集。

实验设置

数据集和token化。用 Pile 数据集 Gao et al. (2020) 的子集对模型进行预训练，并使用 CodeGen token化器 Nijkamp et al. (2023) token化。此数据集包含 300B 个tokens。具体而言，用 一个 30B tokens子集来训练1B 规模的模型，利用 100B tokens子集来训练 3B 和 6B 规模的模型。

模型超参。在下表中列出不同模型架构的超参数设置。这里，Natt 和 Nffd 分别表示每个注意层和每个前馈层的专家数。在模型中，用 GeLU Hendrycks & Gimpel (2016) 激活函数。在注意块中使用分组查询注意 (GQA) Ainslie et al. (2023)。对 1B 模型使用 2 个共享K-V头，对 3B 和 6B 模型使用 4 个共享K-V头。

训练细节。用 AdamW 优化器 Loshchilov & Hutter (2017) 训练模型，学习率为 3 × 10-4。训练包括余弦学习率进度，1B 模型的预热需要 10 亿个 tokens，3B 和 6B 模型的预热需要 20 亿个 tokens。在整个训练过程中，使用 0.01 的恒定权重衰减和 1.0 的裁剪梯度。1B 模型的批次大小为 50 万个 tokens，3B 和 6B 模型的批次大小为 200 万个 tokens，序列长度为 2048 个 tokens。为了优化训练，用完全分片数据并行（FSDP） Zhao et al. (2023); Rajbhandari et al. (2020) 和激活检查点 Korthikanti et al. (2023)。在 8 个 H100 80GB GPU 上，1B 模型的训练时间为 24 小时；在 32 个 H100 GPU 上，3B 模型的训练时间为 64 小时；6B 模型的训练时间为 124 小时。互信息损失权重（公式 6 中的 α）列于下表中。

评估设置

基线。将方法与两个基线进行比较。A. 密集模型。对于DS-MoE 模型的每个实例，都会训练一个类似的密集模型。该模型的参数大小设计为与DS-MoE 模型相似。密集模型和DS-MoE 模型之间的唯一参数差异来自路由器功能。B. 稀疏 MoE。用传统的稀疏梯度传播训练 MoE 模型，以匹配 Dense-1B 和 Dense-3B 模型的性能。Top-K 设置为 2，并使用切换损失Fedus (2022) 来实现路由器的负载平衡。对于稀疏 MoE 块的实现，采用 dMoE Gale (2023)。这些基线模型旨在突出传统稀疏 MoE 训练方法中发现的参数效率低下问题。

下游任务。在一系列下游任务中评估模型，包括常识推理和问答。这些包括对 PiQA Bisk（2020）、HellaSwag Zellers（2019）、WinoGrande Sakaguchi（2021）、SciQ Welbl（2017）和 Arc Clark（2018）的评估。此外，还在 Wikitext 数据集 Merity（2016）上测量并报告模型的困惑度。对于所有这些评估，利用 LM 评估工具 Gao（2023）来确保测试方法的一致性和可靠性。
如下表是基模型的性能评估比较：

如下表是推理速度的比较：