用于 NLG 模型的专家混合

在本教程中，我们介绍了如何将 DeepSpeed 专家混合 (MoE) 应用于 NLG 模型，这将训练成本降低了 5 倍，并将 MoE 模型大小减少了 3 倍（详细信息请参阅我们的博客）。我们使用 Megatron-LM 框架中的类似 GPT-3 的模型作为示例。在阅读本教程之前，我们建议您先阅读有关专家混合和Megatron-LM GPT 预训练的教程。

1. 安装

您需要安装 DeepSpeed v0.6.0 或更高版本才能使用 MoE 功能。用于 NLG 模型的 MoE 示例位于 Megatron-DeepSpeed 存储库的 MoE 文件夹下。

2. 训练 NLG+MoE 模型

2.1. 模型的更改

为了将 MoE 应用于 GPT 样式模型，我们在 Megatron 框架中进行了一些更改，主要是在 megatron/model/ 中，我们向模型中添加了 MoE 层。

2.2. 预训练标准 MoE 模型

我们在 examples_deepspeed/MoE 下提供了示例训练脚本，我们使用这些脚本在我们的博客中执行实验。对于标准 MoE 模型，有一些新的超参数

--num-experts：每个 MoE 层的专家数量。在我们的实验中，我们将其设置为 128。较多的专家数量往往可以提供更好的收敛性，但收益递减。

--moe-expert-parallel-size：MoE 专家并行度的程度。换句话说，每个 GPU 上将有 num-experts/moe-expert-parallel-size 个专家。因此，--moe-expert-parallel-size 不应超过 GPU 数量和 --num-experts。

--moe-loss-coeff：将 MoE 损失添加到模型损失中的缩放系数。在我们的实验中，我们发现 0.01 是一个不错的设置。

--moe-train-capacity-factor、--moe-eval-capacity-factor、--moe-min-capacity：这些配置决定了单个专家可以处理多少个标记。较大的数字可能导致更好的收敛性，但也会导致训练速度变慢，因为不同专家之间的负载会更加不平衡。

--disable-moe-token-dropping：这将完全消除对单个专家可以处理多少个标记的限制。出于与上述相同的原因，我们仅建议在推理/评估期间使用此功能。

2.3. 预训练 PR-MoE 模型

PR-MoE 是一种新设计的 MoE 模型，代表金字塔残差 MoE，与标准 MoE 相比，它将参数效率提高了高达 3 倍。有关更多详细信息，请参阅我们的博客。我们在 examples_deepspeed/MoE 下提供了示例训练脚本。与标准 MoE 相比，PR-MoE 模型有一些不同的超参数

--num-experts：为了启用金字塔 MoE，您需要提供一个列表而不是提供单个数字，该列表的长度与 MoE 层的数量相同。我们建议在模型的后阶段（靠近输出）使用更多专家。

--mlp-type：从 [standard, residual] 中选择。当它是残差时，将启用残差 MoE。

除了上面针对标准 MoE 和 PR-MoE 的新超参数之外，对于 NLG+MoE 模型，我们发现降低学习率并增加学习率衰减持续时间与基本密集模型相比很有帮助。我们的调整细节可以在示例训练脚本中找到。

关于训练数据，我们无法发布我们的内部数据，但 Megatron-LM 预训练的任何公共数据都可以直接用于训练 MoE 模型（需要注意的是，它可能无法提供与我们的实验中完全相同的模型质量）。例如，我们评估了 The Pile 数据集（pile.eleuther.ai、github.com/EleutherAI/the-pile）的密集模型和 MoE 模型。下面的表 1 显示，此公共数据提供了与我们的内部数据类似的评估结果。

模型大小	LAMBADA：完成预测	PIQA：常识推理	BoolQ：阅读理解	RACE-h：阅读理解	TriviaQA：问答	WebQs：问答
密集 NLG
350M，内部数据	0.5203	0.6931	0.5364	0.3177	0.0321	0.0157
350M，公共 Pile	0.5106	0.6589	0.5933	0.3196	0.0257	0.0064
标准 MoE NLG
350M+MoE-128，内部数据	0.6270	0.7459	0.6046	0.3560	0.1658	0.0517
350M+MoE-128，公共 Pile	0.6128	0.7323	0.6040	0.3349	0.1111	0.0335
PR-MoE NLG
350M+MoE-128，内部数据	0.6365	0.7399	0.5988	0.3569	0.1630	0.0473
PR-MoE + MoS NLG
350M+MoE-128，内部数据	0.6346	0.7334	0.5807	0.3483	0.1369	0.0522

表 1：不同密集和 MoE NLG 模型的零样本评估结果（最后六列）。所有零样本评估结果均使用准确率指标。

2.4. 使用缩减模型大小训练 MoS

MoS 代表学生混合，是一种基于分阶段蒸馏的技术，用于压缩大型 MoE 模型。MoS 进一步将模型大小减少了 12.5%，与标准 MoE 相比，与 PR-MoE 结合使用时，模型大小减少了 3.7 倍。缩减的模型大小有助于降低推理期间的延迟和成本。要训练 MoS 模型，需要指定一些其他参数。我们将以 PR-MoE 为例

--mos：这将通过知识蒸馏启用学生混合。

--load-teacher：这指定了教师模型检查点的路径。这是使用 MoS 的必填参数，教师模型检查点可以通过训练标准 MoE 或 PR-MoE 获得。

num-layers-teacher、--hidden-size-teacher、--hidden-size-teacher、--num-experts-teacher：除了教师模型检查点路径之外，我们还需要指定教师模型的模型架构，例如其层数、隐藏维度大小以及每个 MoE 层的专家数量。在 PR-MoE 的情况下，我们需要也为教师模型提供专家列表，我们在其中从教师模型中删除了一些专家层。

除了上述新参数之外，我们观察到在整个训练过程中使用教师 PR-MoE 可能会对最终学生模型的准确性产生不利影响。在我们的实验中，我们使用了一种分阶段蒸馏方法，在训练过程中尽早停止蒸馏（例如，在 400K 步之后），并在训练的剩余时间内仅针对标准语言建模损失执行优化。

我们在 examples_deepspeed/MoE 下提供了示例训练脚本。我们的参数设置详细信息可以在示例训练脚本中找到。MoS 的性能结果可以在我们的博文和我们的论文中看到。