Megatron-LM 分布式执行调研

Megatron-LM 分布式执行调研

Created by: strint Created time: May 31, 2023 6:02 PM

调研作者 strint(https://github.com/strint)，和 BBuf(https://github.com/BBuf), Chengpeng(https://github.com/CPFLAME), Juncheng(https://github.com/liujuncheng), Wenxiao(https://github.com/leaves-zwx) 讨论得到。

本文在 strint.github.io(https://strint.github.io/) 继续完善，目的是跟进Megatron-LM相关进展。

参考资料

• Megatron-LM1 之模型并行（调研的模型并行部分参考的这篇Paper）. https://arxiv.org/abs/1909.08053 . 2020.
  • video https://developer.nvidia.com/gtc/2020/video/s21496
  • Megatron-LM GTC 2020. s21496-megatron-lm-training-multi-billion-parameter-language-models-using-model-parallelism.pdf
  • LiMu. Megatrom-LM paper reading. https://www.youtube.com/watch?v=mk8FuiDmf0I .
• Megatron-LM2 之流水并行和3D并行（调研的模型并行和混合并行部分参考的这篇Paper）. https://arxiv.org/abs/2104.04473 . 2021.
• Megatron-LM 测评环境 Selene Supercomputer 解析 s31700-janneblomqvist-understanding-selene_1616715331011001YQKN.pdf Accelerating-AI-at-Scale_Julie_Prethvi_updated051421.pdf
• DGX A100 配置（超赞的介绍）. https://www.microway.com/hpc-tech-tips/dgx-a100-review-throughput-and-hardware-summary/ .
• 集合通信通信量分析.
  • https://zhuanlan.zhihu.com/p/504957661
  • https://www.cs.utexas.edu/~pingali/CSE392/2011sp/lectures/Conc_Comp.pdf

进一步参考

• Megatron-LM 技术博客.
• Scaling Language Model Training to a Trillion Parameters Using Megatron(https://www.notion.so/Scaling-Language-Model-Training-to-a-Trillion-Parameters-Using-Megatron-6e07c974235b4ecd9f4958579ff37d97?pvs=21)
• Bloom Megatron 训练日志. https://www.cnblogs.com/Matrix_Yao/p/17238627.html .
• Megarton-LM 性能分析. Dive Deep into the Performance Model of GPT-3 Training on Megatron-LM_ Storage, Computation, and Communication.pdf
• Megatron-Turing NLG 530B. https://developer.nvidia.com/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/
  • Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B. https://arxiv.org/abs/2201.11990 . 2022.
• Reducing Activation Recomputation in Large Transformer Models. https://arxiv.org/abs/2205.05198 . 2022.
• FriendliAI. Accelerating LLM Training with Memory-Balanced Pipeline Parallelism. https://medium.com/friendliai/accelerating-llm-training-with-memory-balanced-pipeline-parallelism-3fdfb6ec2c80 . 202307

概要

之前主要的大模型训练方式是数据并行，Megatron-LM 比较成熟的支持 LLM 的模型并行和流水并行。

模型并行部分，主要是提供了 transformer 的模型并行的高效实现。结合数据并行和模型并行，Megtron-LM 把 GPT-2 模型扩展到了8B、 512 GPUs的规模，扩展效率为 75%+.

流水并行部分，Megatron-LM 组合了数据、模型、流水并行，将 GPT 模型扩展到了1000B、3072 GPUs的规模，且 MFU 为 52%。利用了 1F1B 的调度策略缩短 Activation 的生命周期，降低显存峰值；也利用了 Re-Materialization 技术（通常也叫 activation-checkpointing）缩短 Activation 的生命周期，来降低显存峰值。另外流水并行部分提出了 interleaved 1F1B 调度策略，减少流水线气泡，在显存开销不变的情况下，吞吐提高 10%.

Megetron-LM 的文章也介绍和如何设置混合并行的最佳实践。这些技术旨在减少通信并使GPU计算保持繁忙。减少通信、减小显存读写，把任务转换成一个计算受限型任务。

关键词：Transformer，模型并行，流水并行

模型并行 Tensor Parallelism

也叫 Intra-Layer (Tensor) Parallelism.

Tensor Parallelism 下的矩阵乘法：Row Parallel Linear Layer

分块矩阵乘法：

OneFlow 中 SBP 表示：

Tensor Parallelism 下的矩阵乘法：Column Parallel Linaer Layer（这里的图错了，请以开头的原文链接为准）

分块矩阵乘法：

OneFlow 中SBP表示：

Tensor Parallelism 下的通信操作

g 和 f 分别代表该位置需要插入的操作。比如对于 Row Parallel, f 位置，前向需要插入split 操作，而后向需要插入 all-gather 操作。

Tensor Parallelism 的开销

Flops 代表计算量；Bandwidth 代表 Memory access operations；Intensity 代表 Flops/Bandwidth，显示计算操作和内存搬运操作的比例。

Normal 即常规模式，矩阵乘法实际需要执行

次 的(n次乘法 + n次加法):

• Flops = （n次乘法 + n次加法）*

次 =

• Bandwidth = (

)【一个 n*n 矩阵的读或写】 * 2【2是 float16数据大小】) * 3【对应读 X、读A，写 Y】=

• Intensity = Flops / Bandwidth

Column Parallel 模式下（单 device 上的）：

• Flops = （n次乘法 + n次加法）* (

/ p)次 =

• Bandwidth = （

【X 这个 nn 矩阵的读】+

【 A 这个

矩阵的读】+

【 Y 这个

矩阵的写】 ) 2【2是 float16数据大小】)

• Intensity = Flops / Bandwidth

可以看到并行执行的 Bandwidth 需求增加，要考虑 Intensity 的 tradeoff：

• 选择合适的数据切分方式，减少需要聚合数据的同步点，这样可以减少 Bandwidth；
• 数据的切分要使得其聚合操作能用简单的集合通信实现（如 NCCL/PyTorch，它们是高度优化好的）；
• 使用 Volta tensor cores 来做混合精度计算，这样可以减少 Flops；

MLP + GeLU 时，选择 Column Parallel 以避免通信同步

对于 MLP + GeLU 操作，如果选择 Row Parallel，单设备得到的是一个 element-wise partial 的矩阵，而

GeLU(element-wise-partial-M0 + element-wise-partial-M1) ≠ 

GeLU(element-wise-partial-M0) + GeLU(element-wise-partial-M1)

，所以需要引入同步。

选择 Column Parallel 时，后面遇到 GeLU 操作，

 GeLU([split-M0, split-M1]) ==  [GeLU(split-M0), GeLU(split-M1)]

可以不做数据同步。

Fused MLP + GeLU + MLP

第一个 MLP 采用 Column Parallel，第二个 MLP 为了延迟通信同步，可以选择 Row Parallel，这样直到第二个 MLP 执行完之后才需要做一次 all-reduce.

如此就减少了两次通信（在绿色箭头处各减少了一次通信），且红色虚线框内的 MLP + GeLU + MLP 可以被 fuse 成一个算子执行，以减少显存拷贝。

Fused Self-Attention

Self-Attention 的原始计算如下

在 GPT 的 Transformer 中可以分解成如下计算：

• Q 和 K 的 MatMul 计算；
• Scale 运算即除以 K 的维度的开方；
• 为了保证当前 token 只看到之前的 token 而加上的 mask 矩阵；
• SoftMax
• 和 V 的 MatMul 运算

如下图中左侧所示。

对于 Multi-Head Attention，是需要把每个单独的 Attention 得到的结果 Concat起来再继续后面的运算。所以多个 Head 之间的运算是天然可并行的。

如上图所示，蓝色框中代表 Self-Attention 中的 MatMul 运算。类似之前的 Fused MLP，这里也是采用先 Column Parallel 后 Row Parallel 达到 Self-Attention 计算不需要通信同步的效果。

具体的，这里第一个矩阵乘对应 X 转换到 Q、K、V 的 Linear 运算。按 Column Parallel 把Q、K、V的三个参数矩阵分别 Split 了，Split 后的每个分片分别对应一个 Head。这样得到 Multi-Head Attention 算法下的 Q、K、V的计算天然就是符合 Column Parallel 的。以 Q 的生成为例，如下图所示：

第二个矩阵乘是对应的是 multi-head 的 V 的点乘运算执行完成后、再 Concat V之后的那个 Linear 运算。multi-head 的 V 的点乘运算执行完成后的结果对应上图中的

和

，是按列切分的，这时把 Linear 的参数按行切分，就形成了 Row Parallel.

Transformer Layer 中的通信

最终，通过模型并行，一层 transformer layer 的前向、后向一共只需四次 all-reduce 通信，减少了一半通信次数。

以上便是 Megatron-LM 中对 transformer layer 的模型并行处理策略。

GPT-2 中其它模块的并行方式

GPT 算法结构

上文分析了 Transformer Layer 中的 MLP 和 Attention 操作的并行方式，这里补充下 GPT-2 中其它操作的并行方式。

假设输入的 batch-size 为

，输入句子的 token 长度为

。

Input Embedding 和 Output Embedding 共享了参数。设 hidden-size 为

，vocabulary size 为

。 则参数矩阵的 shape 为

。一般

都特别大（万级别），所以选择在

上做切分，设分配大小为

。

那么对于 Input Embedding 运算，就是一个 Row Parallel，

。然后在其后插入一个 all-reduce 操作，把输入转换为 shape 为

的 Embedding。

对于 Output Embedding 运算，其输入是

，要将其转换为

。采用 Column Parallel，这样得到一个按列切分的输出，然后再做一次 all-gather 操作，shape 转换过程为

。但是

很大，这导致 all-gather 操作非常耗时，需要同步

个元素。因此得到

后，延迟通信同步到 Cross entropy loss 之后（训练时使用的是 Cross entropy，在推理时使用的是 Softmax）。

对于 Cross entropy loss 运算，其输入的 shape 是

、输入的 label 的 shape 是

，做点积后输出是

。此时再做 all-reduce 操作得到最终的 Cross entropy loss，而这个 all-reduce 只需要同步

个元素。

对于 Dropout/Layer Norm/Residual connections(即 Add）操作，都没有做并行拆分。比如 Layer Norm 的参数，在每个 GPU 上都是完整的。

在多卡上，因为模型的参数要么是重复的、要么是局部的，所以在做模型参数更新时是不需要额外通信的。

上面的切分方式使得整个网络只需要额外插入几个简单的 all-reduce 操作，所以不需要依赖编译器做复杂通信算子选择。

模型并行的性能测评

Weak Scaling 是衡量并行扩展性的一个指标。假设单任务、单设备的处理时间为

，而

倍的任务、

倍的设备去计算得到的处理时间为

。那么：

如下图所示（来源(https://cvw.cac.cornell.edu/parallel/scale)），理想情况下，在任务增加n倍、设备增加n倍时，处理时间恒定不变，Weak Scaling 为 1，即线性扩展性。实际处理时间会变大，导致 weak_scaling 小于 1。

对于 GPT 的模型并行、单机内扩展测试显示，8卡 Weak scaling 可以达到 77% . 实验的详细配置参见论文(https://arxiv.org/pdf/1909.08053.pdf)实验部分。

而采用节点内模型并行、节点间数据并行，如下图所示，在 512 卡下，仍然可以获得 74% 的 weak scaling.

具体的，相比 Config 1， Config 4 的模型扩大了 7 倍，设备数扩大了8倍相当于任务规模扩大了，而一次前后向计算时间只增加了30% . 证明了模型并行的良好扩展性。

这里补充下硬件配置：

• 32 台 DGX-2H 服务器 (共 512 张 Tesla V100 SXM3 32GB GPUs)；
• 节点内的 GPU 之间有 300 GB/sec 的带宽，基于 NVSwitch；
• 节点间的 GPU 之间有 100 GB/sec 的带宽，基于 InfiniBand，每个节点有 8 个 InfiniBand adapter.

流水并行 Pipeline Parallelism

Pipeline Parallelism（流水并行） 也叫 Inter-Layer (Pipeline) Parallelism. 不同于模型并行中的把参数 tensor 划分到不同的设备，流水并行是把模型分 layer 划分的不同的设备，以达到切分模型的目的。Megatron-LM 对于 GPT 模型的流水并行切分方式如下图所示。

具体的，两个 Transformer Layer 被切分到两个设备上，形成 pipeline，对应绿色的框。而每个 Transformer Layer 内部则是上文的模型并行内容中的 Self Attention 和 MLP 的切分方式。所以流水并行和模型并行是正交使用的。

严格的 Optimizer 语义

为了使流水线可以流转起来，需要把一个 batch 的前后向执行进行切分，切分后的数据对应一个 micro batch. 这样第 n 个 micro batch 的数据在一个 pipeline stage（设备）执行时，第 n + 1 个 micro batch 可以在上一个 pipeline stage（设备）执行。这样就可以形成流水线。

严格的 Optimizer 语义是指，一个 batch 内的所有 micro batch 的数据遇到的模型都是同一个版本的。为了达到这个效果，Megatron-LM 在一个 batch 的结尾引入了流水线 flush，即做一次跨设备的同步，然后再用 Optimizer 更新参数。

这种参数更新的方式也叫做 Gradient Accumulation.

GPipe 流水线和气泡

GPipe 的流水线就实现了严格的 Optimizer 语义。

如下图所示，黑色的竖线对应一次 Pipeline flush. 它之前的区域代表一个 batch 的数据执行的时空图，横向对应时间，纵向代表设备。

具体的，这是一个4级流水线，每级一个 device（对应一个 device 上的运算）。一个 batch 被拆分为8个 micro batch。蓝色代表前向运算，绿色代表后向运算。可以看到 GPipe 是先执行完所有 micro batch 的前向，然后再执行所有 micro batch 的后向，然后在 pipeline flush 处进行一次同步。注意这里假设了后向的计算量是前向的2倍。

上图中灰色的部分就是设备空闲的情况，在本文中被叫做 pipeline bubble(流水线气泡）。显然这些气泡会降低设备的 MFU。

假设一个 batch 中 micro batch 的数量为

，pipeline 的 stage 的个数为

. 一个 micro batch 的整个前向、整个后向的执行时间分别为

和

. 则上图中在前向存在

个

的气泡，在后向存在

个

的 气泡，所以一个迭代的气泡

:

而一个迭代理想的处理时间，即没有气泡的处理时间

:

这样气泡占比 Bubble time fraction:

这样为了降低气泡量，就需要

. 但是每个 micro bath 前向的 Activation 都需要暂存直到依赖它的后向计算完成，这样 micro batch 数量过多，会导致显存占用增加过多。后面对此提供了一种改进策略：1F1B.

1F1B

为了降低流水线的显存，PipeDream-Flush(https://arxiv.org/abs/2006.09503) 中提出了一种缩短 Activation 生命周期的方法，叫 one forward pass one backward pass（简称 1F1B）。

如下图中的上半部分所示，流水线中最长驻留了

个未完成的 micro batch. 而 1F1B 则限制其最多驻留流水线深度

个未完成的 micro batch，如此形成了下图中的下半部分的流水线。这个流水线的特点是一个迭代的时间没有变化，但是

，所以驻留的未完成的 micro batch极大减少，减少了显存峰值。

Interleaved 1F1B（交错的1F1B）

在 1F1B 中还是存在明显的气泡，Interleaved 1F1B 可以如下图所示，用更少的时间完成一轮迭代（看黑色线的变化）。

其核心的想法源于让 micro batch （micro batch size 更小）更多以减少气泡。方法是让一个 device 虚拟成

个 device，从计算 1个连续的 layer 段（有

个 layer）变成计算

个不连续的 layer 段（每段 layer 数量为

）。比如之前 1F1B 时 device 1 负责 layer 1~4，device 2 负责 5~8，在 Interleaved 1F1B 下 device 1 负责 layer 1~2 和 9~10，device 2 负责 3~4 和 11~12，这样可以让流水线中每个 stage 更小，因而下个 stage 的等待时间更短，气泡更小。需要注意的是，

需要是

的整数倍。

此时完成

个 layer 段中一个的前向、后向时间分别为

和

, 流水线的气泡

:

注：因为气泡的分布不规律，这个公式的结论还存疑，但是看流水线中的气泡数量确实如此。

如此得到的气泡占比 Bubble time fraction:

这样气泡占比就减少到了

。但是流水线之间的通信量也增加了

倍。后文会讨论如何使用 InfiniBand networking cards 来减少这部分新增通信带来的影响。

下图显示了对 Interleaved 调度的测评。模型是 GPT-3 175 billion 参数，可以看到在不同的 batch size 下，总是 Interleaved 调度的吞吐更高（依赖于 scatter/gather 通信优化，否则在大 batch size 下Interleaved 调度有负面影响）。不过随着 batch size 变大，Interleaved 调度的优势在减少，一方面是batch size 变大气泡会减少（见后文相关分析），另一方面是batch size 变大导致 pipeline stage 间通信变得显著，而 Interleaved 调度的pipeline stage 间的通信次数是Noe-interleaved 调度的多倍。

3D混合并行的最佳实践

在一定的 GPU 数量和 batch size 的情况下，可以采用不同的数据并行、模型并行、流水并行配置策略（简称3D并行）。不同的并行策略需要在这些因素间做权衡：显存占用、设备使用率、通信量。

一个3D并行的GPU分组示意图

后文使用如下标记：

: 3D 并行维度.

代表流水并行大小，

代表模型并行大小，

代表数据并行大小.

: 总共的 GPU 数量. 要求

.

: Global batch size.

: Microbatch size.

: 一个流水线要处理的 batch size 大小， 等于

.

: 一个 batch 在每个 pipeline 的 microbatch 的数量.

: 输入的 sequence length.

: hidden size.

模型并行和流行并行的组合选择

这里先假设数据并行度为1，即

. 那么

. 之前已经分析，流水并行中会产生气泡，气泡占比为：

在

固定，即

固定的情况下，随着

增加，气泡占比会降低。因此提高模型并行度是有利的。

但是提高模型并行度会增加通信量。

具体的，对于一个 pipeline stage，里面包括多个 Transformer layer，一个 microbatch 在流水并行的多个 pipeline stage 间的通信量是

（考虑前向、后向各一次），采用 point-to-point 通信，且和 Transformer layer 数量无关 。

但是模型并行时，一个 Transformer layer，一个 microbatch，需要4个 all-reduce 通信（参考上文模型并行部分）。一次 all-reduce 在单设备的通信量是：

如果一个 pipeline stage 内有

个 Transformer layer，则一个 pipeline stage 内的单设备通信量：

当

最小为2时，结果为

，它是pipeline间的通信量的

倍；且

约小，通信量越大。

考虑机器间的带宽 （IB网络，200G/s 量级）是比机器内的带宽（NVLink，600G/s）小很多的。所以为了减少气泡占比可以增加

，但是

超过单机设备数量时，通信开销就会显著增加。

如下图所示，固定数据并行大小为1，64 个 A100，使用 GPT 161 billion 参数的模型，可以看到不同的 Batch size 下，在模型并行大小为8的情况下，取得了最优吞吐。因为使用的 DGX A100 机器，单机内 GPU 数量为8. 另外也可以看到模型并行大小过大、流水并行度过大都会带来更低的吞吐。最好让模型并行度尽量大（减少气泡），然后让流水并行度尽量小（减少气泡）。

最佳实践1：当考虑使用模型并行时，如果节点内有

个 GPU，则模型并行的最大并行度是

. 如果这样还放不下模型，再使用流水并行来分 layer 切分模型。

数据并行和其它并行的组合

考虑数据并行和流水并行组合，设

. 则

. 此时流水线气泡占比为：

下图是不同的

组合对应的流水线气泡占比。可以看到固定

时，

越大气泡越小。但是

往往无法无限趋近于

，因为模型参数和 Activation 显存占用会超过单设备的显存限制。

下图固定模型并行大小为1，64 个 A100，使用GPT 5.9 billion 参数的模型。可以看到流行并行数量越小，数据并行度越大，吞吐越高。所以流水并行的大小满足模型显存需求就好，尽量提高数据并行的大小以提高吞吐。

总的来看，增加数据并行度，可以提高吞吐（减少流水线气泡）。而数据并行需要的通信主要是后向汇总梯度时的 all-reduce 操作。它的通信量和

成正比，因此

越大时，通信量越小。因此，在一定的 GPU 数量下，在显存足够模型占用时，数据并行度越高越好。

Global batch size 的影响。当我们固定 micro batch size，考虑增加 Global batch size

，

增大，进而流水线气泡占比减小。而数据并行中需要的梯度 all-reduce 时间没有变化但是因为 micro batch 数量变多，因此 all-reduce 通信占比变低。因此增大 Global batch size

也可以增加吞吐。如下图所示，使用标准的1F1B，一个 pipeline stage 总是对应 3 层 Transformer layer，固定 micro batch size 为1，数据并行大小为1，模型并行大小为8，可以看到在不同的流水并行数量下，总是 batch size 大的配置吞吐更高：

考虑数据并行和模型并行的组合。模型并行组内的，每个 micro batch，每个 transformer layer 内，前后向一共需要4次昂贵的 all-reduce 通信。而数据并行组内，只需要在每个 batch对参数的梯度做一次 all-reduce 操作。另外模型并行时，如果参数矩阵比较小，矩阵乘法规模比较小，也无法达到 GPU 的峰值计算效率。

下图固定流水并行大小为1，64 个 A100，使用GPT 5.9 billion 参数的模型。可以看到模型并行数量越小，数据并行度越大，吞吐越高。所以模型并行的大小满足模型显存需求就好，尽量提高数据并行的大小以提高吞吐。

在一定的 GPU 数量限制下，尽量提高数据并行的大小，不过模型并行度的大小受最大 Global batch size 的限制，因为 Global batch size 过大会导致 loss 不收敛。

最佳实践2：如果模型比较大，需要先组合模型并行和流水并行，

的组合用来满足模型和模型相关的数据的显存需求，但是要让

尽量小。之后使用数据并行来扩展训练规模（扩大数据并行度、扩大 Global batch size）。

Microbatch 的大小

Micro batch size

的大小也会影响模型训练的吞吐，如下图所示，增大

可以带来1.3倍的吞吐提升。

参考上文中 1F1B 和其中气泡的计算，在固定的并行策略

和 Global batch size

的情况下，忽略通信开销，一个迭代的时间为：

如果增大

，

（即一个 pipeline 内的 micro bath 数量）会减少，但是一个 micro batch 的前后向执行总时间

会增加。所以一个 iter 执行的时间和

是非线性的关系，但是可以找到一个极值点。如下图所示，使用64卡 A100，使用 GPT 1 billion 参数的模型，模型并行大小8，流水并行大小8，在两种不同的 Global Batch size 下，最优的 micro batch size 都是 4：

最佳实践3：在一定的并行策略

和 Global batch size

的情况下，micro batch size

对吞吐和显存占用都有影响。（最优的 micro batch size 需要通过实验试验出来.）

Activation Checkpointing

Activation Checkpointing 也叫做 activation recomputation. 目的是减少显存峰值，但是会增加后向的计算时间。具体的，在前向计算时，只保暂存部分 layer 的输出的 activation；在后向计算时，利用这些暂存的 activation 再做一次前向计算以得到后向计算需要的 activation。

参考 1F1B 策略，里面假设后向的计算时间是前向的2倍，那么使用 Activation Checkpointing 后，大致后向计算时间会增加50%，一个迭代的时间会增加 33%.

虽然 Activation Checkpointing 会增加一轮迭代的时间，但是降低显存占用后可以扩大 Batch Size（可以增大吞吐）。如下图所示，黄色线代表Activation Checkpointing关闭，蓝色线代表打开，可以看到扩大 Batch Size 带来的吞吐提升超过了Activation Checkpointing带来的负面影响，可以获得高达两倍的吞吐：

混合并行的实现要点

pipeline stage 间的通信优化：scatter-gather

论文中环境使用的时 DGX A100，每个节点有 8 个 GPU，节点内的8卡用 NVLink 互联。而节点间采用 InfiniBand(IB) 相连，但是这个连接只是点对点的，节点 A 的 i 卡只和节点 B 的 i 卡这样的成对的卡之间联通。这给节点间的通信带来了比较大的限制。

为了适应这种情况，对于 pipeline stage 间的 activation 通信，论文选择传输完整的 activation，这样就可以利用 IB 的点对点通信完成。

如下图中(a) 所示，GPU1 和 GPU2 是一个节点，GPU3 和 GPU4 组成一个节点；浅色的竖条代表上一个 pipeline stage 中的 layer，深色代表下一个 pipeline stage 中的 layer；横条代表 pipeline stage 之间传输的 activation。常规的办法是把 activation 聚合成完整的再通过 IB 传输到下个 pipeline stage.

上图中(a) 的做法，在模型很大，需要8卡模型并行时，就需要传输8份完整的 activation，冗余很大。论文采用了一种 scatter/gather 通信优化方法来减少冗余。具体的就是把 activation 拆分，每个 IB 连接的 GPU 对之间只传输一部分，然后在下个 pipeline stage 使用 activation 前，使用机内的 NVLink 聚合 activation 再使用。因为 NVLink 远快于 IB，这样可以解决 IB 带来的瓶颈，把 IB 的传输量从

减少到了

。

如下图所示，在使用 Interleaved 调度的情况下，175 billion 的 GPT-3 模型，scatter-gather通信优化可以获得约11%的吞吐提升。

计算优化

计算上的优化是针对特定模型（GPT/Bert）做的。主要有三类。

第一类优化是修改数据的 layout。这样一方面可以避免读写密集型的 transpose 操作，另一方面可以使用 strided batched GEMM kernels。具体的论文把数据的 layout 从

变成了

，

分别是 batch、sequence、attention-head、hidden-size.

第二类优化是 kernel fusion. 把 transformer layer 中的 bias + dropout + add 做了 fuse，把 bias + GeLU 也做了 fuse，他们都是 element-wise 的操作，这里采用了 PyTorch JIT 来实现。

第三类优化是自定义 kernel，目的也是做 kernel fuse. 具体的，实现了 scale + mask + softmax 的 fusion。

对于 175 billion 的 GPT-3 模型，这些优化可以提升 19%的吞吐。对于 530 billion 的 GPT 模型，这些优化可以提升 11%的吞吐。

混合并行的运行

集群配置

集群采用的是 Selene Supercomputer，参见参考资料中的“Megatron-LM 测评环境 Selene Supercomputer 解析”。Selene 的配置如下图，有 560 个 DGX A100，每个里面有8卡共4480 个 A100-80G 的 GPU 组成，850 个 200 G IB 交换机，14 PB 的全 NVME（SSD）的共享的文件存储空间。

集群由互联、可扩展的 POD 组成，POD 内可以容纳多达 140 个互联的 DGX A100，集群结构如下图所示：

单节点是DGX A100，其配置如下所示。

DGX A100 配置表

DGX A100 拥有8卡 NVIDIA 80-GB A100 GPU，共 640G 的 GPU 显存空间。单GPU的 fp16 峰值吞吐是 312 teraFLOP/s。8卡之间采用 NVLink 和6个 NVSwitch 互联，GPU 之间双向总带宽是 600G/s（A800被阉割到400G/s)，单向带宽是 300G/s。有2个64核的AMD Rome 7742 CPU，共128核，2TB 的内存，8个3.84TB的 NVME 共 30TB 的 SSD 文件存储空间（给超大模型 Checkpointing 的 save/load 提供保障）。

DGX A100的网络包括8个连接其它节点的 IB，2个连接存储设备的IB。

DGX A100 结构

DGX A100 结构2，来源(https://www.microway.com/hpc-tech-tips/dgx-a100-review-throughput-and-hardware-summary/)

更多关于 DGX A100 的配置和测评参见：

• https://www.microway.com/hpc-tech-tips/dgx-a100-review-throughput-and-hardware-summary/

端到端的性能评估

对于 GPT 模型，设 vocabulary size 为

，sequence 长度为

，hidden size 为

，Transformer 层数为

，则模型的总参数量

:

对于GPT模型，主要的浮点运算是矩阵乘法。batch size 为

，总卡数为

，则每一轮迭代的 FLOPs

，这里的 FLOPs 考虑的是浮点运算操作数量，和精度无关，另外这个计算也考虑了 Activation recomputation：

的计算参见Megatron-LM2之流水并行的论文的 Appeddix 部分。

下表基于这个一轮迭代的 FLOPs 计算公式，然后统计一轮迭代的端到端时间（包括 data load，optimizer step，communication，logging），可以计算出FLOPs/s。下表中就是不同模型配置下的 FLOPs/s 统计，可以看到增加GPU数量，单卡的吞吐越来越高，实现了超线性扩展：

在训练时，另外一个需求是预估训练时间。上面的

有了一轮迭代的吞吐量的预估，再知道需要的总迭代次数，就知道了训练完整个模型需要的 FLOPs。这里假设训练的 token 量为

，则训练迭代次数

。这样需要需要的总 FLOPs 为：

对于总 FLOPs 和 参数量，考虑

，

，

，所以可以得到下面的近似值

然后上表1中提供了一个单卡吞吐量 FLOPs/s 的经验数值

。设有

卡，则可以得到需要的训练时间：

采用这个公式，GPT-3 有

参数，有

tokens，在

个 GPU 上，batch size 为 1536 时，每个 GPU 的

teraFLOP/s，计算可得训练时间为 34 天。

另外，如果要分析GPT模型的参数量，可以进一步参考这篇文章《分析transformer模型的参数量、计算量、中间激活、KV cache》(https://zhuanlan.zhihu.com/p/624740065)。

和 ZeRO-3 的比较

论文中和纯的 ZeRO3 做的对比，保持了 Global batch size 相同，设备数量相同，比较单设备的FLOPs/s。结果如下图：

可以看到 Megatron-LM 的混合并行会比数据并行+ZeRO-3的吞吐高、线性扩展性好。

ZeRO-2中证明了再数据并行组内做梯度切分和聚合带来的通信量和标准数据并行的梯度聚合通信量是相同的，即梯度 reduce-scatter 加参数的 all-gahter 等于梯度的 all-reduce 通信量。

但是在 ZeRO-3时，为了切分模型，在前向还要再增加一次参数的 all-gahter 操作。以 Transformer layer 的前向为例，这些模型参数的发生的 all-gahter 操作的频率（4次，4个参数矩阵），会比模型并行采用的all-reduce操作频率（1次，一个 Activation）高，这看起来是ZeRO-3的一个缺陷（定性的）。

不过因为 ZeRO 可以和混合并行组合使用，比如在数据并行组内增加 ZeRO 并行。所以这部分的结论是不太严谨的，需要更准确的实验和分析。

相关术语

FLOPs

Floating Point Operations Per Second.

FLOPS 等于单核每个时钟周期多少浮点运算数 * 单核每秒多少个时钟周期 * 核心数量.

参考：

• Transformer FLOPs 推导. https://zhuanlan.zhihu.com/p/646905171 .

MFU(Model FLOPs Utilization) and HFU(Hardware FLOPs Utilization)

Model FLOPs(MF) are the floating point operations required to perform a single forward and backward pass (single iteration) regardless of the implementations and hardware limitations. As a result, model FLOPs are hardware and implementation independent and only depend on the underlying model.

The hardware FLOPs(HF) represent the floating point operations that are actually performed on the hardware per iteration.

The MFU and HFU are defined as model and hardware FLOPs per second divided by the accelerator theoretical peak FLOPs per second.

参考：

• nanoGPT 中的 MFU 计算. https://github.com/karpathy/nanoGPT/blob/eba36e84649f3c6d840a93092cb779a260544d08/model.py#L289
• Reducing Activation Recomputation in Large Transformer Models 文章中的 MFU 计算. https://arxiv.org/pdf/2205.05198.pdf