Zero Bubble Pipeline Parallelism (2401.10241v1)

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.

• The paper demonstrates a novel scheduling method that splits backward computation to eliminate pipeline bubbles in distributed training.
• It introduces both handcrafted and automated scheduling techniques using integer linear programming and heuristic strategies.
• Experimental results show up to a 31% increase in throughput with optimized memory usage compared to traditional 1F1B methods.

Zero Bubble Pipeline Parallelism

Introduction

The paper "Zero Bubble Pipeline Parallelism" addresses a major challenge in distributed model training: the inefficiency caused by pipeline bubbles in pipeline parallelism (PP). As the complexity and size of neural networks increase, distributed training involving multiple GPUs has become essential. Traditional methods like data parallelism (DP) and model parallelism, including tensor parallelism (TP) and PP, have been optimized to various extents. However, pipeline bubbles—idle times during execution due to dependencies between stages—persist as a bottleneck. This research proposes a novel scheduling strategy that eliminates these bubbles while maintaining synchronous training semantics.

Key Concepts and Methodology

The main innovation of this paper is a scheduling strategy that eliminates pipeline bubbles. The strategy involves splitting the backward computation into two parts: one for computing gradients with respect to the input, and the other for parameters. This split allows the researchers to design novel pipeline schedules that are highly efficient. The paper introduces both handcrafted and automated scheduling techniques tailored to specific model configurations and memory limitations.

Handcrafted Schedules: The authors initially present two handcrafted pipeline schedules. The first one, \zbh{1}, maintains similar peak memory usage as conventional 1F1B schedules but with reduced bubble sizes. The second one, \zbh{2}, achieves zero bubbles but with higher memory consumption.

Figure 1: Handcrafted pipeline schedules, top: \zbh{1}; bottom: \zbh{2}.

Automatic Scheduling: The paper further develops an automatic scheduling algorithm, leveraging integer linear programming and heuristic methods to optimize scheduling under realistic conditions. This algorithm considers the execution time and memory constraints, achieving near-zero bubble rates while maintaining computational efficiency.

Experimental Results and Implications

The method was evaluated against baseline methods such as 1F1B and 1F1B-I, showing significant improvement in throughput. For instance, the proposed \zb{2} schedule outperforms traditional 1F1B scheduling by up to 31% when memory limits are relaxed. The experiments were conducted on a variety of model sizes, demonstrating the scalability of the approach.

Figure 2: ZB-V schedule.

The zero bubble schedules present clear benefits in terms of throughput and memory usage. The "Zero Bubble V" (ZB-V) schedule, depicted in Figure 2, achieves comparable throughput to \zb{2} but uses significantly less memory, making it a practical choice for resource-constrained environments.

Figure 3: The relation between memory limit and bubble rate for ZB-V, compared with the heuristic method.

Conclusion

The introduction of a zero-bubble scheduling strategy marks a substantial advancement in pipeline parallelism. By eliminating pipeline bubbles, this method optimizes the use of GPUs in distributed training settings, ultimately reducing training times and costs. The research not only provides a theoretical basis for bubble-free pipeline scheduling but also demonstrates practical implementations that can be adopted in existing frameworks, as evidenced by their work with the Megatron-LM repository. Future developments may focus on integrating these techniques with other parallelism strategies, targeting more complex architectures and further improving training efficiency on a large scale.