TLDR: InfiniPipe introduces Elastic Pipeline Parallelism (EPP) to efficiently train Large Language Models (LLMs) with long and varied contexts. It addresses challenges like high communication overhead and memory consumption by adaptively combining batch-level and token-level pipeline parallelism. InfiniPipe intelligently splits long sequences and packs short ones into “heterogeneous micro-batches” and co-optimizes pipeline scheduling with a novel stage-aware adaptive gradient checkpointing mechanism. This results in significant speedups (up to 1.69x) over current state-of-the-art systems.
Training Large Language Models (LLMs) to understand and generate longer contexts is a critical area of research, enabling these powerful AI models to handle more complex tasks. However, this ‘long-context training’ comes with significant computational challenges, particularly in distributed computing environments.
Existing methods often struggle with efficiency. For instance, ‘sequence parallelism’ can incur substantial communication overhead, especially when data needs to be moved between different computing nodes. ‘Pipeline parallelism’ (PP) offers a way to reduce this communication, but its effectiveness depends heavily on how data is divided, or its ‘granularity’.
The paper ‘Data-Centric Elastic Pipeline Parallelism for Efficient Long-Context LLM Training’ introduces a novel solution called Elastic Pipeline Parallelism (EPP) and a distributed training system named InfiniPipe. This new approach aims to overcome the limitations of current methods by intelligently adapting to both the available computing resources and the varied nature of real-world datasets.
One of the core problems EPP addresses is the trade-off in PP granularity. ‘Batch-level PP’ groups multiple input samples, which can lead to high memory consumption when dealing with very long sequences. Conversely, ‘token-level PP’ splits sequences into smaller slices, reducing memory but potentially under-utilizing hardware if not managed carefully. Furthermore, real-world datasets often have a skewed distribution of sequence lengths – many short sequences and a few very long ones – which can lead to imbalanced workloads and inefficient scheduling with static PP methods.
EPP is designed with two key characteristics: adaptability and hybridization. Adaptability means it can dynamically choose the right PP granularity based on the specific workload and hardware resources. Hybridization allows it to orchestrate both batch-level and token-level PP. This means InfiniPipe can intelligently split long sequences to keep memory usage in check, while simultaneously packing shorter sequences together to ensure that computing hardware is used efficiently. This results in ‘heterogeneous micro-batches’ – data chunks of varying sizes and structures – which require a dynamic and workload-aware pipeline schedule.
Also Read:
- Optimizing Large Language Model Training in Mobile Edge Networks with CollaPipe
- PPSD: Boosting LLM Inference Speed with Pipelined Self-Speculative Decoding
How InfiniPipe Works
The InfiniPipe system is built on several key techniques:
-
An Effective Cost Model: This model accurately estimates the computation, communication, and memory footprint for these heterogeneous micro-batches, even considering the impact of gradient checkpointing (a technique to save memory by recomputing some data during training).
-
A Workload-Balanced and Resource-Aware Sequence Processor: This component takes the varied-length sequences from the dataset and organizes them into workload-balanced ‘chunks’. It uses a two-phase process: splitting long sequences and packing short ones. This ensures that the workload is evenly distributed and that the chunk granularity is appropriate for the available resources.
-
A Chunk Scheduler with Co-Optimization: This is where InfiniPipe truly shines. It jointly optimizes the pipeline schedule and the gradient checkpointing configuration. It uses a novel mechanism called ‘stage-aware chunk-level adaptive checkpointing’. This adaptive approach means that checkpointing is applied strategically, reducing unnecessary recomputation while still managing memory effectively across different stages of the pipeline and for different types of data chunks.
The researchers, Shiju Wang, Yujie Wang, Ao Sun, Fangcheng Fu, Zijian Zhu, Bin Cui, Xu Han, and Kaisheng Ma, conducted extensive experiments to evaluate InfiniPipe. Their findings demonstrate that InfiniPipe consistently outperforms state-of-the-art distributed training systems. It achieved a speedup of up to 1.69 times compared to FlexSP, 2.07 times compared to DeepSpeed, and 2.60 times compared to Megatron-LM. These significant performance gains highlight InfiniPipe’s ability to reduce communication overhead and improve overall training throughput, especially for larger models and longer context lengths.
InfiniPipe represents a significant step forward in making long-context LLM training more efficient and accessible, paving the way for even more capable and versatile AI models. For more technical details, you can read the full research paper here.
