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.aigithub.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 的性能结果可以在我们的博客文章和我们的论文中查看。

更新时间: