SPPO:Efficient Long-sequence LLM Training via Adaptive Sequence Pipeline Parallel Offloading

2025-03-13

AI Summary

To address GPU memory and computational bottlenecks in training long-sequence large language models (LLMs), this paper proposes an adaptive sequence-based pipeline parallel offloading framework. Methodologically, it introduces the first sequence-aware CPU/GPU heterogeneous offloading mechanism, integrated with two-level activation lifetime management, a heuristic pipeline scheduler, and dynamic sequence sharding—jointly optimizing memory footprint and computational efficiency. The core innovations are sequence-aware offloading and reuse-oriented sequence partitioning, which overcome scalability limitations of conventional pipeline parallelism for long sequences. Experiments demonstrate successful training of a 7B-parameter model on 4M-token sequences using 128 A100 GPUs, achieving 3.38× higher throughput than Megatron-LM and DeepSpeed. This significantly enhances the practicality and scalability of LLM training for extremely long contexts.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

In recent years, Large Language Models (LLMs) have exhibited remarkable capabilities, driving advancements in real-world applications. However, training LLMs on increasingly long input sequences imposes significant challenges due to high GPU memory and computational demands. Existing solutions face two key limitations: (1) memory reduction techniques, such as activation recomputation and CPU offloading, compromise training efficiency; (2) distributed parallelism strategies require excessive GPU resources, limiting the scalability of input sequence length. To address these gaps, we propose Adaptive Sequence Pipeline Parallel Offloading (SPPO), a novel LLM training framework that optimizes memory and computational resource efficiency for long-sequence training. SPPO introduces adaptive offloading, leveraging sequence-aware offloading, and two-level activation management to reduce GPU memory consumption without degrading the training efficiency. Additionally, SPPO develops an adaptive pipeline scheduling approach with a heuristic solver and multiplexed sequence partitioning to improve computational resource efficiency. Experimental results demonstrate that SPPO achieves up to 3.38x throughput improvement over Megatron-LM and DeepSpeed, realizing efficient training of a 7B LLM with sequence lengths of up to 4M tokens on only 128 A100 GPUs.

Problem

• Optimizes memory and computational efficiency for long-sequence LLM training.
• Reduces GPU memory consumption without compromising training efficiency.
• Improves computational resource efficiency with adaptive pipeline scheduling.

Innovation

• Adaptive offloading reduces GPU memory usage
• Two-level activation management enhances efficiency
• Heuristic solver optimizes pipeline scheduling