流水线并行#

NOTE

重点

1. 画流程图，分析 Bubble 占用时间
2. 和 batch norm 的联系

首先是模型并行，就是将模型（尽可能均匀地）拆成几份放到不同的 GPU 上，然后依次执行 forward 和 backward，如下：

实际有效时间占比：

• $K$：GPU 数量
• $t_{f},t_{b}$：前向/后向传播时间

利用率很低。并且空间复杂度为：

• $N$：训练数据大小
• $L$：模型层数
• $d$：模型宽度

为了解决 GPU 利用率低的问题，引入流水线并行：主要使用 micro batch 技术，将一个 batch（现称为 mini batch）拆分为 $M$ 个 micro batch，如下

同上推导

Gpipe 中证明当 $M\geq 4K$ 时利用率就很高了。

但是 Batch Normalization 会有影响，Gpipe 中的处理方法是：训练时使用 micro batch 中的均值和方差，测试时使用统计的 mini batch 的平均和方差。Layer norm 则不受影响。

再解决 GPU 内存问题，使用 re-materalization 技术，即用时间换空间，几乎不存中间结果，当 backward 时再重新算一遍 forward。

GPU 峰值时刻，空间复杂度为：

数据并行#

核心是把模型拷贝几份到不同的 GPU 上，然后划分数据分别输入，计算梯度后聚合并分发到每一块 GPU 上。

张量并行#

针对 transformer 提出的方式，对权重 $W$ 进行横向/纵向切分，放到不同的 GPU 上，对 MLP、attention、embedding、cross-entropy 层进行处理。

设输入为 $X$，形状 $(b,s,h)$，权重 $W$，形状 $(h,h')$，$Y=XW$，形状 $(b,s,h')$

1. 按行切分：
   1. forward：将 $X$ 纵向分为 $X_{1},X_{2}\dots$，形状 $(b,s,h' / N)$，权重横向分为 $W_{1},W_{2},\dots$，形状 $(h / N, h')$，则 $Y = Y_{1}+Y_{2}+\cdots=X_{1}W_{1}+X_{2}W_{2}+\cdots$
   2. backward：将 $\frac{ \partial L }{ \partial Y }$ 同时广播到各个 GPU 上，分别计算出 $\frac{ \partial L }{ \partial W_{i} }$ 之后进行 concatenating 即可
2. 按列切分：
   1. forward：将 $W$ 纵向切分为 $W_{1},W_{2},\dots$，形状 $(h, h' / N)$，则 $Y=\mathrm{concat}[XW_{1},XW_{2},\dots]$
   2. backward：

举例：

• 对 MLP，如 $Y=\mathrm{GELU}(XA)B$，可将 $A$ 按列拆分，将 $B$ 按行拆分。forward 中在 $g$ 处进行一次 All-Reduce，backward 过程中在 $f$ 处进行一次 All-Reduce，每次通讯量 $2\Phi=2*b*s*h$.

• 对 self-attention，有类似的处理方法，将 attention 层按列拆分，linear 层按行拆分，$Z=\mathrm{Dropout}(\mathrm{Self\_Attention}(X)B)$，仍然是 forward 和 backward 各一次通讯，总通讯量 $4\Phi$

• 对于 embedding encoding，由 word embedding 和 positional embedding 构成，形状 $(v,h)$ 和 $(\mathrm{max\_s},h)$。positional embedding 对短文本的模型可忽略不计。word embedding 被拆分为 $n$ 份存在各个 GPU 上，若无法查询则赋值 $0$，然后进行一次 All-Reduce 即可
• 对于 embedding decoding，和输入层共用一个 embedding 即可。注意如果 embedding 的 encoding 和 decoding 阶段在不同的 GPU 上，则应该在参数更新前进行一次 All-Reduce。
• 对于 cross entropy，由于如果在 embedding decoding 之后对结果 $Y_{1},Y_{2},\dots$ 进行一次 All-Gather 得到最终结果 $Y$，那么通讯量为 $b*s*v$，当词表过大时，开销会很大，可以进行以下优化，使通讯量降为 $N+b*s$