Google GPipe 阅读笔记

背景

  1. 现在很多的成果说明越深的神经网络表现越好,如 BERT、BigGAN 等工作;

  2. 更深的网络意味着更多的参数,在过去的几年中,模型的大小涨了 36 倍,但 GPU 的内存只从 2014 年 K40 的 12G 涨到了 2018 年 V100 的 32G。

  3. 现代的数据集也越来越大,如 OpenImages 等数据库中有着高分辨率的图像(包含更多的信息),这使得网络的 Activation Map 所占的内存变得更大。如果要利用好这些巨大的数据集(减少 Over-Fitting、提高 Performance),则需要打破加速器的内存限制。

方法

结合 Model Parallelism 与 Data Parallelism

有两种不同的方式来并行地执行模型。第一种是将模型的不同部分放在不同的加速器上:计算的时候将一个加速器上输出的 Activation Map 传输到另外一个加速器上;另外一种方式是数据并行,每个加速器上都存储整个网络,每次将当前的 Batch 分为若干份,然后每个加速器算其中的若干份,最后再 CPU 上加起来拿到最后的梯度。

我们可以结合这两种运用多个加速器做并行的方法,对于 Model Parallelism 来说,我们仍然将网络分成若干部分,将每个部分放在不同的加速器上,从 Data Parallism 来说,我们将一个 Batch 拆分成若干个 Micro-Batch,然后通过 Pipeline 的方式来进行运算。

具体来说,我们先给网络的第 $i$ 层 $f_i$ 一个计算代价估计 $c_i$,GPipe 会将整个网络分成 $K$ 份(通过一些 Heuristic 的方式来分层,文中用的方法是每次都尽可能减少每部分与每部分之间的计算代价差距),第$i$个加速器只负责对第 $i$ 层的计算。每个 Mini-Batch,将其分成 $T$ 份($T$ 是预先给定的),对于 Forward 过程,等第k个加速器完成对第 $t$ 个 Micro-Batch 计算的时候,第 $k+1$ 个加速器将开始处理第 $t$ 个 Micro-Batch,与此同时,第$k$个加速器也会开始第 $t+1$ 个 Micro-Batch 的计算(如下图)。

$F\_{a, b}$表示第$a$个加速器(第$a$层)处理第$b$个 Micro-Batch

Re-Computation

为了减少存储 Activation Map 的空间,每个加速器只在其“边界”上存储 Activation Map,在 Backward 的时候,每个加速器将重新去计算在这个加速器上每一层的 Activation Map。

评价

虽然这样的 Pipeline 会造成一段时间的空闲(即 Bubble Overhead),但是实际上在实验中发现,这部分所造成的空闲非常小,可能是在我们可以在 Backward 之前就开始进行 Re-Computation。同时真正最耗时间的地方在于在切分网络的时候,每一层的计算量并不是均匀的。

引用

[1] Huang, Yanping, et al. “GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.” arXiv preprint arXiv:1811.06965 (2018).