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: 这些配置决定了一个专家可以处理多少个 token。更大的值可能带来更好的收敛性，但也会导致训练速度变慢，因为不同专家上的负载会更不均衡。

--disable-moe-token-dropping: 这将完全解除单个专家可以处理多少个 token 的限制。出于同样的原因，我们只建议在推理/评估时使用此选项。

2.3. 预训练PR-MoE模型

PR-MoE 是一种新设计的 MoE 模型，代表 Pyramid-Residual-MoE，与标准 MoE 相比，它将参数效率提高了3倍。请参阅我们的博客了解更多详情。我们提供了 examples_deepspeed/MoE 下的训练脚本示例。与标准 MoE 相比，PR-MoE 模型有一些不同的超参数：

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

--mlp-type: 可选择 [standard, residual]。当选择 residual 时，将启用 Residual-MoE。

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

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

模型尺寸	LAMBADA: 补全预测	PIQA: 常识推理	BoolQ: 阅读理解	RACE-h: 阅读理解	TriviaQA: 问答	WebQs: 问答
密集NLG
3.5亿参数, 内部数据	0.5203	0.6931	0.5364	0.3177	0.0321	0.0157
3.5亿参数, 公开Pile数据集	0.5106	0.6589	0.5933	0.3196	0.0257	0.0064
标准MoE NLG
3.5亿参数+MoE-128, 内部数据	0.6270	0.7459	0.6046	0.3560	0.1658	0.0517
3.5亿参数+MoE-128, 公开Pile数据集	0.6128	0.7323	0.6040	0.3349	0.1111	0.0335
PR-MoE NLG
3.5亿参数+MoE-128, 内部数据	0.6365	0.7399	0.5988	0.3569	0.1630	0.0473
PR-MoE + MoS NLG
3.5亿参数+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%，与 PR-MoE 结合使用时，相对于标准 MoE，模型尺寸可缩小高达3.7倍。减小的模型尺寸有助于降低推理过程中的延迟和成本。要训练 MoS 模型，需要指定一些额外的参数。我们将以 PR-MoE 为例：

--mos: 这将通过知识蒸馏启用学生混合模型 (Mixture-of-Students)。

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

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

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

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