Data-Centric Elastic Pipeline Parallelism for Efficient Long-Context LLM Training

TL;DR

This work tackles the memory and efficiency challenges of long-context LLM training by proposing Elastic Pipeline Parallelism (EPP), a hybrid approach that mixes batch-level and token-level pipeline parallelism to adapt to resource and workload heterogeneity. It introduces InfiniPipe, a distributed system built around a cost model, a workload-aware sequence processor, and a co-optimized chunk scheduler that jointly optimizes pipeline schedule and stage-aware gradient checkpointing via MILP. Key contributions include an extended cost model for varied-length chunks, a workload-balanced sequence processor, a MILP-based checkpointing optimizer, and a scalable runtime that achieves up to 1.69x speedups over state-of-the-art baselines. The approach reduces inter-node communication and pipeline bubbles while handling highly skewed sequence length distributions, enabling efficient ultra-long context training on heterogeneous hardware. Overall, the work provides a practical path to scalable long-context LLM training on realistic datasets with mixed sequence lengths.

Abstract

Long context training is crucial for LLM's context extension. Existing schemes, such as sequence parallelism, incur substantial communication overhead. Pipeline parallelism (PP) reduces this cost, but its effectiveness hinges on partitioning granularity. Batch-level PP dividing input samples exhibits high memory consumption in long-context scenario, whereas token-level PP splitting sequences into slices alleviates memory overhead but may incur hardware under-utilization. This trade-off motivates adaptively selecting PP granularity to match resource and workload characteristics. Moreover, sequence length distribution of the real-world dataset exhibits skewness, posing a challenge on PP's workload balance and efficient scheduling. Current static PP scheduling methods overlook the variance of sequence length, leading to suboptimal performance. In this paper, we propose Elastic Pipeline Parallelism (EPP) that orchestrates token-level PP and batch-level PP to adapt to resource and workload heterogeneity. We build InfiniPipe, a distributed training system that unleashes the potential of EPP via (1) a resource-aware and workload-balanced sequence processor that splits long sequences and packs short ones; and (2) a co-optimization methodology that jointly optimizes pipeline schedule and gradient checkpointing via a mechanism named stage-aware chunk-level adaptive checkpointing. Comprehensive experiments demonstrate that InfiniPipe achieves a 1.69x speedup over state-of-the-art systems.

Figures (13)

• Figure 1: (a) The bottom illustrates DAPPLE ($N_{prefill}$ 1) and Seq1F1B's schedules, where sequences are divided uniformly into $N_{prefill}$ slices, forming homogeneous micro-batches. The upper presents the profiled memory footprint to train GPT-7B on 8 A800 GPUs with a 16K context. Statistics are simulated for DAPPLE due to the OOM error. (b) Heterogeneous micro-batches with B packed from short sequences and the others split from long sequences, requiring a dynamic pipeline schedule. The imbalanced workload introduces pipeline bubbles.
• Figure 2: Statistics of sequences grouped by length intervals. The upper subgraph presents the sample and token distribution, while the bottom one denotes the computation FLOPS distribution.
• Figure 3: Illustration of sequence packing and padding's difference in attention mask and activation arrangement.
• Figure 4: InfiniPipe System Overview
• Figure 5: Illustration of Chunk (heterogeneous micro-batch in EPP). Slices from the same sequence are colored the same. (b) Two split chunks and a tail slice are generated. (c) The tail slice is packed with short sequences to form a hybrid chunk.