摘要
arXiv:2504.14519v1 宣告类型: cross
摘要:管道并行ism(PP)是一种关键的技术,用于训练大型语言模型(LLMs),因为它能够在相对较低的通信开销下缓解模型状态带来的内存压力。然而,在长上下文场景中,现有的管道并行方法无法解决显著的激活内存压力,主要原因在于多个微批量累积的激活导致的峰值内存消耗。此外,这些方法不可避免地引入了大量的管道泡,进一步妨碍了效率。
为了应对这些挑战,我们提出了SlimPipe,这是一种新颖的细粒度管道并行方法,结合了均匀序列分割和一次前向一次后向(1F1B)调度。它将多个微批量的累积激活减少到仅一个,并将其划分为多个片段。尽管这些片段均匀分区,但由于因果注意力机制,计算成本在片段之间是不均匀的。我们开发了一种复杂的工作负载再分配技术来解决这一负载不平衡问题。SlimPipe 同时实现了(1)近乎零的内存开销和(2)最小的管道泡。通过使用多种模型架构、上下文窗口大小和SlimPipe特定配置的全面测试,证明了SlimPipe的有效性。例如,在Llama 70B模型中,与最先进的方法相比,对于512K的上下文长度,SlimPipe 显著提高了模型浮点运算利用率(MFU)至高达1.57倍。更为值得注意的是,在2048K的上下文长度下,它在256块NVIDIA Hopper 80GB GPU 上维持了超过45%的利用率,而其他方法则由于内存约束要么显著性能下降,要么完全失败。