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

2024/06/26 by Diandian Gu, Gu, Diandian, Peng Sun +25 · 5 citations
Computer Science · #Natural Language Processing Techniques #Handwritten Text Recognition Techniques #Speech Recognition and Synthesis

paper · pdf · doi:10.48550/arxiv.2406.18485

Abstract

Efficiently training LLMs with long sequences is important yet challenged by the massive computation and memory requirements. Sequence parallelism has been proposed to tackle these problems, but existing methods suffer from scalability or efficiency issues. We propose LoongTrain, a novel system to efficiently train LLMs with long sequences at scale. The core of LoongTrain is the 2D-Attention mechanism, which combines both head-parallel and context-parallel techniques to break the scalability constraints while maintaining efficiency. We introduce Double-Ring-Attention and analyze the performance of device placement strategies to further speed up training. We implement LoongTrain with the hybrid ZeRO and Selective Checkpoint++ techniques. Experiment results show that LoongTrain outperforms state-of-the-art baselines, i.e., DeepSpeed-Ulysses and Megatron Context Parallelism, in both end-to-end training speed and scalability, and improves Model FLOPs Utilization (MFU) by up to 2.88x.

Cited by

• Libra: Taming Attention Workload Skew in Long-Context LLM Training with Bounded Sequence Pool
• Optimizing Long-context LLM Serving via Fine-grained Sequence Parallelism
• MTraining: Distributed Dynamic Sparse Attention for Efficient Ultra-Long Context Training
• Syncopate: Efficient Multi-GPU AI Kernels via Automatic Chunk-Centric Compute-Communication Overlap
• BurstEngine: an Efficient Distributed Framework for Training Transformers on Extremely Long Sequences of over 1M Tokens

Related

• SPPO:Efficient Long-sequence LLM Training via Adaptive Sequence Pipeline Parallel Offloading
• DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
• DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
• ByteScale: Efficient Scaling of LLM Training with a 2048K Context Length on More Than 12,000 GPUs
• Linear Attention Sequence Parallelism