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