论文标题：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

论文链接：https://arxiv.org/abs/1909.08053

论文来源：NVIDIA

一、概述

随着自然语言处理领域预训练语言模型的规模变得越来越大，它们超过了现代处理器的内存限制，需要额外的内存管理技术，如激活检查点(activation checkpointing)。一些广泛使用的优化算法如Adam需要额外的内存来存储其中的动量和其他优化器状态，这降低了可以有效训练的模型大小。几种模型并行方法通过划分模型来克服这个限制，这样权重及其相关的优化器状态就不需要同时驻留在处理器上。例如，GPipe和Mesh-Tensorflow提供了不同种类的模型并行框架。但是，它们需要重写模型，并依赖于仍在开发中的自定义编译器和框架。

在这项工作中，我们使用简单高效的层内模型并行(intra-layer model-parallelism)来实现模型并行。我们利用transformer基础语言模型中的固有结构来实现一个简单的模型并行，它可以在PyTorch中高效训练，而无需自定义C++代码或编译器。这种方法与GPipe等基于流水线的模型并行方法是正交的(可以同时使用，相互独立而不冲突)。

为了证明我们方法的可扩展性，我们建立了一个baseline，在单个NVIDIA V100 32GB GPU上训练了一个12亿参数的模型，保持39 TeraFLOPs的计算速度。这是DGX-2H服务器中单GPU配置的理论峰值浮点运算能力的30%，因此是一个很强的基线。通过在512个GPU上以8路模型并行将模型扩展到83亿参数，我们实现了高达每秒15.1PetaFLOPs的持续计算速度。与单GPU情况相比，这代表了76%的扩展效率。下图显示了更详细的扩展结果。

拓展效率

为了分析模型大小扩展对准确率的影响，我们训练了自回归的GPT-2语言模型以及自编码的BERT双向transformer，并在几个下游任务上对其进行评估。我们发现，随着模型大小的增加，现有的BERT架构会导致模型退化。我们通过重新排列transformer层中的层标准化和残差连接来克服这一挑战，结果表明，进行这一改变之后，下游任务的开发集结果随着模型大小的增加单调提升。此外，我们的模型在WikiText103、LAMBADA上的闭包式预测准确率以及RACE阅读理解数据集上都取得了SOTA的测试集结果。

二、背景

有两种主要的范式可以将深度神经网络的训练扩展到多个硬件加速器中:(1)数据并行，其中minibatch的训练被划分到多个worker中;(2)模型并行，其将模型的内存使用和计算分布到多个worker中。通过按可用worker的数量成比例增加minibatch大小(即weak scaling)，可以观察到训练数据吞吐量的近乎线性的扩展。但是，大批量训练会使优化过程更复杂，可能导致准确率降低或收敛时间延长，反而抵消了增加训练吞吐量带来的好处。进一步的研究开发了各种技术来缓解这些影响，降低大型神经网络的训练时间。为了进一步扩展训练规模，一些并行工作将数据并行和激活检查点结合起来:在前向传播中重新计算而不是存储激活，以减少内存需求。

然而，这些技术在其可以处理的问题大小方面存在一个根本的局限:模型必须完全能够在一个worker上进行处理。随着BERT和GPT-2等语言模型大小和复杂度的增加，神经网络已经接近了现代硬件加速器的内存容量。这个问题的一个解决方案是采用参数共享来减少模型的内存占用，但这限制了模型的总体容量。我们的方法是利用模型并行将模型划分到多个加速器上。这不仅减轻了内存压力，而且独立于微批量大小增加了并行度。

在模型并行中，还有两种进一步的范式:逐层流水线并行(layer-wise pipeline parallelism)和更通用的分布式张量计算(distributed tensor computation)。在流水线模型并行中，一组操作首先在一个设备上执行，然后将输出传递到流水线中的下一个设备，在下一个设备上执行不同的另一组操作。一些方法与流水线并行结合使用参数服务器。然而，这些方法存在一致性问题。TensorFlow中的GPipe框架通过使用同步梯度下降来解决这种一致性问题。这种方法需要额外的逻辑来处理这些通信和计算操作的高效流水线，并受到减少效率的流水线bubble的影响，或者对优化器本身的更改会影响准确性。

分布式张量计算是一种正交的、更通用的方法，它将张量操作划分到多个设备上以加速计算或增加模型大小。编排这种并行计算的深度学习框架FlexFlow提供了一种选择最佳并行策略的方法。最近，Mesh-TensorFlow在TensorFlow中引入了一种指定分布式张量计算的通用类的语言。并行维度由终端用户在这一语言中指定，生成的图由适当的集体原语编译。我们利用类似于Mesh-TensorFlow中的见解，并利用transformer注意力头的并行计算来并行化我们的transformer模型。但是，我们没有实现一个用于模型并行的框架和编译器，而是仅对现有的PyTorch transformer实现进行了一些有针对性的修改。我们的方法很简单，不需要任何新的编译器或代码重写，可以通过插入几个简单的基元来完全实现，如下一节所述。

三、方法

我们利用transformer网络的结构来创建一个简单的模型并行实现，只需要添加几个同步原语。一个transformer层由一个自注意力块和一个两层多层感知器(MLP)组成，如下图所示。我们在这两个块中分别引入模型并行。

模型架构

首先详细说明MLP块。该块的第一部分是一个GEMM(General Matrix Multiplication)，后面是一个GELU非线性层:

并行化GEMM的一种方法是按行切分权重矩阵，按列切分输入:

，

这将得到。由于GeLU是非线性函数，，这种切分方式需要在GeLU函数之前进行同步。

另一种切分方法是按列切分。这种切分方式允许独立地对每个切分后的GEMM的输出应用GeLU非线性函数:

这种方法的优点是去除了一个同步点。因此，我们采用这种列并行方式切分第一个GEMM，直接将GeLU层的输出作为第二个GEMM(这个GEMM以按行切分的方式并行)的输入，而不需要任何通信，如下图(a)所示。

并行方法

第二个GEMM的输出通过dropout层之前先跨GPU进行reduce。这种方法将MLP块中的两个GEMM切分到不同的GPU上，前向传播中只需要一个all-reduce操作(g操作符)，反向传播中也只需要一个all-reduce(f操作符)。这两个操作符互为共轭，可以通过PyTorch中的几行代码实现。例如，f操作符的实现如下:

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

如上图3(b)所示，对于自注意力块，我们利用了多头注意力操作中固有的并行性，以列并行的方式切分key(K)、query(Q)和value(V)对应的GEMM，这样每个注意力头对应的矩阵乘法在一个GPU上本地计算。这允许我们在GPU之间切分每个注意力头的参数和工作量，并且不需要任何直接的通信就可以完成自注意力。在自注意力之后的输出线性层的GEMM(自注意力之后)沿其行进行并行化，直接获取并行注意力层的输出，而不需要GPU之间的通信。MLP和自注意力这两种块的并行方法都融合了两组GEMM，消除了中间的同步点，从而获得了更好的扩展性。这使我们能够使用仅两个all-reduce在前向路径中完成transformer层中的所有GEMM的计算，并在反向路径中也使用两个all-reduce(如下图所示)。

通信操作

Transformer语言模型的输出嵌入维度为隐层维度乘以词汇表大小。由于词汇表的大小是万级的(例如，GPT-2使用的词汇表大小为50,257)，将输出嵌入GEMM并行化是有益的。但是，在transformer语言模型中，输出嵌入层与输入嵌入共享权重，这需要同时修改这两者。我们沿词典维度切分输入嵌入权重矩阵为 (列向)。由于每个切分部分现在只包含嵌入表的一部分，在输入嵌入之后需要一个all-reduce(g操作符)。对于输出嵌入，一种方法是执行并行GEMM 以获得logits，添加一个all-gather ，并将结果发送到交叉熵损失函数。但是，在这种情况下，all-gather将通信个元素(其中是batch大小，是序列长度)，由于词汇表大小很大，这会产生巨大的通信量。为了减小通信量，我们将并行GEMM 的输出与交叉熵损失函数融合，这将维度降低到。通信标量损失而不是logits极大地减少了通信量，这极大地提高了我们的模型并行方法的效率。

我们的模型并行方法的很大一部分可以归纳为针对减少通信并保持GPU计算限度的技术。对于dropout、层标准化、残差连接等计算，我们选择跨GPU重复计算，而不是让一个GPU计算部分然后广播结果到其他GPU。具体来说，我们在每个GPU上维护层标准化参数的副本，并在模型并行区域的输出上运行dropout和残差连接，然后将其馈送作为下一个模型并行区域的输入。对于模型的优化，我们允许每个模型并行worker优化自己的参数集合。由于所有的值要么是本地的(比如GEMM的参数)，要么被复制在多个GPU上(比如层标准化的参数)，所以在这种形式中不需要通信更新的参数值。

总而言之，我们上述的方法实现起来很监督，只需要在前向和反向传播中添加几个额外的all-reduce操作。它不需要编译器，与GPipe等方法提倡的流水线模型并行是正交互补的。