Optimized pipeline优化

当前1F1B流水线并行典型配置采用相同的micro-batch-size(mbs),这种对流水线气泡时间并不友好,存在进一步优化的可能性。

将流水线并行的过程数学建模成整数规划问题,根据profiling得到的通信和计算时间,最小化端到端耗时,求解整数规划问题以得到mbs的最优配置,在保持micro-batch数量和global batch size不变的情况下,合理设置每个micro-batch-size以减少bubble time,达成性能提升。整体方案包含以下特性:

  1. 小步快跑:mbs逐渐增大,各个流水线(device/node)快速启动。
  2. 中间空隙:增强流水线并行对前向反向计算抖动的鲁棒性。
  3. 缓速结束:mbs逐渐减小,避免流水线停止时气泡过大。
图1 通过最优化mbs减少气泡时间

代码已合入AscendSpeed,当前在LLaMA-7B的Atlas 训练系列产品上性能提升3%~6.6%,Atlas A2 训练系列产品上性能提升3.5%~4.2%,PP越大性能收益越高。