随着AI模型的规模越来越大，分布式训练技术越来越被广泛使用。现行的分布式训练方法主要包含两个部分：数据并行（Data Parallel）和模型并行（Model Parallel）。数据并行是将模型完整拷贝到多张显卡中，对批次数据进行并行计算，适合规模小而数据多的训练场景；而模型并行适合超大规模参数的模型训练，将模型不同的部分分别加载到不同的显卡中，依次计算得出结果。
Megratron是NVIDIA提出的一种分布式训练大规模语言模型的架构，针对Transformer进行了专门的优化，主要采用的是模型并行的方案。这篇文章将描述幻方AI对于NVIDIA Megatron在萤火二号平台上运行的一些实验，以及与我们目前的方法的对比。
模型：GPT
代码：GitHub – NVIDIA/Megatron-LM: Ongoing research training transformer models at scale
环境：幻方萤火二号，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 = 1, PP = 4 | TP = 2, PP = 2 | TP = 4, PP = 1 | |
|---|---|---|---|
| time (s) | 4.43 | 5.41 | 9.54 |
| percentage of peak | 9.4% | 7.6% | 4.3% |
其中，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）组合的计算效率会有较大的差别。
| layer, hidden | 20, 3872 | 32, 3072 | 52, 2400 | 80, 1920 | 80, 2048 |
|---|---|---|---|---|---|
| PFLOP | 16.4 | 16.5 | 16.5 | 16.5 | 18.6 |
| Params (B) | 3.99 | 3.94 | 3.84 | 3.74 | 4.24 |
| Megatron (s/iter) | 5.79 | 4.34 | 7.48 | 8.07 | 5.90 |
| Ours (s/iter) | 8.47 | 6.31 | 8.00 | 8.11 | 7.76 |
| Megatron / Ours | 68.4% | 68.8% | 93.5% | 99.5% | 76.0% |
在一些特殊的形状下，（也就是GPU恰好效率不高的形状下），我们的方法和Megatron还是很接近的…在最优的形状下，差距还是很明显。猜测Megatron的论文里选用这个形状也是因为它的执行效率最高。背后的原因可能是某种尺寸的矩阵的GEMM效率最高。由此进一步推测，tensor parallel也能由于类似的原因实现优化。
Mixed Precision
Megatron默认是使用mixed precision（fp16）进行计算的。而且在前面的图里也能看到A100对于FP16的算力是TF32的两倍。我们做了一个简单的实验对比fp16与tf32在Megatron上的性能差异：
| FP16 | TF32 | |
|---|---|---|
| s/iter | 4.43 | 7.53 |
TF32的时间开销增加了约70%，这个结果是符合预期的。
但是：
- 我们在过去的测试中，发现开关AMP并没有产生太大的性能差距
- Megatron在TF32下的迭代时间与我们的方法有些接近
因此，一个怀疑是，尽管我们在代码中加入了torch.cuda.amp，它并没有真正生效，又或者是使用的方法并不是最恰当的。这可能是我们的方法和Megatron的性能差距的一个来源。当然，也有可能是Megatron完全使用了fp16（而不是mixed precision），这个问题还有待检查。
更新2022 Feb 17：经过检查代码，Megatron 的 —fp16 应该是启用了 mixed precision 的。在 optimizer.py 中可以看到对 fp32 参数的维护。
Pipeline Parallel vs DDP
然而，另一个问题在于，当总GPU数量固定为n时，p 越小，data parallel的进程数 d=n/p 就越多。如果节点间的数据交换需要通过一个中心交换机的话，数据传输的效率可能成为一个问题。
在Megatron中，tensor parallel与pipeline parallel都是将大模型装进显存的方法。但前面的实验已经验证了，在萤火二号的平台上，需要大量带宽的tensor parallel并不合适。那么，理论上，只需要选取最小的 ppp 使得模型能够装入显存就可以了。
接下来是一个关于pipeline的简单实验。为了能够将模型装入单个GPU，我们将hidden size缩小到2048。
| pipeline parallel | 1 | 2 | 4 | 8 | 16 |
|---|---|---|---|---|---|
| Megatron (s/iter) | 2.84 | 3.25 | 3.19 | 3.01 | 3.40 |
| Ours (s/iter) | OOM | 2.80 | 4.40 | 7.69 | not supported |
| Megatron / Ours | 116% | 72.5% | 39.1% |
对于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只能承载流水线中的一步，我们用我们自己的实现做个实验。
| pipeline stages | 4 | 8 | 16 | 32 |
|---|---|---|---|---|
| forward | 0.63 | 0.64 | 0.66 | 0.70 |
| backward | 4.53 | 4.19 | 3.83 | 3.59 |
| total | 5.50 | 5.11 | 4.74 | 4.53 |
可以看到，随着流水线变深，时间开销缩减到了原来的 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≫p，拆散了更好。
Micro Batch Size
在Megatron上，一个经验性的实验是：
| micro batch size | 2 | 4 | 8 |
|---|---|---|---|
| time (s) | 4.79 | 4.43 | 5.69 |
Checkpoint Layers
由于 c的取值理论上并不会影响训练效率，最终不超过显存容量即可。根据Megatron论文中的建议，c取1或2是比较合适的值。
在Megatron上的实验验证了checkpoint的粒度并不会对效率产生显著的影响，浮动约为8%：
| checkpoint layers | 1 | 2 | 4 | 8 |
|---|---|---|---|---|
| total | 4.39 | 4.43 | 4.74 | 4.55 |
在我们的方法上得到的结果类似，浮动仅为4%：
| checkpoint layers | 1 | 2 | 4 | 8 | 16 | 32 |
|---|---|---|---|---|---|---|
| forward | 0.60 | 0.59 | 0.60 | 0.60 | 0.60 | 0.59 |
| backward | 5.64 | 5.62 | 5.38 | 5.38 | 5.40 | 5.42 |
| total | 6.41 | 6.37 | 6.15 | 6.14 | 6.16 | 6.18 |
Fused Kernels
Megatron的作者提供了一些fused kernels来提高计算效率。通过将bias-gelu和masked-softmax两个fusion，模型能获得约8.5%的提升，单次迭代时间从4.43s提高到4.08s。
总结
pipeline优化。pytorch自带的pipeline还是有比较大的改进空间的，尤其是在发生跨numa的时候，与Megatron的实现产生了非常大的差距（而且pytoch的pipe不支持跨节点）。好的pipeline实现是训练大模型必不可少的一部分，可以考虑像hfreduce一样实现一个hfpipe。
关注硬件特性。不同的 (layers, hidden size) 的组合，即使具有相同的理论计算量，实际执行起来的效率差异可以高达85+%。背后的原因应该是矩阵形状的不同。尽管在设计一个模型的时候，训练效率并不是第一考量，但是我们可以通过轻微的调整来追求更高的训练效率。与此同时，另一个可行的做法是通过tensor parallel将矩阵变为最合适的形状。
当硬件架构、带宽已经确定后，tensor parallel能否被高效执行是存疑的。但是依旧可能存在某种专门的实现，使tensor parallel的传输开销能够小于pipeline parallel的bubble开销，来达到Megatron中所宣称的执行效率。