AMDP: Asynchronous Multi-Directional Pipeline Parallelism for Large-Scale Models Training
Abstract
Pipeline parallelism is essential for large-scale model training, but existing asynchronous approaches often degrade convergence due to parameter mismatch between forward and backward passes. We propose Asynchronous Multi-Directional Pipeline parallelism (AMDP) to mitigate this issue while sustaining high utilization. AMDP limits the first stage of each pipeline to process at most two minibatches before backpropagation, bounding the number of parameter updates between forward and backward passes. To alleviate the resulting pipeline bubbles, AMDP launches multiple concurrent pipelines and adapts their number according to pipeline depth. In addition, AMDP accumulates gradients across minibatches and applies them in a single update, ensuring that only a bounded number of minibatches experience parameter mismatch, limited to within one optimization step. Experiments on GPT- and BERT-style models demonstrate that AMDP significantly accelerates training while preserving convergence.
Lay Summary
Training large language models like GPT-3 requires massive computational resources, making single-GPU training impractical due to memory constraints and extremely long training times. While pipeline parallelism distributes the model across multiple GPUs to reduce memory pressure, existing approaches face trade-offs: synchronous methods suffer from pipeline bubbles that idle hardware, whereas asynchronous methods introduce parameter mismatch that harms convergence. We propose Asynchronous Multi-Directional Pipeline parallelism (AMDP) to address these issues. AMDP uses concurrent pipelines with a scheduling strategy that limits parameter mismatch to at most one step. It combines gradient accumulation to reduce communication and integrates memory-efficient ZeRO optimization to avoid redundant storage. Extensive experiments show that AMDP increases training throughput by up to 17% over state-of-the-art baselines with minimal extra memory, significantly accelerating large-scale model training while maintaining convergence quality.