Paper: Zero Bubble (Almost) Pipeline Parallelism (ICLR-2024)

Reviewer comment (Experiments):
The method still requires a large enough number of microbatches, e.g., m >= 2p to achieve an amortized zero bubble PP, and how it performs under m < p is rarely discussed. Experiments (at least a roofline analysis) with m < p can make the methods more sound.

Author response:
By the nature of PP, $m < p$ would cause extremely large bubble rate ($\ge 50%$). Previous PP strategies reduce bubble rate by increasing m, for example, in Megatron-LM, all configurations have $m \ge 8p$, while still having a notable bubble rate. Therefore we consider $m < p$ as a very uncommon setting for PP. We want to highlight that compared to previous methods, ZB-2p requires a much smaller number of microbatches to achieve “near zero bubble” (typically $m=3p$ is enough).

Reviewer comment (Novelty):
There are other methods that can achieve zero bubble PP like PipeDream and PipeMare, which are more strictly flushless (thus bubble-free) methods, though they have altered the math of backprop and pose potential accuracy loss or incur more weight memory. Therefore, it is a bit of an overclaim that this work is the first to achieve pipeline bubbles. Discussions on these previous related works are also encouraged.

Author response:
We were aware that asynchronous PP achieves zero bubbles, and intended to briefly touch this point and state that ours is the first to achieve zero bubble under synchronous setting. Somehow we forgot to do so, thanks for pointing this out, we added this discussion in the updated version.

Reviewer comment (Writing):
"Near zero bubble" is probably more accurate; emphasis on "zero bubble" actually overshadows the solid intuition mentioned as the first bullet point in Strengths.

Author response:
To avoid over-claiming, we changed the title to “Near Zero Bubble Pipeline Parallelism” in our updated pdf.

Reviewer comment (Experiments):
Even though authors proposed two algorithms to schedule unbalanced TF/TB/TW, the experiments are not designed to show their advantages. For example, with sentence-size = 1024, TF/TB/TW are almost equal (24h >> 8s in table 1) for all model sizes. Experiments or roofline analysis with 8s >> 24h (or any models with significantly different T) could be helpful to evaluate the efficacy of these algorithms.

Author response:
- We show our profiled value of $T_F/T_B/T_W$ across different settings in Appendix D. From the table, we can see $T_F=T_B$, but $T_W$ is significantly different, with $T_W/T_B$ varies from 0.516 to 0.756.
- Our algos don’t rely on any assumption on the value of $T_F/T_B/T_W$. The reason why the heuristic works well is because of the flexibility of W and the pattern of “1F-1B-1W”.
- Actually by dividing W into a finer granularity, e.g. each single weight gradient calculation, it’s easy and straightforward to fill the bubble because the FLOPs of each weight gradient calculation $\ll T_F/T_B$. In practice, we just found that scheduling based on grouping W is already good enough. So we didn’t bring this complexity.

Reviewer comment (Experiments):
When authors talk about peak memory, I assume they refer to max of peak memory of each worker, rather than the entire pipeline. If it is the latter case, even ZB-H1 should have a higher peak memory than 1F1B based on figure 3. Can you make this more explicit in the texts?

Author response:
Yes, the peak memory in paper refers to the maximum value of peak memory of each worker. Thanks for the suggestion, we edited the PDF to make it more clear.

Reviewer comment (Writing):
What's a "local search" in 3.1?

Author response:
The “local search” means that we have a hyperparameter to control whether to execute this strategy or not, and we perform a grid search to find the best hyperparameters. We have 3 hyperparameters forming 2^3 combinations, which is a small space, therefore we simply scan through all combinations. The word “local search” is indeed misleading, we rephrased for better clarity in our paper.

Reviewer comment (Experiments):
What's the runtime of ILP, and how much does it improve over HEURISTIC? Can you do an ablation with HEURISTIC algorithm only?

Author response:
As ILP has non-polynomial time complexity, in experiments we limit the runtime of ILP to 2 minutes, initialized with result from heuristic algorithm. For most cases in our experiments ILP shows no significant improvements within the time limit because the heuristic algorithm is already good enough as shown in Table 5. As a result we've also moved the ILP to the Appendix.

Reviewer comment (Theory):
My understanding is section 4 BYPASSING OPTIMIZER SYNCHRONIZATIONS only works for INF/NAN. How does it work for grad-norm when each stage needs the global reduced norm? If the optimizer depends on partially reduced state, how can it provide the same numerical result as baselines like 1F1B?

Author response:
For grad-norm clipping, the clipping is only triggered when the norm exceeds a certain threshold, otherwise it has no effect. So our strategy is: we greedily update each stage **assuming that clipping will not get triggered**, while at the same time we propagate a partially reduced norm from the first stage to the last stage and finally get the globally reduced norm. If it exceeds the clipping threshold, we inform each stage to rollback and **redo the clipping** like how we deal with NAN. The above strategy is based on the observation that in a stable training gradient clipping is rarely triggered, like INF/NAN also rarely happens.
We clarified the process in our updated paper.

Reviewer comment (Experiments):
Can you try to recompute the activation gradient (not activation), e.g., redo B during W to save some peak memory in H2/2P cases?

Author response:
Thanks for the insightful suggestion. We respond in the following aspects.
- We would want to highlight ZB-1P/ZB-H1 here. If the 2x activation memory of ZB-2p is a concern, we would prefer ZB-1p over recomputation, because recomputation needs O(#micro-batch) extra computation while ZB-1p only contains O(#pipeline-stage) bubbles.
- The peak memory actually happens in the first stage before the first B, where we save a lot of activations for backward pass, so “redo B during W” won’t reduce the peak memory.

Reviewer comment (Evaluation):
The performance of the method developed in this paper is limited.

Author response:
Despite the memory concern, in an apple to apple comparison where our memory usage is similar to the baselines, our method ZB-1p improves the throughput by 9.5% on average (min 4%, max 15%, row 2 and 3 in table 4). With larger models (14.6B & 28.3B), the average throughput improvement is 11.7%. We think these numbers are already significant considering the cost of GPU. Furthermore, when memory allows, our method pushes the throughput almost to the theoretical upper limit (red bar in figure 5 and bubble rates of ZB-2p in Table 5). 

Another advantage of ZB-2p is that it can achieve near zero bubble with a smaller number of micro-batch (typically m=3p is enough), for 1F1B in Megatron the setting is m>=8p yet still has notable bubbles. This means more micro-batch can be partitioned over Data Parallelism (DP) dimension which brings a better scalability for the training of LLM. Given a fixed PP size and a fixed global batch size, we can reduce #micro-batch of PP to a relatively small number (e.g. 3p) and increase the size of DP, thus shortening the training time by using more devices. When integrated with ZeRO, we can further partition the parameter memory (including optimizer states) to compensate for the extra activation memory introduced by ZB-2p.

Reviewer comment (Evaluation):
The solution mentioned in this paper is likely to increase the memory usage of pipeline parallelism training.

Author response:
We want to highlight the automatic scheduling algorithm which provides the flexibility for users to trade-off between memory usage and bubble rate by themselves. Both ZB-1p and ZB-2p are special cases generated from this algorithm with different memory limits. Although ZB-2p indeed needs a doubled memory, ZB-1p still significantly outperforms 1F1B under a similar memory constraint.

Reviewer comment (Experiments):
The pipeline strategy developed in this paper will increase the memory usage of pipeline parallelism training, which is likely to decrease the size of the micro-batch; a small micro-batch will decrease the efficiency of the pipeline parallelism strategy.

Author response:
We double the size of each microbatch for 1F1B and compare its throughput with ZB-2p. The experimental results show that ZB-2p also holds a higher throughput than 1F1B by about 14% on average (min 4%, max 22%), even with a half size for each microbatch. Empirically, a larger batch size increases the utilization rate of the GPU and thus improves the efficiency. However, it is less of a concern for large models because the hidden dimension is big enough to saturate device utilization. Based on this consideration and our experimental results, we believe ZB-2p is preferred over increasing the batch size for 1F1B in LLM training. For more details of the experiments, please see the Appendix F in our updated pdf.

Reviewer comment (Presentation):
The writing of this paper could be improved, for example, by using different colors to represent B in Figure 3 and Backward in Figure 2.

Author response:
We adjusted accordingly in the updated PDF.

Reviewer comment (Novelty):
The title of this paper is too large, and the reviewer does not think the method in this paper is zero bubble.

Author response:
It is true that the bubble rate can never go to zero because practically there’s always a difference in computation time of F/B/W, introducing gaps in between. As Reviewer xbKE also suggested “Near Zero Bubble” could be more accurate and also well supported by our experimental evidence, we updated the title accordingly in the PDF.

Reviewer comment (Evaluation):
The algorithm can achieve zero bubble pipeline parallelism but with memory usage nearly doubled. This will limit the application of this method for larger LLM given larger models will use larger memory.

Author response:
For the concern of application to LLM, we want to highlight the following points.
* We want to highlight the automatic scheduling algorithm which provides the flexibility for users to trade-off between memory usage and bubble rate by themselves. Both ZB-1p and ZB-2p are special cases generated from this algorithm with different memory limits. Although ZB-2p indeed needs a doubled memory, ZB-1p still significantly outperforms 1F1B under a similar memory constraint.
* We double the size of each microbatch for 1F1B and compare its throughput with ZB-2p. The experimental results show that ZB-2p also holds a better performance even with a half size for each microbatch. Empirically, a larger batch size increases the utilization rate of the GPU and thus improves the efficiency. However, it is less of a concern for large models because the hidden dimension is big enough to saturate device utilization. Based on this consideration and our experimental results, we believe ZB-2p is preferred over increasing the batch size for 1F1B in LLM training. For more details of the experiments, please see the Appendix F in our updated pdf.
* Another advantage of ZB-2p is that it can achieve near zero bubble with a smaller number of micro-batch (typically m=3p is enough), for 1F1B in Megatron the setting is m>=8p yet still has notable bubbles. This means in ZB-2p, more micro-batch can be partitioned over Data Parallelism (DP) dimension which brings a better scalability for the training of LLM. Given a fixed PP size and a fixed global batch size, we can reduce #micro-batch of PP to a relatively small number (e.g. 3p) and increase the size of DP, thus shortening the training time by using more devices. When integrated with ZeRO, we can further partition the parameter memory (including optimizer states) to compensate for the extra activation memory introduced by ZB-2p.

Reviewer comment (Evaluation):
For Figure 5, sometimes ZB-1p is better than 1F1B-I, and sometime not. Can you explain the reason in the paper?

Author response:
From Table 5, we can find that 1F1B-I has lower bubble rates for models of 1.5B, 6.2B and 28.3B, and a bit higher bubble rate for the 28.3B model. However, 1F1B-I is generally under-performed than theoretical efficiency in our experiments, especially in multi-node setups. We guess the reason is that 1F1B-I has more communication and computation passes than both 1F1B and ZB-1p, so it is more likely to be affected by the fluctuations in the time cost of each pass. For multi-node setups, the time cost of inter-node communication is usually larger than intra-node communication, which also decreases the implemental efficiency of 1F1B-I because we don’t take this difference into account.

Reviewer comment (Reproducibility):
My main concern relates to breaking the synchronization of optimizer step because it complicates the synchronous training semantics, and adoption requires close interaction of rollback and model checkpointing.

Author response:
We guess the “model checkpointing” you mentioned here means saving previous model parameters/optimizer states and restoring them when rollback.
- For efficiency, we recommend the in-place optimizer rollback proposed in Appendix C, which does not require checkpointing or extra copy of parameters and should be more efficient.
- For simplicity, from our study, gradient clipping and NAN are mainly activated during the initial stage of training, which we believe is more a choice for training stability than model performance. Therefore, instead of post validation, there’re indeed other simpler choices that may achieve the same. For example, clipping the gradients locally. Since in this work we would like to focus on the systems aspect, we did not explore the numerical/algorithmic choices, but just align it with our baselines.

Reviewer comment (Experiments):
Evaluating the performance of synchronous optimizer step as way of understanding the trade-off between the simplicity and throughput.

Author response:
Thanks for the valuable suggestion. We add some new experiments in Appendix E, which shows that the synchronous optimizer under-performs our post-validation strategy by about 8.8% on average.

Reviewer comment (Experiments):
Another concern, or perhaps a question, is whether the more memory efficient schemes (1F1B* and ZB-1p) be evaluated with larger micro-batch sizes than ZB-2p; this would help to confirm that the evaluation is not biased towards ZB-2p.

Author response:
We add experiments comparing ZB-2p and 1F1B (with double micro-batch size) under the same memory consumption. The results show that ZB-2p also has a higher throughput than 1F1B by about 14% on average (min 4%, max 22%). Please find related data in Appendix F in the updated manuscript.