并行训练-流水线

简述

并行训练主要有三种策略:

1. 数据并行训练加速比最高，但要求每个设备上都备份一份模型，显存占用比较高，但缺点是通信量大。
2. 张量并行，通信量比较高，适合在机器内做模型并行。
3. 流水线并行，训练设备容易出现空闲状态，加速效率没有DP高；但能减少通信边界支持更多的层数，适合在机器间使用。

流水线并行

Micro-batch(Gpipe)

将网络结构进行纵向拆分, 每张卡训练其中的几层. 下图把网络拆成4层. 如果是按照纯粹的mini-batch训练, 每层之间是纯串行的. 后面的卡会始终等待前面的卡. 所以引入了micro-batch的概念. 把mini-batch进行更细粒度的拆分, 这样在完成batch0的fp之后, 卡0可以进行batch1的fp, 卡1就能开始batch0的fp. 从而提高并行度.

存在的问题:

1. 存在bubble_time: 每张卡的空闲时间 = (stage_num - 1) * (fp_time + bp_time)

\[\frac{(stageNum - 1)(tf + tp)}{(stageNum - 1)(tf + tp) + microNum(tf + tp)} = \frac{stageNum - 1}{microNum + stageNum - 1}

\]

实际应用中 当mico-batch个数大于stageNum的4倍时, 可以忽略bubble_time

2. 显存浪费: 当进行stage3的micro-batch 3时, 还需要保存前面所有mico-batch的fp中间结果用于bp.
3. 在每个mini-batch之间无法并行. 因为下一个minibatch需要等当前所有的micro-batch更新完参数

PipeDream（非交错式1F1B DeepSpeed）

在每个micro-batch fp完成之后立刻优先进行bp. 这样可以把当前batch的中间变量释放掉, bp完成后更新本机参数, 但这种方式存在参数更新冲突, 机器1和机器2使用的参数不一样, 机器1的batch5只用了 batch1反向后更新的参数, 但机器2的使用了batch2的, PipeDream通过多版本参数cache的思想来解决这个问题

为啥worker1需要保存4个版本参数, 而worker4只需要1个呢? 这里的版本数和同一个batch fp和bp的间隔决定的. 如果我跑完fp后, 中间有其他batch更新的bp. 那就需要把这些bp结果给缓存起来, 不然就会导致fp和bp使用的不是同一份参数. 可以看到worker1的batch5 中间间隔了2,3,4 3次bp, 再加上它本身. 就得保存4份...这种方法对显存极度不友好, 所以有了下面的flush方式

1F1B-flush

对比上面的F-then-B的方式, 1F1B优先bp计算. 每个micro-batch完成后直接释放掉了对应micro-batch的计算中间值.

只需要保存1份w, 在固定micro-batch个数后进行一次flush, 同步所有worker的权重使其保持同一个版本.

另外在stage3中 batch1 fp时, 因为batch0已经算完了. 所以可以直接复用batch0的显存不用重新分配.

[!NOTE]

这里有个疑问..越底层的stage需要缓存的中间值其实越多, 这种造成存储不均匀的问题怎么解决? 通过stage切分不同大小参数的方式么

1F1B-flush(交错式, megatron)

这个方案有个新的概念, virtual_pipeline, 方案要求一个小批次中的微批次数量是管道并行大小（流水线中的设备数量）的整数倍

按之前非交错式的方法. 一共有8层, worker1如果是1/2层, worker2是3/4层..worker4是7/8层, 每个worker计算连续的层

那么virtual_pipeline如果是2的话, 会把每个worker进一步拆分, worker1变成了计算1/5层, worker2: 2/6层..类推, 相当于通过把每个worker从单一流水线拆成了virtual_pipeline个流水线.

• 在之前的1F1B模式里, 因为每个机器计算是有先后顺序的, worker2的通信接收worker1的fp结果必须等worker1的fp完成.
• 而在交错式设计里, worker2计算的是2/6层, 当他计算2的时候, 可以同步从worker1拿上一个batch的5层结果, 算完2后的理想状态就是直接算5. 能更好的把通信隐藏起来.

总结这个方案的优点:

• 相邻的计算与通信操作无依赖关系, 可以加速并行执行
• 发起通信操作时，通信的对端通常已经准备好了要通信的数据，通信操作不需要额外的等待时间。