DeepSpeed-MoE 用于 NLG：将语言模型的训练成本降低 5 倍

基于自回归 Transformer 的自然语言生成（在本文的其余部分中称为 NLG）模型可以为广泛的语言任务提供令人信服的解决方案，从文档摘要、标题生成、问答到各种编程语言中的代码生成。由于这些模型的普遍适用性，提高其质量一直是学术界和工业界共同关注的焦点。

NLG 的质量随着模型规模的增加而提高。然而，今天我们正在接近当前一代硬件所能做到的极限。Megatron-Turing NLG 530B 模型在 NVIDIA Selene 超级计算机上使用超过 2K 个 A100 GPU 训练了 3 个月，消耗了超过 300 万个 GPU 小时。模型规模再增加 3 到 5 倍在合理的时间范围内是不可行的。鉴于训练最先进的 NLG 模型所需的巨额计算资源，一个自然而然的问题是：“是否可以在不增加计算成本的情况下对模型质量进行非凡改进？”或者等效地，“是否可以使用 3 到 5 倍的资源来生成质量相似的模型？”

最近的一些工作，例如 GShard 和 Switch Transformers 表明，混合专家 (MoE) 模型结构显着降低了基于 Transformer 的编码器-解码器模型的大模型训练成本。MoE 模型包含一组稀疏门控专家。在训练和推理期间，只有这些专家的子集被激活以用于每个输入标记。因此，模型可以扩展到数十亿个参数，而不会成比例地增加计算量。尽管 MoE 显示出有希望的结果，但 MoE 对计算密集型 NLG 家族模型的有效性在很大程度上仍然未知。

鉴于训练 NLG 家族模型的巨大计算和能源需求，我们探索了 MoE 带来的降低其训练成本的机会。我们证明 MoE 可以应用于 NLG 家族模型，以显着提高其模型质量，而训练成本相同。或者，它可以实现 5 倍的训练成本降低，以实现与密集型 NLG 模型相同的模型质量。例如，通过应用 MoE，我们在训练 1.3B 参数密集型模型的成本下实现了 6.7B 参数密集型 NLG 模型的模型质量，这得益于 MoE 的稀疏结构。

假设缩放保持有效，这些结果有可能从成本方面彻底改变大型模型训练格局。例如，一个万亿参数密集型模型可以以 200B 参数（例如 GPT-3）大小的密集模型的成本进行训练，这将转化为数百万美元的训练成本降低和能源节省（Brown 等人，2020，语言模型是少样本学习者）。

基于 MoE 的 NLG 模型架构

为了创建一个基于 MoE 的 NLG 模型，我们研究了类似 GPT 的基于 Transformer 的 NLG 模型。为了在合理的时间范围内完成训练，选择了以下模型：350M（24 层，1024 个隐藏大小，16 个注意力头）、1.3B（24 层，2048 个隐藏大小，16 个注意力头）和 6.7B（32 层，4096 个隐藏大小，32 个注意力头）。我们使用“350M+MoE-128”来表示一个使用 350M 密集模型作为基础模型并在每隔一个前馈层添加 128 个专家的 MoE 模型。也就是说，对于 350M+MoE-128 和 1.3B+MoE-128，总共有 12 个 MoE 层。

我们使用一个门控函数来激活 MoE 层中每个标记的专家子集。具体来说，在我们的实验中，只选择了前 1 名专家。因此，在训练和推理期间，我们的 MoE 模型将与密集部分具有相同的参数数量被激活以用于每个标记。例如，我们的 1.3B+MoE-128 将仅激活每个标记的 1.3B 参数，每个标记的训练计算量将类似于 1.3B 密集模型。

MoE 训练基础设施和数据集

我们使用 DeepSpeed 在 128 个 A100 GPU 上对上述模型的密集版本和 MoE 版本进行了预训练。DeepSpeed 使用数据并行和专家并行训练的组合来有效地扩展 MoE 模型训练。

我们使用了与 MT-NLG 博客 中描述的相同的训练数据。为了公平比较，我们使用 300B 个标记来训练密集模型和 MoE 模型。

MoE 提高了 NLG 模型的质量

图 1 显示了模型的 MoE 版本的验证损失显着优于其密集对应部分。此外，请注意 MoE 模型的验证损失 350M+MoE-128 与基础更大 4 倍的 1.3B 密集模型的验证损失相当。这也适用于 1.3B+MoE-128 与基础更大 5 倍的 6.7B 密集模型的比较。此外，模型质量不仅在验证损失方面相当，而且在各种 6 个 Zero-shot 评估任务中也相当，如表 1 所示，证明这些模型实际上具有非常相似的模型质量。

图 1：具有不同模型尺寸的密集和 MoE NLG 模型的标记级验证损失曲线。

模型大小	LAMBADA：完成预测	PIQA：常识推理	BoolQ：阅读理解	RACE-h：阅读理解	TriviaQA：问答	WebQs：问答
密集 NLG
350M	0.5203	0.6931	0.5364	0.3177	0.0321	0.0157
1.3B	0.6365	0.7339	0.6339	0.3560	0.1005	0.0325
6.7B	0.7194	0.7671	0.6703	0.3742	0.2347	0.0512
MoE NLG
350M+MoE-128 (13B)	0.6270	0.7459	0.6046	0.3560	0.1658	0.0517
1.3B+MoE-128 (52B)	0.6984	0.7671	0.6492	0.3809	0.3129	0.0719

表 1：不同密集和 MoE NLG 模型的 Zero-shot 评估结果（最后六列）。所有 Zero-shot 评估结果都使用准确度指标。

以 5 倍的训练成本实现相同质量

正如我们从上面的结果中看到的那样，向 NLG 模型添加具有 128 个专家的 MoE 显着提高了 NLG 模型的质量。但是，这些专家不会改变模型的计算需求，因为每个标记只由单个专家处理。因此，具有相同基础的密集模型及其相应的 MoE 模型的计算需求相似。

更具体地说，1.3B+MoE-128 模型训练需要与 1.3B 密集模型大致相同的计算操作，同时提供更好的模型质量。此外，我们的结果表明，通过应用 MoE，我们可以在 1.3B 参数密集型模型的训练成本下实现 6.7B 参数密集型模型的模型质量，从而有效地将训练计算量减少了 5 倍。

这种计算成本降低可以通过利用高效的 DeepSpeed MoE 训练系统直接转化为吞吐量增加、训练时间和训练成本降低。表 2 显示了 1.3B+MoE-128 模型的训练吞吐量与 128 个 NVIDIA A100 GPU 上的 6.7B 密集模型的比较。

	每秒训练样本	吞吐量增加 / 成本降低
6.7B 密集	70	1x
1.3B+MoE-128	372	5x

表 2：比较基于 MoE 的模型与密集模型的训练吞吐量（在 128 个 A100 GPU 上），这两个模型都可以实现相同的模型质量。

用于推理的 MoE

MoE 的训练成本降低并非免费的，它需要增加实现与密集模型相同模型质量所需的总参数数量。1.3B+MoE-128 的参数数量（52B）大约是 6.7B 密集模型的 8 倍。那么，这是否意味着推理会比密集模型慢 8 倍，因为推理通常受读取所有模型参数所需的时间限制，特别是对于小批量大小？

不尽然。请注意，在 1.3B+MoE-128 模型中，每个标记由每个 MoE 层的唯一专家处理，用于处理标记的总参数数量仅为 1.3B。理论上，由于计算量和参数读取量减少了 5 倍，这可能会导致比质量等效的密集 6.7B 模型更快的推理速度。然而，在现实中，推理期间批次中的标记数量通常大于 1。推理、长序列长度或非单元批量大小可能需要加载所有专家，与质量等效的密集模型相比，加载的总参数数量增加了 8 倍。因此，即使每个标记使用的参数和发生的计算量与质量等效的密集模型相比很小，但使用 MoE 实现良好的推理性能仍然具有挑战性。

尽管如此，我们相信可以通过不同的并行形式来利用巨大的内存带宽，通过跨大量设备进行扩展来加速 MoE 推理，使其与质量等效的密集模型相当或更快，用于扩展推理场景，并创造机会使基于 MoE 的模型除了训练外，在推理方面也具有成本效益。

结论与发布

我们证明了基于 MoE 的模型可以应用于 NLG 任务，与 GPT-3 和 MT-NLG 530B 等密集型、基于自回归 Transformer 的模型相比，将训练成本降低了 5 倍。通过基于 MoE 的低成本训练，我们希望使高质量的语言模型能够被广泛的用户群体使用，即使他们拥有有限的计算资源。

为此，我们发布了 用于训练基于 MoE 的 NLG 模型的端到端管道，以及 特定示例脚本 和 教程，以帮助您开始使用我们的管道。我们期待着这种技术可能带来的应用和创新，并将它们带给深度学习社区。

致谢

这项工作是与 Turing 团队的 Brandon Norick、Zhun Liu、Xia Song 以及 Z-Code 团队的 Young Jin Kim、Alex Muzio、Hany Hassan Awadalla 合作完成的。我们还感谢 Luis Vargas、Umesh Madan、Gopi Kumar、Andrey Proskurin 和 Mikhail Parakhin 的持续支持和指导。