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

Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (arxiv.org)

概要

megatorn时特定针对语言模型

最近在大预言模型上的工作展示了训练大的Transformer模型提高了自然语言处理应用的技术水平。然而，非常大的模型相当难以训练由于内存限制。在这个工作中，我们提出了我们的训练大的Transformer模型的技术并实现了一个简单有效的intra-layer模型并行方法，使得训练具有十亿个参数的Transformer模型成为可能。我们方法不需要编译器或者库文件的改变，是正交和互补的管道模型并行，可以通过在原生pytorch中插入一些通信操作来完全实现。 我们使用512个GPU通过聚合基于transformer的模型到多达83亿哥参数来说明这种方法。相比于一个能支撑39%TeraFLOPs的单GPU标准，我们在整个应用集群中保持15.1PetaFLOPs（每秒千万亿次浮点运算）,76%的扩展效率。为了演示大语言模型是可以进一步提高最先进技术水平（SOTA），我们训练了83亿个参数的transformer语言模型类似GPT2，和39亿的的参数类似BERT。我们表明，随着模型的增长，在类bert模型中，仔细注意归一化层的位置对于提高性能至关重要。使用GPT2模型我们达到了SOTA在WikiText103（（10.8，而SOTA困惑度为15.8））和LAMBADA（66.5%而SOTA准确率为63.5）数据集上的结果。我们的BERT模型在RACE数据集（90.9%对比SOTA的89.4%的准确率）达到了SOTA的结果。

1 介绍

自然语言处理（NLP）在某种程度上快速发展，这是由于可用计算的增加和数据集大小。富裕的计算和数据集，使得通过无监督的预训练训练更大的语言模型成为可能。经验证据表明，大语言模型是对于NLP任务如文章补，问题回答和自然语言推理全非常有用。通过微调这些下游自然语言任务上的预训练模型，人们可以达到最先进的结果，如同展示的一样。

随着这些模型变大，它们超过了通用处理器的内存限制，需要额外的内存管理技术比如activation checkpointing。广泛被使用的优化算法比如ADAM需要每个参数额外的内存来存储动量和其他的优化状态，这减少了可以有效训练的模型的小大小。模型并行的几种方法客服了这一限制，它们对模型进行分区，让权重和相关的优化器状态不需要在处理器上同时存在。例如，GPipe和Mesh-Tensorflow提供了不同程度的模型框架。然而，它需要重写模型，并且依赖于仍在开发中的定制编译器和框架。

在这个工作中，我们使用层内模型并行实现了简单和有效的模型并行方法。我们利用transformer中基于语言模型的内在结构制作了可以在pytorch中有效训练的简单的模型并行实现，它不需要定制C++代码或者编译器。这种方法与Gpipe等方法所提倡的基于管道的模型性是正交的。

为了演示我们方法的可扩展性，我们通过在一个英伟达V100 32G GPU上训练一个有12亿参数的模型，建立了一个基线，它维持了39TeraFLOPs。这是如DGX-2H服务配置的单个GPU理论峰值的30%，因此是一个强有力的基线。在512个GPU上使用8种模型并行方式扩展模型到83亿个参数，我们达到了整个应用程序中持续每秒15.1PetaFLOPs。与单个GPU相比，这是76%的扩展效率。如图1展示了更详细的扩展结果。

图1：模型（蓝)和模型+数据（绿)并行FLOPS随GPU的数量的变化。模型并行（蓝)：高达8路模型并行弱缩放，每个GPU大概有10亿个参数。模型+数据并行（绿)：接近于模型并行和64路数据并行的配置。

为了分析模型尺寸缩放对于精度的影响，我们训练了left-to-right GPT-2语言模型以及BERT双向transformers并且在几个下游任务上评估他们。我们表明，随着大小的增加，现有的BERT架构会导致模型退化。我们通过重新排列transformer中的层归一化和残差连接，表明通过这种改变，开发集上下游任务的结果随着模型大小的增加而单调改善。此外，我们表明我们的模型在WikiText103上达到了测试集最先进(SOTA)的结果，在LAMBADA上达到了完形风格的预测精度，以及阅读理解RACE数据集。

简单来说，我们的贡献如下。

• 我们通过只对现有的Pytorchtransformer实现进行针对性的修改，实现了简单有效的模型并行方法。
• 我们进行了对于我们模型和数据并行技术进行深入经验的分析，使用512个GPU演示76%的扩展效率。
• 我们表明，在类似BERT的模型中，层归一化的放置位置对随着规模增长而实现更高的准确性至关重要。
• 我们演示了缩放的模型大可以提高GPT2和BERT模型的精确度（研究了多达3.9B个模型参数）。
• 我们展示我们的模型在测试数据集上得到了最先进的结果：WikiText103上的困惑度
• (10.8 ppl)， LAMBADA上的准确率(66.5%)，RACE上的准确率(90.9%)
• 我们开源了我们的代码于训练评估管道https://github.com/NVIDIA/Megatron-LM

2 背景和挑战

2.1 自然语言模型预训练

预训练大语言模型已经变成了NLP研究工具箱中不可或缺的部分。充分利用大的预料预训练以学习健壮的语言自然表现是活跃的研究领域，这种情况已经持续了十年。早期的预训练和传递语言的神经表示演示了，与从头开始学习词嵌入相比，预训练词嵌入改善了下游任务的结果。后来的工作通过学习和转移神经模型来获取单词的上下文推荐了这一领域的研究。最近的拼写工作在这些想法上进一步建立，不仅仅是转移语言模型来提取上下文的单词表示，还可以在下游任务上以端到端方式对语言模型进行微调。通过这些工作，技术水平从仅仅传输词嵌入表发展到传输整个数十亿参数的语言模型。这个方法的进展已经已经不可避免需要硬件系统技术和能够大规模高效运行和满足增长的计算需求的框架。我们的工作瞄准了提供必要的工具以在这趋势中再迈出一步。

2.2 Transformer语言模型和多头注意力

当前的NPL工作趋势向着使用Transformer模型因为它强大的准确和有效的计算。原始的Transformer被构想为一个机器翻译架构，它使用编码器和解码器将输入序列转换位另一个输出序列。然而最近的工作充分利用Transformer进行语言建模比如BERT和GPT2根据需要只是用编码器或者解码器。这个工作探索了解码器架构的GPT2和编码器架构的BERT。

图2：Transformer架构，紫色的块相当于全连接网络。每个蓝色块代表一个被复制N次的变压器层。

bert只有Transformer的编码器而没有解码器

图2展示了我们使用的模型概要图。我们推荐读者参考之前的工作已获得模型详细结构的描述。值得一提的是GPT-2和BERT都使用了GeLU非线性和层归一化到多头注意力和前馈层的输入。然而原始的Transformer使用了ReLU非线性并在输出用了层归一化。

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

弱缩放：是指在保持每个计算节点上处理的数据量不变的情况下，增加计算节点的数量，以评估系统在处理更大规模数据集时的性能和效率。弱缩放通常用于研究并行计算和分布式系统的性能。

弱缩放主要关注的是系统的扩展性，即在增加节点数时，系统的性能是否能按比例提高。例如，假设在两个节点上完成任务的时间为T，如果将节点数增加到四个，期望的时间应为T/2。

强缩放 Strong scaling指的是：每个节点处理的数据减少，节点数增加，保持总数量不变。目标是评估固定数量情况下，增加计算资源时的效率。

将深度神经网络扩展到众多硬件加速器有俩个主要范例：数据并行度，其中一小个批量被分散到多个机器上时，以及模型的并行性，其中被分割到多个机器，以及模型并行性，其中内存的使用模型计算分布在多个机器之间。通过按比例的增加可用机器的最小批大小（即弱缩放），可以观察到训练数据吞吐量接近线性缩放。然而大批训练给优化带来了复杂性，可能导致了准确性和收敛时间变长，增加了训练吞吐量的好处。进一步的研究已经开发出缓解这些影响，并降低了大型神经网络的训练时间。为了进一步扩大训练规模，并行工作组合数据的并行性与激活检查点相结合。在向后传递中重新计算激活，而不将它们存储在向前传递中，以减少内存需求。

模型的并行度（Model Parallelism）是指在训练或推理过程中，将一个机器学习模型的不同部分分布到多个计算资源（如CPU、GPU或多个节点）上，以实现更高效的计算和更快的处理速度。并行度在处理大规模模型或大规模数据时尤为重要。

然而这个技术在解决的问题规模上有一个基本的限制：这种模式必须完全适合只有一个机器。对于逐渐增大尺寸和复杂性的语言模型比如BERT和GPT2，神经网络已经接近了现代硬件加速器的内存容量。这个问题的一个方法是使用参数共享以减少模型的内存占用，但是这个限制了整个模型的容量。我们的方法使用利用模型的并行度以在多个加速器之间分割模型。这不止缓和了内存压力，也增加了独立于微批量大小的平行性。

在模型并行中，流水线级数（pipline stage）指的是模型被划分成多少个连续的子部分。

微批次（micro-batching）：流水线中引入了微批次级数，将大的批次数据拆分。这些批次可以在不同的阶段进行处理。通过微批次，每个阶段开始可以处理下一小批数据，而不必等待前一个阶段完全处理完完整的大批次。

流水线冒泡是指在流水线并行计算中，由于某些阶段的任务等待数据或资源，导致其他阶段无法充分利用计算资源的情况。这种现象会造成流水线中的“空闲”时间段，降低整体系统的效率。

在模型并行性中有俩个进一步的范例：层级流水线并行性和更一般的张量计算。在流水线并行性中，操作组在一个设备上执行然后输出传递给流水线的下一个设备在那里执行另一组操作。一些方法将参数服务器与流水线进行合并使用。然而这些都存在不一致的问题。TensorFlow的GPipe框架通过使用同步梯度下降克服了这一问题。这种方法需要额外的逻辑来处理这些通信和计算操作的高效流水线化，并且会受到流水线气泡的影响，导致效率降低，或者需要对优化器本身进行修改从而影响准确性。

分布式tensor计算是一种正交的并且更一般的方法，它划分tensor操作跨越多个设备以加速计算或者增加模型大小。FlexFlow一个协调这种并行计算的深度学习框架，提供选择最佳化并行策略的方法。最近 Mesh-TensorFlow引入了一种语言，用于在tensorFlopw中指定分布式tensor计算的一般类别。并行维度在语言中由最终用户指定，生成的计算图用适当集群的操作原语进行编译。我们利用了与mesh-TensorFlow中相似的见解，并且通过计算Transformer的注意力头来实现对Transformer模型的并行化。然而，不是实现一个框架爱和编译器来实现模型的并行性，我们对一个现有的pytorch转换器实现做了一定针对性的修改。我们的方法很简单，不需要任何编译器或重写，就可以完全通过增加几个简单的集群进行实现。将会在下一章节进行描述。

3 模型并行

懒得翻译了，总结一下

张量并行是一种并行优化策略。当模型规模大到单个GPU无法处理的时候。张量并行的思想是已将一个张量（神经网络中的参数或数据）按维度拆分到多个设备上进行并行的处理，以提高计算效率

数据并行（Data Parallelism）：每个 GPU 上有相同的模型副本，但不同 GPU 处理不同的数据 mini-batch，每个 GPU 独立计算梯度，最终将梯度聚合。

张量并行（Tensor Parallelism）：不同 GPU 负责相同计算中的不同张量部分，并行计算同一个 mini-batch 的不同部分，最终的计算结果需要通过设备间通信来整合。

模型并行（Model Parallelism）：核心思想是将模型的不同部分（而不是数据或张量的不同部分）拆分到多个设备（通常是GPU）上进行运算

megatron使用的方法其实是张量并行，但在论文中他把这称之为模型并行。因为这篇文章使用了这种表达方式使得之后的文章也会搞混张量和模型并行。

MLP

假设mlp的模块输入是一个x，他的行数为批量大小b乘以序列长度L列是隐藏层宽度也就是K（一般为俩个全连接层）

输入X假设第一个隐藏层权重为A的话，就做矩阵乘法然后是激活函数（relu或者Gelu）。

$\sigma(XA)B=Y$。B在这个公式中充当了激活值映射到输出空间的角色是一个矩阵或者向量。

接下来是A。A的行数是K，列数是K'隐藏层的大小。可以把A和X俩个矩阵一个横切或竖切进行拆分。另一个必须保证是完整的矩阵。我们这里对A进行拆分，也就是说每个GPU上都存有X。（理论上可以一个横切一个竖切。如果最后将这俩个矩阵加起来，能得到最后的结果但是这里的加操作就会浪费带宽。因此不选择这种方案。）如果有俩个gpu它们会分别得到一半完整的结果。然后直接对B算矩阵乘法。原来的矩阵是竖切的，所以B这个矩阵需要横切（这样只需要存储一半的B）。算出来Y后进行相加就得到了完整的Y的结果。

进行allreduce之后所有的GPU拿到了完整的Y的结果。这样每个GPU可以直接进行下一次的运算，而不需要copy各个gpu上的矩阵。中间不需要任何通讯。

自注意力

为了简单起见我们假设输入的批量为1，因为批量不管为多少，它总是只和自己做计算的。

设矩阵X。矩阵行数是l列数为k。自注意力中x会被复制为三份。分别对应的是Q，K，V矩阵。其中h是头的个数。投影之后Q它的维度就会成为：

长为k/h高为l。K和V同理。然后对Q和K的投影做自注意力分数（点积+softmax）。得到的权重和V的投影相加。就会得到输出长为k/h高为l的矩阵。

如果是多个头的话最终会得到$l\times k$的 矩阵，在之后对$k\times k$进行乘法得到Y（$l\times k$）。

所以可以每个头放一个GPU。最后allreduce相加最后输出的结果。

输入输出层

输入层为批量大小乘以序列长度$b\times l$然后是embedding层。它的大小字典的大小v（一般是几万）列数是k（嵌入维度），然后进入字典查找词向量。查找之后的结果为$b\times\ l \times k$的矩阵。

可以把输入x复制到俩个gpu上面，词表切一半。然后查找之后相加就可以输出。

而输出层是反过来的。输入是$b\times l \times k$的矩阵。输入向量来查找词语。然后输出$b\times l \times v$。

这里每个gpu各自算各自的指数。每个gpu把自己那个按行求和。然后求和后的行进行allreduce。这样大大减小了通讯量。

其中所有的allreduce操作都是可以自动求导的。因为它只是一个加法操作。其中每一块的通讯量都是一个固定值都是：

O（$b\times l\times k\times n$）b为小批量大小，l序列长度，乘以k隐藏层大小，n就是层数。

如果做的是数据并行的话就是O($k^2\times n$)所以从通讯量的角度来说，它们是相似的。（k平均是几千到一万）

张量并行好处：每个gpu维护的模型更小了，坏处：保证头的个数和隐藏层k的大小能够被gpu数量整除。此外它没有办法做到数据计算和数据通讯的并行。（数据并行不存在这个问题）

但是数据并行可能无法保存特别大的模型。