Megatron-LM：使用模型并行训练数十亿参数的语言模型

摘要

最近在语言建模方面的工作表明，训练大型Transformer模型能够推动自然语言处理应用的技术前沿。然而，由于内存限制，训练非常大的模型可能相当困难。在这项工作中，我们展示了训练极大Transformer模型的技术，并实现了一种简单、高效的层内模型并行方法，使得训练具有数十亿参数的Transformer模型成为可能。我们的方法不需要新的编译器或库的更改，与管道模型并行正交且互补，并且可以通过在原生PyTorch中插入少量通信操作来完全实现。我们通过在512个GPU上收敛基于Transformer的模型，参数高达83亿，来阐述这一方法。我们在整个应用中维持了15.1 PetaFLOPs的计算量，与强大的单GPU基线相比，具有76%的扩展效率，该基线维持了39 TeraFLOPs，这是峰值FLOPs的30%。为了证明大型语言模型可以进一步推动技术前沿（SOTA），我们训练了一个类似于GPT-2的83亿参数Transformer语言模型和一个类似于BERT的39亿参数模型。我们展示了在BERT类模型中，随着模型规模的增大，对层归一化放置的仔细关注对于实现性能提升至关重要。使用GPT-2模型，我们在WikiText103（10.8对比SOTA困惑度15.8）和LAMBADA（66.5%对比SOTA准确率63.2%）数据集上取得了SOTA结果。我们的BERT模型在RACE数据集上取得了SOTA结果（90.9%对比SOTA准确率89.4%）。

1. 引言

自然语言处理（NLP）的快速发展部分归功于可用计算资源和数据集规模的增加。计算和数据的丰富使得通过无监督预训练（Devlin等，2018；Radford等，2019）训练越来越大的语言模型成为可能。经验证据表明，更大的语言模型对于NLP任务（如文章补全、问答和自然语言推理）来说显著更有用（Lan等，2019；Raffel等，2019）。通过在自然语言下游任务上微调这些预训练的语言模型，可以取得技术前沿的结果，如最近的工作所示（Devlin等，2018；Peters等，2018；Howard & Ruder，2018；Radford等，2018；2017；Ramachandran等，2016；Liu等，2019b；Dai等，2019；Yang等，2019；Liu等，2019a；Lan等，2019）。

随着这些模型变得越来越大，它们超出了现代处理器的内存限制，并需要额外的内存管理技术，如激活检查点（Chen等，2016）。广泛使用的优化算法（如ADAM）需要额外的内存来存储动量和其他优化器状态，这减少了可以有效训练的模型的大小。几种模型并行的方法通过分区模型来克服这一限制，使得权重及其相关的优化器状态不需要同时驻留在处理器上。例如，GPipe（Huang等，2018）和Mesh-Tensorflow（Shazeer等，2018）提供了不同类型的模型并行框架。然而，它们需要重写模型，并依赖于仍在开发中的自定义编译器和框架。

在这项工作中，我们实现了一种简单高效的模型并行方法，使用层内模型并行。我们利用基于Transformer的语言模型的固有结构，实现了一个简单的模型并行实现，该实现在PyTorch中高效训练，无需自定义C++代码或编译器。这种方法与基于管道的模型并行（如GPipe（Huang等，2018）所倡导的）正交。

为了证明我们方法的可扩展性，我们通过在单个NVIDIA V100 32GB GPU上训练一个12亿参数的模型来建立基线，该模型维持了39 TeraFLOPs的计算量。这是在DGX-2H服务器中配置的单GPU理论峰值FLOPS的30%，因此是一个强有力的基线。将模型扩展到83亿参数，在512个GPU上使用8路模型并行，我们在整个应用中实现了高达15.1 PetaFLOPs的持续计算量。与单GPU情况相比，这是76%的扩展效率。图1显示了更详细的扩展结果。
在这里插入图片描述

为了分析模型规模扩展对准确性的影响，我们训练了从左到右的GPT-2（Radford等，2019）语言模型以及BERT（Devlin等，2018）双向Transformer，并在几个下游任务上进行了评估。我们展示了现有的BERT架构随着规模的增加会导致模型性能下降。我们通过重新排列Transformer层中的层归一化和残差连接来克服这一挑战，并展示了随着模型规模的增加，开发集上的下游任务结果单调提高。此外，我们展示了我们的模型在WikiText103上的测试集技术前沿（SOTA）结果，LAMBADA上的填空式预测准确率，以及阅读理解RACE数据集上的结果。

总之，我们的贡献如下：

我们通过对现有的PyTorch Transformer实现进行少量有针对性的修改，实现了一种简单高效的模型并行方法。
我们对我们的模型和数据并行技术进行了深入的实证分析，并展示了使用512个GPU时高达76%的扩展效率。
我们展示了在BERT类模型中，随着模型规模的增大，对层归一化放置的仔细关注对于实现准确性提升至关重要。
我们证明了模型规模的扩展可以提高GPT-2（研究至83亿参数）和BERT（研究至39亿参数）模型的准确性。
我们展示了我们的模型在测试集上取得了技术前沿的结果：WikiText103上的困惑度（10.8 ppl），LAMBADA上的准确率（66.5%），以及RACE上的准确率（90.9%）。
我们在https://github.com/NVIDIA/Megatron-LM开源了我们的代码以及训练和评估流程。

2. 背景与挑战

2.1. 神经语言模型预训练

预训练的语言模型已成为NLP研究人员工具包中不可或缺的一部分。利用大型语料库预训练来学习鲁棒的语言神经表示是一个活跃的研究领域，已经跨越了过去十年。预训练和转移语言神经表示的早期例子表明，与从头学习的词嵌入表相比，预训练的词嵌入表提高了下游任务的结果（Mikolov等，2013；Pennington等，2014；Turian等，2010）。后来的工作通过学习并转移捕捉词上下文表示的神经模型，推动了这一领域的研究（Melamud等，2016；McCann等，2017；Peters等，2018；Radford等，2017；2019）。最近的并行工作（Ramachandran等，2016；Howard & Ruder，2018；Radford等，2018；Devlin等，2018；Liu等，2019b；Dai等，2019；Yang等，2019；Liu等，2019a；Lan等，2019）进一步基于这些思想，不仅转移语言模型以提取上下文词表示，还通过在下游任务上端到端微调语言模型。通过这些工作，技术前沿已经从仅转移词嵌入表发展到转移整个数十亿参数的语言模型。这一方法的进步需要硬件、系统技术和框架能够在大规模上高效运行，并满足日益增长的计算需求。我们的工作旨在提供必要的工具，以在这一趋势中迈出另一步。

2.2. Transformer语言模型与多头注意力机制

当前NLP领域的工作趋势倾向于使用Transformer模型（Vaswani等，2017），因为它们具有卓越的准确性和计算效率。最初的Transformer设计是作为一种机器翻译架构，使用编码器和解码器两部分将输入序列转换为另一个输出序列。然而，最近利用Transformer进行语言建模的工作，如BERT（Devlin等，2018）和GPT-2（Radford等，2019），根据需求仅使用编码器或解码器。本研究探讨了两种架构：解码器架构GPT-2和编码器架构BERT。
在这里插入图片描述

图2展示了我们使用的模型的示意图。我们建议读者参考先前的工作以获取模型架构的详细描述（Vaswani等，2017；Devlin等，2018；Radford等，2019）。值得注意的是，GPT-2和BERT都使用GeLU（Hendrycks & Gimpel，2016）非线性和层归一化（Ba等，2016）应用于多头注意力和前馈层的输入，而原始Transformer（Vaswani等，2017）使用ReLU非线性并将层归一化应用于输出。

2.3. 深度学习中的数据并行和模型并行

扩展深度神经网络训练到多个硬件加速器有两个核心范式：数据并行（Valiant，1990），其中训练小批量数据在多个工作节点间分割；模型并行，其中模型的内存使用和计算分布在多个工作节点上。通过将小批量大小与可用工作节点数成比例增加（即弱扩展），可以观察到训练数据吞吐量的近线性扩展。然而，大批量训练引入了优化过程中的复杂性，可能导致准确性降低或收敛时间延长，抵消了训练吞吐量增加的好处（Keskar等，2017）。进一步的研究（Goyal等，2017；You等，2017；2019）已经开发了技术来减轻这些影响并缩短大型神经网络的训练时间。为了进一步扩展训练，并行工作（Chen等，2016）将数据并行与激活检查点结合：在反向传播中重新计算激活而不在正向传播中存储它们以减少内存需求。

然而，这些技术在它们可以处理的问题大小上有一个基本限制：模型必须完全适合一个工作节点。随着BERT和GPT-2等语言模型的规模和复杂性增加，神经网络已经接近现代硬件加速器的内存容量。解决这个问题的一个方案是采用参数共享来减少模型的内存占用（Lan等，2019），但这限制了模型的整体容量。我们的方法是利用模型并行将模型分割到多个加速器上。这不仅缓解了内存压力，还增加了与微批量大小无关的并行量。

在模型并行中，还有两个进一步的范式：层间管道并行和更一般的分布式张量计算。在管道模型并行中，一组操作在一个设备上执行，然后输出传递到管道中的下一个设备，在那里执行另一组操作。一些方法（Harlap等，2018；Chen等，2018）将参数服务器（Li等，2014）与管道并行结合使用。然而，这些方法存在一致性问题。TensorFlow的GPipe框架（Huang等，2018）通过使用同步梯度下降克服了这个一致性问题。这种方法需要额外的逻辑来处理这些通信和计算操作的高效管道化，并且受到减少效率的管道泡沫或影响准确性的优化器本身变化的影响。

分布式张量计算是一种正交且更一般的方法，它将张量操作分割到多个设备上以加速计算或增加模型大小。FlexFlow（Jia等，2018），一个协调这种并行计算的深度学习框架，提供了一种选择最佳并行化策略的方法。最近，Mesh-TensorFlow（Shazeer等，2018）引入了一种语言，用于在TensorFlow（Abadi等，2015）中指定一类通用的分布式张量计算。最终用户在语言中指定并行维度，并使用适当的集体原语编译结果图。我们利用与Mesh-TensorFlow中类似的见解，并利用Transformer注意力头的并行性来并行化我们的Transformer模型。然而，我们不是实现一个模型并行的框架和编译器，而是仅对现有的PyTorch Transformer实现进行少量有针对性的修改。我们的方法简单，不需要任何新的编译器或代码重写，并且可以通过插入一些简单的原语完全实现，如下一节所述。

3 模型并行的Transformer

我们利用Transformer网络的结构，通过添加一些同步原语创建了一个简单的模型并行实现。一个Transformer层由一个自注意力块和一个两层多层感知器（MLP）组成，如图2所示。我们在这两个块中分别引入模型并行。

我们首先详细说明MLP块。块的第一部分是一个GEMM，随后是一个GeLU非线性：
$Y=\mathrm{GeLU}(XA)\qquad{(1)}$

并行化GEMM的一种选择是沿着行分割权重矩阵 $A$ ，并沿着列分割输入 $X$ ，如下所示：
$X=[X_1,X_2],\:A=\begin{bmatrix}A_1\\A_2\end{bmatrix}\qquad{(2)}$
这种分割将导致 $Y =$ GeLU $X_1A_1+$ $X_{2}A_{2}).$ 由于GeLU是一个非线性函数，GeLU $X_1A_1+$ $X_{2}A_{2})\neq$ GeLU $X_1A_{1})+$ GeLU $X_{2}A_{2})$ ，这种方法在GeLU函数之前需要一个同步点。

另一种选择是沿着列分割 $A$ ，即 $A=[A_1,A_2].$ 这种分割允许GeLU非线性独立地应用于每个分割后的GEMM的输出：

$[Y_1,Y_2]=[\mathrm{GeLU}(XA_1),\mathrm{GeLU}(XA_2)] \qquad{(3)}$
这种方法的优势在于它移除了一个同步点。因此，我们以这种列并行的方式对第一个GEMM进行分割，并沿着行分割第二个GEMM，使其直接接收GeLU层的输出，而无需任何通信，如图3a所示。第二个GEMM的输出随后在传递到dropout层之前跨GPU进行归约。这种方法将MLP块中的两个GEMM分割到多个GPU上，并且在前向传播中仅需要一个all-reduce操作（g操作符），在反向传播中也仅需要一个all-reduce操作（f操作符）。这两个操作符互为共轭，可以在PyTorch中用几行代码实现。例如，f操作符的实现如下所示：

在这里插入图片描述

如图3b所示，对于自注意力块，我们利用了多头注意力操作中的固有并行性，以列并行的方式分割与键（K）、查询（Q）和值（V）相关的GEMM，使得每个注意力头对应的矩阵乘法在单个GPU上本地完成。这使得我们能够将每个注意力头的参数和工作负载分配到多个GPU上，并且不需要任何即时通信来完成自注意力操作。随后的输出线性层（自注意力之后）的GEMM沿其行并行化，并直接接收并行注意力层的输出，而无需GPU之间的通信。这种针对MLP和自注意力层的方法将两组GEMM融合在一起，消除了中间的同步点，从而实现了更好的扩展性。这使得我们能够在简单Transformer层中仅使用前向路径中的两个all-reduce操作和反向路径中的两个all-reduce操作来执行所有GEMM（见图4）。

在这里插入图片描述

Transformer语言模型具有一个输出嵌入层，其维度为隐藏大小（H）乘以词汇表大小（v）。由于现代语言模型的词汇表大小通常在数万个词元（例如，GPT-2使用了50,257的词汇表大小），因此并行化输出嵌入层的GEMM是有益的。然而，在Transformer语言模型中，输出嵌入层与输入嵌入层共享权重，因此需要对两者进行修改。我们沿词汇表维度对输入嵌入权重矩阵 $E_H×v$ 进行并行化分割 $E = [E 1; E 2]$ （按列分割）。由于每个分区现在只包含嵌入表的一部分，因此在输入嵌入之后需要一个all-reduce操作（g操作符）。对于输出嵌入，一种方法是执行并行GEMM $[Y 1; Y 2] = [XE 1; XE 2]$ 以获得logits，然后添加一个all-gather操作 $Y = a ll - g a t h er ([Y 1; Y 2])$ ，并将结果发送到交叉熵损失函数。然而，在这种情况下，all-gather将通信 $b \times s \times v$ 个元素（b是批量大小，s是序列长度），由于词汇表大小较大，通信量非常大。为了减少通信量，我们将并行GEMM [Y1; Y2]的输出与交叉熵损失融合，从而将维度减少到 $b \times s$ 。通信标量损失而不是logits大大减少了通信量，提高了我们模型并行方法的效率。

我们的模型并行方法在很大程度上可以被描述为旨在减少通信并保持GPU计算密集型的技术。我们选择在GPU之间复制计算，而不是让一个GPU计算部分dropout、层归一化或残差连接并将结果广播到其他GPU。具体来说，我们在每个GPU上维护层归一化参数的副本，并在模型并行区域的输出上运行dropout和残差连接，然后将这些张量作为输入传递给下一个模型并行区域。为了优化模型，我们允许每个模型并行工作节点优化其自己的参数集。由于所有值要么是本地值，要么在GPU上复制，因此在这种公式中不需要通信更新的参数值。

我们在附录B中提供了关于混合模型和数据并行性以及处理随机数生成的更多详细信息以供参考。总之，我们上述的方法实现简单，仅需在前向和反向传播中添加一些额外的all-reduce操作。它不需要编译器，并且与（Huang等，2018）等倡导的管道模型并行正交且互补。

4 实验设置

预训练的语言理解模型是自然语言处理和语言理解中的核心任务。语言建模有几种不同的形式。在本研究中，我们主要关注GPT-2（Radford等，2019），一种基于从左到右生成的Transformer语言模型，以及BERT（Devlin等，2018），一种基于语言模型掩码的双向Transformer模型。我们将在以下部分解释这些模型的配置，并建议读者参考原始论文以获取更多详细信息。

4.1. 训练数据集

为了收集一个具有长期依赖性的多样化大型训练集，我们整合了几个最大的语言建模数据集。我们创建了一个包含Wikipedia（Devlin等，2018）、CC-Stories（Trinh & Le，2018）、RealNews（Zellers等，2019）和OpenWebtext（Radford等，2019）的聚合数据集。为了避免训练集数据泄露到下游任务中，我们移除了WikiText103测试集（Merity等，2016）中存在的Wikipedia文章。我们还从CC-Stories语料库中移除了由预处理引入的不必要的换行符。对于BERT模型，我们在训练数据集中包含了BooksCorpus（Zhu等，2015），然而，由于该数据集与LAMBADA任务重叠，因此在GPT-2训练中排除了该数据集。

我们合并了所有数据集，然后从聚合数据集中过滤掉所有内容长度小于128个词元的文档。由于相似的内容可能在聚合数据集中多次出现，我们使用局部敏感哈希（LSH）对Jaccard相似度大于0.7的内容进行去重。最终的聚合语料库包含174 GB的去重文本。

4.2. 训练优化和超参数

为了高效训练我们的模型，我们利用混合精度训练和动态损失缩放来充分利用V100的Tensor Cores（Micikevicius等，2017；NVIDIA 2018）。我们首先使用简单的正态分布 $W\sim\mathcal{N}(0,0.02)$ 初始化权重 $W$ 。然后，在残差层之前立即将权重缩放 $\frac1{\sqrt{2N}}$ ，其中N是由自注意力和MLP块组成的Transformer层数。对于优化器，我们使用Adam（Kingma & Ba，2014）并加入权重衰减（Loshchilov & Hutter，2019） $\lambda$ = 0.01。此外，我们使用全局梯度范数裁剪为1.0来提高训练大模型的稳定性。在所有情况下，dropout设置为0.1。最后，为了更好地管理内存占用，我们在每个Transformer层之后使用激活检查点（Chen等，2016）。

对于GPT-2模型，所有训练均使用1024个子词单位的序列，批量大小为512，进行300k次迭代。我们的学习率为1.5e-4，在前3k次迭代中使用预热期，然后在剩余的297k次迭代中遵循单周期余弦衰减。我们将最小学习率停止在1e-5。

对于BERT模型，我们主要遵循（Lan等，2019）中描述的训练过程。我们使用原始的BERT词典，词汇表大小为30,522。此外，我们按照（Lan等，2019）的建议，将下一句预测头替换为句子顺序预测，并使用（Joshi等，2019）的整词n-gram掩码。在所有情况下，我们将批量大小设置为1024，并使用1.0e-4的学习率，在10,000次迭代中预热，并在200万次迭代中线性衰减。其他训练参数保持与（Devlin等，2018）相同。

5 实验

我们所有的实验最多使用32台DGX-2H服务器（共512个Tesla V100 SXM3 32GB GPU）。我们的基础设施针对多节点深度学习应用进行了优化，通过NVSwitch在服务器内部的GPU之间提供300 GB/s的带宽，并使用每台服务器8个InfiniBand适配器在服务器之间提供100 GB/s的互连带宽。

5.1. 扩展性分析

为了测试我们实现的扩展性，我们考虑了具有四组参数的GPT-2模型，详见表1。为了在自注意力层中保持一致的GEMM大小，每个注意力头的隐藏大小保持恒定在96，同时改变头数和层数以获得从10亿到80亿参数的配置。12亿参数的配置适合单个GPU，而80亿参数的模型需要8路模型并行（8个GPU）。原始词汇表大小为50,257，然而，为了在logit层中获得高效的GEMM，每个GPU的词汇表大小最好是128的倍数。由于我们研究最多8路模型并行，我们将词汇表填充为128 × 8 = 1024的倍数，结果填充后的词汇表大小为51,200。我们研究了模型并行和模型+数据并行的扩展性。对于模型并行扩展，所有配置使用固定的批量大小8。数据并行扩展对于训练许多通常使用更大全局批量大小的最先进模型是必要的。为此，对于模型+数据并行的情况，我们将所有实验的全局批量大小固定为512，对应于64路数据并行。

5.1.1. 模型和数据并行

在本节中，我们将展示模型并行和模型+数据并行情况下关于模型参数的弱扩展性。弱扩展通常通过扩展批量大小来完成，然而，这种方法不适用于无法在单个GPU上训练的大型模型，并且对于大批量大小会导致训练收敛性下降。相比之下，我们在这里使用弱扩展来训练否则无法训练的大型模型。所有扩展数字的基线是表1中第一个配置（12亿参数）在单个GPU上运行。这是一个强有力的基线，因为它在整个训练过程中实现了39 TeraFLOPS，这是DGX-2H服务器中单个GPU理论峰值FLOPS的30%。
在这里插入图片描述

在这里插入图片描述

图5显示了模型并行和模型+数据并行的扩展值。我们在这两种设置中观察到了优秀的扩展数字。例如，83亿参数的8路（8 GPU）模型并行案例实现了77%的线性扩展。模型+数据并行需要进一步的梯度通信，因此扩展数字略有下降。然而，即使是在512个GPU上运行的最大配置（83亿参数），我们也实现了相对于强单GPU基线配置（12亿参数）线性扩展的74%扩展。进一步的扩展分析见附录D。

5.2. 使用GPT-2的语言建模结果

为了证明大型语言模型可以进一步推动技术前沿，我们考虑训练表2中列出的规模和配置的GPT-2模型。355M模型在规模和配置上与BERT-Large模型（Devlin等，2018）相当。2.5B模型比之前最大的GPT-2模型更大，而8.3B模型则是迄今为止训练的最大的从左到右的Transformer语言模型（据我们所知）。为了训练和评估我们的语言模型，我们使用了第4节中描述的程序。表2还列出了完成一个epoch所需的时间，相当于68,507次迭代。例如，对于512个GPU上的8.3B模型，每个epoch大约需要两天时间。与表1中用于扩展研究的配置相比，2.5B模型相同，8.3B模型有24个注意力头而不是32个，而355M模型比之前看到的任何模型都小得多，但仍然使用64个GPU进行训练，因此每个epoch的时间大大减少。

在这里插入图片描述

图6显示了验证困惑度随迭代次数的变化。随着模型规模的增加，验证困惑度降低，8.3B模型的验证困惑度达到9.27。我们在表3中报告了训练模型在LAMBADA和WikiText103数据集上的零样本评估结果。有关评估方法的更多详细信息，请参阅附录E。我们观察到模型规模增加的趋势也导致WikiText103上的困惑度降低和LAMBADA上的填空准确率提高。我们的8.3B模型在WikiText103测试集上实现了技术前沿的困惑度，经过适当调整后的困惑度为10.81。在LAMBADA任务中，8.3B模型以66.51%的准确率同样超越了之前的填空准确率结果。我们在附录C中包含了从83亿参数模型生成的样本。最近，微软的研究人员与NVIDIA合作，使用Megatron训练了一个170亿参数的GPT-2模型，称为Turing-NLG（微软，2020），并表明随着模型规模的扩大，准确率进一步提高，凸显了更大模型的价值。

在这里插入图片描述

为了确保我们没有在测试集中训练任何数据，我们按照之前的工作（Radford等，2019）计算了测试集中8-gram与训练集的重叠百分比。WikiText103测试集最多有10.8%的重叠，而LAMBADA测试集（Paperno等，2016）最多有1.4%的重叠。需要注意的是，WikiText103测试集与WikiText103训练集本身已经有9.09%的重叠（Radford等，2019）。由于这些结果与之前的工作一致，我们确信测试数据中的文档没有被无意中包含在训练数据中。

5.3. 使用BERT的双向Transformer结果

在本节中，我们将我们的方法应用于BERT风格的Transformer模型，并研究模型规模对多个下游任务的影响。先前的工作（Lan等，2019）发现，将模型规模增加到超过336M参数的BERT-large会导致意外的模型性能下降。为了解决这种性能下降问题，该工作的作者（Lan等，2019）引入了参数共享，并表明他们的模型比原始BERT模型具有更好的扩展性。

我们进一步研究了这一现象，并通过实验证明，重新排列层归一化和残差连接的顺序（如图7所示）对于使BERT风格模型的规模超越BERT-Large至关重要。图7中的架构（b）消除了使用原始BERT架构（a）时观察到的不稳定性，并且具有更低的训练损失。据我们所知，我们是第一个报告这种变化能够训练更大BERT模型的研究。

在这里插入图片描述

使用图7(b)中的架构变化，我们考虑了三种不同的情况，详见表4。336M模型的规模与BERT-large相同。1.3B模型与之前显示的BERT-xlarge配置相同，该配置在之前的研究中表现不如336M的BERT-large模型（Lan等，2019）。我们通过增加隐藏层大小和层数进一步扩展了BERT模型，达到了3.9B参数的规模。在所有情况下，每个注意力头的隐藏层大小保持恒定在64。336M和1.3B模型训练了200万次迭代，而3.9B模型训练了150万次迭代，并且仍在训练中。

在一个3%的保留集上，336M、1.3B和3.9B模型的验证集困惑度分别为1.58、1.30和1.16，随着模型规模的增加，困惑度单调下降。我们在多个下游任务上对训练好的模型进行了微调，包括GLUE基准测试（Wang等，2019）中的MNLI和QQP，斯坦福问答数据集（Rajpurkar等，2016；2018）中的SQuAD 1.1和SQuAD 2.0，以及阅读理解RACE数据集（Lai等，2017）。对于微调，我们遵循与（Liu等，2019b）相同的程序。我们首先对批量大小和学习率进行超参数调优。一旦获得最佳值，我们在5个不同的随机种子初始化下报告开发集结果的中位数。每个模型和任务使用的超参数在附录A中提供。

表5显示了MNLI、QQP、SQuAD 1.1和SQuAD 2.0的开发集结果以及RACE的测试集结果。对于RACE的测试集结果，我们首先使用开发集找到在5个随机种子下给出中位分数的检查点，并在测试集上报告该检查点的结果。我们还报告了SQuAD开发集和RACE测试集的5路集成结果。从表5中我们观察到：（a）随着模型规模的增加，下游任务的性能在所有情况下都有所提高；（b）与其他基于BERT的模型相比，我们的3.9B模型在开发集上建立了技术前沿的结果；（c）我们的3.9B模型在RACE测试集上实现了单模型和集成模型的技术前沿结果。

6. 与未来工作

在这项工作中，我们通过仅对现有PyTorch Transformer实现进行少量修改，成功实现了模型并行，从而超越了传统单GPU训练模型的限制。我们在512个NVIDIA V100 GPU上高效训练了基于Transformer的模型，参数高达83亿，并实现了在整个应用中持续15.1 PetaFLOPs的计算量。我们还展示了对于BERT模型，随着模型规模的增加，仔细关注层归一化的位置对于提高准确性至关重要。我们研究了模型规模对下游任务准确性的影响，并在下游任务上取得了显著优于以往的结果，为WikiText103、LAMBADA和RACE数据集建立了新的技术前沿。最后，我们开源了我们的代码，以便未来的工作能够利用模型并行的Transformer。

未来工作有多个方向。继续扩大预训练的规模是一个有前景的研究方向，这将进一步测试现有的深度学习硬件和软件。为了实现这一点，需要改进优化器的效率和内存占用。此外，训练一个超过160亿参数的模型将需要比DGX-2H服务器中16个GPU更多的内存。对于这样的模型，混合层内和层间模型并行以及节点间模型并行将更为合适。其他三个研究方向包括：（a）预训练不同的模型家族（如XLNet、T5）；（b）评估大型模型在更困难和多样化的下游任务（如生成式问答、摘要生成和对话）上的性能；（c）使用知识蒸馏从这些大型预训练教师模型中训练小型学生模型。

A. BERT微调超参数

表6展示了在微调过程中每个模型和任务所使用的超参数。

B. 模型并行补充材料

在本节中，我们将进一步介绍混合模型和数据并行的细节，以及随机数生成的处理方法。

B.1. 混合模型与数据并行

模型并行与数据并行是正交的，因此我们可以同时使用两者来在合理的时间内训练大型模型。图8展示了混合模型和数据并行的GPU分组。同一服务器内的两个或多个GPU组成模型并行组（例如图8中的GPU 1到8），并在这些GPU上分布一个模型实例。
在这里插入图片描述

其余的GPU（可能位于同一服务器内，但更常见的是位于其他服务器中）运行额外的模型并行组。每个模型并行组中位置相同的GPU（例如图8中的GPU 1、9、…、505）组成数据并行组，因此数据并行组中的所有GPU持有相同的模型参数。在反向传播过程中，我们并行运行多个梯度all-reduce操作，以减少每个独立数据并行组内的权重梯度。所需GPU的总数是模型并行组和数据并行组数量的乘积。例如，对于83亿参数的模型，我们每个模型并行组使用8个GPU，并采用64路数据并行，总共使用512个GPU。所有通信均通过PyTorch调用NCCL实现。每个模型并行组内的GPU在组内执行all-reduce操作。对于数据并行，每个all-reduce操作在每个模型并行组中的一个GPU之间进行。

B.2. 模型并行随机数生成

利用随机数生成的技术（如dropout）是现代深度学习训练的基石。Transformer在模型并行区域之外的残差连接之前以及模型并行区域内的自注意力块中都有dropout层。由于一些dropout层位于模型并行区域内，而另一些则不在，我们需要谨慎处理随机数生成，以确保dropout正常工作。为了在模型并行工作节点之间同步残差连接的dropout，我们在训练开始时使用相同的种子初始化随机数生成器。这导致所有模型并行工作节点上的dropout模式相同。然而，模型并行区域内的dropout应为每个工作节点生成不同的随机模式，以实现整个操作的随机性。为此，我们为模型并行区域内的dropout维护一个独立的随机数生成器。该随机数生成器为每个模型并行工作节点单独设置种子。
在这里插入图片描述

C. 文本示例

以下是由Megatron-LM使用上下文提示生成的一些文本示例。部分文本被截短。

上下文：美国两党政治家之间有一股强大的力量将他们团结在一起：大型科技公司。联邦和州一级的民主党和共和党人正在联合起来审查硅谷巨头的权力，并可能对其进行约束。纽约州民主党总检察长莱蒂西亚·詹姆斯（Letitia James）周五宣布，八个州（四名民主党人和四名共和党人）以及哥伦比亚特区的总检察长已开始对Facebook进行反垄断调查。

Megatron-LM生成的文本：詹姆斯女士表示，这项调查可能会导致诉讼，这是州和联邦法院一系列挑战科技公司市场主导地位的高调监管行动中的最新一起。司法部也在调查Facebook的数据实践，以及该公司是否违反了2011年与联邦贸易委员会达成的一项同意令，该同意令要求公司在与第三方共享用户个人信息之前获得用户的明确同意。

是什么可能促使两党联盟起诉Facebook？目前尚不清楚是否有任何单一原因。但立法者和监管机构对该公司最近的争议越来越感到担忧，这些争议包括影响超过5000万用户的数据泄露、平台上虚假新闻的传播，以及2016年大选期间俄罗斯支持账户的宣传活动。

目前也不清楚民主党总检察长们对提起诉讼的认真程度。虽然总检察长们在新闻稿中宣布了他们的调查，但他们表示尚未就是否提起诉讼做出决定。过去，一些总检察长将精力集中在试图从科技公司中达成和解，而不是提起诉讼。

詹姆斯女士在一份声明中表示：“作为调查和起诉医疗欺诈、价格操纵以及其他违反我们州反垄断法行为的总检察长，我们致力于在科技公司未能保护消费者和市场竞争力时追究其责任。”

这项调查最初由《纽约时报》报道，正值主导科技行业的公司在全球范围内面临反垄断监管机构的审查之际。谷歌、Twitter、Facebook和亚马逊在美国和欧洲面临了一系列反垄断调查和集体诉讼……

上下文：花，有时被称为花朵或花蕾，是开花植物（被子植物门，也称为被子植物）中的生殖结构。花的生物学功能是通过提供精子和卵子结合的机制来实现繁殖。

Megatron-LM生成的文本：花的结构组成部分包括花瓣和萼片（生殖器官）、花托（形成花的基部）以及花萼（柱头）。花是一种复合结构，由萼片和花瓣组成，萼片由三个融合的花被片组成，花瓣由两个融合的花瓣组成，尽管这一规则有一些例外。三个萼片融合形成一个称为花萼的管状结构，这种结构能够吸引传粉昆虫。花瓣由两个融合的萼片（外层的）组成，它们在顶端打开并向内折叠（抱合），形成花瓣。四个萼片和花瓣结合在一起形成花冠，这是花的最外层结构，负责携带花粉。

花的结构具有重要的功能意义，包括保护子房免受昆虫侵害以及在花的雄性和雌性部分之间传递花粉。花的结构还被植物用来吸引传粉者，例如通过花瓣的形状和雄蕊的颜色。

花的结构在植物之间差异很大。最简单的花是不分枝的，仅由两个萼片和一个花瓣组成。其他例子由多个花瓣连接在一起组成，最内层的萼片完全融合形成一个称为花萼的管状结构。花瓣通常具有颜色以吸引昆虫，因为昆虫只能看到特定波长的光。萼片通常也有颜色，但颜色范围通常更有限。

一些花具有额外的结构，例如雄蕊、花药和子房。某些花中雄蕊和花药的存在用于异花授粉。子房是产生雌配子的器官，也是花胚胎中最早发育的结构。子房是雌花的生殖系统，由子房壁、胚珠和输卵管组成。

花的大小和形状各异，从微小到巨大，通常颜色鲜艳。花可以被视为更大结构的一部分，例如花序、花序梗、苞片、佛焰苞或心皮。

“花”一词通常用于指外部可见的花结构，但许多……

上下文：OpenAI由于担心恶意使用，没有发布完整的GPT-2模型，但他们发布了一个较小版本的模型，其规模与原始GPT（1.17亿参数）相当，并在新的更大数据集上进行了训练。尽管这个较小版本不如大型模型强大，但它仍然具有一定的语言生成能力。

Megatron-LM生成的文本：2016年12月，DeepMind发布了一个更大的模型GPT-3，其参数超过15亿。DeepMind没有计划发布该模型，因此问题仍然在于其他AI研究人员是否能够复制并在此基础上进行开发。

其他语言模型使用了不同的技术构建。2016年，哥伦比亚大学、多伦多大学和Google DeepMind的研究团队展示了基于神经网络的模型（包括深度神经网络）可以学习语言的语法规则。Google DeepMind的一个团队证明，这些规则可以在网络中表示为一种决策树，其中树的节点是句子中的单词，而单词的规则由树末端的决策节点表示。网络会遍历网络，根据是否遵循规则来选择句子中的下一个单词。2017年，DeepMind团队发布了一篇论文，详细介绍了他们模型的训练和性能。

该模型存在许多问题。其中一个最大的问题是语法规则的模糊性。例如，短语“I came down”可能意味着我从楼梯上下来，或者我从天上下来。在英语中，还有许多其他模糊句子的例子。该模型无法区分这些模糊句子，因此它只能根据已知规则对整个句子进行预测。这使得模型容易出现过拟合现象。为了避免这种情况，模型必须使用大量数据进行训练，以便模型能够在数据中找到有用的模式，并利用这些模式帮助进行预测。这种方法也容易出错，因为它依赖于训练数据来帮助做出正确的预测，而数据可能被错误标记，导致模型生成一组糟糕的预测结果。

在GPT-2的情况下，模型每种短语类型（主谓宾）只有大约100个训练句子。此外，不同语言对不同类型的短语有不同的规则。在英语中，主谓宾结构相当普遍，但其他语言有不同的规则，这使得训练GPT更加复杂。最后，规则数量与模型性能之间存在固有的权衡……

上下文：第36届国际机器学习会议（ICML 2019）将于2019年6月10日至6月15日在美国加州长滩举行。会议将包括一天的教程（6月10日），随后是三天的正式会议（6月11日至13日），最后是两天的研讨会（6月14日至15日）。

Megatron-LM生成的文本：会议范围广泛，旨在为从业者和研究人员探讨机器学习领域的挑战与机遇。会议议程将涵盖从强化学习到计算生物学、从人工智能到深度学习等不同机器学习主题。它是全球探索和展示机器学习进展的首要论坛。

ICML 2019将提供一系列令人印象深刻的会议，这些会议将吸引从业者、研究人员和学生的兴趣。我们将涵盖的一些主题包括：

强化学习
软计算
计算生物学
卷积网络
深度学习
统计学习理论
统计学习机器
自然语言处理
生物学
奖项

自第一届会议以来，ICML不断发展壮大。如今，它已成为展示机器学习领域最前沿研究的首要论坛。该会议被广泛认为是展示机器学习最新研究的卓越平台，也是从业者、研究人员和学生交流思想和见解的场所。

该会议是一个同行评审的学术平台，其声望和影响力来自于演讲的卓越性和受邀演讲者的质量。超过50%的演讲内容由专家委员会根据其内容选出，这清楚地表明会议对研究和质量的高标准保持了极高的门槛。

上下文：俄罗斯国防部周二表示，救援人员找到了周末坠入黑海的俄罗斯飞机的飞行记录仪。周日早上，俄罗斯军方图-154飞机从俄罗斯南部城市索契起飞两分钟后坠毁，机上92人全部遇难。

Megatron-LM生成的文本：一名儿童在彼得罗巴甫洛夫斯克-堪察加斯基当地政府大楼内的临时停尸房举起一位已故亲属的肖像。周日早上，俄罗斯国防部一架载有92人的图-154飞机从俄罗斯南部城市索契起飞后不久坠入黑海。（尤里·科切特科夫/EPA）

据国防部称，飞机在超过1000英尺的深度被发现，记录仪已损坏。声明称：“国防部决定将记录仪移交给国家间航空委员会，该委员会是空难的主要调查机构。”由于强风和汹涌的海浪，寻找飞机黑匣子的工作受到阻碍。调查人员已对坠机事件展开刑事调查，官员称事故可能是由技术故障引起的。

俄罗斯总统弗拉基米尔·普京宣布周三为遇难者哀悼日。

D. 进一步扩展性分析

在本节中，我们研究了注意力头数量对扩展结果的影响，并展示了我们12亿参数模型的强扩展性结果。

D.1. 注意力头与扩展性

本节研究注意力头数量对模型并行扩展性的影响。为此，我们考虑83亿参数的配置，采用8路模型并行，并将注意力头数量从16个增加到32个。结果如表7所示。随着注意力头数量的增加，自注意力层中的一些GEMM操作变得更小，同时自注意力softmax中的元素数量增加。这导致扩展效率略有下降。未来的研究应注意这一超参数，以设计出平衡模型速度和准确性的大型Transformer模型。

D.2. 强扩展性

我们的模型并行性主要是为了训练比单个GPU内存容量更大的模型，但它也可以在不增加批量大小的情况下加速较小模型的训练。为了测量这种加速效果，我们训练了一个固定的12亿参数模型。我们使用每次迭代8个样本的固定批量大小，并通过模型并行性增加GPU数量。结果列于表8中。使用两个GPU使训练速度提高了64%。在此之上，随着每GPU计算量的减少以及内存带宽和通信开销开始占主导地位，收益逐渐递减。
在这里插入图片描述

E. 使用WikiText103和LAMBADA评估语言模型

在本节中，我们详细介绍了对WikiText103数据集（Merity等，2016）的评估方法，以及LAMBADA数据集（Paperno等，2016）上的填空式预测准确率。

E.1. WikiText103困惑度

WikiText103困惑度是自基准数据集创建以来过去几年中被广泛研究的评估标准。困惑度是语料库平均交叉熵的指数化（Mikolov等，2011）。这使其成为语言模型的自然评估指标，因为语言模型代表了整个句子或文本的概率分布。
$PPL=\exp(-\frac{1}{T_o}\sum_{t}^{T}\log P(t|0:t-1))\qquad{(4)}$
为了计算（4）中的困惑度，我们根据子词词汇表对WikiText103测试语料库进行分词，并对每个词元[0; T]的交叉熵损失求和。然后，我们通过原始分词方案中的词元数量 $T_o$ 对交叉熵损失进行归一化。WikiText103测试语料库已经预先分词为词级词元，之前的工作使用这些词元来计算困惑度。为了在与之前工作相同的水平上评估我们模型的困惑度，我们必须根据原始词元数量 $T_o$ 而不是实际输入到模型中的分词数据中的词元数量T进行归一化。这种预分词还会在文本中引入训练数据中不存在的伪影。为了缓解这种分布不匹配问题，我们首先使用可逆的去分词器对WikiText103测试数据集进行预处理，以删除与标点符号和空格相关的各种伪影。 $T_o$ 的值在此预处理之前计算。对于WikiText103的测试集， $T_o = 245566，T = 270329$ 。

我们还必须对困惑度计算进行一项针对Transformer的进一步修改。与基于RNN的语言模型不同，Transformer在固定窗口大小的输入上运行。因此，它们无法完全计算 $P (t ∣0 : t - 1)$ ，而只能计算 $P (t ∣ t - w : t - 1)$ ，其中w是我们的上下文大小：1024个词元。然而，为数据集中的每个词元计算此值是非常昂贵的，因为我们必须计算大约T次大小为w的上下文评估。为了高效评估我们的模型，我们采用了一种称为重叠评估的折中方法，其中我们每次将滑动窗口推进一定的重叠o，并且仅计算与窗口最后o个词元对应的交叉熵损失。在我们的实验中，我们使用重叠o为32，并以这种方式计算所有滑动窗口的损失。

E.2. LAMBADA填空准确率

处理长期上下文的能力对于最先进的语言模型至关重要，并且是长文本生成和基于文档的问答等问题的必要前提。像LAMBADA这样的填空式数据集旨在衡量模型在这些长期上下文中操作和推理的能力。填空式阅读理解使用词元上下文 $x = x_{1:t}$ ，其中一个词元xj被掩码；模型的目标是正确预测缺失的第j个词元的值。为了准确预测缺失的词元，模型需要深入理解周围的上下文以及在此上下文中如何使用语言。LAMBADA通过构建4-5个句子的示例来测试生成式从左到右的语言模型，其中上下文xt中的最后一个词被掩码。

我们的模型使用子词单元，因此对于LAMBADA评估，我们使用原始的、未经处理的LAMBADA数据集，并要求我们的模型预测组成词元的多个子词词元。我们使用教师强制，并且只有当所有输出预测都正确时才认为答案正确。这种公式等同于原始的词元预测任务。