知识站

Sequence Parallelism序列并行（一）

· 3 min read

原始论文 Reducing Activation Recomputation in Large Transformer Models

序列并行

序列并行是在张量并行的基础上进行的进一步深度优化，旨在减少“中间值”带来的显存占用(“中间值”是反向传播所必需的。如果不保存这些中间值，在反向传播过程中就必须重新执行前向计算来生成它们，这会显著增加训练的时间开销)。^1

关于 Transformer 各层的显存占用分析，请参考我的文章：

关于张量并行，请参考我的文章：

“中间值”显存占用分析

符号约定

符号 含义 符号 含义
aa number of attention heads pp pipeline parallel size
bb microbatch size ss sequence length
hh hidden dimension size tt tensor parallel size
LL number of transformer layers vv vocabulary size

未采用并行机制的 Transformer 架构

架构如下图，我们主要关注图中灰色部分的 Transformer Layer，因为它会重复 LL 层，占显存开销的大头。

假设采用 FP16 / BF16 精度，下面我们对中间值的显存占用进行拆解分析。

Attention Block

中间值 大小
Q/K/V projection 共享输入 2sbh2sbh
QK⊤QK^\top 需要存 Q 和 K 4sbh4sbh
Softmax 输出 2as2b2as^2b
Softmax dropout mask as2bas^2b
Attention over V: dropout 输出 2as2b2as^2b
V 2sbh2sbh
Output linear projection 输入 2sbh2sbh
Attention dropout mask sbhsbh

中间值合计占用：

Attention=11sbh+5as2b \text{Attention} = 11sbh + 5as^2b

MLP Block

项 大小
第一个 linear 输入 2sbh2sbh
第二个 linear 输入 8sbh8sbh
GeLU 输入 8sbh8sbh
MLP dropout mask sbhsbh

中间值合计占用：

MLP=11sbh+5as2b \text{MLP} = 11sbh + 5as^2b

LayerNorm

两个 LayerNorm，每个存输入 2sbh2sbh，所以

LayerNorms=4sbh\text{LayerNorms}=4sbh

合计

Total=sbh(34+5ash)\text{Total}=sbh(34+\frac{5as}{h})

采用张量并行机制的 Transformer 架构

采用张量并行机制的 Transformer 架构如下图：

引入张量并行度为 tt，进行中间值占用的显存分析：

Attention Block

项 原大小 Tensor Parallel 后
Q/K/V projection 共享输入 2sbh2sbh 不切，仍是 2sbh2sbh
QK⊤QK^\top 需要存 Q 和 K 4sbh4sbh 切，变成 4sbht\frac{4sbh}{t}
Softmax 输出 2as2b2as^2b 切，变成 2as2bt\frac{2as^2b}{t}
Softmax dropout mask as2bas^2b 切，变成 as2bt\frac{as^2b}{t}
Attention over V: dropout 输出 2as2b2as^2b 切，变成 2as2bt\frac{2as^2b}{t}
VV 2sbh2sbh 切，变成 2sbht\frac{2sbh}{t}
Output linear projection 输入 2sbh2sbh 切，变成 2sbht\frac{2sbh}{t}
Attention dropout mask sbhsbh 不切，仍是 sbhsbh

因此 attention 部分变成：

Attention=3sbh+8sbht+5as2bt\text{Attention} = 3sbh + \frac{8sbh}{t} + \frac{5as^2b}{t}

MLP Block

项 原大小 Tensor Parallel 后
第一个 linear 输入 2sbh2sbh 不切，仍是 2sbh2sbh
第二个 linear 输入 8sbh8sbh 切，变成 8sbht\frac{8sbh}{t}
GeLU 输入 8sbh8sbh 切，变成 8sbht\frac{8sbh}{t}
MLP dropout mask sbhsbh 不切，仍是 sbhsbh

因此 MLP 部分变成：

MLP=3sbh+16sbht\text{MLP}=3sbh+\frac{16sbh}{t}

LayerNorm

张量并行未应用到 LayerNorm，故存储的中间值不变，即两个 LayerNorm，每个存输入 2sbh2sbh：

LayerNorms=4sbh\text{LayerNorms}=4sbh

合计

Total=sbh(10+24t+5asht)\text{Total}=sbh(10+\frac{24}{t}+\frac{5as}{ht})

采用序列并行 + 张量并行的 Transformer 架构

序列并行进一步将公式中的 10sbh10sbh 项并行化。

其具体实现方式如下图。需要说明的是，序列并行并非独立运作，而是与张量并行协同工作的。

我们以已经张量并行的 MLP 块和它前面的 LayerNorm 和 它后面的 Dropout 层为例，如下图所示，其中有两个 GPU：

[Y1s,Y2s]=LayerNorm([X1s,X2s]),Y=g(Y1s,Y2s),[Z1h,Z2h]=[GeLU(YA1c),GeLU(YA2c)],W1=Z1hB1randW2=Z2hB2r,[W1s,W2s]=gˉ(W1,W2),[V1s,V2s]=[Dropout(W1s),Dropout(W2s)] \begin{align*} [Y_1^s, Y_2^s] &= \text{LayerNorm}([X_1^s, X_2^s]), \cr Y &= g(Y_1^s, Y_2^s), \cr [Z_1^h, Z_2^h] &= [\text{GeLU}(Y A_1^c), \text{GeLU}(Y A_2^c)], \cr W_1 &= Z_1^h B_1^r \quad \text{and} \quad W_2 = Z_2^h B_2^r, \cr [W_1^s, W_2^s] &= \bar{g}(W_1, W_2), \cr [V_1^s, V_2^s] &= [\text{Dropout}(W_1^s), \text{Dropout}(W_2^s)] \end{align*} \

此处的一般维度设定如下：

  • 输入：X:s×b×hX: s \times b \times h
  • 投影矩阵：A:[h,4h]A: [h, 4h]，B:[4h,h]B: [4h, h]

符号说明：

  • Y1sY_1^s：将 YY 按 ss 方向切分成两份，其中第一份（存放在第一个 GPU 上）记为 Y1sY_1^s。
  • A1cA_1^c：将 AA 按列方向切分成两份，其中第一份（存放在第一个 GPU 上）记为 A1cA_1^c。
  • B1rB_1^r：将 BB 按行方向切分成两份，其中第一份（存放在第一个 GPU 上）记为 B1rB_1^r。

这些符号不必刻意记忆。关键是，序列并行的可行性在于：LayerNorm 和 Dropout 按 ss 方向（即序列维度）切分后，并不会影响其计算结果。中间涉及的其他切分操作，只是为了配合 LayerNorm 和 Dropout 按 ss 维度切分而做的必要调整。按照维度动手推一遍即可理解。

为了完成整个操作，我们需要设定通信算子 gg 和 g‾\overline{g}。其中：

  • gg: forward 时为 all-gather; 根据通信算子和伴随通信算子[1]的性质，快速得到 backward 时为 reduce-scatter
  • g‾\overline{g}: forward 时为 reduce-scatter; 根据通信算子和伴随通信算子的性质，快速得到 backward 时为 all-gather

因此，对于图 5 所示的“张量并行 + 序列并行”架构，每个 Transformer 层在每轮前向与反向传播中，共需要执行 4 次 all-gather 和 4 次 reduce-scatter。

而在图 4 的纯张量并行架构中，每层每轮前向与反向传播总共需要 4 次 all-reduce。

乍看之下，通信操作似乎变多了，但由于一次 all-reduce 的通信量相当于一次 all-gather 与一次 reduce-scatter 之和[2]，因此总的通信量实际上并未增加。

此外，还需注意一点：在 MLP 块中，完整的 YY 仍然在每个 GPU 上都有备份。为解决这一问题，可将 YY 也按 ss 维度切分为 Y1Y_1、Y2Y_2。不过，在计算 GeLU\text{GeLU} 之前，仍需执行一次 all-gather 操作。本文的做法是将这次 all-gather 与反向传播中计算 YY 的梯度（该梯度的计算仅需 AA 即可完成）进行重叠，从而降低通信延迟[3]。

Attention 块与 MLP 块类似，这里就不赘述了，论文中其实也只是以 MLP 块举例。

合计

采用“序列并行+张量并行”后，Transformer 架构中中间值的显存占用变为：

Total=sbh(10t+24t+5asht)=sbht(34+5ash)\text{Total}=sbh(\frac{10}{t}+\frac{24}{t}+\frac{5as}{ht})=\frac{sbh}{t}(34+\frac{5as}{h})

后记采用激活重计算

上述的显存占用公式中，其实最后一项:

sbht(5ash)\frac{sbh}{t}(5\frac{as}{h})

占用的显存显著大于前面的常数项。为了量化这一点，让我们考虑一下 GPT-3 和 MT-NLG 模型。对于 GPT-3，a=96a = 96，s=2048s = 2048，h=12288h = 12288，因此 5ash=80\frac{5as}{h} = 80。对于 MT-NLG，a=128a = 128，s=2048s = 2048，h=20480h = 20480，所以 5ash=64\frac{5as}{h} = 64。

所以，可以将 sbht(5ash)\frac{sbh}{t}(5\frac{as}{h}) 进行激活重计算[4]，从而达到显存开销和训练时间开销的一种权衡。

评论