详解MegatronLM序列模型并行训练(Sequence Parallel)_megatron-lm

1. 背景介绍

MegatronLM的第三篇论文【Reducing Activation Recomputation in Large Transformer Models】是2022年出的。在大模型训练过程中显存占用过大往往成为瓶颈，一般会通过recomputation重计算的方式降低显存占用，但会带来额外的计算代价。这篇论文提出了两种方法，分别是sequece parallel和selective activation recomputation，这两种方法和Tensor并行是可以相结合的，可以有效减少不必要的计算量。

下图中绿色部分表示不同模型中需要用于保存activation需要的显存大小，蓝色部分表示不同模型中需要用于保存parameter和optimizer state需要的显存大小。红色线表示A100的显存大小80G。

2. Pipeline Parallel详细介绍

2.1 估算Transformer Activation Memory大小

以Transformer结构为例估算Activation Memory大小，这里的Activation定义是指前向和反向梯度计算中创建的所有tensor。按这个定义来说，计算不包含模型参数大小和优化器中状态大小，但是包含dropout op用到的mask tensor。

一个Transformer块中由一个Attention块和一个MLP块组成，中间通过两个LayerNorm层进行连接。在Transformer中用到的参数表示如下：

Attention模块的计算公式如下：

A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V

$$\begin{gather*} Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V \end{gather*}$$ Attention(Q,K,V)=softmax(dk​ ​QKT​)V​

对于Attention块来说，输入的element个数为sbh个，每个element以16-bit的浮点数(也就是2 bytes)来进行存储的话，对应输入的element大小为2sbh bytes，后续计算默认都是按bytes为单位进行计算。

Attention块中包含一个self-attention块、一个linear线性映射层和attention dropout层。对于linear线性映射层来说需要保存输入的Activation大小为2sbh, 对于attention dropout层需要mask的大小为sbh(对于一个元素的mask只用1个bytes即可)，对于self-attention块的Activation Memory的计算有以下几块：

• $Query(Q),Key(K),Value(V)$矩阵相乘：输入input是共享的，元素个数为sbh个，总大小是 2sbh bytes。
• $Q{K}^T$ 矩阵相乘：需要分别创建保存 $Q$ 和 $K$ 的矩阵，每个矩阵元素总大小为 2sbh bytes, 总共大小为 4sbh bytes。如下图以b=1, s=2, h=6为例，输入$X$元素个数为1 * s * h = 12个，计算完后$Q$ 和 $K$ 的矩阵中元素个数各有 1 * s * h = 12个，总元素大小为2 * 2 * b * s * h = 48 bytes。

• softmax的输出总的元素大小为$2a{s}^{2}b$ bytes, 分别计算每个Head头的$Q_n\times K_n$ 的乘积。计算公式如下, 图中计算以b=1, s=2, h=6, a=2为例：
  

• 在softmax后还有dropout的mask层大小，mask矩阵的大小与softmax的输出一样，元素个数都是 $a{s}^{2}b$ 个，但mask单个元素的大小只用1 bytes即可，总的大小为 $a{s}^{2}b$ bytes

• softmax的输出也会用于反向的计算，需要缓存下来，对应大小也是 $2a{s}^{2}b$

• $V$ 矩阵的大小之前没有统计，和 $Q$、$K$矩阵一样，大小也是2sbh bytes

综上，Attention Block总的大小为 11sbh + 5as^2b bytes。

MLP的Activation大小计算：MLP中有两层线性layer，分别存储输入矩阵大小为 $2sbh$ bytes和 $8sbh$ bytes；GeLU的反向也需要对输入进行缓存，大小为 $8sbh$ bytes; dropout层需要 sbh bytes; 总大小为 19sbh。

LayerNorm的Activation大小计算：每个LayerNorm层的输入需要 $2sbh$ 大小，有两个LayerNorm层，总大小为 4sbh bytes.

最终transformer网络中一层(含Attention/MLP/LayerNorm)的Activation总的大小为：

A c t i v a t i o n M e m o r y P e r L a y e r = s b h ( 34 + 5 a s h )

$$\begin{gather} ActivationMemoryPerLayer = sbh \left( 34 + 5 \frac{as}{h} \right) \end{gather}$$ ActivationMemoryPerLayer=sbh(34+5has​)​​

注意: 这里公式(1)计算的Activation总和是在没有应用模型并行策略的前提下进行的。

2.2 Tensor Parallel的Activation Memory计算

如下图，在Tensor模型并行中只在Attention和MLP两个地方进行了并行计算，对于Attention(Q/K/V)和MLP(Linear Layer)的输入并没有并行操作。图中 $f$ 和 $\overline{f}$ 互为共轭(conjugate)，$f$ 在前向时不做操作，反向时执行all-reduce; $\overline{f}$ 在前向时执行all-reduce, 反向时不做操作。

参虑上Tensor并行的话(Tensor并行度为 $t$)，并行部分有MLP的Linear部分($18sbh$ bytes)和Attention的QKV部分($6sbh$ bytes)， ActivationMemoryPerLayer相比公式(1)中的值降为：
A c t i v a t i o n M e m o r y P e r L a y e r = s b h ( 10 + 24 t + 5 a s h t )

$$\begin{gather} ActivationMemoryPerLayer = sbh \left( 10 + \frac{24}{t} + 5 \frac{as}{ht} \right) \end{gather}$$ ActivationMemoryPerLayer=sbh(10+t24​+5htas​)​​

2.2 Sequence Parallel

在Tensor模型并行基础上提出了Sequence Parallel，对于非Tensor模型并行的部分在sequence维度都是相互独立的，所以可以在sequence维度上进行拆分(即sequence parallel)。拆分后如下图，$f$ 和 $\overline{f}$ 替换为 $g$ 和 $\overline{g}$， $g$ 和 $\overline{g}$ 也是共轭的，$g$ 在前向是all-gather通信，反向是reduce-scatter通信；$\overline{g}$在前向是reduce-scatter, 反向是all-gather通信。

接下来以MLP为例，详细说明拆分步骤。MLP层由两个Linear层组成，对应的计算公式如下, 其中 $X$ 的大小为 $s\times b\times h$ ; $A$ 和 $B$ 是Linear的权重weight矩阵，大小为 $h\times 4h$ 和 $4h\times h$。

Y = L a y e r N o r m ( X ) Z = G e L U ( Y A ) W = Z B V = D r o p o u t ( W )

$$\begin{gather*} \begin{aligned} Y &= LayerNorm(X) \\ Z &= GeLU(YA) \\ W &= ZB \\ V &= Dropout(W) \\ \end{aligned} \end{gather*}$$ YZWV​=LayerNorm(X)=GeLU(YA)=ZB=Dropout(W)​​

如下图，切分时说明如下：

1. 对 $X$ 按sequence维度切分，$X=\left[X_1^s,X_2^s\right]$，LayerNorm的结果 $Y=\left[Y_1^s,Y_2^s\right]$；
2. 由于接下来的GeLU不是线性的，所以要进行all-gather操作，计算 $Z=GeLU(YA)$；
3. 对 $A$ 进行列切分的tensor并行，得到结果 $Y{A}_1^c$ 和 $Y{A}_2^c$
4. 对 $B$ 进行行切分的tensor并行，得到结果 $Z_1^h{B}_1^r$ 和 $Z_2^h{B}_2^r$
5. 得到 $W_1$ 和 $W_2$ 后进行累加操作(reduce-scatter)

对应的计算公式如下：

[ Y 1 s , Y 2 s ] = L a y e r N o r m ( [ X 1 s , X 2 s ] ) Y = g ( Y 1 s , Y 2 s ) [ Z 1 h , Z 2 h ] = [ G e L U ( Y A 1 c ) , G e L U ( Y A 2 c ) ] W 1 = Z 1 h B 1 r W 2 = Z 2 h B 2 r [ W 1 s , W 2 s ] = g ‾ ( W 1 , W 2 ) [ V 1 s , V 2 s ] = [ D r o p o u t ( W 1 s ) , D r o p o u t ( W 2 s ) ]

$$\begin{gather} \begin{aligned} \left[ Y^s_1, Y^s_2 \right] &= LayerNorm([X^s_1, X^s_2]) \\ Y &= g(Y^s_1, Y^s_2) \\ \left[ Z^h_1, Z^h_2 \right] &= [GeLU(YA^c_1), GeLU(YA^c_2)] \\ W_1 &= Z^h_1 B^r_1 \\ W_2 &= Z^h_2 B^r_2 \\ \left[ W^s_1, W^s_2 \right] &= \overline{g}(W_1, W_2) \\ \left[ V^s_1, V^s_2 \right] &= [Dropout(W^s_1), Dropout(W^s_2)] \\ \end{aligned} \end{gather}$$ [Y1s​,Y2s​]Y[Z1h​,Z2h​]W1​W2​[W1s​,W2s​][V1s​,V2s​]​=LayerNorm([X1s​,X2s​])=g(Y1s​,Y2s​)=[GeLU(YA1c​),GeLU(YA2c​)]=Z1h​B1r​=Z2h​B2r​=g​(W1​,W2​)=[Dropout(W1s​),Dropout(W2s​)]​​​

Tensor并行在一次前向和后向总共有4次的all-reduce操作，在Sequence并行一次前向和后向总共有4次all-gather和4次reduce-scatter操作。ring all-reduce 执行过程中有两步，先是一个reduce-scatter然后跟着一个all-gather，Sequence并行相比没有引入更多的通信代价。一个使用reduce-scatter和all-gather实现all-reduce的Python代码示例如下：

import torch
import torch.distributed as dist

# 初始化进程组
dist.init_process_group(backend='gloo')

# 获取进程组中的进程数和当前进程的排名
world_size = dist.get_world_size()
rank = dist.get_rank()

# 定义输入和输出张量
x = torch.tensor([1, 2, 3, 4])
result = torch.zeros_like(x)

# 使用 reduce_scatter 将每个进程的输入张量的部分和归约到每个进程的输出张量
dist.reduce_scatter(input_list=[x], output=result)

# 使用 all_gather 将每个进程的输出张量收集到所有进程中
output_list = [torch.zeros_like(result) for _ in range(world_size)]
dist.all_gather(output_list, result)

# 在每个进程上打印结果
print(f"Process {rank}: {output_list}")

# 清理资源
dist.destroy_process_group()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

通过使用sequence parallel和tensor parallel以后，ActivationMemoryPerLayer相比公式(2)的值再次减少，相比公式(1)相当于对所有的ActivationMemory进行Tensor并行, 即 $\frac{ActivationMemoryPerLayer}{t}$：

A c t i v a t i o n M e m o r y P e r L a y e r = s b h ( 10 t + 24 t + 5 a s h t ) = s b h t ( 34 + 5 a s h )

$$\begin{gather} \begin{aligned} ActivationMemoryPerLayer &= sbh \left( \frac{10}{t} + \frac{24}{t} + 5 \frac{as}{ht} \right) \\ &= \frac{sbh}{t} \left( 34 + 5 \frac{as}{h} \right) \\ \end{aligned} \end{gather}$$ ActivationMemoryPerLayer​=sbh(t10​+t24​+5htas​)=tsbh​(34+5has​)​​​

2.3 Pipeline Parallel

加上Pipeline Parallel后，对具有 $L$ 层的layer的transformer来说，Pipeline Parallel并行度为 $p$, 对应会分为 $\frac{L}{p}$ 组(即stage个数)。以PipeDream中的1F1B调度为例，要完成初始化的话，第1个stage必须处理完 $p$ 个micro-batch，让其他stage至少有1个micro-batch在处理，也就是要缓存 $p$ 个micro-batch的activation。由于每个stage都有 $\frac{L}{p}$ 个Layer，一共需要 $p\times \frac{L}{p}=L$ 个layer的activation信息，对应总的计算如下：

T o t a l A c t i v a t i o n M e m o r y = s b h L t ( 34 + 5 a s h )

$$\begin{gather} TotalActivationMemory = \frac{sbhL}{t} \left( 34 + 5 \frac{as}{h} \right) \\ \end{gather}$$ TotalActivationMemory=tsbhL​(34+5has​)​​

当然这里的公式(5)的ActivationMemory的计算没有加上EmbeddingLayer、最后的LayerNorm和输出的OutputLayer。加上这三部分的结果会略大于公式(5), 但以22B参数模型来说只增加0.01%的大小，这部分可忽略，证明请参考原论文。未计算部分如下图红色部分：

3. 可选Activation重计算介绍(Selective Activation Recomputation)

在后向过程中通过重计算方式重新计算前向结果来节省显存大小，这种方式文中称为full activation recomputation，以transformer为例会增加30%~40%的计算量。Selective的方式主要思路是选择 FLOPs 计算量小，且activation占用大的算子进行重计算，这里的 FLOPs 的衡量标准是GEMM的计算量大小。以公式(5)为例，针对大模型来说 $5as/h>34$, 如果重计算这部分layer的话可以减少快一半的activation大小。对于GPT-3来说，这种方式可以减少70%的activation显存大小，同时只增加了2.7%的 $FLOPs$ 计算量。采用Selective Activation Recomputation后，公式(5)的结果可以减少为：

T o t a l   r e q u i r e d   m e m o r y = 34 s b h L t

$$\begin{gather*} Total\ required\ memory = 34 \frac{sbhL}{t} \\ \end{gather*}$$ Total required memory=34tsbhL​​

以下是不同方法组合下Activation Memory占用情况：

4. 参考

• Reducing Activation Recomputation in Large Transformer Models
• The Illustrated Transformer