GPipe - efficient training of giant neural networks using pipeline parallelism

神经网络基础

![[nano-MLP]]
如图是一个基础的神经网络结构, 我们这里首先要明确几个函数:

每一层内部的计算是 $h_i := \sigma(z_i)$ , $z_i := W_i h_{i-1} + b_i$
我们最终的输出函数是 $L :=$ loss(sample; forward)
forward: 指的是从输入样本数据 $h_0$ 到最终输出 $\sigma(z_2)$ 的所有计算过程
backward: 从 $L$ 开始不断计算每一层中的 $\frac{\partial L}{\partial W_i}$ 和 $\frac{\partial L}{\partial b_i}$ 的方向梯度, 然后用这个来不断更新 $W_i,\ b_i$

现在来拆解一下反向传播的过程: 基本根据的是莱布尼茨链式法则 (chain-rule)

$\frac{\partial L}{\partial W_2}= \frac{\partial L}{\partial z_2}\frac{\partial z_2}{\partial W_2} =\frac{\partial L}{\partial z_2}h_1^T$ , 这里的 $h_1^T$ 的转置暂时不细看, 理解为 $h_1$ 就行了;
$\frac{\partial L}{\partial b_2}= \frac{\partial L}{\partial z_2}\frac{\partial z_2}{\partial b_2} =\frac{\partial L}{\partial z_2}$ , 因为 $z_2 = h_2 + W_2h_1$ 对 $b_2$ 求导就是 1
这里唯一需要的是 $\frac{\partial L}{\partial z_2}$ 这个值, 并且这个值本身是可计算的, 也就是我们只需要对 $L=$ loss(Sample; $\sigma(z_2)$ ) 这个函数是已知的
我们对这个 $\frac{\partial L}{\partial z_2}$ 暂时记录成 grad_2

这就完成了第一个模块的计算; 然后接着对前一个模块进行传递:

$\frac{\partial L}{\partial W_1}= \frac{\partial L}{\partial z_1}\frac{\partial z_1}{\partial W_1} =\frac{\partial L}{\partial z_1}h_0^T$ , 这里的 $\frac{\partial L}{\partial z_1} =\frac{\partial L}{\partial h_1} \frac{\partial h_1}{\partial z_1}$ , $h_1 := \sigma(z_1)$ 的导数可看作已知, $\frac{\partial L}{\partial h_1}$ 就是 $\frac{\partial L}{\partial z_2} \frac{\partial z_2}{\partial h_1} =$ grad_2 $\cdot W_2$ , 这一条其实是依赖于下游的模块的梯度反向传播的
总结来讲, 对中间每一层在 backward propagation 过程的独立计算, 需要传入的参数就是 $\frac{\partial L}{\partial z_{\text{local}}} := \frac{\partial L}{\partial z_{\text{local+1}}} \cdot W_{\text{local+1}}$ , 这就是梯度的反向传播过程
$\frac{\partial L}{\partial b_1}= \frac{\partial L}{\partial z_1}\frac{\partial z_1}{\partial b_1} =\frac{\partial L}{\partial z_1}$

论文定位

GPipe 的核心目标不是单纯 “把模型切到多张 GPU 上”, 而是提出一种 通用, 可靠, 可扩展的 pipeline parallelism 训练抽象, 让能够表示成 sequence of layers 的深度网络可以被拆分到多个 accelerator 上训练, 同时尽量保持:

训练语义一致
显存占用可控
吞吐随设备数扩展
通信开销较低

论文解决的核心问题

随着神经网络越来越大, 单个 GPU / TPU 常常无法容纳完整模型及其训练所需的激活, 梯度, 优化器状态;已有的一些 model parallel 方法往往:

很依赖具体模型结构
很依赖特定任务
需要复杂的手工并行实现
训练语义可能随并行方式变化

GPipe 试图回答的问题是:

能不能把大模型切到多设备上训练?
能不能尽量保持训练语义和普通 mini-batch SGD 一致?
能不能把这种并行方式做成一种通用接口, 而不是某个模型的 ad-hoc trick?

insight

把模型统一抽象成 layer sequence

整个网络由 $L$ 层组成, 每层包含:

forward computation function $f_i$
parameters $w_i$

于是整个模型可以表示为:

$h_1 = f_1(x)$
$h_2 = f_2(h_1)$
$\dots$

把层序列切成多个连续 partition / stage

给定 partition 数 $K$ , GPipe 将 $L$ 层切成 $K$ 个连续的 cell 或 stage:
$p_0, p_1, \dots, p_{K-1}$
每个 stage 包含一段连续层; 例如一个 8 层模型在 $K=4$ 时可以切成:

stage 0: layers 1-2
stage 1: layers 3-4
stage 2: layers 5-6
stage 3: layers 7-8

第 $k$ 个 stage 的复合 forward function 记为:
$F_k = f_j \circ f_{j-1} \circ \dots \circ f_i$
其 backward function 记为 $B_k$ ;

mini-batch 与 micro-batch

项目	mini-batch	micro-batch
定义	优化意义上的一批样本	执行调度意义上的一小块样本
作用	一次参数更新真正对应的数据量	为了填满 pipeline 而把一个 mini-batch 再拆开
是否直接决定一次 update	是	否
与梯度的关系	所有样本的梯度累积后, 统一做一次 update	每个 micro-batch 只贡献 mini-batch 梯度的一部分
在 PP 中的角色	优化单位	流水线中的流动单位
例子	mini-batch size = 128	将 128 个样本切成 4 个 micro-batch, 每个 32 个样本

3. 为什么 PP 下不能把整个 mini-batch 一次性塞进去

如果 8 个样本整体塞入两段 pipeline:

stage 0 先处理全部 8 个
stage 1 等待
stage 1 再处理全部 8 个
stage 0 空闲

这几乎退化成朴素 model parallelism, 利用率低;
如果拆成 4 个 micro-batch, 每个 2 个样本, 则:

时刻 1: stage 0 跑 mb0
时刻 2: stage 0 跑 mb1, stage 1 跑 mb0
时刻 3: stage 0 跑 mb2, stage 1 跑 mb1
时刻 4: stage 0 跑 mb3, stage 1 跑 mb2
时刻 5: stage 1 跑 mb3

结论: micro-batch 是为了提高 pipeline 利用率, 并降低单次显存压力;

GPipe 的训练语义: synchronous mini-batch GD

GPipe 强调的一点是:

micro-batch 只是执行单元
mini-batch 才是优化单元

具体流程是:

取一个 mini-batch
把它拆成多个 micro-batches
micro-batches 在多个 stage 上流水线执行
每个 micro-batch 分别做 forward / backward
各 micro-batch 产生自己的梯度贡献
在 mini-batch 结束时统一累积梯度
只做一次参数更新

所以不是:

每个 micro-batch 算完就立刻 update
而是:
所有 micro-batch 共用同一版参数
最后统一做 single synchronous gradient update

这保证了训练语义尽量接近普通 mini-batch SGD, 避免了 weight staleness;

为什么普通反向传播很吃内存

反向传播在第 $i$ 层计算时, 不仅需要上游传回来的梯度, 还需要该层 forward 时的 activation, 例如:

$h_{i-1}$
$z_i$
$h_i$
例如: $\frac{\partial L}{\partial W_i} = \frac{\partial L}{\partial z_i} h_{i-1}^T$ 这里必须用到前向时的输入 $h_{i-1}$ ;
又例如:
$\frac{\partial L}{\partial z_i} = \frac{\partial L}{\partial h_i} \odot \sigma'(z_i)$ 这里必须用到前向时的 $z_i$ 或 $h_i$ ; 如果不重算, 这些 activation 就必须在 forward 时全部缓存下来;

activation 空间复杂度

设:

batch size 为 $N$
网络层数为 $L$
每层 activation 大小量级近似相同
那么:
一层 activation memory 量级大约与 $N$ 成正比
共 $L$ 层都要缓存

所以 activation memory 粗略量级为: $O(N \times L)$
这说明:

batch 越大, activation 内存越大
层数越多, activation 内存越大

re-materialization 的作用

GPipe 为了减少 activation memory, 引入 re-materialization:

forward 时只保存 partition boundary 的 activation
backward 时重新计算 stage 内部的 forward 结果
本质是: 用额外计算, 换更低显存;

四个核心优化点

Re-materialization: 降低 activation memory

GPipe forward 时每个 accelerator 只保存 partition boundary activation, 不保存所有中间层结果; backward 时重算该 partition 的 forward;
论文给出的 peak activation memory 量级为: $O\left(N + \frac{L}{K} \times \frac{N}{M}\right)$

其中:

$K$ 是 partition 数
$M$ 是 micro-batch 数
$N/M$ 是 micro-batch size
$L/K$ 是每个 partition 的层数

直观理解:

partition 越多, 每个 stage 的层越少
micro-batch 越多, 每次处理的数据越少
activation 内存显著下降

Bubble overhead: 吞吐损失来自流水线填充/排空

因为 pipeline 一开始没有被填满, 结束时也要排空, 所以会产生空闲时间, 称为 bubble overhead;

论文给出的量级为: $O\left(\frac{K-1}{M+K-1}\right)$
结论:

$K$ 越大, bubble 越明显
$M$ 越大, bubble 占比越小

经验结论:
当 $M \ge 4K$ 时, bubble overhead 通常很小;

Communication overhead 低

GPipe 的通信只发生在 partition boundary:

forward 时传 boundary activation
backward 时传 boundary gradient

因此通信是: 相邻 stage 间的点对点通信，而不是频繁全局同步
这意味着 PP 对高速互联依赖较弱;

Load imbalance: 现实性能瓶颈

理论分析常假设 partitions 负载均衡, 但现实中不同层在:

memory
flops
compute time
上常常不均衡; 如果某个 stage 明显更慢, 它会拖累整个 pipeline; 因此，更好的 partitioning algorithm 是后续重要研究方向;

DP 与 PP 的区别

项目	DP (Data Parallelism)	PP (Pipeline Parallelism)
并行切分对象	数据	模型
每张 GPU 持有什么	完整模型副本	模型的一部分层 / 一个 stage
数据如何流动	不同 GPU 各自处理不同数据子集	同一个样本或 micro-batch 需要按 stage 顺序流过多张 GPU
forward / backward 如何执行	每张 GPU 都独立执行完整网络的 forward/backward	每张 GPU 只执行自己那部分层的 forward/backward
并行的核心方式	多个完整模型副本同时处理不同数据	不同 micro-batches 在不同 stage 上重叠执行
梯度如何计算	每张 GPU 都能算一份整网梯度	每张 GPU 只算自己那部分参数的梯度
梯度如何汇总	最后做 all-reduce / average	通过 stage 间接力反传形成整网梯度
本质	切数据, 不切模型	切模型, 不切完整副本
典型优势	实现简单, 容易扩展数据吞吐	能让单卡放不下的大模型分布到多卡上训练
典型限制	单卡必须能放下完整模型	有流水线 bubble, stage 间需要传 activation / gradient

梯度视角对比

视角	DP	PP
单张卡能否独立算完整梯度	能	不能
单张卡计算的梯度范围	整个模型的参数梯度	自己负责那部分参数的梯度
梯度形成方式	各卡各自算完整梯度, 再同步	各 stage 局部反传, 拼成整网梯度

2. 梯度是怎么同步的

在某个 PP replica 内部:

后半 stage 先反传
算出自己的参数梯度
再把边界梯度传给前半 stage
前半 stage 继续反传
在不同 DP replicas 之间:
对应 stage 的参数副本做梯度同步
例如前半模型的副本彼此同步, 后半模型的副本彼此同步
所以:
DP 同步是按 stage 对齐进行的;

系统化刻画 PP 的 tradeoff

GPipe 不只提出"切模型", 还明确把以下问题系统化:

activation memory
bubble / throughput
boundary communication
partition balance
这为后续大量并行训练系统论文奠定了分析框架;

[!论文主流观点总结]

如果按论文主线浓缩, GPipe 的主流观点可以总结为:

大模型持续扩展需要新的训练基础设施

Model parallelism 不应只是特定模型技巧, 而应成为通用抽象

PP 的关键不只是切层, 而是 micro-batch pipeline

执行优化不能破坏训练语义, 应该保持 synchronous mini-batch GD

真正高效的 PP 需要同时处理 memory, throughput, communication 和 load balance

PP 可以并且应该与 DP 组合, 成为更大规模训练的组成模块