Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

文章标题: Megatron-LM: 使用模型并行化训练数十亿参数的语言模型
作者/机构: Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro (1: NVIDIA, 2: 斯坦福大学)

A1 主要贡献

本文提出了一种简单高效的层内模型并行方法，以解决因内存限制而难以训练超大规模Transformer模型的核心问题。研究目标是实现能够训练数十亿参数Transformer模型的可扩展方案，且该方案无需新的编译器或库更改，易于在现有的PyTorch框架中实现。

核心创新点如下：
* 简单高效的模型并行方法：本文实现了一种层内模型并行方法，仅通过对现有PyTorch Transformer实现进行少量针对性修改即可完成。该方法与流水线模型并行正交且互补。
* 卓越的可扩展性：通过对模型和数据并行技术进行深入的实证分析，本文展示了在使用512个GPU时高达76%的扩展效率，在整个应用中实现了15.1 PetaFLOPs的持续计算性能。
* BERT模型扩展的关键发现：本文发现，对于BERT这类模型，仔细调整层归一化（Layer Normalization）的位置对于模型规模增大时性能的提升至关重要。
* 模型规模与性能的关系：本文证明了扩大模型规模可以显著提升GPT-2（研究至83亿参数）和BERT（研究至39亿参数）模型的准确率。
* 刷新SOTA记录：本文训练的模型在多个基准测试中取得了SOTA（state-of-the-art）成果，包括WikiText103数据集的困惑度（10.8 ppl）、LAMBADA数据集的完形填空准确率（66.5%）以及RACE数据集的阅读理解准确率（90.9%）。
* 代码开源：本文开源了其代码、训练和评估流程，以促进社区的进一步研究。

A3 背景知识与挑战

神经语言模型预训练

预训练语言模型的演进。预训练语言模型已成为自然语言处理（NLP）研究人员不可或缺的工具。利用大规模语料库进行预训练以学习语言的鲁棒神经表示，是过去十年中一个活跃的研究领域。早期的工作，如预训练词嵌入表【29, Distributed representations of words and phrases and their compositionality, 2013】【33, Glove: Global vectors for word representation, 2014】【45, Word representations: A simple and general method for semi-supervised learning, 2010】，证明了这能改善下游任务的效果。后续工作通过学习和迁移能够捕捉词语上下文表示的神经模型，进一步推动了该领域的发展【24, context2vec: Learning generic context embedding with bidirectional lstm, 2016】【23, Learned in translation: Contextualized word vectors, 2017】【34, Deep contextualized word representations, 2018】【35, Learning to generate reviews and discovering sentiment, 2017】【37, Better language models and their implications, 2019】。近期的并行工作【41, Unsupervised pretraining for sequence to sequence learning, 2016】【10, Fine-tuned language models for text classification, 2018】【36, Improving language understanding by generative pre-training, 2018】【6, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】【21, Roberta: A robustly optimized BERT pretraining approach, 2019b】【5, Transformer-xl: Attentive language models beyond a fixed-length context, 2019】【48, Xlnet: Generalized autoregressive pretraining for language understanding, 2019】【20, Multi-task deep neural networks for natural language understanding, 2019a】【18, Albert: A lite bert for self-supervised learning of language representations, 2019】不仅迁移语言模型以提取上下文词表示，还在下游任务上对语言模型进行端到端微调。这一系列方法的发展，从迁移词嵌入表演进到迁移数十亿参数的整个语言模型，对硬件、系统技术和框架的规模化高效运行能力提出了越来越高的计算需求。本研究旨在提供必要的工具，以推动这一趋势继续向前发展。

Transformer语言模型和多头注意力

Transformer架构。当前NLP领域的研究趋势是使用Transformer模型【47, Attention is all you need, 2017】，因其具有优越的准确性和计算效率。原始的Transformer架构是为机器翻译设计的，包含编码器（Encoder）和解码器（Decoder）两部分，用于将输入序列转换为输出序列。然而，近期用于语言建模的Transformer模型，如BERT【6, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】和GPT-2【37, Better language models and their implications, 2019】，根据其需求仅使用编码器或解码器。本工作同时探索了基于解码器的GPT-2架构和基于编码器的BERT架构。

模型架构细节。图2展示了我们使用的模型示意图。关于模型架构的详细描述，读者可以参考先前的工作【47, Attention is all you need, 2017】【6, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】【37, Better language models and their implications, 2019】。值得一提的是，GPT-2和BERT都在多头注意力和前馈网络的输入端使用GeLU非线性激活函数【9, Bridging nonlinearities and stochastic regularizers with gaussian error linear units, 2016】和层归一化（Layer Normalization）【2, Layernorm, 2016】，而原始的Transformer模型【47, Attention is all you need, 2017】则使用ReLU非线性激活函数，并将层归一化应用于输出端。

深度学习中的数据并行与模型并行

数据并行。将深度神经网络训练扩展到众多硬件加速器上有两种核心范式：数据并行【46, A bridging model for parallel computation, 1990】和模型并行。在数据并行中，一个训练小批量（minibatch）被分割到多个工作节点上。通过将小批量大小与可用工作节点数量成比例增加（即弱扩展），可以观察到训练数据吞吐量近似线性增长。然而，大批量训练会给优化过程带来复杂性，可能导致准确率下降或收敛时间延长，从而抵消了训练吞吐量增加带来的好处【14, On large- batch training for deep learning: Generalization gap and sharp minima, 2017】。后续研究【7, Accurate, large minibatch SGD: training imagenet in 1 hour, 2017】【49, Large batch training of convolutional networks, 2017】【50, Large batch optimization for deep learning: Training bert in 76 minutes, 2019】已开发出相应技术来缓解这些影响，并缩短大型神经网络的训练时间。为了进一步扩展训练规模，一些并行工作【4, Training deep nets with sublinear memory cost, 2016】将数据并行与激活检查点（activation checkpointing）相结合：在反向传播中重新计算激活值，而不是在前向传播中存储它们，以减少内存需求。

模型并行。然而，这些技术存在一个根本限制：模型必须能完全装入单个工作节点。随着BERT和GPT-2等语言模型的规模和复杂性日益增加，神经网络已接近现代硬件加速器的内存容量上限。一个解决方案是采用参数共享来减少模型的内存占用【18, Albert: A lite bert for self-supervised learning of language representations, 2019】，但这限制了模型的整体容量。我们的方法是利用模型并行将模型分割到多个加速器上。这不仅减轻了内存压力，还增加了与微批量大小（microbatch size）无关的并行度。

模型并行的两种范式。在模型并行中，又存在两种范式：层级流水线并行（layer-wise pipeline parallelism）和更通用的分布式张量计算（distributed tensor computation）。在流水线模型并行中，一组操作在一个设备上执行，然后其输出被传递到流水线中的下一个设备，由后者执行另一组操作。一些方法【8, Pipedream: Fast and efficient pipeline parallel dnn training, 2018】【3, Efficient and robust parallel dnn training through model parallelism on multi-gpu platform, 2018】结合了参数服务器【19, Scaling distributed machine learning with the parameter server, 2014】和流水线并行，但存在不一致性问题。GPipe框架【11, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2018】通过使用同步梯度下降克服了这个问题，但这需要额外的逻辑来处理通信和计算操作的高效流水线，并且会产生降低效率的“流水线气泡”（pipeline bubbles），或者需要改变优化器本身，从而影响准确性。

分布式张量计算。分布式张量计算是另一种正交且更通用的方法，它将一个张量操作划分到多个设备上，以加速计算或增加模型大小。FlexFlow【12, Beyond data and model parallelism for deep neural networks, 2018】是一个协调此类并行计算的深度学习框架，它提供了一种选择最佳并行策略的方法。最近，Mesh-TensorFlow【42, Mesh-TensorFlow: Deep learning for supercomputers, 2018】在TensorFlow【1, TensorFlow: Large-scale machine learning on heterogeneous systems, 2015】中引入了一种语言，用于指定一类通用的分布式张量计算。并行维度由终端用户在语言中指定，最终的计算图会被编译并插入适当的集体通信原语（collective primitives）。我们利用了与Mesh-TensorFlow相似的见解，并利用Transformer注意力头计算中的并行性来并行化我们的Transformer模型。然而，我们没有实现一个模型并行的框架和编译器，而是仅对现有的PyTorch Transformer实现进行了少量针对性的修改。我们的方法很简单，不需要任何新的编译器或代码重写，并且可以通过插入几个简单的原语来完全实现，具体将在下一节中描述。

A2 方法细节

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

MLP模块的并行化。我们首先详细说明MLP模块。该模块的第一部分是一个GEMM（通用矩阵乘法），后面跟着一个GeLU非线性激活函数：

class f(torch.autograd.Function):
    def forward(ctx, x):
        return x
    def backward(ctx, gradient):
        all_reduce(gradient)
        return gradient

代码1. f算子的实现。g算子与f类似，其forward函数执行all-reduce，backward函数是恒等映射。

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

输入输出嵌入层的并行化。Transformer语言模型的输出嵌入层维度为隐藏层大小（H）乘以词汇表大小（v）。由于现代语言模型的词汇表大小通常在数万级别（例如，GPT-2使用了50257的词汇表大小），并行化输出嵌入层的GEMM是有益的。然而，在Transformer语言模型中，输出嵌入层与输入嵌入层共享权重，这要求对两者都进行修改。我们沿词汇表维度（列式）将输入嵌入权重矩阵$E_{H \times v}$并行化为$E = [E_1, E_2]$。由于每个分区现在只包含部分嵌入表，因此在输入嵌入之后需要一个all-reduce操作（g算子）。对于输出嵌入，一种方法是执行并行的GEMM $[Y_1, Y_2] = [XE_1, XE_2]$来获得logits，然后进行一次all-gather操作$Y = \text{all-gather}([Y_1, Y_2])$，并将结果发送到交叉熵损失函数。然而，在这种情况下，all-gather操作将通信$b \times s \times v$个元素（b是批量大小，s是序列长度），由于词汇表大小巨大，通信量会非常大。为了减少通信量，我们将并行GEMM的输出$[Y_1, Y_2]$与交叉熵损失函数融合，这将维度降低到$b \times s$。通信标量损失而不是logits，极大地减少了通信量，从而提高了我们模型并行方法的效率。

通信优化与实现总结。我们的模型并行方法很大程度上可以归结为旨在减少通信并保持GPU计算密集的技术。我们没有让一个GPU计算部分dropout、层归一化或残差连接，然后将结果广播给其他GPU，而是选择在所有GPU上重复计算。具体来说，我们在每个GPU上维护层归一化参数的副本，并接收模型并行区域的输出，在这些张量上执行dropout和残差连接，然后将它们作为输入送入下一个模型并行区域。为了优化模型，我们允许每个模型并行的工作节点优化自己的一组参数。由于所有值要么是本地的，要么是在GPU上复制的，因此在这种设计中不需要通信更新后的参数值。关于混合模型和数据并行以及处理随机数生成的更多细节，我们在附录B中提供以供参考。总而言之，我们上述的方法实现简单，只需在前向和反向传播中添加几个额外的all-reduce操作。它不需要编译器，并且与诸如【11, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2018】等方法所倡导的流水线模型并行是正交且互补的。

A4 实验环境

训练数据集

• 数据集构成: 本研究使用了一个聚合数据集，包含维基百科(Wikipedia)【6, Bert: Pre-training of deep bidirectional transformers for language understanding, 2018】、CC-Stories【43, A simple method for commonsense reasoning, 2018】、RealNews【51, Defending against neural fake news, 2019】和OpenWebtext【37, Better language models and their implications, 2019】。为避免数据泄露，移除了WikiText103测试集中的维基百科文章。对于BERT模型，额外加入了BooksCorpus【52, Aligning books and movies: Towards story-like visual explanations by watching movies and reading books, 2015】。
• 数据规模与预处理: 数据集总大小为174GB。所有文档经过筛选，移除了长度小于128个token的文档，并使用局部敏感哈希（LSH）对Jaccard相似度大于0.7的内容进行去重。

训练优化与超参数

• 优化技术: 使用混合精度训练和动态损失缩放以利用V100的Tensor Cores【26, Mixed precision training, 2017】【30, Mixed precision training: Choosing a scaling factor, 2018】。权重初始化采用$W \sim N(0, 0.02)$的正态分布，残差层前的权重会按$1/\sqrt{2N}$进行缩放（N为Transformer层数）。优化器采用带权重衰减（λ=0.01）的Adam【16, Adam: A method for stochastic optimization, 2014】【22, Decoupled weight decay regularization, 2019】。使用全局梯度范数裁剪（值为1.0）以提高训练稳定性。所有模型均使用0.1的dropout率，并采用激活检查点（activation checkpointing）【4, Training deep nets with sublinear memory cost, 2016】来节省内存。
• GPT-2超参数: 序列长度1024，批量大小512，总计训练30万次迭代。学习率为1.5e-4，前3000次迭代进行预热（warmup），之后采用单周期余弦衰减，最小学习率为1e-5。
• BERT超参数: 遵循【18, Albert: A lite bert for self-supervised learning of language representations, 2019】的训练流程，使用原始BERT词典（词汇量30522），用句子顺序预测（sentence order prediction）替代下一句预测（next sentence prediction），并使用n-gram全词掩码（whole word n-gram masking）【13, Spanbert: Improving pre-training by representing and predicting spans, 2019】。批量大小1024，学习率1.0e-4，预热1万次迭代后线性衰减，总计训练200万次迭代。

硬件与软件配置

• 硬件平台: 实验使用了多达32台DGX-2H服务器，共计512块Tesla V100 SXM3 32GB GPU。服务器内部GPU间通过NVSwitch连接，带宽达300GB/s；服务器之间通过每台8个InfiniBand适配器连接，带宽达100GB/s。
• 软件平台: 所有实现均基于PyTorch，通信操作通过Python调用NCCL库完成。

A4 实验结果

扩展性分析

• 实验内容: 为了测试实现的扩展性，研究者对GPT-2模型进行了弱扩展（weak scaling）分析。基线是一个12亿参数模型在单块V100 GPU上的性能，该基线达到了39 TeraFLOPs，是DGX-2H服务器单GPU理论峰值性能的30%，是一个非常强的基线。实验评估了模型并行和模型+数据并行两种场景下的性能。模型参数配置见表1。
• 实验结果:
  • 模型并行: 83亿参数的模型在8路模型并行下（8个GPU），实现了77%的线性扩展效率。
  • 模型+数据并行: 即使在最大的配置（83亿参数，512个GPU，8路模型并行+64路数据并行）下，相对于单GPU基线，也实现了74%的线性扩展效率。
• 分析结论: 本文提出的模型并行方法具有出色的扩展性，能够有效利用大规模GPU集群来训练巨型模型。
  

  

  
  

  

  
  

  

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

• 实验内容: 训练了三种不同规模的GPT-2模型（3.55亿、25亿、83亿参数，见表2），并在WikiText103和LAMBADA数据集上进行零样本（zero-shot）评估。
• 实验结果:
  • 随着模型规模的增加，验证集困惑度（perplexity）单调下降，83亿参数模型收敛到9.27（图6）。
  • 在WikiText103测试集上，83亿参数模型取得了10.81的SOTA困惑度。
  • 在LAMBADA测试集上，83亿参数模型取得了66.51%的SOTA完形填空准确率（表3）。
• 分析结论: 扩大模型规模能显著提升语言模型的性能，并能在多个基准测试上刷新SOTA记录。最近，微软与NVIDIA合作使用Megatron训练了一个170亿参数的GPT-2模型Turing-NLG【27, Turing-nlg: A 17-billion-parameter language model by microsoft, 2020】，进一步证实了模型规模越大，准确率越高的趋势。
  

  

  
  

  

  
  

  

使用BERT的双向Transformer结果

• 实验内容: 研究发现，当BERT模型规模超过BERT-Large时会出现性能退化。通过调整Transformer层中层归一化和残差连接的顺序（如图7b所示），成功解决了这一问题。基于此修改，训练了三种规模的BERT模型（3.36亿、13亿、39亿参数，见表4），并在MNLI、QQP、SQuAD和RACE等下游任务上进行了微调和评估。
• 实验结果:
  • 修改后的架构（图7b）能够稳定地训练更大的BERT模型，并获得更低的训练损失（图7）。
  • 随着模型规模的增加，所有下游任务的性能都得到了单调提升（表5）。
  • 39亿参数的模型在多个开发集上取得了SOTA性能，并在RACE测试集上取得了单模型（90.9%）和集成模型（91.7%）的SOTA准确率。
• 分析结论: 对于BERT类模型，层归一化和残差连接的布局是实现更大规模模型稳定训练和性能提升的关键。通过正确的架构调整，扩大BERT模型的规模能持续带来下游任务性能的提升。
  

  

  
  

  

  
  

  

A5 结论

本文成功地通过实现一种仅需对现有PyTorch Transformer实现进行少量修改的模型并行方法，克服了传统单模型单GPU训练的局限性。研究团队在512个NVIDIA V100 GPU上，使用8路模型并行高效地训练了高达83亿参数的Transformer模型，并在整个应用中实现了15.1 PetaFLOPs的持续计算性能。研究还发现，对于BERT模型，仔细调整层归一化的位置对于随着模型规模增大而提升准确率至关重要。通过研究模型规模对下游任务准确率的影响，本文在WikiText103、LAMBADA和RACE数据集上取得了远超以往的成果，并建立了新的SOTA记录。最后，本文开源了代码以支持未来的研究工作。

未来工作方向包括：
1. 继续扩大预训练规模：这将进一步考验现有的深度学习软硬件。
2. 优化器改进：需要提高优化器的效率和减小其内存占用。
3. 混合并行策略：对于超过160亿参数的模型，需要结合层内、层间模型并行以及跨节点模型并行。
4. 探索其他模型：将该方法应用于不同的模型家族（如XLNet, T5）。
5. 评估更复杂的任务：在更困难和多样化的下游任务（如生成式问答、摘要、对话）上评估大模型的性能。
6. 知识蒸馏：使用这些大型预训练教师模型来训练小型的学生模型。

A6 附录

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。所有通信都是通过Python调用NCCL在PyTorch中实现的。每个模型并行组内的GPU在该组内的所有GPU之间执行all-reduce。对于数据并行，每个all-reduce操作都由来自每个模型并行组的一个GPU参与。

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

Dropout的随机数处理。利用随机数生成的技术，如dropout，是现代深度学习训练的主要组成部分。Transformer模型在模型并行区域之外、残差连接之前，以及在模型并行区域内部的自注意力模块中都有dropout层。由于一些dropout层在模型并行区域内，而另一些则不在，我们需要仔细处理随机数生成以确保dropout正常工作。为了在模型并行工作节点之间同步残差连接的dropout，我们在训练开始时使用相同的种子来初始化随机数生成器。这使得所有模型并行工作节点的dropout模式保持一致。然而，模型并行区域内的dropout应该为每个工作节点产生不同的随机模式，以在整个操作中实现随机性。为了实现这一点，我们为模型并行区域内的dropout维护一个独立的随机数生成器。这个随机数生成器对每个模型并行工作节点使用唯一的种子进行初始化。

C. 文本样本

（此部分为模型生成的示例文本，非技术内容，故此处省略）

D. 进一步的扩展性分析

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

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

D.2. 强扩展性

模型并行对小模型的加速效果。我们的模型并行主要设计用于训练那些超出单个GPU内存容量的大型模型，但它也可以在不增加批量大小的情况下加速小型模型的训练。为了衡量这种加速效果，我们训练一个固定的12亿参数模型。我们使用固定的每轮迭代8个样本的批量大小，并利用模型并行增加GPU数量。结果列于表8。使用两个GPU使训练速度提高了64%。超过这个数量后，我们观察到收益递减，因为每个GPU的计算量减少，内存带宽和通信开销开始占据主导地位。

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

E.1. WikiText103困惑度

困惑度计算方法。WikiText103困惑度是评估语言模型的一个常用指标。困惑度是语料库平均交叉熵的指数【28, Empirical evaluation and combination of advanced language modeling techniques, 2011】。这使其成为语言模型的自然评估指标。我们按如下公式计算困惑度：

Transformer的评估调整。由于Transformer模型在固定窗口大小的输入上操作，我们采用一种称为“重叠评估”的折中方法。我们每次将滑动窗口移动一个重叠量o，并只计算窗口最后o个token对应的交叉熵损失。在我们的实验中，我们使用32的重叠量o。

E.2. LAMBADA完形填空准确率

评估长期依赖能力。像LAMBADA这样的完形填空式数据集旨在衡量模型处理和理解长距离上下文的能力。该任务要求模型预测一段4-5句话的上下文中的最后一个词。我们的模型使用子词单元，因此在LAMBADA评估中，我们使用原始未处理的数据集，并要求模型预测构成目标单词的多个子词token。我们使用teacher forcing，并且只有当所有预测输出都正确时才认为答案正确。这种形式等同于原始的词级别预测任务。