Seq1F1B: Efficient Sequence-Level Pipeline Parallelism for Large Language Model Training

本文介绍了大型语言模型（LLMs）的分布式训练策略，其中管道并行在其中起着至关重要的作用。随着LLMs的训练序列长度扩展到32k甚至128k，当前的管道并行方法面临着严重的瓶颈，包括高内存占用和大量的管道气泡，极大地阻碍了模型的可扩展性和训练吞吐量。为了提高内存效率和训练吞吐量，本文介绍了一种名为Seq1F1B的高效的序列级一前一后（1F1B）管道调度方法，专门用于在长序列上训练LLMs。Seq1F1B将批级可调度单元分解为更细的序列级单元，减小了气泡大小和内存占用。考虑到如果序列被均匀分割，Seq1F1B可能会产生轻微的额外气泡，因此我们设计了一种计算策略来分割输入序列并减轻这种副作用。与竞争的管道基准方法（如Megatron 1F1B管道并行）相比，我们的方法在更少的内存占用下实现了更高的训练吞吐量。值得注意的是，Seq1F1B使用64个NVIDIA A100 GPU在长度为64k的序列上高效地训练了一个具有30B参数的LLM，而现有方法无法达到这一成就。我们的源代码基于Megatron-LM，现在可在https://github.com/MayDomine/Seq1F1B.git上获得。

论文旨在解决大型语言模型在训练过程中遇到的瓶颈问题，特别是在面对长度为32k或128k的序列时，存在内存占用和管道阻塞等问题。同时，论文还试图提高模型的可扩展性和训练吞吐量。

论文提出了一种名为Seq1F1B的序列级别的一前一后（1F1B）管道调度方法，将批级别的可调度单元细分为更细的序列级别单元，从而减少了管道阻塞和内存占用。此外，论文还设计了一种计算策略来划分输入序列，以减轻由于均匀拆分序列而产生的额外阻塞。

论文的实验表明，相比竞争基线方法，如Megatron 1F1B管道并行方法，Seq1F1B具有更高的训练吞吐量和更少的内存占用。此外，Seq1F1B能够在64个NVIDIA A100 GPU上训练具有30B参数的LLM，而无需重新计算策略，这是现有方法无法实现的。论文基于Megatron-LM的源代码，已在github上开源。

在这个领域中，最近还有一些相关研究，如《MegatronLM: Training Multi-Billion Parameter Language Models Using Model Parallelism》和《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》等。