May, 2024

可控内存的管道并行

TL;DR通过重复构建块的方式,将流水线调度分解为生命周期的概念,并发现现有调度方案在内存利用方面存在问题。为了解决这个问题,引入了一系列具有可控激活内存的高效构建块,可以在不影响效率的情况下将峰值激活内存减少到 1F1B 的 1/2 甚至 1/3,同时几乎没有流水线气泡,以及在吞吐量方面表现优于 1F1B。在实践场景中对混合并行化超参数进行网格搜索时,相较于 1F1B 基准,我们提出的方法在大型语言模型上实现了 16% 的吞吐量提升。