dense并行训练1-流水线并行

dense · 浏览次数 : 4

小编点评

并行训练主要有三种策略:数据并行训练、张量并行和流水线并行。数据并行训练加速比最高,但显存占用较高,通信量大;张量并行适合在机器内做模型并行,通信量较高;流水线并行能减少通信边界,支持更多层数,但在某些情况下加速效率不如数据并行。 流水线并行Micro-batch(Gpipe)将网络结构纵向拆分,每张卡训练其中的几层。通过更细粒度的拆分,提高并行度。但存在bubble_time和显存浪费问题。 为解决这些问题,提出了1F1B-flush方案,优先进行BP计算,每个micro-batch完成后直接释放对应计算中间值。此外,还有交错式设计,如virtual_pipeline方案,将每个worker从单一流水线拆分成多个流水线,使通信隐藏在计算过程中。 总之,这些并行训练策略在不同方面有各自的优势和局限性,需要根据具体任务和硬件条件选择合适的策略。

正文

并行训练-流水线

简述

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

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

流水线并行

Micro-batch(Gpipe)

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

image-20240626203004792

存在的问题:

  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

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

PipeDream(非交错式1F1B DeepSpeed)

image-20240626191344318

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

image-20240626210948570

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

1F1B-flush

image-20240626201932709

对比上面的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)

image-20240626213836419

这个方案有个新的概念, 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. 能更好的把通信隐藏起来.

总结这个方案的优点:

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

与dense并行训练1-流水线并行相似的内容: