[NLP]深入理解 Megatron-LM

一. 导读

NVIDIA Megatron-LM 是一个基于 PyTorch 的分布式训练框架，用来训练基于Transformer的大型语言模型。Megatron-LM 综合应用了数据并行（Data Parallelism），张量并行（Tensor Parallelism）和流水线并行（Pipeline Parallelism）来复现 GPT-3.

在自然语言处理（NLP）领域，大型模型能够提供更精准和强大的语义理解与推理能力。随着计算资源的普及和数据集的增大，模型参数的数量呈指数级增长。然而，训练这样规模庞大的模型面临着一些挑战：

1. 显存限制： 即便是目前最大的GPU主内存也难以容纳这些模型的参数。举例来说，一个1750亿参数的GPT-3模型需要约700GB的参数空间，对应的梯度约为700GB，而优化器状态还需额外的1400GB，总计需求高达2.8TB。
2. 计算挑战： 即使我们设法将模型适应单个GPU（例如通过在主机内存和设备内存之间进行参数交换），模型所需的大量计算操作也会导致训练时间大幅延长。举个例子，使用一块NVIDIA V100 GPU来训练拥有1750亿参数的GPT-3模型，大约需要耗时288年。
3. 并行策略挑战： 不同的并行策略对应不同的通信模式和通信量，这也是一个需要考虑的挑战。

二 . 并行策略和方法简介

1 数据并行

数据并行模式会在每个worker之上复制一份模型，这样每个worker都有一个完整模型的副本。输入数据集是分片的，一个训练的小批量数据将在多个worker之间分割；worker定期汇总它们的梯度，以确保所有worker看到一个一致的权重版本。对于无法放进单个worker的大型模型，人们可以在模型之中较小的分片上使用数据并行。

数据并行扩展通常效果很好，但有两个限制：

• a）超过某一个点之后，每个GPU的batch size变得太小，这降低了GPU的利用率，增加了通信成本；
• b）可使用的最大设备数就是batch size，限制了可用于训练的加速器数量。

2 模型并行

人们会使用一些内存管理技术，如激活检查点（activation checkpointing）来克服数据并行的这种限制，也会使用模型并行来对模型进行分区来解决这两个挑战，使得权重及其关联的优化器状态不需要同时驻留在处理器上。

模型并行模式会让一个模型的内存和计算分布在多个worker之间，以此来解决一个模型在一张卡上无法容纳的问题，其解决方法是把模型放到多个设备之上。

模型并行分为两种：流水线并行和张量并行，就是把模型切分的方式。

• 流水线并行（pipeline model parallel）是把模型不同的层放到不同设备之上，比如前面几层放到一个设备之上，中间几层放到另外一个设备上，最后几层放到第三个设备之上。
• 张量并行则是层内分割，把某一个层做切分，放置到不同设备之上，也可以理解为把矩阵运算分配到不同的设备之上，比如把某个矩阵乘法切分成为多个矩阵乘法放到不同设备之上。

具体如下图，上面是层间并行（流水线并行），纵向切一刀，前面三层给第一个GPU，后面三层给第二个GPU。下面是层内并行（tensor并行），横向切一刀，每个张量分成两块，分到不同GPU之上。

这两种模型切分方式是可以同时存在的，实现正交和互补的效果。

从另一个角度看看，两种切分同时存在，是正交和互补的（orthogonal and complimentary）。

通信分析

对于模型并行的通信状况。

流水线并行：通信在流水线阶段相邻的切分点之上，通信类型是P2P通信，单词通信数据量较少但是比较频繁，而且因为流水线的特点，会产生GPU空闲时间，这里称为流水线气泡（Bubble）。

比如下图之中，上方是原始流水线，下面是流水线并行，中间给出了 Bubble 位置。

 

张量并行：通信发生在每层的前向传播和后向传播过程之中，通信类型是all-reduce，不但单次通信数据量大，并且通信频繁。

张量并行一般都在同一个机器之上，所以通过 NVLink 来进行加速，对于流水线并行，一般通过 Infiniband 交换机进行连接。

2.1 张量并行

张量模型并行化（tensor model parallelism）将每个transformer 层内的矩阵乘法被分割到多个GPU上，虽然这种方法在NVIDIA DGX A100服务器（有8个80GB-A100 GPU）上对规模不超过200亿个参数的模型效果很好，但对更大的模型就会出现问题。因为较大的模型需要在多个multi-GPU服务器上分割，这导致了两个问题。

• 张量并行所需的all-reduce通信需要通过服务器间的链接，这比multi-GPU服务器内的高带宽NVLink要慢；
• 高度的模型并行会产生很多小矩阵乘法（GEMMs），这可能会降低GPU的利用率。

2.2 流水线并行

流水线模型并行化是另一项支持大型模型训练的技术。在流水线并行之中，一个模型的各层会在多个GPU上做切分。一个批次（batch）被分割成较小的微批（micro-batches），并在这些微批上进行流水线式执行。

通过流水线并行，一个模型的层被分散到多个设备上。当用于具有相同transformer块重复的模型时，每个设备可以被分配相同数量的transformer层。Megatron不考虑更多的非对称模型架构，在这种架构下，层的分配到流水线阶段是比较困难的。在流水线模型并行中，训练会在一个设备上执行一组操作，然后将输出传递到流水线中下一个设备，下一个设备将执行另一组不同操作。

原始的流水线并行会有这样的问题：一个输入在后向传递中看到的权重更新并不是其前向传递中所对应的。所以，流水线方案需要确保输入在前向和后向传播中看到一致的权重版本，以实现明确的同步权重更新语义。

模型的层可以用各种方式分配给worker，并且对于输入的前向计算和后向计算使用不同的schedule。层的分配策略和调度策略导致了不同的性能权衡。无论哪种调度策略，为了保持严格的优化器语义，优化器操作步骤需要跨设备同步，这样，在每个批次结束时需要进行流水线刷新来完成微批执行操作（同时没有新的微批被注入）。Megatron-LM引入了定期流水线刷新。

在每个批次的开始和结束时，设备是空闲的。我们把这个空闲时间称为流水线bubble，并希望它尽可能的小。根据注入流水线的微批数量（micro-batches），多达50%的时间可能被用于刷新流水线。微批数量与流水线深度（size）的比例越大，流水线刷新所花费的时间就越少。因此，为了实现高效率，通常需要较大的batch size。

一些方法将参数服务器与流水线并行使用。然而，这些都存在不一致的问题。TensorFlow的GPipe框架通过使用同步梯度下降克服了这种不一致性问题。然而，这种方法需要额外的逻辑来处理这些通信和计算操作流水线，并且会遇到降低效率的流水线气泡，或者对优化器本身的更改会影响准确性。

某些异步和bounded-staleness方法，如PipeMare、PipeDream和PipeDream-2BW完全取消了刷新，但这样会放松了权重更新语义。Megatron会在未来的工作中考虑这些方案。

2.3 技术组合

用户可以使用多种技术来训练大型模型，每种技术都涉及不同的权衡考量。此外，这些技术也可以结合使用。然而，技术的结合可能导致复杂的相互作用，特别是在系统拓扑方面的设计，不仅需要根据算法特点对模型进行合理切割，还需要在软硬件一体的系统架构设计中进行推敲，以实现良好的性能。因此，以下问题显得尤为重要：

如何组合并行技术，以在保留严格的优化器语义的同时，在给定的批量大小下最大限度地提高大型模型的训练吞吐量？

Megatron-LM的开发人员演示了一种名为PTD-P的技术，它结合了流水线、张量和数据并行。这种技术在1000个GPU上训练大型语言模型，以良好的计算性能（达到峰值设备吞吐量的52%）。 PTD-P利用跨多GPU服务器的流水线并行、多GPU服务器内的张量并行和数据并行的组合，利用了在同一服务器和跨服务器的GPU之间具有高带宽链接的优化集群环境，能够训练具有一万亿参数的模型，并具备良好的扩展性。

这种技术示范了如何在大规模分布式系统中充分发挥不同并行技术的优势，以实现高效的大型模型训练。

要实现这种规模化的吞吐量，需要在多个方面进行创新和精心设计：

1. 高效的核实现： 关键是实现高效的核（kernel），使大部分计算操作成为计算绑定而不是内存绑定。这意味着计算任务能够更快地完成，从而提高整体的计算效率。
2. 智能的计算图分割： 针对设备上的计算图进行智能的分割，以减少通过网络传输的数据量。通过将计算分散到多个设备上，不仅减少了数据传输的成本，还可以限制设备的空闲时间，从而提高整体的计算效率。
3. 通信优化和高速硬件利用： 在特定领域实施通信优化，利用高速硬件如先进的GPU，以及在同一服务器内和不同服务器GPU之间使用高带宽链接，可以大幅提升数据传输的速度和效率。这对于分布式系统中的数据交换至关重要。

通过在上述方面进行创新和优化，可以有效地提高大型模型训练的规模化吞吐量，实现更高的训练效率和性能。这需要结合领域专业知识和系统设计，以解决各种挑战并取得成功。

2.4 指导原则

Megatron开发者对不同的并行模式组合以及其之间的影响进行了研究，并总结出了分布式训练的一些指导原则：

1. 并行模式的相互作用： 不同的并行化策略之间以复杂的方式相互影响。并行模式的选择会影响通信量、计算核的效率以及由于流水线刷新（流水线气泡）而导致的worker空闲时间。例如，张量模型并行在多GPU服务器上表现良好，但对于大型模型，最好采用流水线模型并行。
2. 流水线并行的调度影响： 用于流水线并行的调度方式会影响通信量、流水线气泡的大小以及存储激活所需的内存。Megatron提出了一种新的交错调度方式，相比先前的调度方式，它在稍微增加内存占用的基础上，可以提高多达10%的吞吐量。
3. 超参数的影响： 超参数的值，如微批量大小（microbatch size），会影响内存占用、在worker上执行的核效果以及流水线气泡的大小。
4. 通信密集性： 分布式训练是通信密集型的过程。使用较慢的节点间连接或者更多的通信密集型分区会限制性能表现。

综合上述指导原则，Megatron开发者通过深入研究不同并行技术的相互作用，超参数的调优以及通信密集性等因素，为分布式训练提供了更加明确的方向，以实现更高效的大型模型训练吞吐量。

Megatron在训练拥有万亿参数的大型模型时，采用了PTD-P（Pipeline, Tensor, and Data Parallelism）方法，从而实现了高度聚合的吞吐量（502 petaFLOP/s）。

在该方法中，Tensor模型并行用于intra-node transformer层，这使得在基于HGX系统的平台上能够高效运行。同时，Pipeline模型并行则被应用于inter-node transformer层，充分利用了集群中多网卡的设计，提升了模型训练的效率。

除此之外，数据并行也在前述两种并行策略的基础上进行了加强，从而使得训练能够扩展到更大规模，并且实现更快的训练速度。

三 张量模型并行（Tensor Model Parallelism）

3.1 原理

矩阵乘的张量模型并行充分利用矩阵分块乘法的原理。举例来说，要实现如下矩阵乘法Y=X*A，其中X是维度为MxN的输入矩阵，A是维度为NxK的参数矩阵，Y是结果矩阵，维度为MxK。如果参数矩阵A非常大，甚至超出单张卡的显存容量，那么可以把参数矩阵A切分到多张卡上，并通过集合通信汇集结果，保证最终结果在数学计算上等价于单卡计算结果。

什么是GEMM？ 它的英文全称是 GEneral Matrix to Matrix Multiplication (通用矩阵的矩阵乘法)

这里通过 GEMM 来看看如何进行模型并行，这里要进行的是 XA=Y ，对于模型来说，  X是输入， A是权重， Y 是输出。从数学原理的角度来看，对于神经网络中的线性层（Linear层），可以将其看作是将输入矩阵分块进行计算，然后将计算结果合并成输出矩阵。这个过程涉及矩阵乘法和加法操作，其中矩阵乘法涉及到权重矩阵和输入数据之间的乘法，然后再加上偏置向量。

对于非线性层（例如激活函数层），通常不需要进行额外的设计。这些层的计算过程是基于输入数据应用某种非线性函数，例如ReLU（修正线性单元）、Sigmoid、Tanh等。这些函数在数学上是已知的，只需要将输入数据传递给这些函数，然后得到输出。

整体来看，神经网络的计算可以被抽象为一系列的矩阵和向量操作，其中线性层涉及矩阵乘法和加法，而非线性层涉及特定的函数计算。这些操作在深度学习框架中会被高度优化，以提高计算效率和训练速度。

3.2 行并行（Row Parallelism）

我们先看看Row Parallelism，就是把 A 按照行分割成两部分。为了保证运算，同时我们也把 X 按照列来分割为两部分，这里 X1​的最后一个维度等于 A1​ 最前的一个维度，理论上是：

所以，X1​和 A1​ 就可以放到GPU1之上计算，X2​ 和 A2​ 可以放到 GPU2 之上，然后把结果相加。

 我们接下来进行计算。第一步是把图上横向红色箭头和纵向箭头进行点积，得到Y中的绿色。

 

 得出了绿色Y1 与蓝色的 Y2，此时，可以把 Y1,Y2 加起来，得到最终的输出 Y 。

为了满足矩阵乘法规则，输入矩阵X需要按列切分X=[X1 | X2]。同时，将矩阵分块，分别放置在两张卡上，每张卡分别计算Y1=X1*A1，Y2=X2*A2。计算完成后，通过collective通信Allreduce_sum，归约其他卡上的计算结果，可以得到最终的结果矩阵Y。同样，这种切分方式，既可以保证数学上的计算等价性，并解决单卡显存无法容纳，又可以保证单卡通过拆分方式可以装下参数A的问题。

(1) forward

用N来表示GPU的数量。有几块GPU，就把W按行维度切成几份。下图展示了N=2时的切割方式：

W按照行维度切开后，X的维度和它不对齐了，这可怎么做矩阵乘法呢？很简单，再把X“按列切开”就行了，如下图所示：

(2) backward

做完forward，取得预测值Y，进而可计算出损失L，接下来就能做backward了。我们重画一下forward的过程，并在其中加入backward的部分，整体流程图如下：

3.2 列并行（Column Parallelism）

另外一种并行方式Column Parallelism，就是把 A按照列来分割。

最终计算结果如下： 

分别将A1，A2放置在两张卡上。两张卡分别计算Y1=X*A1和Y2=X*A2。计算完成后，通过collective通信AllGather（一种跨GPU卡的通信方式），获取其它卡上的计算结果，拼接在一起得到最终的结果矩阵Y。综上所述，通过将单卡显存无法容纳的矩阵A拆分，放置在两张卡上，并通过多卡间通信，即可得到的最终结果。该结果在数学上与单卡计算结果上完全等价。

（1）forward

按列切分权重后，forward计算图如下：

（2）backward

介绍完了“按行”和“按列”切分权重的方法。在Megatron-LM中，权重的切分操作就是由这两个基础算子组合而成的。接下来，针对Transformer模型，我们依次来看在不同的部分里，Megatron-LM是怎么做切分的。

四 Transformer 张量并行

Megatron-LM 提出了一种针对 Transformer 结构的层内切分方式，可以把每一层的参数切分放在不同 GPU 上。

具体来说，Transformer 结构可以分成三种 block：MLP，self attention layer 以及 embedding layer (& output layer)。下面看看针对这三种 block，Megatron-LM 如何设计切分模式

这里Transformer的模型并行，特指层内切分，即 Tensor Model Parallel。

4.1 Transformer

自从2018年Google的Attention论文推出之后，近年的模型架构都是在 Transformer基础之上完成，模型有多少层，就意味着模型有多少个Transformer块，所以语言模型的计算量主要是Transformer的计算，而Transformer本质上就是大量的矩阵计算，适合GPU并行操作。

Transformers层由一个Masked Multi Self Attention和Feed Forward两部分构成，Feed Forward 部分是一个MLP网络，由多个全连接层构成，每个全连接层是由矩阵乘操作和GeLU激活层或者Dropout构成。

Transformer的每个块包含以下两个主要部分：

1. Masked Multi-Head Self Attention： 这是Transformer中的关键机制，用于建立输入序列中各个位置之间的关系。它涉及多个头（head）的自注意力计算，其中每个头会学习不同的上下文关系。在计算过程中，涉及到大量的矩阵乘法操作，这些操作可以被高效地并行执行。
2. Feed Forward Neural Network： 这部分也被称为“位置前馈网络”。它包含多个全连接层，每个全连接层都涉及矩阵乘法、激活函数（通常是GeLU）和可能的Dropout层。这些操作也是高度并行化的，可以在GPU上迅速执行。

Megatron 的 Feed Forward 是一个两层多层感知器（MLP），第一层是从 H变成4H，第二层是从 4H 变回到 H，所以Transformer具体架构如下，紫色块对应于全连接层。每个蓝色块表示一个被复制N次的transformer层，红色的 x L 代表此蓝色复制 L 次。

4.2 切分Transformer

分布式张量计算是一种正交且更通用的方法，它将张量操作划分到多个设备上，以加速计算或增加模型大小。FlexFlow是一个进行这种并行计算的深度学习框架，并且提供了一种选择最佳并行化策略的方法。最近，Mesh TensorFlow引入了一种语言，用于指定TensorFlow中的一般分布式张量计算。用户在语言中指定并行维度，并使用适当的集合原语编译生成一个计算图。我们采用了Mesh TensorFlow的相似见解，并利用transformer's attention heads 的计算并行性来并行化Transformer模型。然而，Megatron没有实现模型并行性的框架和编译器，而是对现有的PyTorch transformer实现进行了一些有针对性的修改。Megatron的方法很简单，不需要任何新的编译器或代码重写，只是通过插入一些简单的原语来完全实现，

Megatron就是要把 Masked Multi Self Attention 和Feed Forward 都进行切分以并行化，利用Transformers网络的结构，通过添加一些同步原语来创建一个简单的模型并行实现。

4.2.1  切分MLP

上图第一个是 GeLU 操作，第二个是 Dropout操作，具体逻辑如下：

1. MLP的整个输入 X 通过 f 放置到每一块 GPU 之上。

2. 对于第一个全连接层：

• 使用列分割，把权重矩阵切分到两块 GPU 之上，得到 A1,A2。
• 在每一块 GPU 之上进行矩阵乘法得到第一个全连接层的输出 Y1和 Y2​。

3. 对于第二个全连接层：

1. 使用行切分，把权重矩阵切分到两个 GPU 之上，得到 B1,B2​。
2. 前面输出 Y1​ 和 Y2​ 正好满足需求，直接可以和 B 的相关部分（B1​,B2​）做相关计算，不需要通信或者其他操作，就得到了 Z1​,Z2​。分别位于两个GPU之上。

4. Z1,Z2通过 g 做 all-reduce（这是一个同步点），再通过 dropout 得到了最终的输出 Z。

然后在GPU之上，第二个GEMM的输出在传递到dropout层之前进行规约。这种方法将MLP块中的两个GEMM跨GPU进行拆分，并且只需要在前向过程中进行一次 all-reduce 操作（g 操作符）和在后向过程中进行一次 all-reduce 操作（f 操作符）。

这两个操作符是彼此共轭体，只需几行代码就可以在PyTorch中实现。作为示例，f 运算符的实现如下所示：

     f算子的实现。g类似于f，在后向函数中使用identity，在前向函数中使用all-reduce。

 其中，GELU是激活函数，A和B分别为两个线性层。在Transformer里，一般设h' = 4h。假设现在有N块GPU，我们要把MLP层的权重拆到上面做计算，要怎么拆分呢？Megatron提供的拆分办法如下：

为什么我们对A采用列切割，对B采用行切割呢？这样设计的原因是，我们尽量保证各GPU上的计算相互独立，减少通讯量。对A来说，需要做一次GELU的计算，而GELU函数是非线形的，它的性质如下：

也就意味着，如果对A采用行切割，我们必须在做GELU前，做一次AllReduce，这样就会产生额外通讯量。但是如果对A采用列切割，那每块GPU就可以继续独立计算了。一旦确认好A做列切割，那么也就相应定好B需要做行切割了。

Transformer中的FFN结构均包含两层全连接（FC）层，即存在两个矩阵乘，这两个矩阵乘分别采用上述两种切分方式，如下图所示。对第一个FC层的参数矩阵按列切块，对第二个FC层参数矩阵按行切块。这样第一个FC层的输出恰好满足第二个FC层数据输入要求（按列切分），因此可以省去第一个FC层后的AllGather通信操作。

MLP层的通讯量分析

分析可知，MLP层做forward时产生一次AllReduce，做backward时产生一次AllReduce。AllReduce的过程分为两个阶段，Reduce-Scatter和All-Gather，每个阶段的通讯量都相等。现在设每个阶段的通讯量为 Φ ，则一次AllReduce产生的通讯量为 2Φ 。MLP层的总通讯量为 4Φ 。

根据上面的计算图，我们也易知， Φ=b∗s∗h

4.2.2 切分Self-Attention层

现在，我们来看稍微复杂一点的self-attention层切割方式（Transformer中Encode和Decoder之间还有做cross-attention，但计算逻辑和self-attention一致，因此这里只拿self-attention举例）。
 

attention的多头计算简直是为张量模型并行量身定做的，因为每个头上都可以独立计算，最后再将结果concat起来。也就是说，可以把每个头的参数放到一块GPU上。

如下图所示。

• 首先，对于自我注意力块，Megatron 利用了多头注意力操作中固有的并行性，以列并行方式对与键（K）、查询（Q）和值（V）相关联的GEMM进行分区，从而在一个GPU上本地完成与每个注意力头对应的矩阵乘法。这使我们能够在GPU中分割每个attention head参数和工作负载，每个GPU得到了部分输出。
• 其次，对于后续的全连接层，因为每个GPU之上有了部分输出，所以对于权重矩阵B就按行切分，与输入的 Y1,Y2 进行直接计算，然后通过 g 之中的 all-reduce 操作和Dropout 得到最终结果 Z。

       具有模型并行性的transformer块。f和g是共轭的。f在前向传播中使用一个identity运算符，在后向传播之中使用了all reduce，而g在前向传播之中使用了all reduce，在后向传播中使用了identity运算符。

       对三个参数矩阵Q，K，V，按照“列切割”，每个头放到一块GPU上，做并行计算。对线性层B，按照“行切割”。切割的方式和MLP层基本一致，其forward与backward原理也一致，这里不再赘述。
      最后，在实际应用中，并不一定按照一个head占用一块GPU来切割权重，我们也可以一个多个head占用一块GPU，这依然不会改变单块GPU上独立计算的目的。所以实际设计时，我们尽量保证head总数能被GPU个数整除。

一个典型的模型并行版本的 Multiheads Attention 如下代码所示：

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5
 
        self.world_size = get_tensor_model_parallel_world_size() # 模型并行路数
        self.dim_per_partition = divide(dim, self.world_size)
        self.dim_per_attention_head = divide(dim, num_heads)
        self.num_heads_per_partition = divide(num_heads, self.world_size)
 
        self.qkv = ColumnParallelLinear(dim, dim * 3, gather_output=False) # gather_output=False，延迟 gather 的时间，和后面的 RowParallelLinear 组合起来
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = RowParallelLinear(dim, dim, input_is_parallel=True) # input_is_parallel=True，输入已经是按列切分的了
        self.proj_drop = nn.Dropout(proj_drop)
 
    def forward(self, x):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, self.num_heads_per_partition, 3, self.dim_per_attention_head).permute(3, 0, 2, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
 
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
 
        x = (attn @ v).transpose(1, 2).reshape(B, N, self.dim_per_partition)
        x = self.proj_drop(self.proj(x))
        return x

4.2.3 Embedding层

对于Embedding算子，如果总的词表非常大，会导致单卡显存无法容纳Embedding层参数。举例来说，当词表数量是50304，词表表示维度为5120，类型为FP32，那么整层参数需要显存大约为50304*5120*4/1024/1024=982MB，反向梯度同样需要982MB，仅仅存储就需要将近2GB。对于Embedding层的参数，可以按照词的维度切分，即每张卡只存储部分词向量表，然后通过AllReduce汇总各个设备上的部分词向量结果，从而得到完整的词向量结果。

上图描述了单卡Embedding和Embedding两卡张量模型并行的示意图。

在单卡上，执行Embedding操作，bz是batch size大小，Embedding的参数大小为[word_size, hidden_size]，计算得到[bz, hidden_size]张量。下图为Embedding张量模型并行示例，其将Embedding参数沿word_size维度，切分为两块，每块大小为[word_size/2, hidden_size]，分别存储在两个设备上，即每个设备只保留一半的词表。当每张卡查询各自的词表时，如果无法查到，则该词的表示为0，各自设备查询后得到[bz, hidden_size]结果张量，最后通过AllReduce_Sum通信，跨设备求和，得到完整的全量结果，可以看出，这里的输出结果和单卡执行的结果一致。

输入层Embedding

Embedding层一般由两个部分组成：

• word embedding：维度(v, h)，其中v表示词表大小。
• positional embedding：维度(max_s, h)，其中max_s表示模型允许的最大序列长度。

对positional embedding来说，max_s本身不会太长，因此每个GPU上都拷贝一份，对显存的压力也不会太大。但是对word embedding来说，词表的大小就很客观了，因此需要把word embedding拆分到各个GPU上，具体的做法如下：

详细说明下这张图。对于输入X，进入word embedding的过程，就是等于用token的序号去word embedding中查找对应词向量的过程。例如，输入数据为[0, 212, 57, 9]，数据中的每一个元素代表词序号，我们要做的就是去word embedding中的0，212，57，9行去把相应的词向量找出来。

假设词表中有300个词，现在我们将word embedding拆分到两块GPU上，第一块GPU维护词表[0, 150)，第二块GPU维护词表[150, 299)。当输入X去GPU上查找时，能找到的词，就正常返回词向量，找到不到就把词向量中的全部全素都置0。按此方式查找完毕后，每块GPU上的数据做一次AllReduce，就能得到最终的输入。

对于 embedding layer，在 V（词库大小）维度上进行切分。输入 tokens 分别在各个 GPU 上“查字典”。如果该 GPU 维护了这个 token 对应的向量，那么记录该向量；如果没有，结果就是 0.

最后把各个 GPU 的结果进行 All Reduce，每个 GPU 就都得到了所有 tokens 的向量表示

例子中，第一块GPU的查找结果为[ok, 0, ok, ok]，第二块为[0, ok, 0, 0]，两个向量一相加，变为[ok, ok, ok, ok]

通信成本：依旧是 b*l*k

输出层Embedding

语言模型中，output layer 一般与 embedding layer 共享参数。

输出层中，同样有一个word embedding，把输入再映射回词表里，得到每一个位置的词。一般来说，输入层和输出层共用一个word embeding。其计算过程如下：

必须时刻保证输入层和输出层共用一套word embedding。

在backward的过程中，我们在输出层时会对word embedding计算一次梯度，在输入层中还会对word embedding计算一次梯度。在用梯度做word embedding权重更新时，我们必须保证用两次梯度的总和进行更新。

值得注意的是，为了节约通讯成本，我们不对 Y 直接进行 All Reduce，因为 V 太大了，通常能够达到几万。由于之后要计算 softmax 得到概率，所以先在每个 GPU 上计算指数，并且按照行加和，对该结果进行 All Reduce。

这一步的通信成本只有 b*l

当模型的输入层到输入层都在一块GPU上时（即流水线并行深度=1），我们不必担心这点（实践中大部分用Megatron做并行的项目也是这么做的）。但若模型输入层和输出层在不同的GPU上时，我们就要保证在权重更新前，两块GPU上的word embedding梯度做了一次AllReduce。

4.2.5 CrossEntropy层

分类网络最后一层一般会选用softmax和cross_entropy算子来计算交叉熵损失。如果类别数量非常大，会导致单卡显存无法存储和计算logit矩阵。针对这一类算子，可以按照类别数维度切分，同时通过中间结果通信，得到最终的全局的交叉熵损失。

回顾输出层过完embedding后的样子：

正常来说，我们需要对Y1和Y2做一次All-Gather，把它们concat起来形成Y，然后对Y的每一行做softmax，就可得到对于当前位置来说，每个词出现的概率。接着，再用此概率和真值组做cross-entropy即可。

但是All-Gather会产生额外的通讯量 b∗s∗v 。当词表v很大时，这个通讯开销也不容忽视。针对这种情况，可以做如下优化：

• 每块GPU上，我们可以先按行求和，得到各自GPU上的GPU_sum(e)
• 将每块GPU上结果做AllReduce，得到每行最终的sum(e)，也就softmax中的分母。此时的通讯量为 b∗s
• 在每块GPU上，即可计算各自维护部分的e/sum(e)，将其与真值做cross-entropy，得到每行的loss，按行加总起来以后得到GPU上scalar Loss。
• 将GPU上的scalar Loss做AllReduce，得到总Loss。此时通讯量为N。

这样，我们把原先的通讯量从 b∗s∗v 大大降至 b∗s+N 。

4.2.5 通信成本

来自线性层（在 self attention 层之后）输出的后续GEMM会沿着其行实施并行化，并直接获取并行注意力层的输出，而不需要GPU之间的通信。这种用于MLP和自我注意层的方法融合了两个GEMM组，消除了中间的同步点，并导致更好的伸缩性。这使我们能够在一个简单的transformer层中执行所有GEMM，只需在正向路径中使用两个all-reduce，在反向路径中使用两个all-reduce。

类比于MLP层，self-attention层在forward中做一次AllReduce，在backward中做一次AllReduce。总通讯量也是 4Φ ，其中 Φ=b∗s∗h

 Transformer语言模型输出了一个嵌入，其维数为隐藏大小（H）乘以词汇量大小（v）。由于现代语言模型的词汇量约为数万个（例如，GPT-2使用的词汇量为50257），因此将嵌入GEMM的输出并行化是非常有益的。然而，在transformer语言模型中，想让输出嵌入层与输入嵌入层共享权重，需要对两者进行修改。

我们沿着词汇表维度 E=[E1，E2]（按列）对输入嵌入权重矩阵E(h*v)​进行并行化。因为每个分区现在只包含嵌入表的一部分，所以在输入嵌入之后需要一个all-reduce（g操作符）。对于输出嵌入，一种方法是执行并行 GEMM[Y1，Y2]=[XE1，XE2]以获得logit，然后添加一个all-gather Y=all−gather([Y1，Y2])，并将结果发送到交叉熵损失函数。但是，在这种情况下，由于词汇表的很大，all-gather 将传递b×s×v个元素（b是batch size，s是序列长度）。为了减小通信规模，我们将并行GEMM[Y1，Y2]的输出与交叉熵损失进行融合，从而将维数降低到b×s.

需要特别考虑的是:  由于前向和后向传播中每层都有两个 all reduce，因此 TP 需要设备间有非常快速的互联。因此，除非有一个非常快的网络，否则不建议跨多个节点进行 TP。

训练 BLOOM 的硬件配置中，节点间的速度比 PCIe 慢很多。实际上，如果节点有 4 个 GPU，则最高 TP 度设为 4 比较好。如果需要 TP 度为 8，则需要使用至少有 8 个 GPU 的节点。

Megatron-LM 针对 Transformer 结构设计了张量并行。不管是哪一种 block（MLP，自注意力或者是 embedding 层）它都是把层内的参数平均切分到各个 GPU 单元上，并且在每一层都需要通过 All Reduce 获取完整的输出。

可以看出 Megatron-LM 的几个不足之处：

1. 计算和通讯不能并行：必须等待该层的 All Reduce 通讯完成之后，才能开始下一层的计算。对于带宽的要求很高，否则性能会急剧下降
   论文中实验所用的机器，同一机器内 GPU 的通讯带宽能达到 300 GB/s；不同机器间 GPU 通讯带宽也有 100 GB/s. 这是非常豪华的机器配置。
2. 尽管每个 GPU 只保留每层中的一部分参数，但是它要保存每一层完整的输入，这一部分的内存占用也不小
3. 张量的切分方式专为 Transformer 结构设计，具有较强的专用性

五 流水线并行与混合并行

5.1 流水线模型并行（Pipeline Model Parallelism）

目前主流的流水线并行方法包括了两种：Gpipe和PipeDream。与这两者相比，Megatron中的流水线并行实现略有不同，它采用了Virtual Pipeline的方法。简而言之，传统的流水线并行通常会在一个设备上放置几个模块，通过在计算强度和通信强度之间取得平衡来提高效率。然而，虚拟流水线则采取相反的策略。在设备数量不变的前提下，它将流水线阶段进一步细分，以承载更多的通信量，从而降低空闲时间的比率，以缩短每个步骤的执行时间。

5.2 混合并行设置（Mix Parallelism）

参考的 Megatron 的论文，先对使用的符号做一个说明。

5.2.1  张量和流水线模型并行 （Tensor and Pipeline Model Parallelism） 

不同GPU之间通信量也受 p和 t 的影响。流水线模型并行具有开销更小的点对点通信；另一方面，张量模型并行性使用更消耗带宽的all-reduce通信（向前和向后传递中各有两个all-reduce操作）。

• 使用流水线并行，在每对连续设备（向前或向后传播）之间为每个微批次（micro-batch）执行的通信总量为 b⋅s⋅h ， s 是序列长度， h是隐藏大小（hidden size）。
• 使用张量模型并行，每个层前向传播和后向传播中，总大小 b⋅s⋅h 的张量需要在 t个模型副本之中 all-reduce 两次。

因此，这里看到张量模型并行性增加了设备之间的通信量。当 t大于单个节点中的GPU数量时，在较慢的节点间链路上执行张量模型并行是不合算的。

所以，当考虑不同形式的模型并行时，当使用 g台-GPU服务器，通常应该把张量模型并行度控制在 g 之内，然后使用流水线并行来跨服务器扩展到更大的模型。

 5.2.2 数据和模型并行 （Data and Model Parallelism）

数据并行和流水线并行

 数据并行和张量并行

 5.2.3 Microbatch Size

 

5.2.4 对比

Tensor vs Pipeline Parallelism.

在节点内部（如DGX A100服务器），张量模型的并行性表现最佳，因为这可以降低通信量。另一方面，流水线模型并行性采用更经济的点对点通信方式，可以跨节点执行，而不会受到整个计算的限制。然而，流水线并行性可能会在流水线“气泡”中消耗大量时间，因此应该限制流水线的总数，以确保流水线中微批次（micro-batches）的数量是流水线深度的合理倍数。

因此，当张量模型的并行大小等于单个节点中GPU的数量（例如DGX A100有8个GPU的节点）时，性能会达到峰值。这一结果表明，仅使用张量模型的并行性（如Megatron V1）或仅使用流水线模型的并行性（如PipeDream），都无法与将这两种技术结合使用时的性能相媲美。

Pipeline vs Data Parallelism.

通过实验观察发现，针对每个批次大小（batch size），随着流水线并行规模的增加，吞吐量逐渐减少。因此，流水线模型并行的主要应用场景是支持不适合单个处理单元的大型模型训练，而数据并行则适用于扩大训练规模。

Tensor vs Data Parallelism.

接下来我们来看看数据并行性和张量模型并行性对性能的影响。在处理较大批次量和微批次为1的情况下，数据并行通信并不频繁；而张量模型并行则需要对批次中的每个微批次进行全对全（all-to-all）通信。这种全对全通信在张量模型并行中占据主导地位，对整个端到端的训练时间产生影响，尤其是当通信需要跨多个GPU节点进行时。此外，随着张量模型并行规模的增加，每个GPU上执行的较小矩阵乘法也降低了每个GPU的利用率。

需要注意的是，虽然数据并行可以有效地扩展训练，但不能仅凭数据并行来处理训练批次受限的大型模型，原因如下：a）内存容量不足，b）数据并行的扩展受限（例如，GPT-3的训练批次为1536，因此数据并行仅支持最多1536个GPU的并行化；然而，该模型的训练涉及约10000个GPU）。

六  张量模型并行 + 数据并行

在实际应用中，常常将模型并行与多种并行方法结合在一起。对Transformer类的模型，采用最经典方法是张量模型并行 + 数据并行，并在数据并行中引入ZeRO做显存优化。具体的架构如下：

在 Megatron 中，一般在单机内使用模型并行，在机器之间使用数据并行。

以下图为例，一共8张卡，2路数据并行，4路模型并行，则在数据并行组（data parallel group） 内部卡的序号从0到3，一共两组，在模型并行组（model parallel group）内部卡的序号从0到1，一共四组。在通信初始化的时候，这些卡号和组别就需要被定义清楚。在训练中，一开始的参数对齐和每轮迭代时的梯度对齐都需要指定对应组别。

在同一数据并行组内，与单纯的数据并行方式相比，几乎没有发生变化。以PyTorch为例，数据并行模型一般采用 DistributedDataParallel，有一种方式是将其参数 process_group 替换成数据并行组。

# model 的 ddp 要传入 data_parallel_group 作为它的 process_group：在初始化 broadcast 阶段和 average_gradient 阶段起作用
model = DistributedDataParallel(model, process_group=data_parallel_group)

下面解释一下同一模型并行组的行为。 首先，为了保证在分布式训练中不同卡上的参数在初始化时保持同步，除了上面隐含在 DDP 中的数据并行组里的参数 broadcast，需要额外在模型并行组内做一次非并行layer的参数 broadcast。其次，在同一个模型并行组内，处理的数据是同一份，这不仅要求在 data sampler 里保证，为了防止输入到网络中的数据有不同的变形，一般也在 forward 之前，在组内对数据做 broadcast。这样，在非并行的 layer，比如 patch embedding，组内执行的是完全一致的行为，在并行的 layer，比如包含 parallel attention 和 parallel mlp 的 parallel transformer layer，组内处理的仍然是同一份数据，但是参数已经切片，所以得到的是部分结果，再通过通信组合到一起。

Megatron 将上面所提到的功能合在一起写在了一个类似于 DistributedDataParallel的 class 里面。

其中，node表示一台机器，一般在同一台机器的GPU间做张量模型并行。在不同的机器上做数据并行。图中颜色相同的部分，为一个数据并行组。凭直觉，我们可以知道这么设计大概率和两种并行方式的通讯量有关。具体来说，它与TP和DP模式下每一层的通讯量有关，也与TP和DP的backward计算方式有关。

TP与DP计算backward的方式
TP在从上一层往下一层做backward的过程中，所有GPU间需要做一次AllReduce的。例如下图：

而对DP(数据并行)来说，本层算完梯度以后，就正常把本层的梯度发出去，和属于一个DP组的GPU做AllReduce，同时继续往下一层做backward。下一层也是同理。也就是在DP组中，下一层不依赖上一层的梯度聚合结果。因此在DP组中对带宽的要求就没那么高了。所以可以放到机器间做DP。例如下图：

七 总结

模型并行方法旨在减少通信和控制GPU计算范围的。不是让一个GPU计算dropout、layer normalization或 residual connection，并将结果广播给其他GPU，而是选择跨GPU复制计算。

模型并行性与数据并行性是正交的，因此我们可以同时使用二者在来训练大型模型。下图显示了一组用于混合模型并行和数据并行性的GPU。

• 一个模型需要占据8张卡，模型被复制了64分，一共启动了512个即成。
• 模型并行。同一服务器内的多个GPU形成模型并行组（model parallel group），例如图中的GPU 1到8，并包含分布在这些GPU上的模型实例。其余的GPU可能位于同一台服务器内，也可能位于其他服务器中，它们运行其他模型并行组。每个模型并行组内的GPU执行组内所有GPU之间的all-reduce。
• 数据并行。在每个模型并行组中具有相同位置的GPU（例如图中的GPU 1，9，…，505）形成数据并行组（data parallel group），即，具有相同模型参数的进程被分配到同一个数据并行组之中。对于数据并行，每个all-reduce操作在每个模型并行组中一个GPU之上执行。
• 所有通信都是通过pytorch调用NCCL来实现的。

在反向传播过程中，我们并行运行多个梯度all-reduce操作，以规约每个不同数据并行组中的权重梯度。所需GPU的总数是模型和数据并行组数量的乘积。

Megatron使用了PTD-P（节点间流水线并行、节点内张量并行和数据并行）在训练具有万亿参数的大型模型时候达到了高聚合吞吐量（502 petaFLOP/s）。

• Tensor模型并行被用于intra-node transformer 层，这样在HGX based系统上高效运行。
• Pipeline 模型并行被用于inter-node transformer 层，其可以有效利用集群中多网卡设计。
• 数据并行则在前两者基础之上进行加持，使得训练可以扩展到更大规模和更快的速度。

深入理解 Megatron-LM（1）基础知识 - 知乎 (zhihu.com)

深入理解 Megatron-LM（2）原理介绍 - 知乎 (zhihu.com)

[源码解析] 模型并行分布式训练Megatron (1) --- 论文&基础 - 掘金 (juejin.cn)

图解大模型训练之：张量模型并行(TP)，Megatron-LM - 知乎 (zhihu.com)

分布式并行计算——张量并行 - 知乎 (zhihu.com)

张量模型并行详解 | 深度学习分布式训练专题 (paddlepaddle.org.cn)

大规模训练之 transformer 中的张量模型并行 - 知乎 (zhihu.com)