Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer

TL;DR

The paper addresses the challenge of training LLMs with ultra-long context under limited hardware. It introduces the Fully Pipelined Distributed Transformer (FPDT), a memory-hierarchy-aware, fully pipelined approach that offloads and prefetches data between host and GPU memories to dramatically extend trainable sequence lengths. FPDT demonstrates up to a 16x increase in context length on the same hardware (e.g., training an $8\mathrm{B}$ model with $2\mathrm{M}$ tokens on 4 GPUs) while achieving MFU over 55% and maintaining model-agnostic applicability across GPT and Llama architectures. The results show substantial practical impact for long-context tasks in NLP and computational biology, enabling efficient scaling without proportionally increasing hardware budgets.

Abstract

Large Language Models (LLMs) with long context capabilities are integral to complex tasks in natural language processing and computational biology, such as text generation and protein sequence analysis. However, training LLMs directly on extremely long contexts demands considerable GPU resources and increased memory, leading to higher costs and greater complexity. Alternative approaches that introduce long context capabilities via downstream finetuning or adaptations impose significant design limitations. In this paper, we propose Fully Pipelined Distributed Transformer (FPDT) for efficiently training long-context LLMs with extreme hardware efficiency. For GPT and Llama models, we achieve a 16x increase in sequence length that can be trained on the same hardware compared to current state-of-the-art solutions. With our dedicated sequence chunk pipeline design, we can now train 8B LLM with 2 million sequence length on only 4 GPUs, while also maintaining over 55% of MFU. Our proposed FPDT is agnostic to existing training techniques and is proven to work efficiently across different LLM models.

Figures (14)

• Figure 1: Comparison of end-to-end training Model FLOPs Utilization (MFU) and maximum context length per GPU supported. We show 3 model sizes, i.e. 2.7B, 13B, and 70B.
• Figure 2: DeepSpeed Ulysses with distributed attention.
• Figure 3: DeepSpeed Ulysses works orthogonally with ZeRO-3.
• Figure 4: The design of distributed attention with offloading
• Figure 5: The design of distributed attention with fetching and offloading. We follow the online attention policy to update attention output.