LoongTrain: Efficient Training of Long-Sequence LLMs with Head-Context Parallelism

高效地训练具有长序列的LLM模型对于应对大量的计算和内存需求至关重要，但是现有的序列并行方法存在可扩展性或效率问题。我们提出了一种新的系统LoongTrain，以高效地训练具有长序列的LLM模型为目标。LoongTrain的核心是2D-Attention机制，它结合了头部并行和上下文并行技术，打破了可扩展性限制，同时保持了效率。我们引入了双环形注意力，并分析了设备放置策略的性能，以进一步加快训练速度。我们使用混合ZeRO和Selective Checkpoint++技术实现了LoongTrain。实验结果表明，LoongTrain在端到端训练速度和可扩展性方面均优于最先进的基线模型，即DeepSpeed-Ulysses和Megatron Context Parallelism，并将模型FLOPs利用率提高了最多2.88倍。

该论文旨在解决训练长序列的LLMs所需的大量计算和内存需求的问题，并提出一种新的系统LoongTrain来实现高效训练。

LoongTrain的核心是2D-Attention机制，结合了头并行和上下文并行技术，以打破可扩展性限制同时保持效率。同时，引入了Double-Ring-Attention并分析了设备放置策略的性能，以进一步提高训练速度。

论文在实验中使用了混合ZeRO和Selective Checkpoint++技术，并展示了LoongTrain在训练速度和可扩展性方面优于DeepSpeed-Ulysses和Megatron Context Parallelism等现有技术，并将模型FLOPs利用率提高了2.88倍。

在这个领域中的相关研究包括DeepSpeed、Megatron和GShard等。