论文解读：GPipe

背景：

增大模型规模通常是提高模型效果的有效方案，但是更多的模型参数，更复杂的模型计算逻辑，导致模型的训练过程需要更多的内存。同时计算设备 (GPU) 的内存资源往往非常非常有限，单台设备不足以支持超大规模模型的训练，而将模型分布在多台设备上则是一种可能的解决方案。

简介：

GPipe 使用模型并行方案，将模型切分成一连串 stage，每个 stage 放在独立的设备（GPU/TPU）上，实现对超大规模模型的支持，同时利用 Pipeline 的方案，提高了模型并行模式下的设备利用率。最终 GPipe 通过更大规模的模型和更大的 batch_size，在图像和 NLP 的模型上都得到更好的模型效果。

方案设计

1. 划分 stage

如图 1.(a) 左图所示，将模型划分为连续的几个 stage，每个 stage 各自对应一个设备。这样就使得模型的大小可以突破单个设备内存的大小，因为一台设备只需要能够容纳部分模型的参数和计算。

2. 划分 Micro-Batch 和 Pipeline

图 1.(b) 为一般的模型并行的运算模式，在每个时间点只有一台设备在处理计算逻辑，完成计算后将结果发送给下一台设备，其中设备空闲的时间称为 Bubble。GPipe 将 mini-batch 进一步划分成更小的 micro-batch，同时利用 pipipline 方案，每次处理一个 micro-batch 的数据，得到结果后，将该 micro-batch 的结果发送给下游设备，同时开始处理后一个 micro-batch 的数据，通过这套方案减小设备中的 Bubble。

GPipe 实验发现，当增加 micro_batch 数目的时候，设备上的 bubble 可忽略 (一定条件下)，一方面是由于如图 1.(c) 所示的的运行模式导致的，另一方面是由于 Re-Materializaition 方案中，计算梯度过程中的前向计算部分可以不依赖前一部分数据的计算结果 (即在图 1.(c) 中的 Bubble 阶段执行。)

3. Re-Materialization

基于现有的训练框架，在前向计算时，训练框架会记录每一个算子的计算结果，用于 backward 时的梯度计算。在很多模型中，这部分的内存消耗可能大于模型参数，是限制模型大小（以及 batch-size）的一个重要原因。

Re-Materialization 具体是指在前向计算过程中，GPipe 只记录 stage 划分处的输出，在计算梯度时，GPipe 会重新执行前向计算逻辑，从而得到各个算子的前向结果，然后再计算梯度结果。

Re-Materialization 的好处是大大减少了内存需求，可以有效增大训练的 batch_size，比如正常的训练模式下 (忽略输出所占的内存，这部分一般占比较小)，内存能够支持 mini-batch=16，在 Gpipe 的模式下，可以将 micro-batch 设置为 16, 这样实际的 batch_size 就扩大了 min-batch/micor_batch 倍。