AI因你而升温，记得加个星标哦！
随着大模型参数量的爆炸性增长，其所需内存也呈爆炸性增长，最现实的问题就是单块显卡装不下模型，所以我们需要进行分布式训练。
演进路线
- 数据并行Data Parallelism：一台机器可以装下模型，所以将同一个模型同时部署在多台机器，用多份数据分开训练
- 流水线并行Pipeline Parallelism：一台机器装不下模型，但模型的一层或多层一台设备装得下，所以同一个模型按层拆分到不同机器进行训练
- 张量并行Tensor Parallelism：模型的一层都装不下了，所以同一个模型层内拆分开训练
本文我们主要介绍流水线并行，对PipeDream和GPipe两个最知名的方法进行介绍，张量并行以及Megatron和ZeRO会陆续介绍。不了解分布式训练的同学介意先阅读这两篇文章：
GPipe
神经网络是在设计之初就是按照层来组合的，每一层都有指定格式的输入和输出，并在层内部进行运算，因此可以按照神经网络的层将一个完整的网络拆分到多个设备上进行并行计算。
**图a：**假设神经网络一共有4层，将4层分配到4个设备上，其中前向传播用F表示，反向传播用B表示
**图b：**在GPU0上做完一次forward，然后将GPU0上最后一层的输入传给GPU1，继续做forward，直到四块GPU都做完forward后，再依次backward。等把四块GPU上的backward全部做完后，最后一个时刻统一更新每一层的梯度。这样做随之会带来两个问题：
其一是GPU利用度不够。上图纵轴表示机器，横轴表示每个时间步，字母F与B的下标表示第几次训练。可以看到在1-8一次训练的每个时间步中，每个时间步都只有一台机器在运行。那么当模型越大，所需GPU的数量越多（K块）时，机器的空闲时间的复杂度为 O ( K − 1 K ) O(\frac{K-1}{K}) O(KK−1)，即k越大空闲时间越多，机器的空闲时间的复杂度称为Bubble。
其二是中间结果占据大量内存，因为在做backward计算梯度的过程中，每一层的中间结果都需要储存。
**图c：**第一个下标表示GPU编号，第二个下标表示micro-batch编号。为解决以上问题，Gpipe在模型并行的基础上，进一步引入数据并行的办法，即把原先的数据再划分成若干个batch，送入GPU进行训练。未划分前的数据，叫mini-batch，在mini-batch上再划分的数据，叫micro-batch。
因为micro-batch很小，所以一台GPU上的一层网络计算一个micro-batch耗时就很短，可以及时把结果送到接下来的GPU机器。假设我们将mini-batch划分为M个，则流水线并行下，机器的空闲时间复杂度为 O ( K − 1 K + M − 1 ) O(\frac{K-1}{K+M-1}) O(K+M−1K−1)，Gpipe通过实验证明，当 M > = 4 K M>=4K M>=4K时，bubble产生的空转时间占比对最终训练时长影响是微小的，可以忽略不计。
在解决了GPU空置问题后，就要解决GPU的内存问题了。对此，Gpipe采用了一种非常简单粗暴但有效的办法：用时间换空间，在论文里，这种方法被命名为re-materalization，后人也称其为active checkpoint。具体的做法是每块GPU只保存来自上一块的最后一层输入z（一块GPU上可能有几层网络），其余的中间结果不保存，等到backward的时候再由保存下来的z进行重新计算。
最后需要提一点，在micro-batch的划分下，我们在计算Batch Normalization时会有影响。Gpipe的方法是，在训练时计算和运用的是micro-batch里的均值和方差，但同时持续追踪全部mini-batch的移动平均和方差，以便在测试阶段进行使用。Layer Normalization则不受影响。
PipeDream
GPipe的思想是算完全部micro-batch后再开始反向传播计算，通过切分micro-batch的方法降低Bubble时间，但是Microbatch的尺寸过小可能会导致通信效率降低，从而影响整体性能。而PipeDream的优化方案是无需等待全部micro-batch算完，即每台机器算完一次micro-batch就立刻计算这台机器的梯度。
为了解决Bubble问题，作者提出了 1F1B（1次Forward，1次Backward）的调度机制，在起始阶段允许执行多个Mini Batch的Forward，稳定后就保持Forward和Backward的交替执行，这样可以保证GPU在稳定状态下没有空闲，并且始终继续学习。
仔细观察会发现1F1B存在一个问题，在 Backward 后立即更新模型参数会引入两种类型的参数不一致：
- 同一个Mini Batch的Forward和Backward可能使用不同的参数进行计算。比如，以Machine 1为例，Forward 5使用的是 Mini Batch 1 更新后的参数；而Backward 5使用的是Mini Batch 1,2,3,4共同更新后的参数
- 同一个Mini Batch在不同机器上使用的参数版本不一致。还是以Mini Batch 5为例，Machine 1上的 Forward 5 使用的是Mini Batch 1更新后的参数；Machine 2上的Forward 5使用的是Mini Batch 1和 Mini Batch 2更新后的参数
为了解决参数不一致的问题，PipeDream引入了Weight Stashing和Vertical Sync：
- Weight Stashing：为每个正在计算的Mini Batch都保存一份参数。Forward计算时，每个设备都使用最新的权重参数计算输入的 Mini Batch，并将这个参数保存，直到当前设备上对应的 Backward 计算完成，即Forward用哪个Weights算的，Backward就还用它。这样可以解决上述的第一个参数不一致问题，但无法解决第二个
- Vertical Sync：每个 Mini Batch 开始计算时都使用最新版本的权重参数，且参数的版本号会伴随该 Mini Batch 数据的整个生命周期，在所有机器都使用同一版本的参数，从而实现每台机器的参数一致性。举个例子，我们在算第Batch 5（图中深蓝5）时，Machine1使用的weight是更新了一次的，它算完得到的output扔给Machine2的时候，Machine2也用只更新了一次的weight算（就好像浅绿色的2没有用），算完得到的output扔给Machine3，Machine3也用只更新了一次的算（就好像浅绿色的2与3没有用）。这样可以解决上述的第二个参数不一致问题
不过作者坦言，Vertical Sync优化提升效果不太大，最重要的还是Weight Stashing优化策略，在实际使用中也只采用了Weight Stashing方案。
总结
GPipe需要等所有的microbatch前向传播完成后，才会开始反向传播。PipeDream则是当一个microbatch的前向传播完成后，立即进入反向传播阶段。理论上，反向传播完成后就可以丢弃掉对应microbatch缓存的激活。由于PipeDream的反向传播完成的要比GPipe早，因此也会减少显存的需求。
GPipe与PipeDream主要差别是在梯度更新上，Gpipe是最后同步一次更新的，而PipeDream是异步的。异步方法更进一步降低了GPU的空转时间比。虽然PipeDream设计更精妙些，但是Gpipe因为其"够用"、稳定和浅显易懂，更受大众欢迎（torch的PP接口就基于Gpipe）。
