Efficient Large-Scale Language Model Training on GPU ClustersUsing Megatron-LM

1 INTRODUCTION

在这篇文章中展示了如何将 tensor ，pipeline， data 并行组合，扩展到数千个GPU上。
提出了一个新的交错流水线调度，可以提升10%的吞吐量。propose a novel interleaved pipelining schedule that can improve throughput by 10+% with memory foot-print comparable to existing approaches
训练transformer 语言模型等这类模型，模型规模超大，挑战：

不能将模型直接放入GPU内存中训练，既是GPU内存非常大：
即使可以将模型装入GPU，太多次数的操作也会导致难以忍受的训练时间

利用数据并行进行规模化，通常表现很好但是存在两方面的限制：
a) 除了一个节点的情况，每个GPU的batch size 太小会将带GPU的利用率，增加通信的损耗
b) 可用设备的最大数是batch的size，限制了加速器的数量。
一些模型被提出来解决这两个挑战

tensor (intra-layer) model parallelism,层内并行模型。transformer每一层内的矩阵乘被切分到多态GPU上。在更大的模型上表现不好；
更大的模型被被割到多个· 多GPU的服务器上，导致两个问题：

allreduce实现tensor并行，需要经过服务器间的链路（ inter-server links ），要比在一个多GPU服务器内的高带宽 NVLink 慢得多。
高程度得模型并行，会造成小矩阵相乘，可能会降低GPU利用率（为什么？？）（GEMMs？小矩阵的运算没有大矩阵高效，降低了GPU利用率？）

流水线模型并行，将模型的各层分不到CPU上。将一个batch 分成多个更小的 microbatches，通过流水线执行。
为了达到高效，需要更大的batch size

本文中介绍了一种新的流水线调度方式，可以提升小的batch size的效率。

本文解决如下问题：

在保证严格的优化器语义的同时，如何结合并行化技术来最大化给定 batch size的大模型训练吞吐效率。（strict optimizer semantics？？？）
How should parallelism techniques be combined to maximize the training throughput of large models given a batch size while retaining strict optimizer semantics?

PTD-P：将 pipeline，tensor，data parallesim 进行结合的技术。

研究了不同组件之间相互作用对吞吐量的影响。
在这些研究的基础上，我们提供了如何配置分布式训练的指导原则：

Different forms of parallelism interact in non-trivial ways:不同方式的并行以非平凡的方式交互。 tensor并行最好在多GPU的sever内部使用，才更高效。
不同的并行化策略会影响：通信量、 kernel的计算效率、worker的空闲时间、 pipeline 气泡。提出交错式的流水线调度，寄生吞吐量，限制内存占用。
超参数，例如micro batch的大小会影响内存占用、计算效率、流水线气泡大小
分布式训练是通信密集型的，如果 inter-node 通信慢，或者在通信密集处分区隔断，将严重影响效率

2 并行化模型

PTD-P ：我们将管道模型并行性和张量模型并行性（组合如图2所示）与数据并行性相结合，并称之为 PTD-P

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-neZO7MOh-1669727514425)(assets/markdown-img-paste-20220712113741693.png)]

2.1 数据并行化

每个工作节点都有完整的模型，输入的数据被切分，worker 定期聚合梯度，保证所有worker 看到版本一致的权重。
对于太大的模型，无法适配到单个工作节点上。

2.2 流水线并行模型

模型的层被分割到多个设备上。
当应用到有相同的transformer模块重复的模型时，每个设备分得相同数量得transformer 层。我们不考虑更多非对称得模型结构，分配策略会更难。

一个batch 被分成更小得 microbatches，然后在 microbatches 上执行流水线。

流水线模式需要保证：流水线方案需要确保输入在前向和后向传递中看到一致的权重版本，保证是严格优化的。

为准确保留严格优化器语义（strict optimizer semantics），我们引入了流水线周期性刷新，以便优化步骤跨设备同步。

在设备空闲时，空闲时间被称为流水线气泡(pipeline bubble)，要使得它足够小。

Asychromous （异步）和 bounded-staleness （边界陈旧）的方式比如：PipeMare，PipeDream 和 PipeDram-2BW 完全消除了 flushes ，但是松弛了权重更新语义。

有几种可能的 microbatch 跨设备前推和后推调度策略。对 bubble size，通信，内存占用进行 trade off

2.2.1 Default Schedule

GPipe 第一次执行将一个batch的所有 microbatches 向前传递，然后再将所有的 microbatches 向后传递。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-80mfGana-1669727514427)(https://raw.githubusercontent.com/novaCoder-zrk/Picture-Bed/master/img[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-YyQKcTFF-1669727514773)(…assetsmarkdown-img-paste-20220709213408346.png)]].png)

每个decive可以运行多个层的集合，称为 model chunck
对一个设备，气泡时间
$t_{pb} = (p-1)*(t_f + t_b)$

理想的处理时间（不包含气泡的时间）
$t_{id} = m*(t_f + t_b)$
m 是一个batch 分成的 microbatch 的数量

$\,time \, fraction(pipeline \,bubble\, size) = \frac{t_{pb}}{t_{id}} = \frac{p-1}{m}$

要使的气泡时间分数小，使得 m 远大于 p，对于大的m会导致搞得内存占用率，因为需要将中间激活值保存在内存中，在 m microbatch的生命中周期中，都需要保存

PipeDream-Flush
20221129101651
因为，在处理一个batch时，如果让所用的microbatch 都处于处理中的状态(in-flight microbatches)，内存占用会很大，因此需要限制管道中还没处理完的microbatch的数量。

PipeDream-Flush 调度，首先进入热身步骤，worker进行一定次数的向前传递，没有一下子将所有的forward都进行完，等有microbtach可以进行backward的时候立刻进行backward，
20221129102507
这个调度限制了处理中的 microbatch的数量，限制在了管道深度depth，而不是batch中microbatch数。

在经过了热身阶段后，每个worker进入了一种稳定状态，向前和向后传递交替进行（1F1B）。
20221129102618

这个新的调度的气泡时间是相同的，但是需要更少的激活被存储，不需要保存 m 个中间激活，只需要depth个。

当m远大于p的时候，PipeDream-Flush 比 Gpipe 内存效率更优。

2.2.2 带交错阶段的调度Schedule with Interleaved Stages

本文提出的新方法
为了减少流水线的气泡，设备可以为层的多个子集进行计算(model chunk)，而不是计算一组连续的层。也就是说一个设备可以执行多股效地chunk，而不是一整个连续的chunk，相当于流水线的细分。

例如，一个机器可以计算4层，设备1 计算1，2，9，10；设备2 计算3，4，11，12；每个设备计算两个模型块，每个模型块有2层。

每个设备承担更多计算阶段，每个阶段的计算任务却更少。

使用1F1B 和交错阶段调度；
要求microbatch的数量，是流水线并行度的整数倍。

20221129103147
将流水线的每个阶段减小，气泡就小了。但是数据之间的传输通信速度又会造成代价。

以前每个device只有一个stage（或者说 model chunk），如今每个device都有 v 个model chunk

因此气泡的时间减少为 $t_{pb}^{int.} = \frac{(p-1)(t_f + t_b)}{v}$

$\,time \, fraction(pipeline \,bubble\, size) = \frac{t_{pb}^{int.}}{t_{id}} = \frac{1}{v}\cdot\frac{p-1}{m}$

将流水线细分，可以提高流水线的并行度，会减少流水线气泡的大小，但代价是增加了机器之间的通信量

2.3 Tensor Model Parallesim

Tensor Model 会将单独的层，划分到多个机器上
Megatron 划分策略：

用于transformer模型层的划分
一个transformer层由一个self-attention块，后跟一个两层的多层感知器（MLP）组成。

多层感知机的划分策略
20221129103948

split A= [A1,A2]
Y = [Y1,Y2] = [GeLU(X A1), GeLU(X A2)]
B矩阵按行划分

multi-head attention 的划分策略

在multi-head attention中，因为固有的并行性，每个GPU可以只负责一个头的attention的计算

3 PERFORMANCE ANALYSIS OF PARALLELIZATION CONFIGURATIONS

并行化配置的性能分析
不同维度的并行化都进行了内存占用，设备利用率和通信量的tradeoff

讨论tradeoff
提出了有关流水线气泡大小的的分析模型
定量描述通信时间行为（communication time behaves），并且提出了通信代价量的模型，但没有提出直接通信时间代价模型，分层通信建模比较困难
据目前所知，这是第一个分析不同维度并行化交互性能的工作。

3.1 符号

（p,t,d）：p 流水线模型并行大小，t tensor模型并行大小，d 数据并行大小
n GPU的数量 p·t·d = n
B: 全局batch size ，input的大小
b: microbatch 的大小
$\frac{1}{b} ·\frac{B}{d}$

3.2 Tensor and Pipeline Model Parallelism

流水线模型，有更廉价的点对点通信； Tensor并行模型使用all-reduce 通信
有n个GPU，不考虑data 并行，data 并行度设为1,d = 1. $t\cdot p = n$

流水线气泡大小：
$\frac{p-1}{m} =\frac{n/t-1}{m}$
m 是确定的，当t增加，流水线的气泡减少。

流水线模型，在每对连续设备间需要进行通信，为每个 microbatch 的通信总量为为 bsh。

s 是序列长度 (sequence length)，s怎么理解？每个样本序列的长度？ h是（hidden size）

大小为 bsh的tensor需要在t个设备之间进行all-reduce，forward和backward个进行一次，每个设备每层的总传输量为 $8bsh(\frac{t-1}{t})$
每个设备有 $l^{stage}$ 层，每个设备总的传输量为 $l^{stage}\cdot(8bsh(\frac{t-1}{t}))$

tensor 并行增加了设备之间的通信量。
当t大于单个接单的GPU数量时，需要节点之间进行通信。更慢的节点之间的链路的影响是很大的。

要点#1
当使用 g-GPUS服务器时，将tensor 并行度控制在g之内，使用pipeline 并行来跨服务器扩展模型

3.3 Data and Model Parallelism

3.3.1 流水线模型并行

$m = B / (d b) = b^{'} / d$
其中， $\frac{B}{b}$ $p = n / (t \cdot d) = n / d$

$\,time \, fraction(pipeline \,bubble\, size) = \frac{p-1}{m} = \frac{n/d - 1}{b'/d} = \frac{n-d}{b'}$

当 d 变大，流水线气泡变小。
虽然将 d增加，会减少气泡，但是将 d增加到 n 是不太现实的，因为一个模型的大小会超出单个GPU的内存

20221129115718

如果数据并行all-reduce 通信，不会随着 d的升高而剧烈增加，比如 ring方法的实现，总的吞吐量会增加
分析 B batch size 增加的影响。 B增加，b’增加， (n - d)/ b’ 减少，因此吞吐量增加。
B增大，data 并行的all-reduece 频率降低，更提高了吞吐率

3.3.2 Data and Tensor Model Parallelism

因为 tensor 模型并行，每个microbatch 都需要执行 all-reduce 通信。
因为数据并行，每个batch 都要进行代价较大的 all-reduce 通信。
GPU负责计算一个层的一部分，如果模型的层不够大划分给一个GPU的矩阵计算太小，GPU的利用率会不高
tensor 的通信代价比 data并行的通信代价大得多。

要点#2
data 并行，只需要在 batch 层面进行 all-reduce。但是tensor 并行，需要在每个micro batch 层面进行all-reduce。
总的模型并行度 M = tp,tensor 和 pipeline 主要是为了将大模型切分以适应GPU的主存，但是数据并行，主要是为了将训练扩展至更多的GPU。

3.4 Microbatch Size

microbatch 的大小 b 同样也可以影响模型训练的吞吐率。

想在给定(p,t,d) 和 batch size B的情况下，确定最优的 microbatch 的大小 b
无论microbatch的大小如何，数据并行的通信量是不变的
给定 $t_f(b) \quad t_b(b)$ ,分别表示了单个microbatch 的大小与向前、向后的传递计算时间。不考虑通信的损耗，用于计算一个 batch的时间为

20221129121216

microbatch的大小，影响计算操作 forward backward的强度，也影响气泡的大小。batch确定，d 确定，b 越大，microbatch的数量越少。 b 越大，forward，backward 计算量更大

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-BHdRNdP8-1669727514434)(https://raw.githubusercontent.com/novaCoder-zrk/Picture-Bed/master/img[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-2lmtSckd-1669727514735)(…assetsmarkdown-img-paste-20220710215139274.png)]].png)
20221129121307

要点#3
对于 microbatch size b的优化，要考虑吞吐率，内存占用，pipeline深度，data 并行度 d，batch大小B

3.5 Activation Recomputation

激活值的重新计算是计算量和内存占用的 tradeoff
只存给定流水线阶段输入的激活值，而不是保存整个集合上的大得多的激活值
在训练相当大的模型时，重新计算激活值，保证内存占用在一个较低的可接受的水平。
激活值检查点的数量不影响吞吐率，但是影响内存占用。
大多数情况下，每一到两 transformer层设置一个检查点最优
其他技术，比如（激活分区）activation partitioning 可与和 tensor model 并行一起使用以降低激活值得内存占用

4 IMPLEMENTATION

实现了PTD-P，作为 Megatron-LM 的代码库
使用PyTorch 构建
使用NCCL 作为设备间的通信
以优化通信和计算为目的

4.1 Communication Optimizations 通信优化

使用A100时，配备了8 个IB网卡，但是在两个sever之间，只能进行一对GPU之间的通信，这就很难充分利用8个卡。

使用 tensor模型和流水线模型，来降低跨界节点通信的额外开销。
每个transformer 层的输出是重复的，因为经过 all-reduce后（如 MLP block 中的g）每个GPU上的参数是一样的，如果8个卡，每一对之间发送的数据是相同的，并且sever之间的传输速度慢，会导致效率降低。

20221129201249

scatter/gather communication optimization

每个卡只需要发送 output的一部分，每个卡发送的数据量相等，在接收端通过NVLink 执行all-gather
将发送端的tensor1 划分成大小相等的块，每个rank 只发送一块到下个设备，通过 InfiniBand card
有 8 个tensor 并行，每个块是原先的 1/8大。在接受端，使用 NVLink 上的 all-gather 来得到完整的tensor。
通过 scatter-gather 通信优化，在两个阶段之间的总的通信量减少到了 $\frac{bsh}{t}$ , t 是tensor模型并行的大小， s是序列长度， h 是隐藏层大小

4.2 计算优化

三种特定模型的优化

首先：更改transformer 层的数据布局，来避免内存占用较多的转置操作(transpose) , 并且可以使用 strided batched GEMM ，将数据布局从[b,s,a,h] 和修改为[s,b,a,h] 。
其中 b,a,s,h, 分别为 batch 大小，序列长度， attention-head ，hidden-size 的维度？？为什么原先的数据布局是这样？？
其次，使用 PyTorCH JIT 为一系列元素级操作生成融合核，将几个kernel 融合在一起
第三，创建了两个自定义核，来实现 scale, mask, and softmax (reduction) 的融合。一个用来支持一般的mask ，另一个用于隐含式因果mask (implicit causal masking)