[Paper Review] Zero Bubble Pipeline Parallelism

Penghui Qi, Xinyi Wan+2 more|arXiv (Cornell University)|Nov 30, 2023

Parallel Computing and Optimization Techniques4 citations

TL;DR

This paper proposes Zero Bubble Pipeline Parallelism (ZB), a novel scheduling strategy that achieves near-zero pipeline bubbles in synchronous training by decoupling backward computation into input gradient (B) and parameter gradient (W) components. By reorganizing the computation graph at a finer granularity and introducing a new schedule (ZB-V), the method reduces bubble rates by up to 75% compared to 1F1B, improving throughput by 23% under similar memory limits and up to 31% when memory is relaxed.

ABSTRACT

Pipeline parallelism is one of the key components for large-scale distributed training, yet its efficiency suffers from pipeline bubbles which were deemed inevitable. In this work, we introduce a scheduling strategy that, to our knowledge, is the first to successfully achieve zero pipeline bubbles under synchronous training semantics. The key idea behind this improvement is to split the backward computation into two parts, one that computes gradient for the input and another that computes for the parameters. Based on this idea, we handcraft novel pipeline schedules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds an optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce a novel technique to bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the 1F1B schedule up to 23% in throughput under a similar memory limit. This number can be further pushed to 31% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. We open sourced our implementation based on the popular Megatron-LM repository on https://github.com/sail-sg/zero-bubble-pipeline-parallelism.

Motivation & Objective

• Address the persistent inefficiency in pipeline parallelism caused by pipeline bubbles under synchronous training semantics.
• Overcome the limitations of existing scheduling strategies like 1F1B, which still suffer from non-zero bubble rates despite improvements.
• Design a scheduling framework that minimizes bubble rates while respecting memory constraints and enabling efficient large-scale model training.
• Enable zero-bubble training without sacrificing optimization semantics or increasing memory usage beyond baseline methods.
• Develop an automated algorithm to generate optimal schedules based on model configuration and memory limits.

Proposed method

• Decompose the backward pass into two independent components: gradient computation with respect to input (B) and with respect to parameters (W), breaking the traditional monolithic backward function.
• Introduce a new pipeline schedule, ZB-V, that reorders B and W computations to eliminate dependencies that cause bubbles, achieving zero bubbles when forward, backward, and parameter gradient times are equal.
• Design a novel synchronization bypass technique during the optimizer step to eliminate idle time and further reduce bubbles.
• Develop an automated scheduling algorithm that finds optimal schedules under given memory and model constraints, minimizing bubble rate.
• Use a computation graph representation at the function level (F, B, W) instead of layer-level abstractions to enable finer-grained scheduling and optimization.
• Implement and open-source the solution based on Megatron-LM, enabling integration into existing large-scale training frameworks.

Experimental results

Research questions

• RQ1 Can pipeline bubbles be eliminated in synchronous pipeline parallelism through a rearchitecting of backward computation?
• RQ2 What scheduling strategy achieves the lowest possible bubble rate while maintaining memory efficiency and compatibility with standard training semantics?
• RQ3 How does decomposing the backward pass into input and parameter gradient components affect pipeline efficiency and memory usage?
• RQ4 To what extent can automated scheduling reduce bubble rates across diverse model sizes and memory budgets?
• RQ5 What is the trade-off between microbatch size and bubble rate, and can it be optimized for maximum throughput?

Key findings

• ZB-V achieves a bubble rate of 0.0274 for a 28.3B model with 32 stages and 256 microbatches, compared to 0.2676 for 1F1B, representing a 89.8% reduction.
• Under similar memory constraints, ZB-V improves throughput by up to 23% over 1F1B, and by up to 31% when memory is relaxed.
• The method reduces bubble rate by up to 75% compared to 1F1B across all evaluated model sizes and configurations.
• ZB-V maintains the same memory footprint as 1F1B while achieving near-zero bubbles, unlike ZB-H2, which requires twice the memory for comparable performance.
• The automated scheduling algorithm successfully generates optimal schedules that outperform heuristic baselines, especially under tight memory limits.
• Increasing microbatch size improves throughput by up to 8% for larger models, but the benefit diminishes for smaller models, indicating a trade-off between microbatch size and bubble rate.