面向 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：不同密集 NLG 和 MoE NLG 模型的零样本评估结果（最后六列）。所有零样本评估结果都使用准确率指标。

2.4. 训练 MoS 并减少模型大小

MoS 代表专家混合，是一种基于分阶段蒸馏的技术，用于压缩大型 MoE 模型。MoS 进一步将模型大小降低了 12.5%，与 PR-MoE 相比，与标准 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 的性能结果。