GPipe & 1F1B

1F1B 由 pipedream 提出。

1f1b 的终极目标是减少activation 的缓存数量,降低显存占用,从而可以训练更大的模型。

理解了这个,一下就能看懂这张图了。每一批数据计算完 forwar的之后,需要马上进行 backward。

image.png

1f1b 相比 gpipe 可以省显存,但是不能提速。但是省下来的显存可以用来提高 mbs 数量,从而提速。

K 是流水线段数,M 是 micro bs。

gpipe bubble rate: $O(\frac{K-1}{K+M-1})$

1f1b 的 bubble rate 也是一样。

朴素的切分流水线的 bubble rate: $O(\frac{K-1}{K})$

其实就是 GPipe 令 M=1。这样好记一些。

Virtual Pipeline

由 megatron-lm 论文提出。也叫交错式调度(interleaved schedule)。

这个可以提速,但是是增加通信量换来的。

image.png

bubble 计算

用这个图更好记。就记 bubble time / idle time。