GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

1、Title
GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
Author: Google AI
来源：NeurIPS 2019
2、Main idea:

• 划分stage
  所有卷积神经网络都可以看作是L个层的堆叠，每一层都由一个 正向计算的函数fi和一组对应的参数wi组成；可以把整个L层的 卷积神经网络划分为K个复合层，每个复合层放到一个worker节 点上；这K个复合层的每一层还是拥有着正向计算函数和一组对 应的参数。模型组成的基本单位是layer，PipeDream将这些layer 划分在不同的stage，不同的的stage部署在不同的机器上。
  将模型划分为连续的几个stage，每个stage各自对应一个设备（一 台加速器）。这样就使得模型的大小可以突破单个设备内存的大 小，因为一台设备只需要能够容纳部分模型的参数和计算；划分stage之后，处理最慢的stage会成为瓶颈。所以应该平均分 配算力。(加速器：可以是GPU，也可以是一台计算设备)

• Mini-batch划分为更小的micro-batch
  mini-batch进一步划分成更小的micro-batch，同时利用流水线方 案，每次处理一个micro-batch的数据，得到结果后，将该 micro-batch的结果发送给下游设备，同时开始处理后一个 micro-batch的数据，通过这种方法减小设备中的Bubble（设备空 闲的时间称为 Bubble）提高设备计算资源的利用率
  

• 使用重计算
  Re-Materialization是指在前向计算过程中，GPipe只记录stage划 分处的输出，在计算梯度时，GPipe会重新执行前向计算逻辑，从 而得到前向结果，然后再计算梯度结果。
  
  如图，反向传播过程中，在计算第四层的梯度信息时，需要先进行一次L3到L4的一次正向计算，得到L4的输出特征信息，再利用L5的梯度信息一起计算L4的梯度。也就是说利用了正向计算的时间换取了存储网络模型参数的内存空间，即以时间换空间的思想。通过减少保存的激活值，进而压缩模型占用空间
  实验证明，当M>=4×K 时，训练中的bubble开销可以忽略不计。

• 梯度累积
  梯度累加本质上就是累加 accumulation_steps 个 batch_size/accumulation_steps 的梯度， 再根据累加的梯度来更新网络参数，以达到真实梯度类似batch_size 的效果。在使用时，需要注意适当的扩大学习率。
  也就是说：
  首先将整个dataset分成多个batch，每个 batch size = 32，且假定 accumulation steps = 8 ；
  因为 batch size = 32，太大了，单机显卡无法跑，于是我们在前向传播的时候以 batch_size = 32 / 8 = 4 来计算梯度；
  这样就再分别将每个batch分成多个batch size 为 4 的小批次，将每个小批次逐一喂给神经网络；
  每个小批次虽然计算梯度，但是在每次反向传播（在反向传播的时候，会将mean_loss也除以8）后，先不进行优化器的迭代更新。
  经过 accumulation steps 个小批次后（即一个batch中的所有小批次），用每个小批次计算梯度的累积和去进行优化器迭代更新参数。
  最后进行梯度清零的操作。
  处理下一个batch。
  这样就跟把 32 batch size 一次性送入模型进行训练效果一样了。
  
  3、Experiment & Conclusion
  最后在两个任务上进行了测试：(i)图像分类:在ImageNet-2012上训练了一个5.57亿参数的AmoebaNet模型，并获得了84.4%的top-1准确率，(ii)多语言神经机器翻译:在一个跨越100多种语言的语料库上训练了一个60亿参数的128层Transformer模型，并获得了比所有双语模型更好的质量。

不足：
假如一个minibatch被分为 n 个micro-batch，则需要缓存 n 份activation(正向计算的结果)。这个 n 是梯度累加的次数，为了尽可能流水，通常这个累加次数都比较大，一般大于两倍 stage 数目。那么即使只缓存少数 Tensor，这种策略依然需要较多显存。