想效仿英伟达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 次迭代练习。大型言语模型的收敛速度显着加速，且收敛完结的验证困惑度比同类小模型的困惑度还低。

本文为机器之心编译，转载请联络本大众号取得授权。

------------------------------------------------