想效仿英伟达50分钟训练 BERT？只有GPU还不够……

选自arXiv

作者：Mohammad Shoeybi 等

机器之心编译

参与： 魔王

此前，机器之心曾经介绍过英伟达的一项研究，打破了 NLP 领域的三项记录：将 BERT 的训练时间缩短到了 53 分钟；将 BERT 的推理时间缩短到了 2.2 毫秒；将 GPT-2 的参数量推向 80 亿（之前 GPT-2 最大为 15 亿参数量）。很多人将此成绩归结为英伟达的优越硬件条件，反正 GPU 有的是。其实不然，英伟达近期的一篇论文公布了该研究中采用的模型并行化方法：层内模型并行化。该方法无需新的编译器或库更改，只需在 PyTorch 中嵌入几个通信操作即可完整实现。

近期关于无监督语言建模的研究证明，训练大型神经语言模型推动了自然语言处理应用中的 SOTA 结果。但是，对于非常大的模型而言，内存限制了实际训练的模型大小。模型并行化使得我们能够训练更大的模型，因为模型并行化可以将参数分割并分配至多个处理器。

英伟达在近期的一项研究中实现了一种简单高效的层内模型并行化方法，可以训练具备数十亿参数的当前最优 transformer 语言模型。该方法无需新的编译器或库更改，它与 pipeline 模型并行正交且互补，只需在 PyTorch 中嵌入几个通信操作即可完整实现。利用该方法，研究者使用 512 个 GPU 收敛了一个具备 83 亿参数的 transformer 语言模型，该模型是目前最大的 transformer 模型，其规模是 BERT 的 24 倍，GPT-2 的 5.6 倍。

图 2 是该模型的示意图：

图 2：GPT-2 Transformer 架构。紫色矩形块表示全连接层，每个蓝色矩形块表示一个 transformer 层（重复 N 次）。

整个应用可保持 15.1 PetaFLOPs/秒的性能、76% 的扩展效率，而强大的单处理器基线方法仅能维持 39 TeraFLOPs/秒的性能，是峰值 FLOPs 的 30%。在 174GB 文本数据上训练该模型，则需以 12 ZettaFLOPs 训练 9.2 天方可实现收敛。将该语言模型应用于 WikiText103 和 LAMBADA 数据集均可实现当前最优结果：它在 WikiText103 数据集上的困惑度达到 10.8，而之前 SOTA 的困惑度为 16.4；它在 LAMBADA 数据集上的准确率可达 66.5%，而之前 SOTA 的准确率为 63.2%。目前，英伟达研究人员已经公开训练和评估代码，以及小型可移植模型的权重。

• 论文地址：https://arxiv.org/abs/1909.08053v1

• 代码地址：https://github.com/NVIDIA/Megatron-LM

研究贡献

英伟达研究人员利用层内模型并行化方法高效训练具备 83 亿参数的 transformer 语言模型。他们利用 transformer 语言模型的内在结构做了一个简单的模型并行实现，该实现可在 PyTorch 中高效训练完成，且无需任何定制 C++ 代码或编译器。该方法与基于 pipeline 的模型并行化是正交的。

为了展示该方法的可扩展性，研究者建立了一个基线：他们在单个 NVIDIA V100 32GB GPU 上训练了一个具备 12 亿参数的模型，整个训练应用维持 39 TeraFLOPs/秒的性能，是单个 GPU 在 DGX-2H 服务器上运行的理论峰值 FLOPS 的 30%，因此这是一个非常强大的基线模型。将该模型扩展至 83 亿参数，并使用 8-way 模型并行化在 512 个 GPU 上进行训练，达到了 15.1 PetaFLOPs/秒的性能。与单个 GPU 的情况相比，它实现了 76% 的扩展效率。在 174 GB 文本数据上收敛该模型需要以 12 ZettaFLOPs 训练 9.2 天。

详细扩展结果见下图 1，随 GPU 数量的增加，其所提供有效计算力的增长差不多接近线性。

图 1：model parallel（蓝色）和 model+data parallel（绿色）的 FLOPS 性能，x 轴为 GPU 数量。

model parallel（蓝色）：8-way 模型并行，每个 GPU 使用约 10 亿参数进行弱扩展（如 2 个 GPU 20 亿，4 个 GPU 40 亿）。model+data parallel（绿色）：与 model parallel 配置类似，不过它还要加上 64-way 数据并行。

研究者在 WikiText103 数据集和 LAMBADA 数据集上分析了该研究训练模型的准确率，发现随着模型规模的增大，模型在 WikiText103 数据集上的困惑度有所下降，在 LAMBADA 数据集上的准确率有所上升，并在这些任务上达到了当前最优结果。

模型并行 transformer

研究者利用 transformer 网络的结构，仅添加几个同步基元（synchronization primitives）即创建出一个简单的模型并行化实现。他们对 transformer 中的自注意力模块和多层感知机（MLP）模块均采用了模型并行化。

图 3：使用模型并行化后的 Transformer 模块。f 和 g 是共轭的，f 在前向传播中是恒等算子，在反向传播中是全规约（all reduce），而 g 在前向传播中是全规约，在反向传播中是恒等算子。

图 4：transformer 层中的通信操作。单个模型并行 transformer 层中的前向和反向传播中共有 4 个通信操作。

混合模型和数据并行化

模型并行与数据并行是正交的，因此我们可以同时使用二者在合理时间内训练大型模型。图 5 展示了混合模型和数据并行的 GPU 分组情况。

图 5：混合模型和数据并行化在 8-way 模型并行和 64-way 数据并行时的 GPU 分组情况。

实验

所有实验均在英伟达 DGX SuperPod4 上实施，研究者使用了多达 32 个 DGX-2H 服务器（共 512 块 Tesla V100 SXM3 32GB GPU）。

为了测试研究中实现的可扩展性，研究人员考虑使用四种参数设置的 GPT-2 模型，如下表所示：

表 1：可扩展性研究中所使用的参数。每个注意力头的隐藏层大小均为 96。

下图 6 展示了 model 和 model+data 并行化的扩展效率。我们可以看到，在两种设置下均出现了不错的扩展效率。

图 6：model parallel (a) 和 model+data parallel (b) 的弱扩展效率，x 轴表示 GPU 数量。

为了研究模型并行扩展中注意力头的作用，研究人员考虑使用 8-way 模型并行来处理 83 亿参数，并将注意力头的数量设置为 16、24、32。结果如下表 2 所示：

表 2：使用 8-way 模型并行处理 83 亿参数时，注意力头数量的作用。

该研究中的模型并行化旨在在单个 GPU 上训练超出其内存承载的模型，以及在不增加批大小的情况下加速小模型的训练。为了衡量加速效果，研究人员使用固定的 12 亿参数训练模型，结果见下表 3：

表 3：使用模型并行化训练 12 亿参数模型得到的加速（批大小保持不变）。

为了证明大型语言模型对 SOTA 结果的推动作用，图 7 展示了模型在验证集上的困惑度，x 轴表示迭代数。

图 7：验证集困惑度。所有语言模型均进行 300k 次迭代训练。大型语言模型的收敛速度明显加快，且收敛实现的验证困惑度比同类小模型的困惑度还低。

本文为机器之心编译，转载请联系本公众号获得授权。

✄------------------------------------------------

加入机器之心（全职记者 / 实习生）：hr@jiqizhixin.com

投稿或寻求报道：content@jiqizhixin.com

广告 & 商务合作：bd@jiqizhixin.com