并行训练主要有三种策略:
将网络结构进行纵向拆分, 每张卡训练其中的几层. 下图把网络拆成4层. 如果是按照纯粹的mini-batch训练, 每层之间是纯串行的. 后面的卡会始终等待前面的卡. 所以引入了micro-batch的概念. 把mini-batch进行更细粒度的拆分, 这样在完成batch0的fp之后, 卡0可以进行batch1的fp, 卡1就能开始batch0的fp. 从而提高并行度.
存在的问题:
实际应用中 当mico-batch个数大于stageNum的4倍时, 可以忽略bubble_time
在每个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方式
对比上面的F-then-B的方式, 1F1B优先bp计算. 每个micro-batch完成后直接释放掉了对应micro-batch的计算中间值.
只需要保存1份w, 在固定micro-batch个数后进行一次flush, 同步所有worker的权重使其保持同一个版本.
另外在stage3中 batch1 fp时, 因为batch0已经算完了. 所以可以直接复用batch0的显存不用重新分配.
[!NOTE]
这里有个疑问..越底层的stage需要缓存的中间值其实越多, 这种造成存储不均匀的问题怎么解决? 通过stage切分不同大小参数的方式么
这个方案有个新的概念, 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个流水线.
总结这个方案的优点: