模型并行 | 大规模语言模型架构 Megatron

随着AI模型的规模越来越大，分布式训练技术越来越被广泛使用。现行的分布式训练方法主要包含两个部分：数据并行（Data Parallel）和模型并行（Model Parallel）。数据并行是将模型完整拷贝到多张显卡中，对批次数据进行并行计算，适合规模小而数据多的训练场景；而模型并行适合超大规模参数的模型训练，将模型不同的部分分别加载到不同的显卡中，依次计算得出结果。

Megratron是NVIDIA提出的一种分布式训练大规模语言模型的架构，针对Transformer进行了专门的优化，主要采用的是模型并行的方案。这篇文章将描述幻方AI对于NVIDIA Megatron在萤火二号平台上运行的一些实验，以及与我们目前的方法的对比。 ​

模型：GPT

代码：https://github.com/NVIDIA/Megatron-LM

环境：幻方萤火二号，16个节点共128张A100（A100-40GB x128）

Megatron简介

Megatron是NVIDIA提出的一种由于分布式训练大规模语言模型的架构，针对Transformer进行了专门的优化（也就是大矩阵乘法）。

第一篇论文发表于2019年9月：Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism，主要提出了通过将矩阵分块提高并行度的方法。

第二篇论文发表于2021年4月：Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM，对于分布式中的一些重要的设计，如tensor parallel、pipeline parallel、micro batch size等进行了一些分析与讨论。同时提出了更加精细的pipeline结构与communication模式。

Megatron作者提供的性能结果如下：

这些测试是基于DGX A100-80GB集群完成的。和萤火二号的测试环境相比，硬件环境上主要存在着如下区别：

• DGX配备了NVLINK、NVSWITCH与IB，使得AllReduce的效率非常高。萤火二号的通讯效率相对更低，尤其是在不使用hfreduce的前提下。
• Megatron使用了A100-80GB而萤火二号使用了A100-40GB。除了显存大小不同之外，二者的内存带宽也不同（2.0T/s vs 1.6T/s）。

除此之外，除去实现细节上的差异，Megatron和我们的方法的主要不同在于

• Megaton支持tensor parallel，并相应地优化了数据传输。
• Megatron对于pipeline机制进行了一些优化（原论文Fig.4）。
• Megatron自行实现了一套DDP框架，而非使用pytorch提供的DDP。
• Megatron自行实现了一些fused kernels，但是为了公平对比被我们disable掉了。

我们关心的问题是，在萤火平台上，Megatron的架构能够达到怎样的训练效率。

我们用于测试的模型配置为：

• hidden size = 3072
• layers = 32
• attention heads = 32
• batch size = 512
• context size = 1024
• vocabulary size = 50264
• micro batch size = 4

总参数量为 3.8 B，单次迭代计算量为 16.53 PFLOP。

Tensor Parallel vs. Pipeline Parallel

在Megatron的论文中，作者的建议是在节点内尽量使用tensor-parallel。pipeline-parallel主要在节点间使用，目的在于增大可用显存，而且应当尽量少。这一观点的原因在于pipeline会不可避免地引入bubble而降低效率。然而，另一方面，tensor-parallel需要更多的AllReduce操作，不适合在带宽较低的设备之间进行。

在这一节中我们测试不同的tensor parallel与pipeline parallel的训练效率。我们的实验观察到了与论文结论相反的现象：

其中，TP = tensor parallel，PP = pipeline parallel。为了避免跨numa，每个进程使用4块GPU。更细致的log显示，无论是forward还是backward，在使用tensor parallel后的时间开销都显著增加。此外，GPU利用率远远低于论文中的结果（48%）。

为了排除数据传输的影响，我们在单卡上跑了一个小模型作为测试，单次迭代需要的计算量约为119.4 TFLOP，时间约为1950ms，那么也只有理论峰值的 19.6%。

Model Shape

第二个有趣的发现是，在总运算量一定的情况下，不同的（layers，hidden size）组合的计算效率会有较大的差别。

在一些特殊的形状下，（也就是GPU恰好效率不高的形状下），我们的方法和Megatron还是很接近的…在最优的形状下，差距还是很明显。猜测Megatron的论文里选用这个形状也是因为它的执行效率最高。背后的原因可能是某种尺寸的矩阵的GEMM效率最高。由此进一步推测，tensor parallel也能由于类似的原因实现优化。

Mixed Precision

Megatron默认是使用mixed precision（fp16）进行计算的。而且在前面的图里也能看到A100对于FP16的算力是TF32的两倍。我们做了一个简单的实验对比fp16与tf32在Megatron上的性能差异：

TF32的时间开销增加了约70%，这个结果是符合预期的。

但是：

1. 我们在过去的测试中，发现开关AMP并没有产生太大的性能差距
2. Megatron在TF32下的迭代时间与我们的方法有些接近

因此，一个怀疑是，尽管我们在代码中加入了torch.cuda.amp，它并没有真正生效，又或者是使用的方法并不是最恰当的。这可能是我们的方法和Megatron的性能差距的一个来源。当然，也有可能是Megatron完全使用了fp16（而不是mixed precision），这个问题还有待检查。

更新2022 Feb 17：经过检查代码，Megatron 的 —fp16 应该是启用了 mixed precision 的。在 optimizer.py 中可以看到对 fp32 参数的维护。

Pipeline Parallel vs DDP

假设micro batch size 为 $b$ ，每个micro batch的forward时间为 $t_f(b)$，backward时间为 $t_b(b)$，pipeline一共分成 $p$ 步，那么pipeline的bubble为 $(p - 1) (t_f(b) + t_b(b))$。

如果mini batch size 为 $B$，那么会被拆分成 $\frac{B}{b}$ 个micro batch，每个GPU的实际计算时间为 $\frac{B}{b} (t_f(b) + t_b(b))$。

这里放张图帮助理解： $p = 3$ 放在3块GPU上，红绿黄代表3个micro-batch。先是流水线化的forward，$t_f(b)$是三个方块；然后是backward，$t_b(b)$是四个方块。

那么bubble与有效计算的时间比例为 $\frac{p - 1}{(B/b)} = \frac{b}{B}(p-1)$。考虑到 $B$ 一般是固定的，$b$ 的问题等下再讨论，结论似乎就是 $p$ 越小越好。

然而，另一个问题在于，当总GPU数量固定为$n$时，$p$ 越小，data parallel的进程数 $d = n / p$ 就越多。如果节点间的数据交换需要通过一个中心交换机的话，数据传输的效率可能成为一个问题。

在Megatron中，tensor parallel与pipeline parallel都是将大模型装进显存的方法。但前面的实验已经验证了，在萤火二号的平台上，需要大量带宽的tensor parallel并不合适。那么，理论上，只需要选取最小的 $p$ 使得模型能够装入显存就可以了。

接下来是一个关于pipeline的简单实验。为了能够将模型装入单个GPU，我们将hidden size缩小到2048。

对于Megatron来说，单卡是最快的，这个符合预期。然而，2卡反而是最慢的，8卡也并没有因为跨numa而显著变慢，16卡也没有因为跨node而显著变慢。作为对比，我们的方法随着pipeline parallel规模的增加，性能下降非常明显。猜测是因为Megatron的pipeline实现更加高效。

补充 2022 Feb 17：在上表中，Megatron 的 pipeline parallel 没有因为跨numa/node而显著性能下降是因为其显卡分配方式本来就是优先跨node的。即使当pipeline深度仅仅为2时，两块显卡也分布在不同的节点上。详见这份文档。

此外，我们的方法在单卡下会OOM而Megatron不会，可能是同样是因为我们没有正确地启用fp16，又或者是我们的实现有一些粗糙（考虑到我们只是调用torch，Megatron也是，按理说不该有很大差别）。

PyTorch自带的pipeline基于RPC实现，无法支持跨节点的pipeline。Megatron中的pipeline应该是作者自行实现的，还有待确认。

另外，上述关于pipeline的分析均基于一个假设：在pipeline的每一步，GPU都能够跑满。然而，实际情况可能并非如此。

先放张示意图：

为了简单，只画了forward部分，backward本质上是一样的。现在假设每个GPU的算力没有跑满，所以它有两行。最上面的图表示一个常规的做法，分成3个micro batch。中间的图表示，如果GPU没有跑满的话，可以拆成6步流水线，每个GPU处理两步，那么效率会明显高于3步流水线。而最下面的图表示了另一种可能：尽管我们拆成了6步流水线，但是每一步都会把GPU跑满，体现为每个时间点GPU都只在处理一步运算。然而，在这种情况下，也并不会比拆成3步流水线的情况差。

由于Megatron的实现已经固定了每个GPU只能承载流水线中的一步，我们用我们自己的实现做个实验。

可以看到，随着流水线变深，时间开销缩减到了原来的 83%。尽管foward小幅变慢，但是backward时间显著缩短。

需要指出的是，类似的提升 并不 总是能够在任何模型结构上都会发生。推测这是由于一些特殊的矩阵形状在GPU上的执行效率不高导致GPU没有跑满。上述结果来自于 layers = 54，hidden size = 1920，attention head = 20。这个实验仅仅提示在某些情况下，这样做可能带来加速。

在这些分析中，我们并没有考虑到pipeline带来的sync等overhead。但是从实验结果来看，这些开销是可以接受的。

实际上，我们真正应该减少的是pipeline所占用的GPU数量，而非总的流水线步骤。当我们观察到GPU并没有跑满的时候，在单个GPU内增加pipeline数量是有可能提高效率的：存在着某一个配置，能够平衡过多的pipeline数量带来的overhead与GPU利用率之间的关系。但是理论计算很难给出这个结果，因为overhead受到大量因素的影响。我们更可能需要经验性的实验去寻找这个平衡。显而易见的是，在节点，或者numa内部去分割pipeline的开销是注定远小于跨节点的pipeline的。

一般而言，总的GPU数量是一定的。在这种情况下，pipeline占用更多的GPU，就意味着其它“完美并行”（无bubble）的方式会占用更少的GPU，比如DDP与tensor parallel。二者中，DDP的传输可以和backward并行，但tensor parallel不行，所以会更加显著地受到带宽的限制。一个经验性的建议是，应该尽量减少单个进程占据的GPU的数量，但是可以考虑适当增加每个GPU内的pipeline数量。最后，我们再来解释一下Megatron的pipeline。

本质上，pipeline的bubble来自于启动时GPU i 要等待GPU i-1 的计算结果。每个GPU执行的时间越久，等待的时间就越久，bubble也会更大。核心的思路在于，如何减少GPU处理每个stage的时间。一方面，我们可以通过减少micro batch size的大小，这个接下来会有讨论。另一方面，我们可以把stage拆得更细致，让每个stage的执行时间更短。这实际上就是Megatron中提出的pipeline的做法。

这里放一张图帮助理解：

假设三个颜色分别代表3个micro batch，总的forward时间为18个格子，backward时间为24个格子。

在默认做法中，我们将计算平均分配到三个GPU上，每个GPU分到6个格子，执行完再传给流水线上的下一个GPU。对应上半个示意图。

另一种方式是，我们将计算拆成6份，每个GPU负责 不连续 的两份，以”interleaved”的方式进行分配。这样，GPU i只需要三个单位的时间就能够执行完毕传递给 GPU i + 1。本质上就是让启动时间尽量短。

这样做带来的额外开销是，数据的传输量增加了。假设每个GPU负责 $v$ 个stage，那么传输量就变为了原来的 $v$ 倍。

Megatron论文里的图片把forward和backward拆散了，看起来不那么直观。注意backward和forward是否并行都不会影响pipeline的执行效率，因为GPU已经跑满了（没跑满的话可以参照前面的讨论）。

Megatron拆开forward和backward的目的在于节约显存占用。如果等待forward全部完成后再进行backward，那么全部的micro batch都要被存下来留着backward时用。但是如果拆散forward与backward，可以理解为对于每一个单独的micro batch都尽可能早地完成forward+backward（而不是等待所有forward都完成后再一起backward），完成backward后就不需要再保留micro batch的数据了，所以节约了显存。

如果等待全部forward完成再进行backward的话，需要保存的数据量正比于micro batch的数量 $m$；如果拆散了的话，需要保存的数据量正比于pipeline的深度 $p$。考虑到 $m \gg p$，拆散了更好。

Micro Batch Size

假设micro batch size 为 $b$ ，实际batch size为 $B$，那么一个mini-batch会被拆分成 $B/b$ 步执行。对于每个micro batch的forward时间为 $t_f(b)$，backward时间为 $t_b(b)$，那么总时间为 $\frac{B}{b}(t_f(b) + t_b(b))$。假设 pipeline一共分成 $p$ 步，那么pipeline的bubble为 $(p - 1) (t_f(b) + t_b(b))$。故每次迭代的总时长为 $(\frac{B}{b} + p - 1) (t_f(b) + t_b(b))$。考虑到GPU的硬件特性，$t_f$与$t_b$均不是$b$的线性函数，因此很难通过理论给出$b$的最优取值。因此，对于不同的模型结构（影响$t_f$与$t_b$）、pipeline划分、batch size的设定，需要通过实验来经验性地确定合适的micro batch size。

在Megatron上，一个经验性的实验是：

Checkpoint Layers

理论上，尽管activation checkpoint会在backward中引入额外的计算，但是能够有效增大batch size，因而降低pipeline bubble的相对占比。综合来讲，是能够提高系统的吞吐量的。checkpoint的粒度并不会影响运行效率，但是会影响内存占用。假设每一层的输入的参数量为 $A$，为了求导而保存的中间结果数据量为 $B$，模型一共有 $l$ 层，我们每 $c$ 层做一次checkpoint，那么整个模型需要存储的checkpoint占用显存为 $\frac{l}{c}A$，backward时重算需要存储的临时变量为 $cB$，总内存开销为 $Bc + lA \frac{1}{c}$，最小值对应 $c = \sqrt{l\frac{A}{B}}$。

由于 $c$ 的取值理论上并不会影响训练效率，最终不超过显存容量即可。根据Megatron论文中的建议，$c$ 取1或2是比较合适的值。

在Megatron上的实验验证了checkpoint的粒度并不会对效率产生显著的影响，浮动约为8%：

在我们的方法上得到的结果类似，浮动仅为4%：

Fused Kernels

Megatron的作者提供了一些fused kernels来提高计算效率。通过将bias-gelu和masked-softmax两个fusion，模型能获得约8.5%的提升，单次迭代时间从4.43s提高到4.08s。

总结

1. pipeline优化。pytorch自带的pipeline还是有比较大的改进空间的，尤其是在发生跨numa的时候，与Megatron的实现产生了非常大的差距（而且pytoch的pipe不支持跨节点）。好的pipeline实现是训练大模型必不可少的一部分，可以考虑像hfreduce一样实现一个hfpipe。

2. 关注硬件特性。不同的 (layers, hidden size) 的组合，即使具有相同的理论计算量，实际执行起来的效率差异可以高达85+%。背后的原因应该是矩阵形状的不同。尽管在设计一个模型的时候，训练效率并不是第一考量，但是我们可以通过轻微的调整来追求更高的训练效率。与此同时，另一个可行的做法是通过tensor parallel将矩阵变为最合适的形状。

3. 当硬件架构、带宽已经确定后，tensor parallel能否被高效执行是存疑的。但是依旧可能存在某种专门的实现，使tensor parallel的传输开销能够小于pipeline parallel的bubble开销，来达到Megatron中所宣称的执行效率。