[NeurIPS 2019] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Introduction

• 作者提出了一种高效且通用的流水线并行方法 GPipie，可以用于对任何可以被拆分为一系列串联层的模型进行模型并行训练，从而可以训练超过单一加速器内存的大模型

The GPipe Library

Algorithm

• Partition Algorithm. 假设模型可以被表示为 $L$ 层的串联形式，每层 $L_i$ 都由前向传播函数 $f_i$ 和参数 $w_i$ 组成。给定划分数 $K$ 以及每层的计算代价 $c_i$，就可以将整个模型的 $L$ 层划分成计算量近似相同的 $K$ 份，分别存放在 $K$ 个加速器上，相邻划分之间可以进行数据传输
• Micro-batches. 如下图 (a) (b) 所示，将模型划分为 4 份，如果采用最简单的模型并行方法，由于前向传播和后向传播函数之间具有相互依赖性，整个模型的执行过程并不能很好地并行。为了解决上述问题，如下图 (c) 所示，GPipe 将 micro-batch 进一步划分为 $M$ 个相同大小的 micro-batches，大大增加了并行效率。在每个 micro-batch 的反向传播结束后，$M$ 个 micro-batches 的梯度累加后进行参数更新。注意到，如果模型中有 BN，那么训练时 BN 是作用在 micro-batch 上的，同时 mini-batch 上统计量的移动平均也会被记录，以便用于后续的模型推理；另外，micro batch size 需要兼顾运算强度和 pipeline bubble，过小会导致运算强度过低，GPU 利用率不高，过大会导致 pipeline bubble 过大，同样会降低 GPU 利用率
  
• 1F1B Pipeline (One Forward pass followed by One Backward pass from PipeDream). 如下图所示 (backward 的计算开销大约是 forward 的两倍)，GPipe 仍然存在比较多的 bubble，并且由于 GPipe 是在所有 micro-batch 的前向过程计算完毕后再进行反向传播，因此如果不使用 checkpointing 的话，就需要缓存所有 micro-batch 的中间激活值，内存开销比较大
  为此，PipeDream 提出 1F1B 策略，每个 micro-batch 前向计算完毕后就立刻开始反向传播，从而有效降低内存开销，允许使用更大的 mivro-batch 从而减少 bubble
  

Performance Optimization

• Re-materialization / Checkpointing. 为了进一步减少内存占用，GPipe 使用了 re-materialization，也就是在前向传播时，每个加速器只存储划分边界上的激活值 (i.e., patition 的第一个层的输入)，patition 的其他层的激活值则不保存，要用的时候再根据存储的划分边界上的激活值重新计算，相当于用时间成本换取存储成本，这样内存占用峰值就减少到了 $O(N+\frac{L}{K}\times \frac{N}{M})$，其中 $N$ 为 batch size，$L$ 为模型总层数，$K$ 为划分数，$M$ 为 micro-batch 数 (i.e., 一整个 mini-batch 在划分边界上的激活值 $O(N)$ 以及一个 micro-batch 在该划分的 $\frac{L}{K}$ 层上的前向传播结果 $O(\frac{L}{K}\times \frac{N}{M})$，该结果用于计算该 micro-batch 的反向传播梯度). 相比之下，如果使用单卡训练，内存峰值占用为 $O(N\times L)$ (计算梯度需要保存所有前向传播激活值)
• Bubble overhead. GPipe 的 bubble time 为 $O(\frac{K-1}{M+K-1})$，也就是上图 (c) 中白色方块数除以彩色方块数 (不包含最后的 update). 实验中作者发现，当 $M\ge 4\times K$ 时，bubble 带来的开销就可以忽略不计，这可能是因为还存在其他的开销，例如 re-materialization 需要重新进行 micro-batch 的前向传播，以及不同加速器之间的通信
• Communication overhead. GPipe 只需要传输每个划分最后的输出激活值，因此通信开销较小
• Imperfect partitioning. 在实际的模型中，模型各个层的内存消耗和计算量通常都是不均衡的，此时不完美的划分就可能会降低训练效率

Performance Analyses

• 下表展示了不同模型在使用单卡 / GPipe 训练时能达到的最大模型规模，可见 GPipe 让训练超大规模模型成为了可能
  
• Normalized training throughput using GPipe with different # of partitions $K$ and different # of micro-batches $M$ on TPUs. 可以注意到，AmoebaNet 相比 Transformer 模型来说更难均匀切分，因此加速效果要差一些
  
• Performance Overhead Breakdown.
  

References

• Huang, Yanping, et al. “Gpipe: Efficient training of giant neural networks using pipeline parallelism.” Advances in neural information processing systems 32 (2019).
• code: https://github.com/KakaoBrain/torchgpipe