NeurlPS '19 | GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

[jasperzhong]

https://papers.nips.cc/paper/2019/file/093f65e080a295f8076b1c5722a46aa2-Paper.pdf

[jasperzhong]

想把pipeline做好还真是件复杂的事情. GPipe简单化处理了，效果还不错.

一共有K个设备. 最naive的做法就是图b. 但这样效率太低了. 所以GPipe提出了图c，把一个mini-batch拆成M个micro-batch. 作者说当 M >= 4 * K 的时候，产生的bubble可以忽略不计. 嘛. 看上去mini-batch可以弄大点.

为什么说这件事比较复杂呢?

1. partition. 很难做得到每个stage时间差不多.
2. communication. stage之间需要传输intermediate output，如何选择正确的partition边界很难.
3. re-computation. 如果还考虑re-computation，那事情更复杂了. BP的时候还要插入FP节点，这个pipeline分析变得很复杂.
4. micro-batch如果太小，对于BN这样的操作不友好.

总之主要难点还是如何做partition. 感觉需要一个搜索.

实验还可以.

这个breakdown不错. bubble overhead才4.1%.

section 6有些惊喜. model parallelism初衷是支持大模型. 但是naive的方法资源利用率太低了. 解决办法通常有两个办法:

1. single program multiple data (SPMD)
2. pipeline parallelism

[jasperzhong]

今天重读了这篇论文，收获不小. 温故而知新.

GPipe的core idea很简单：

• 把L层layer的模型，分成K个partition，分别放到K个机器上
• FW的时候，把一个mini-batch N拆成M个micro-batch，pipeline执行；BW的时候用同样的参数计算gradients；最后将多个micro-batch的gradients w.r.t parameters累加，得到最终的gradient，然后做参数更新.

除此之外，有一个重要的优化，就是recomputation——只保存partition boundaries的output activations. 这样peak activation memory的空间复杂度从O(N L / K)降低到O(N + (L / K) (N / M)），其中N是指保存的output activations，因为每个micro-batch都要保存，所以总的batch数量还是N，而L/K是这个partition的layer层数，N/M是micro-batch size，(L/K)*(N/M)就是在这个partition上做一个micro-batch的output activation size.

其实对比其他parallelism方法，可以发现Pipeline parallelism引入的通信开销是很低的，仅需要传输boundary的activations & gradients w.r.t. activations即可. 而tensor model parallelism以及data parallelism，还需要引入很多all-reduce.

这个实验注意一下. 训练Big Transformer的实验. T(L, H, A). L是层数，H是hidden size，A是#attention heads. 对比T(24, 8192, 16)和T(12, 16384, 32)就会发现，deeper model表现要比wider model更好.

另外和 #167 对比，GPipe是同步更新，其weight update semantics没变. 而Pipedream是异步更新，虽然降低了bubble size，但是却需要引入多份model weights. 其实不太划得来. 因为文中表明了bubble size其实不是个很大的问题.

有个ablation实验吸引了我的注意. 看上去M >> K是个不错的configuration. 也就是说，Pipeline要尽可能多stage（K大），然后尽可能pipelined（M大）.