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.

• The paper presents a novel AMDP framework that strictly limits parameter mismatch to one optimization step by capping stage 0 to two minibatches before backpropagation.
• It employs multi-directional scheduling to fill pipeline bubbles and achieve up to 17% higher throughput compared to prior asynchronous and synchronous methods.
• By integrating a Zero Redundancy Optimizer, AMDP reduces memory footprint and communication overhead, ensuring near-synchronous convergence with scalable performance.

AMDP: Asynchronous Multi-Directional Pipeline Parallelism for Large-Scale Model Training

Motivation and Context

Scaling transformer-based models necessitates distributed training strategies to address prohibitive memory and computational requirements. Pipeline parallelism partitions deep neural networks across multiple GPUs, providing memory efficiency and increased throughput compared to purely data-parallel approaches. However, traditional synchronous pipeline paradigms are hampered by pipeline bubbles—idle periods induced by inter-stage dependencies. Asynchronous approaches, like PipeDream, eliminate bubbles but introduce parameter mismatch between forward and backward passes, leading to degraded convergence, particularly as pipeline depth increases. Existing mitigation techniques, such as parameter stashing and prediction, are limited by delayed gradients or forecasting inaccuracies.

AMDP Design and Structural Advances

AMDP introduces a structurally bounded pipeline-parallel training framework, comprising four core advances:

1. Parameter Mismatch Control:

AMDP enforces that stage 0 processes no more than two minibatches before backpropagation, formally capping the degree of parameter mismatch at one optimization step, independent of pipeline depth or device mapping. This eliminates the linear growth of mismatch inherent to classic asynchronous pipelines.

Figure 1: Parameter mismatch at stage 0 grows with the number of minibatches read before the first backward; limiting to two minimizes mismatch to one step.

2. Multi-Directional Scheduling:

To counterbalance the bubbles resulting from mismatch control, AMDP launches multiple concurrent pipelines with complementary traversal patterns ("multi-directional" scheduling). The number of pipelines (d/2 for pipeline depth d) is chosen analytically, ensuring all idle periods are filled and hardware utilization is maximized.

Figure 2: AMDP bidirectional scheduling with four stages and devices; both trailing and leading bubbles are eliminated, achieving bubble-free execution.

Figure 3: In AMDP for $d=8$, four pipelines with alternate directions fill bubbles efficiently, enabling full device utilization.

3. Gradient Accumulation Updates:

Instead of updating parameters after each backward pass, gradients are accumulated across multiple minibatches, and updates are applied in batch at synchronization points. Only the first $d$ minibatches per window experience mismatch, strictly bounded within one optimization step. This significantly reduces communication overhead without impacting convergence beyond negligible second-order perturbation.

Figure 4: With accumulation threshold equal to pipeline depth $d=4$, each minibatch incurs a one-step mismatch.

4. Zero Redundancy Optimizer (ZeRO):

To avoid the memory penalty of multiple pipelines, optimizer states are sharded across devices, with each stage's optimizer state stored only once. Gradient reduce followed by broadcast replaces traditional all-reduce, maintaining communication efficiency.

Theoretical Guarantees

AMDP's structural controls yield a fundamental improvement in optimization stability for asynchronous pipeline parallelism. It is proven that bounded parameter mismatch (max one step) ensures convergence characteristics equivalent to synchronous SGD or adaptive optimizers, with only second-order $\mathcal{O}(\eta^2)$ perturbation. This topological invariance holds under arbitrary device mappings and multi-node deployments.

Empirical Results

Experiments are conducted on 8 and 16 NVIDIA A800 GPUs, evaluating AMDP against contemporary synchronous and asynchronous pipeline methods using large GPT- and BERT-style models.

Throughput:

AMDP achieves consistently superior throughput, outperforming PipeDream-2BW and XPipe by up to 17% under high pipeline depth and large update batch size. Its advantage increases with depth and batch size due to enhanced bubble mitigation and utilization.

Figure 5: AMDP memory footprint in training a GPT-style model for $d=4$.

Memory Footprint:

AMDP delivers balanced per-GPU memory utilization and only marginally greater peak memory than top-performing baselines, attributed to temporary retention of extra parameters and gradients from multi-directional scheduling. Integration of ZeRO reduces optimizer-state storage by up to 19.4% and improves throughput by 4%.

Figure 6: Loss versus iterations for GPT-style model shows AMDP closely tracks synchronous baselines.

Convergence:

Loss curves and validation perplexity show AMDP maintains near-synchronous convergence. The final training loss is within 0.02 of synchronous Inter-1F1B on GPT, and identical to DAPPLE on BERT. AMDP shortens time to target loss by 23% (GPT) and 22% (BERT) due to higher practical throughput. Validation perplexity exhibits negligible degradation relative to synchronous pipeline.

Ablation:

Removing ZeRO increases memory usage and slows throughput, confirming its crucial role in AMDP's scalability.

Practical and Theoretical Implications

AMDP redefines the trade-offs in pipeline-parallel training for large-scale models. By restricting parameter staleness asynchrony to a single step and systematically maximizing utilization through multi-directional scheduling, AMDP provides high-performance distributed training with robust optimization stability. In practice, it enables more predictable learning dynamics, scalable execution across heterogeneous topologies, and better resource allocation in multi-node clusters. Theoretical foundations are strengthened, showing asynchronous pipeline methods can be made competitive with synchronous ones even as models scale.

Looking forward, AMDP's framework can be aligned with mismatch-aware learning rate adaptation or further optimized pipeline deployment strategies. Its principles are expected to generalize to emerging weight-passing (WeiPipe) and topology-aware (TawPipe) parallelism techniques, especially for long-context transformer-based models.

Conclusion

AMDP introduces an asynchronous, structurally-bounded multi-directional pipeline parallelism regime, achieving state-of-the-art throughput and near-synchronous convergence in large-scale transformer training with minimal memory overhead. Its convergence guarantees and empirical dominance suggest that future distributed training infrastructures should incorporate mismatch-bounded, multi-directional execution with ZeRO-based optimizer state partitioning. Extensions to adaptive scheduling and learning-rate control promise further gains in scalability and robustness for distributed deep learning architectures.