Zero Bubble (Almost) Pipeline Parallelism
Authors: Penghui Qi
Xinyi Wan
Guangxing Huang
Min Lin
Venue: ICLR 2024 poster (2024)
Keywords: Pipeline Parallelism
Zero Bubble
Primary area: infrastructure, software libraries, hardware, etc.

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 15\% in throughput under a similar memory limit. This number can be further pushed to 30\% when the memory constraint is relaxed. We believe our results mark a major step forward in harnessing the true potential of pipeline parallelism. The source code based on Megatron-LM is publicly avaiable at \url{https://github.com/sail-sg/zero-bubble-pipeline-parallelism}.

Published as a conference paper at ICLR 2024
ZERO BUBBLE (ALMOST) PIPELINE PARALLELISM
Penghui Qi", Xinyi Wan', Guangxing Huang & Min Lin
Sea Lab
{qiph, wanxy, huanggx, linmin 8sea.ccmmoommmommmmcmm
ABSTRACT
Pipeline parallelism is one of the key components for large-scale distributed train- ing, 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 seman- ties. 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 com- putes for the parameters. Based on this idea, we handcraft novel pipeline sched- ules that significantly outperform the baseline methods. We further develop an algorithm that automatically finds optimal schedule based on specific model configuration and memory limit. Additionally, to truly achieve zero bubble, we introduce novel technique bypass synchronizations during the optimizer step. Experimental evaluations show that our method outperforms the IFIB schedule up to 15% in throughput under a similar memory limit. This number can be further pushed to 30% when the memory constraint is relaxed. We believe our results mark major step forward harnessing the true potential of pipeline par- allelism. The source code based on Megatron-LM is publicly avaiable at ht htsps: / //gtthu.. com/ sai 1ii1--s9h
ine waapeeeeee-earraaeeeeeeeeeeaaaaa le sm.
INTRODUCTION
The realm of distributed model training has become focal point in the deep learning community, especially with the advent of increasingly large and intricate models. Training these behemoths often requires a vast amount of GPUs interconnected with various topologies. Various parallelism techniques have been proposed for training DNN in the past years. Data parallelism (DP) (Goyal al., 2017; Li et al., 2020) is the default strategy for models of small to moderate sizes due to its simplicity. Beyond model size. it is no longer possible to fit the model parameters in one single GPU. This is when model parallelism comes to the rescue (Harlap et al., 2018; Huang al., 2019; Fan et al., 2021; Zheng et al., 2022). There are two main model parallel schemes, tensor parallelism (TP) and pipeline parallelism (PP). TP splits the matrix multiplication in one layer to several devices, while PP segments the entire model into different stages which can be processed across different devices. Notably, ZeRO (Rajbhandari et al., 2020) provides a strong alternative to model parallelism by sharding parameters across devices, while keeping the simplicity of DP.
Recent research indicates that achieving optimal performance in large-scale training scenarios re- quires non-trivial interaction of DP, TP and PP strategies. In the abundance of interconnection resources, e.g. NVLink between GPUs within one compute node, a hybrid of DP, TP and ZeRO strategies works efficiently. Whereas there numerous empirical evidences Fan et al. (2021); Zheng et al. (2022); Narayanan et al. (2021) showing PP is particularly advantageous for utilizing cross-server connections, especially the scale of thousands of GPUs. This highlights the primary aim of our work: enhancing the efficiency of PP.
Going deeper into the intricacies of PP, the efficiency of its implementation relies heavily on the amount of device idle time referred to as pipeline bubbles. Due to the dependency between lay- ers. bubbles seem inevitable. prominent early work to address this issue GPipe (Huang al., 2019), which attempts to reduce the bubble ratio by increasing the number of concurrent batches in the pipeline. However, direct consequence of this is an increase in peak memory demands.
"Equal Contributors

---

depends W at the layer i, which is usually detrimental for the efficiency of the Published as a conference paper at ICLR 2024
To mitigate this, GPipe discards part of the intermediate activations while recomputing them dur- ing the backward pass. Yet, this approach introduced a computation overhead of around 20% (Fan et al.. 2021). One line of work that improves over GPipe focuses asynchronous PP, including PipeDream (Harlap et al., 2018), PipeMare (Yang et al., 2021). Asynchronous PP is theoretically bubble free, they greatly improve pipeline efficiency, however, the sacrifice of exact optimiza- tion semantics. On the other hand, improvements are also made under synchronous settings. A notable scheduling strategy to address the limitation of GPipe called one-forwarddone-baakiee backward (IFIB). It was first proposed in PipeDream (Harlap et al., 2018) under the asynchronous setting, and later introduced under synchronous settings (Fan al., 2021; Narayanan et al., 2021). 1FIB offers faster memory clearance by early scheduling the backward passes. With the same number of microbatches, yields similar bubble ratios but with a distinct advantage peak memory. Based on IFIB, Narayanan et al. (2021) introduced the IFIB interleaved strategy. By assigning multiple stages to the same device, further reduces the bubble size at the cost of more communication and higher peak memory.
Despite various efforts, to this date the remaining bubbles still pose the largest issue for PP under synchronous training semanties. In this work, spotted the opportunity that PP can be further optimized by representing and scheduling the computation graph at a finer granularity. Classical deep learning frameworks are designed the granularity of layers, whereas modern deep learning compilers use different intermediate representations for optimizations at various levels. (Chen et al., 2018: Roesch al., 2018; Sabne, 2020; Tillet et al., 2019; Lattner et al., 2020). Although a finer granularity always means larger space for searching, it is often impeded by the lack of optimization tools navigate the space. Therefore, choosing suitable granularity is crucial.
Forward
V.L
x
W
P.L
N
7,L
(2)
dar(2)
V,L
7,6
V,LX"
Vy
dz
F
W
Figure 1: Computation Graph for MLP.
Traditionally, neural networks are granularized as stacked layers. There are two functions associated with each layer, forward and backward. In the forward pass, the input transformed into the out- put with the parameterized mapping f(æ, W). The backward pass, crucial for training, involves Correspondingly, they compute the
two computations: fæ W) dy
de and Vw f(x W) dy
de
gradient with respect to the input æ and the layer's parameters W. For convenience, we use single letters B and W to denote these two computations respectively. and denote forward pass (Figure 1). Traditionally. B and W are grouped and provided as a single backward function. This design is conceptually friendly to the user, and it happens to work well for DP, because the communication of the weights' gradient at layer can be overlapped with the backward computation layer 1. However, in PP, this design unnecessarily increases the sequentially dependent computations, i.c. the layer pipeline.
Based split B and W, we present pipeline schedules that greatly improve pipeline efficiency. The remainder this paper is organized as follows: In Section 2, we introduce handcrafted sched- ules based on ideal assumption that the execution times of F, B and W are identical. Subsequently,
2

---

Published as a conference paper at ICLR 2024
in Section 3, remove this assumption and propose automatic scheduling algorithm that works under more realistic conditions. To achieve zero bubble, Section details method that sidesteps the need for synchronization during the optimizer step, yet preserves synchronous training seman- tics. We conclude with empirical evaluations of our methods against baseline methods under diverse settings.
We should note that we do not aim to explore general mixed strategies for large scale distributed training. Instead, specifically target improve the efficiency of pipeline scheduling, supported with apple to apple comparisons with baselines. Our method is orthogonal to DP, TP and ZeRO strategies, and can be used parallel replacement for the PP part large scale training.
HANDCRAFTED PIPELINE SCHEDULES
Based on the key observation that splitting B and W could reduce sequential dependency and thus improve efficiency, we redesign the pipeline starting from the commonly utilized IFIB schedule. As depicted in Figure 2, IFIB initiates with warm-up phase. In this phase, workers conduct varying numbers of forward passes, with each stage typically performing one more forward pass than its immediately subsequent stage. Following the warm-up phase, each worker transits to steady state where they alternately execute one forward pass and one backward pass, ensuring an even workload distribution among stages. In the final phase, each worker processes the backward passes for the outstanding in flighthh microbatches, completing the batch.
In improved version split the backward pass into and W passes, it imperative that F and B from the same microbatch must still remain sequentially dependent across pipeline stages. However, W be flexibly scheduled anywhere after the corresponding B of the same stage. This allows for strategic placement of W to fill the pipeline bubbles. There are many possible schedules that improve over IFIB, trading off differently the bubble size and the memory footprint. We introduce two particularly interesting handcrafted schedules in this section to show the great poten- tial of finer granularity reducing pipeline bubbles (see Figure 3). For the sake of clarity in our initial design, we assume that the time costs for F, B, and W are identical, an assumption shared by earlier studies (Narayanan al., 2021; Huang et al., 2019). However, Section 3, we re-evaluate this assumption to optimize scheduling efficiency in real-world scenarios.
Device
Devic
5
8
Devkc 3
6
8
Budlward
Figure 2: IF1B pipeline schedule.
Dedio
Desko 2
5
6
Devin
3
5
8
Devic 2
Deske 3
Desko4 4
stop
Figure 3: Handerafted pipeline schedules, top: ZB-HI; bottom: ZB-H2
2.1
MEMORY EFFICIENT SCHEDULE
Our first handcrafted schedule, named ZB-HI, ensures that the maximum peak memory usage over all workers doesn't ' exceed that of 1F1B. ZB-H1 generally follows the FFIB schedule, but it adjusts
3

---

Published as a conference paper at ICLR 2024
the starting points of W depending the number of warm-up microbatches. This ensures all work- ers maintain the same number of in-flight microbatches. As a result, as seen in Figure (top). the bubble size is reduced to third of 1FIB's size. This reduction is because B is initiated earlier across all workers compared to 1FIB, and the tail-end bubbles are filled by the later-starting W passes. As W typically uses less memory than B (Table 1), the first worker has the maximum peak memory usage which is consistent with IFIB.
2.2 ZERO BUBBLE SCHEDULE
When we permit a larger memory footprint than IFIB and have sufficient number of microbatches, it's possible to achieve zero bubble schedule, which we label ZB-H2. As illustrated in Fig- ure (bottom), we introduce more passes during the warm-up phase to fill the bubble preceding the initial B. We also reorder the W passes at the tail, which changes the layout from trapezoid into a parallclogram, climinating all the bubbles the pipeline. It is important to highlight that the synchronization between the optimizer steps is removed here, we discuss how this is safely done in Section 4.
2.3
QUANTITATIVE ANALYSES
We use p to denote the number of stages and b denote the size of each microbatch. For transformer architecture, we denote the number of attention heads as a, the sequence length as 8 and the hidden dimension size as h. We use the notations Mp to represent the memory required activations for one B/W pass, and Tp/T'//w to represent the running time for one F/B/W pass. For simplicity, we only do quantitative analyses on transformer architecture (Vaswani et al., 2017), using a typical setting similar to GPT-3 (Brown et al., 2020) where the hidden dimension size inside feedforward is 4h and the dimension size for each attention head h/a.
store
As in Narayanan et al. (2021), we only consider matmul operations when calculating FLOPs because they contribute most of the computations in a transformer layer. For each matmul operation in the forward pass, there are two matmul operations with the same FLOPs in corresponding backward pass (see Figure 1), each of which belongs to either B or W. The approximate formula for calculating the FLOPs of transformer layer is in Table 1. We can see that Tw < T'p TB and TB + Tw 2Tp. We use the same method in Korthikanti et al. (2023) to estimate activations memory required for B. After B completes, it releases some activations not used anymore but keeps some extra gradients (V&L L in Figure 1) for W. The total memory required by W, as in Table 1, is less than B.
Table 1: FLOPs and activations memory required per transformer layer for each pass
Pass
FL.OPs
Activations Memory Required
F
sbh(24h + 4s)
0
sbh(24h + 88)
sb(3h + 5as)
W
32sbh
Without the assumption of Tr Tn Tw. the peak activations memory and bubble size of ZB- HI and 7B-H2 are quantified in Table 2. Notably, the activations memory of worker is (p - + M) + (i Mw for ZB-HI and (2p 2i + 1) + (2i 2) for ZB-H2. As in Table 1. the activations memory required for W is smaller than that for B. Therefore, the peak activations memory is pMg and (2p- MB, for ZB-HI and ZB-H2 respectively.
Table 2: Comparison between IFIB and our handcrafted schedules.
Schedule
Bubble size
Peak activations memory
IFIB
(p- [) + Ta + +w)
PMB
ZB-HI
(p- l) TB -Tw)
PM
ZB-H2
1-]] + +TB -2Iw)
(2p 1)MB
4

---

Published as a conference paper at ICLR 2024
3 AUTOMATIC PIPELINE SCHEDULING
While handcrafted schedules offer simplicity and better comprehensibility, they face several issues in practical applications. For one, scheduling under the assumption that TF = TB Tw introduces unwanted bubbles, especially for models where these values differ significantly. Moreover, com- munication time (denoted as T comm)) required transfer activation/gradien between stages is often ignored in handcrafted schedules, leading to noticeable latencies in the pipeline stream. Finally, striking a balance between minimizing bubble size and adhering memory limit becomes particu- larly challenging when the available memory is insufficient to accommodate enough microbatches for a bubble-free schedule.
To address these challenges and ensure generalization to practical scenarios, we propose algorithms to automatically search the optimal schedule given the number of pipeline stages p, the number of microbatches the activations memory limit Miimit, and the running time estimations Tp. TB. Tw and T Teomm- We design a heuristic strategy, which always generates an optimal or near optimal solution especially when m is large enough. We also systematically formulate the problem Integer Linear Programming (for more details see Appendix G), which can be solved by an off-the-shelf II.P solver (Forrest & Heimee- 2005) when the problem is under certain scale. These two approaches can be combined: first, use the heuristic solution as initialization, and then optimize further with ILP.
3.1 THE HEURISTIC ALGORITHM
We present our heuristic algorithm in the following steps:
In the warm-up phase, within the memory limit, we schedule as many F' passes as possible to minimize the bubble before the first B. The resulting schedule may still have a small bubble (less than 7'p) before the first B if not reaching memory limit, where scheduling another F may delay the following B. We use a binary hyperparameter to control whether to do it or not.
After the warm-up phase, we adhere the pattern where one F and one B are scheduled itera- tively. We insert W to fill the bubble when there is gap larger than Tw. When a bubble occurs but the size is less than Tw, we still insert a W if current bubble makes the largest cumulative bubble size among all stages become larger. We also insert W to recycle some memory when the memory limit is hit. Typically, our hcuristic strategy enters a steady state that follows 1F-1B-1W pattern.
Throughout this process, pipeline stage i is always guaranteed to schedule at least one more F than stage + anytime before used up. When this difference exceeds one, we use another binary hyperparameter to decide whether to skip one F in pipeline stage if it doesn't cause more bubbles. We perform grid search find the best combination of hyperparameters.
In each stage, when F and passes run out, we schedule all the left W passes one by one.
BYPASSING OPTIMIZER SYNCHRONIZATIONS
In
most practices of PP, synchronizations over pipeline stages are usually performed in optimizer step for the sake of numerical robustness. For example, global gradient norm needs to be com- puted for gradient norm clipping (Pascanu al., 2013); a global check for NAN and INF values are performed in the mixed precision settings (Micikevicius et al., 2017); both of them require an all-reduce communication across all stages. However, synchronization at the optimizer step de- stroys the parallelogram (Figure 3) and makes zero bubble impossible. In this section, we propose an alternative mechanism to bypass these synchronizations, while still maintaining synchronous optimization semantics.
In existing implementations, an all-reduce communication is first launched to collect the global states, followed by the optimizer steps which are conditioned the global states. However, we noticed that most of the time the global states have no effects, e.g., the global check for NAN and INF rarely trigger because in robust setting most iterations shouldn' 't have numerical issues; the gradient clipping rate also quite low empirically to justify a synchronization global gradient norm at every iteration.
5

---

Published as a conference paper at ICLR 2024
Based on these observations, propose to replace the before-hand synchronizations with post update validation. The idea is illustrated in Figure 4, at each stage before the optimizer step, a partially reduced global state is received from the previous stage, combined with the current stage local state, and passed on to the next stage. The optimizer step of each stage controlled by the partially reduced state, e.g. skip the update when NAN is spotted the partially reduced gradient norm exceeds the clipping threshold. During the warm-up phase of the next iteration, the fully reduced global state is then propagated back from the last stage to first stage. Upon receiving the global state, each stage performs a validation to decide whether the previous optimizer step is legitimate. If an amendment to the gradient required, rollback will be issued (for more details see Appendix C) and then we redo the optimizer step based on the fully reduced global state.
123
by
propagating from to
Optimizer stop
3
Propagute globally reduced value to each stage
Rollback if validation fails
Figure 4: The post-validation strategy to replace optimizer synchronization.
5 EXPERIMENTS
5.1 SETUP
We base our implementation on the open-source Megatron-LM project (Narayanan et al., 2021) and assess performance using models analogous GPT-3 (Brown al., 2020), as detailed in Table 3. During our experiments, we first conducted specific number of iterations for profiling, collecting empirical measurements for Tr. TB. Tw. and Toomm. After obtaining these values, we fed them into our automatic pipeline scheduling algorithm to determine the optimal schedule. It's worth noting that both the initial and final pipeline stages possess one fewer transformer layer compared to the intermediate stages. This design is to compensate for the extra embedding lookup and loss computations in the initial and final stages so that they won't become bottleneck and cause bubbles to other stages.
Table 3: Models and fixed settings used in experiments
Model
Layers
Attention Heads
Hidden
Sequence Length
Pipelines (GPUs)
Microbatch Size
Number Microbatches
Size
1.5B
22
24
2304
1024
24 / 32
6.2B
30
32
4096
1024
8
24 / 32 64
14.6B
46
40
5120
1024
16
48 / 64 128
28.3B
62
48
6144
1024
32
96 / 128 256
Compared methods:
ZB-1p: Automatically searched schedule with the activation memory limited to pMB, which theoretically has the same peak memory as IFIB.
ZB-2p: Automatically searched schedule with the activation memory limited 2pM B, which is the least amount of memory to empirically achieve close to zero bubble (see Figure 7).
IFIB and IFIB-I: IFIB and interleaved IFIB methods introduced by Harlap et al. (2018) and Narayanan et al. (2021) with implementation from Megatron-LM, For interleaved IFIB, the entire model is divided into sequence of chunks, which cyclically taken by each stage, forming an interleaved pipeline. In our interleaved experiments, we always use the maximum number of chunks ensure least bubble, i.e. each transformer layer serves chunk.

---

Published as a conference paper at ICLR 2024
Our experiments utilize up 32 NVIDIA A100 SXM 80G GPUs distributed across 4 nodes inter- connected by a RoCE RDMA network. The running time of each iteration recorded after several warm-up iterations. Thanks to the reproducibility provided by Megatron-LM implementation, can verify the correctness of ZB--pp and ZB-2p without running models until convergence. We use fixed random seed to initialize the model, record the loss after every iteration for ZB-Ip. ZBB-2p, and IFIB, and then verify that they're re bit-to-bit identical.
5.2
MAIN RESULTS
28-1p
upper bound
28-1p
upper bound
1F1B
28 2p
1F1B
2B 2p
15
4.25
4.00
13
3.75
3.50
11
3.25
24 microbatchhss
32 mecrobatches
64 microbatche
24 macropacches
32 m cr-baachhs
54 micebatches
14.68 Model, 16G
28.38 Model, 32G
2.0
1.1
1.0
0.9
1.6
micrabatches
micropasches
microtatches
Figure 5: Comparison of throughput across different pipeline schedules.
Table 4: Experiment result details
Model
1.5B
6.2B
14.6B
28.3B
Setup
#GPU
8
8
16
32
#Microbatch
24
32
64
24
32
64
48
64
128
96
128
256
Samples
ZB-2p
14.5
14.8
14.9
4.32
4.35
4.39
1.81
1.83
1.85
0.99
1.00
1.00
per GPU
ZB-1p
12.9
13.4
14.2
3.88
4.00
4.20
1.61
1.67
1.76
0.87
0.90
0.96
per second
IFIB
11.8
12.5
13.6
3.50
3.70
4.03
1.40
1.49
1.64
0.76
0.80
0.88
IFIB-1
13.1
13.4
13.9
4.01
4.08
4.19
1.54
1.59
1.66
0.82
0.85
0.90
ZB-2p
59
59
59
70
70
70
51
51
51
74
74
74
Memory
ZB-1p
32
32
42
42
33
33
44
44
44
(GB)
IFIB
30
30
30
39
39
39
32
32
32
43
43
43
IFIB-1
40
40
40
48
48
48
39
39
39
58
58
58
We present the throughput of all methods Figure 5, and leave the additional details for each setup in Table 4. Our experiments demonstrate that ZB-2p consistently outperforms all other methods across various settings. Notably, the throughput of 1F1B, 1FIB-I and ZB-1p show a strong positive correlation with the number of microbatches. In contrast. ZB-2p maintains the efficiency even with fewer microbatches. This is because the bubble rate ZB-2p has almost reached zero (Table 5), and its throughput is already close the upper bound. Here the upper bound is roughly estimated (for more details see Section 5.3).
by multiplying the throughput of 1FIB and I-bubble of IFIB ZB-2p comes at the cost of a higher memory
As mentioned before, the improved efficiency consumption compared to the 1FIB baseline. We also compare ZB-2p with IFIB under the same memory consumption in Appendix F, and the experimental results also show that ZB-2p achieves a higher throughput even with half microbatch size compared 1F1B.

---

Published as a conference paper at ICLR 2024
In contrast, ZB- 1p is designed have peak memory cost similar to the 1FIB baseline. It shows comparable throughput to 1FIB-I the 8 GPUs setups. In multi-node setups where communication bandwidth is more of bottleneck, ZB-Ip clearly outperforms IFIB-I. highlighting its advantage in reducing pipeline bubbles without incurring extra communication cost.
In most of settings set number of microbatches larger than number of stages because they'r re more common use cases of pipeline parallelism. However we conducted experiments listed in Appendix H for m S p cases which shows 20% to 30% improvements with a similar memory consumption.
5.3 EFFICIENCY OF AUTOMATIC SCHEDULING
Table 5: Bubble rates of IFIB, IFIB-I, I, ZB-HI,
H2.
ZB- 1p. ZB-2p under different settings.
Model
#Stage (p)
#Microbatch (m)
IF1B
IF1B-I
ZB-H1
ZB-H2
ZB-Ip
ZB-2p
24
0.2431
0.1055
0.1585
0.1083
0.1585
0.0433
1,5B
32
0.1985
0.0818
0.1242
0.0837
0.1242
0.0039
64
0.1240
0.0443
0.0674
0.0444
0.0674
0.0026
24
0.2347
0.0808
0.1323
0.0698
0.1323
0.0029
6.2B
32
0.1898
0.0628
0.1045
0.0559
0.1045
0.0022
64
0.1091
0.0320
0.0554
0.0294
0.0554
0.0010
48
0,2552
0.1104
0.1397
0.0672
0.1397
0.0066
14.6B
16
64
0.2082
0.0852
0.1088
0.0516
0.1088
0.0054
128
0.1251
0.0445
0.0576
0.0266
0.0576
0.0028
96
0.2646
0.1493
0.1421
0.0641
0.1421
0.0038
28.3B
32
128
0,2168
0.1164
0,1
0.0490
0.1106
0.0029
256
0.1352
0.0624
0.0594
0.0257
0.0594
0.0018
Devies 0
Device
Device
Device 6
Device 7
Devies $
Device
Device
Device
Device
Device 15
Device 0
Device I
Device 2
Device
Device
Device 7
Device 8
Device 9
Device 10
Devise
Device
Device
Time
W
Optimizer Step
Figure 6: A schedule produced by ZB-2p (top) and its profiled execution process (bottom).
We study the efficiency of the schedules generated from our automatic scheduling algorithm. The same setups as our main experiments are used, however, since our purpose is to study the efficiency of the automatic scheduling algorithm, the numbers here based on theoretical caleulations instead of real experiments. To quantify the efficiency of a pipeline schedule, we introduce the concept of bubble rate, which is calculated as (cost m(Tr + TB Tw Tw))/cost The cost here is defined as
8

---

Published as a conference paper at ICLR 2024
the largest execution time of all stages, calculated for each schedule using profiled Tp, TB, Tw and Teomm values. The m(TF + Tp + Tw) is the optimal execution time when all communications are overlapped with computations and hence bubbles in the pipeline.
The bubble rates for different schedules are presented in Table 5, We include the handcrafted sched- ules ZB-H1 and ZB-H2 baselines to the automatically searched schedules. In most of the settings, ZB-2p produces a bubble rate of less than 1%, which is the best among all schedules. In contrast, ZB-H2 consistently performs worse than ZB-2p. This provides strong evidence that our automatic scheduling algorithm adapts better to realistic scenarios by using more accurate estimates of Tr, TB. Tw and Toumm- On the contrary, this improvement is not observed in ZB-lp vs ZB-HI, hy- pothetically because the memory limit becomes the dominate factor. Notably, all of our methods significantly outperform 1FIB.
We also plot ZB-2p and its profiled real execution on 16 GPUs to provide direct visual evidence that it is truly a zero bubble schedule. As shown in Figure 6, the automatically generated ZB-2p schedule has almost no bubble. The profiled execution has slightly more bubbles but retains good overall alignment.
1.58 Model, p
6,28 Model, p
14.68 Model, p 16
28.38 Model, p 32
#hhrrobattches 24
#Mcrrubrcchee 24 #Mcrubucchee MMcrrbacches
#thirobbatches 48 #ttursbatchss #Miccrsbttcches
#Mirobbwcches 95
#Mturrbbttulss
*Microbatches
AMicrobacches
0.05
6.00
2.0pff
3.0p14 1.0p%
2.0
3.0pM, 1.5pM
2,0pMr
3.0pMy 1.0p9f
2.0pM
3.0pM
Mose
Mas
Mu.i
M
Figure 7: The relation between memory limit and bubble rate using our heuristic algorithm.
5.4 MEMORY LIMIT
To better understand the effect of memory limit, we study the relationship of the bubble rate to Minit- We run our heuristic algorithm with series of Misi and plot them in Figure 7. Initially, the bubble rate shows a close-to-linear decreasing trend as we increase the value of Mimit- Theoretically, the curve should plateau around (p-1)(T'++2aa 2TTmmm) TpTe MB. Empirically, we find 2p a good threshold for achieving close to zero bubble rate when Tp R TB and Teomm is relatively small. Beyond the inflection point, although a sufficiently large memory limit does result in a theoretically zero bubble rate, in general the cost outweighs the gain. For more details see Appendix B.
T'
6 CONCLUSION AND DISCUSSION
In this work, we introduced a novel strategy to improve the efficiency of pipeline parallelism by splicting the activation gradient and parameter gradient in backward computation, and we design an automatic pipeline scheduling algorithm that can minimize the pipeline bubble rate under different memory budgets. The schedules produced by this algorithm consistently outperform IFIB and even achieve close to zero bubble rate. Empirically, achieving zero bubble requires approximately twice the activation memory compared 1F1B, which raises concerns about out of memory issues. According to Appendix F, we believe it is worth trading some memory for a zero bubble pipeline schedule in the training of large models. Strategies like ZeRO, tensor parallelism accommodate the increased memory need. Another advantage of zero bubble schedule is that it can achieve optimal efficiency with a smaller number of mierobatches (typically 3p is enough), which means more microbatches can be partitioned over data parallelism dimension. This brings better scalability for the training of large models.
be used to
REFERENCES
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1 2020.

---

Published as a conference paper at ICLR 2024
Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng. Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. {TVM}: An automated (End-to-Endy op- timizing compiler for deep learning. In 13th USENIX Symposium Operating Systems Design and Implementation (OSDI 18), pp- 578-594, 2018.
Shiqing Fan, Yi Rong, Chen Meng, Zongyan Cao, Siyu Wang, Zhen Zheng, Chuan Wu, Guoping Long. Jun Yang, Lixue Xia, et al. Dapple: A pipelined data parallel approach for training large models. In Proceedings of the 26th ACM SIGPIAN Symposium on Principles and Practice of Parallel Programming, pp. 431-445, 2021.
John Forrest and Robin Lougee- Lougee-Heimer. cbe user guide. In Emerging theory, methods, and appli- cations, pp. 257-277. INFORMS, 2005.
Heimer.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint arXiv:1706.000000 70670067077777 2017.
Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. Pipedream: Fast and efficient pipeline parallel dnn training. arXiv preprint arXiv:1806.03321 2018.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hyoukloong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mo- hammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.
Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. Mlir: A compiler infras- tructure for the end of moore's law. arXiv preprint arXiv:2002. 1 2020.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. Pytorch distributed: Experiences on accelerating data parallel training. arXiv preprint arXiv:2006. 15704, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled arXiv:171..05101, 2017.
weight decay regularization. arXiv preprint
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Gareia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprini arXiv:1710.0377. 2017.
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al. Efficient large-scale language model training on gpu clusters using megatron-lm. In Proceed- ings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15, 2021.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318. Pmlr, 2013.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. SC20: International Conference for High Perfor- mance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Jared Roesch, Steven Lyubomirsky, Logan Weber, Josh Pollock, Marisa Kirisame, Tianqi Chen, and Zachary Tatlock. Relay: A new ir for machine learning frameworks. In Proceedings of the 2nd ACM SIGPLAN international workshop on machine learning and programming languages, pp- 58 68, 2018.
10

---

Published as a conference paper at ICLR 2024
Amit Sabne. Xla: Compiling machine learning for peak performance, 2020.
Philippe Tillet, Hsiang- Tsung Kung, and David Cox. Triton: intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Bowen Yang, Jian Zhang, Jonathan Li, Christopher Ré, Christopher Aberger, and Christopher De Sa. Pipemare: Asynchronous pipeline parallel dnn training. Proceedings of Machine Learning and Systems, 3:269-2966 2021.
Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing. al. Alpa: Automating inter-and {Intra- Operator} parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pp. 559-578, 2022.

---

Published as a conference paper at ICLR 2024
A OVERLAP COMMUNICATION IN DATA PARALLELISM
When data parallelism taken into consideration, an all-reduce communication launched to col- lect gradients before optimizer step. Generally, such communication is poorly overlapped with com- putation pass, resulting in a latency especially when the communication bandwidth is limited, As in Figure 3, usually number of W passes scheduled at the tail of iteration. For each W pass, it consists of several independent computations calculating gradients for different parameters. As in Figure 8, We reorder all of these computations to cluster those calculating the gradients for the same parameter, thus achieving the optimal overlapping between computation and communication.
AR
AR
AR
AR
(e)
The schedule grouped by W
schadule grouped by
Figure 8: Comparison between the original schedule grouped by W with poor overlapping (top) and the reordered schedule grouped by parameters with optimal overlapping (bottom). The number i represents the computation belongs to i-th W, and different colors represent computations for different paramters.
B THE MEMORY LIMIT FOR AUTOMATIC SCHEDULING ALGORITHM
The relation between memory limit and bubble rate is highly affected by the bubbles preceding the first B in the initial stage. For the first microbatch, the forward pass needs to go through from the initial stage to final stage, and the backward pass reverses this process until it eventually goes back the initial stage. The total time for the first mierobatch from start to complete takes least p(T'F + TB) + 2(p 1)Tcmmm and it can not be squeezed due to the dependency chains, We denote the number F passes as k(2 1) and the bubble size { 0), preceding the first pass in the initial stage. Then we have:
Mumit 2 KM
(1)
82 P(TF+TBb) TB) 2(p 1)Tcomm KTF TB (p 1)(T'B 2Tommm + (p kYTF
The lower bound of Miimnt is in proportion to k (see Formula 1), and 3 is inversely proportional to (see Formula 2). When increasing and keeping k (p-1)(Tn 2T +T +pff B decreases lin-
TF
early, meanwhile the lower bound of Mimi increases linearly. When k: = ()--))(T++) + 27 + reaches its minimum value without delaying and its value is less than Tr, with peak activation memory at least (p-1))(T ( +:7p M- Beyond this point, further reducing pipeline bubbles to zero is not easy. This is because there is small bubble less than Tç each stage (see Figure 6), and scheduling another F will delay the starting time of B thus causing more requirements previous stages. Theoretically, another p F passes are required in the initial stage to fully elim- inate bubbles preceding the first B for all stages (see Figure 9), which also means total activation memory usage at least (P-1))Tee-- +2T (2p 1)Te
TF
F in
[Mp.
Tv
Devize
Deoter
Device
Dester
Device
Figure 9: Zero bubble schedule 1.5B model with 32 microbatches.
12

---

Published as a conference paper at ICLR 2024
C IN-PLACE OPTIMIZER ROLLBACK
When need rollback optimizer step, typical method is to store historic version of param- eters and optimizer states, and revert to this historic version when needed. However, this method is memory inefficient and lots of copy operations formance. For most optimizers, we notice that the step function is arithmetically reversible. Under this observation, propose novel technique perform in-place optimizer rollback, which avoids allocating extra memory and requires extra computations only when the rollback is performed. As in Algorithm 1, we show how rollback the step function for Adam W optimizer (Loshchilov & Hutter, 2017).
needed, which definitely hurts the training per-
Algorithm In- in-place rollback for AdamW
Optimizer States:
7(rr), 31, 8y(btas) e (epsilon), A(weight decay), 3: m (first moment), v second moment), 0 (parameters),
2:
4: t(time stamp)-
5: function STEP(G)
b In-place step
6:
=+1 m = f.m + (1 31)g
7:
8:
32 (1 82)g
9: m' = mmm11 (1-31) B))
o' =v((1 88)
0 -0-7h0 -A0 7m' /LvV'/VV +
12: end function
function ROLLBACK(g)
In-place rollback
14: m' = m/ 3))
o' ===(1 B5)
16:
= (0 + ym /(v1' + (
7A)
17:
(m (1 31)g)/80 /3
18:
(v (1 32)g²) //2
19:
== -
20:
end function
D PROFILED TIME IN EXPERIMENTS
In experiments, record the profiled time of Tr, TB, Tw, and Teomm in ZB-2p across different settings. These values are then used to calculate bubble rates for all the methods considered in Section 5.3 and 5.4. These values be found in Table 6.
Table 6: Profiled time of TF. TB, Tw. and Teomm-
Model
#Stage (p)
#Microbatch (m))
Tr
Tn
Tw
T Toomm
24
18.522
18,086
9.337
0.601
1.5B
8
32
18.513
18.086
9.331
0.626
64
18.546
18.097
9.321
0.762
24
29.718
29.444
19,927
0.527
6.2B
8
32
29.802
29,428
19.530
0.577
64
29.935
29.621
19.388
0.535
48
11.347
11.248
8.132
0.377
14.6B
16
64
11.307
11.254
8.101
0.379
128
11.325
11.308
8.109
0.378
96
10.419
10.207
7.715
0,408
28.3B
32
128
10.408
10.204
7.703
0.408
256
10.402
10.248
7.698
0.460
13

---

Published as a conference paper at ICLR 2024
E ABLATION STUDY ON OPTIMIZER POST- VALIDAALI STRATEGY
In this section, we provide an ablation study on the effectiveness of the optimizer post-validation strategy. The study compares the throughput of ZB-2p under two conditions: with post-validation and with all-reduce synchronization. According to the experimental results in Table 7, the synchro- nized version of ZB-2p demonstrates performance decrease of approximately 8% compared to 7B-2p 2p with optimizer post-validation.
Table 7: Throughput (Samples per GPU per second) comparison between ZB-2p and synchronized ZB-2p
Model
#Stage (p)
#Microbatch (m)
Post-validation
All-reduce synchronization
1,5B
24
14.5
13.11
6.2B
8
24
4.32
4.00
14.6B
48
1.81
1.68
28.3B
32
96
0.99
0.91
F COMPARE ZB-2p WITH 1F1B UNDER THE SAME MEMORY CONSUMPTION
Under the same memory consumption, we double the size of each microbatch for IFIB and ZB-Ip and compare their throughput with ZB-2 in Table 8. The experimental results show that ZB-2p also holds a better performance even with a half microbatch size compared to IFIB. Empirically, a larger batch size increases the utilization rate of GPU and thus improves the efficiency. However, it less of a concern for large models because the hidden dimension is large enough to saturate device utilization. Based on this consideration and our experimental results, we believe ZB-2p more preferred than increasing the batch size for IFIB. In some experiments where the device utilization less saturated and m/p is relatively large, ZB-1p with a doubled microbatch size may slightly outperform than ZB-2p.
G ILP FORMULATION
Any pass in the pipeline can be uniquely indexed by (i,j,c) c), where i E (1,2,.... p) indexes the stage, E {1. 2, ......,,, m indexes the microbatch, and c E {F, B, W} denotes the specific pass of the microbatch. We define the variable (i.j.e) as the time cost and E(s.j.c) as the ending time of pass. We introduce A A.M e) to denote the memory increment incurred by the pass (i.j, c). For example, A.M....F) = MB because the forward pass leads to net increase of MB of activation stored for the backward pass. AMG...)
Mw MB which removes the memory stored for B while adding those required by W, and AM(,,.,) = -Mw. Finally, the variable that we want to search is the ordering of the passes in the schedule, for which we introduce the variable O (s...c)++(1.) which is an indicator whether the pass index by (i,j, c) is scheduled before (1.j', c.).
E {0, 1).
min max E(G.m. E(G.1,)) + T(....)) O.E
s.t.
Eúj 2 E(i-1557) + Toomm + T(i.j.)
(4)
F(i.,B) 2 F(1415.8)) + Teumm + -T(s..B)
(5)
F(s.,) 2 Euss'.et le) + ((ss..) -O(...e)
(6)
Mimit 2 AMGg.e + Σ AMMGs.. (15.))))
j.c
Overall, the optimization target (3) is to minimize the time spent by the longest stage. Constraints (4) and (5) add the sequential dependency requirements the F and B passes of the same microbatch in adjacent stages. Additionally, (6) adds the dependency constraint imposed by our decision of the scheduling order. Finally, (7) limits the peak activations memory to be below Minin
14

---

Published as a conference paper at ICLR 2024
Table 8: Comparison between IFIB, ZB-Ip and ZB-2pp under the same memory consumption.
Model
p
m
Samples per GPU per second
Memory(GBB)
Schedule
12
12.0
57
IFIB
24
12
13.0
61
ZB-1p
6
14.5
59
ZB-2p
12
12.6
57
IFTB
1.5B
8
32
12
13.6
61
ZB-1p
6
14.8
59
ZB-2p
12
13.8
57
IFIB
64
12
14.4
61
ZB-1p
6
14.9
59
ZB-2p
6
3.56
66
IFIB
24
6
3.95
71
ZB-Ip
3
4.32
70
ZB-2p
6
3.76
66
IFTB
6.2B
8
32
6
4.05
71
ZB-1p
3
4.35
70
ZB-2p
6
4.09
66
IFIB
64
6
4.24
71
ZB-1p
3
4.39
70
ZB-2p
2
1.53
50
IFIB
48
2
1.73
51
ZB-lp
1
1.81
51
ZB-2p
2
1.62
50
IFTB
14.6B
16
32
2
1.79
51
ZB-1p
1.83
51
ZB-2p
2
1.78
30
IFIB
128
2
1.89
51
ZB-1p
1
1.85
51
ZB-2p
2
0.81
72
IFIB
96
2
0.93
74
ZB-1p
1
0.99
74
ZB-2p
2
0.85
72
IFIB
28.3B
32
32
2
0.96
74
ZB-1p
1.00
74
ZB-2p
2
0.94
72
IFIB
256
2
1.02
74
ZB-Ip
1
1.00
74
ZB-2p
H COMPARE ZB METHODS WITH 1F1B ON SMALL NUMBER OF MICROBATCHES
By nature PP when the number of microbatches m is less then number of stages there'll be large bubble rate. However zerobubble methods still boost performance under these settings by approximately 20% to 30%. In rough analysis ignoring communication and assuming < and Tw < TB, an IFIB iteration takes (m p 1) (IF TB + Tw) to complete, while ZB iteration takes (m + p 1) * (Т +TB) Tw. The experiment result is shown in Table 9. Noticeably when < p ZB- 1p and ZB-2p are essentially the same and consumes similar memory as IFIB.
15

---

Published as a conference paper at ICLR 2024
Table 9: Comparison between IFIB and ZB-2p on small number of microbatches.
Model
p
Samples per GPU per second
Memory(GB)
Schedule
3.56
11
IFIB
2
6
4.25
12
ZB-2p
5.74
18
1FTB
1.5B
8
4
6
6.92
19
ZB-2p
8.26
29
IFIB
8
6
9.90
34
ZB-2p
3
1.04
21
IFIB
2
3
1.33
21
ZB-2p
3
1.69
28
IFIB
6.2B
8
4
3
2.16
29
ZB-2p
3
2.45
39
IFTB
8
3
3.07
44
ZB-2p
1
0.39
19
IFIB
4
0,52
20
ZB-2p
0.65
24
IFIB
14.6B
16
8
0.85
24
7B-2p
0.95
32
IFIB
16
1.25
33
ZB-2p
16


Reviews:
[{"review_id": "ORu62cRJu6", "reviewer": "Reviewer_hNFq", "summary": "The paper presents a novel approach, called ZeroBubble, for eliminating bubbles in pipeline parallelism in order to training efficiency. There are two key aspects to the proposal: (1) splitting the backward pass into the B and W components to increase scheduling flexibility, and (2) eliminating synchronization of optimizer step with a rollback mechanism to preserve semantics. The paper also presents various analysis to illustrate the memory consumption effects of different ZeroBubble strategies. Overall, the evaluation results show that ZeroBubble provides up 30% throughput improvement over 1F1B schedules for Megatron GPT models.", "questions": "See weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "I think splitting the backward pass into B (activation gradient computation) and W (weight gradient computation) as a way of improving scheduling flexibility is a nice touch of creativity. \n\nThe paper includes ample analytical details that foster intuition and understanding of the relevant memory and TFLOPs consideration of pipeline parallelism schedules. \n\nThe authors did a really great job of writing and organizing the paper.", "weaknesses": "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.  Moreover, I think there are few design alternatives that the paper failed to explore or discuss. \n1. Evaluating the performance of synchronous optimizer step as way of understanding the trade-off between the simplicity and throughput. \n2. Scheduling W before B, which is less memory-efficient but allows staggering the optimizer step of the stages while preserving the synchronous semantics. \n\nAnother 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OVttDPdQNE", "reviewer": "Reviewer_Uec3", "summary": "The manuscript titled \"ZERO BUBBLE PIPELINE PARALLELISM\" presents a novel strategy for achieving pipeline parallelism with the primary goal of eliminating pipeline bubbles. This approach includes the development of an automatic scheduling algorithm, which is then compared with established methods like 1F1B and interleaved 1F1B. The key contributions of the paper are as follows:\n\n1. It introduces a pipeline parallelism strategy designed to eliminate pipeline bubbles, assuming that forward, backward, and weight gradient calculations all take the same amount of time.\n\n2. It presents an automatic scheduling algorithm that consistently outperforms traditional methods, with the highlight being the achievement of a zero bubble rate, indicating optimal computational resource utilization. However, this comes at the cost of nearly doubling memory usage.\n\n3. The paper delves into the relationship between memory constraints and bubble rates, providing insights into the trade-offs between memory requirements and scheduling efficiency.", "questions": "For Figure 5, sometimes ZB-1p is better than 1F1B-I, and sometime not. Can you explain the reason in the paper?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper introduces a novel strategy for pipeline parallelism aimed at completely eliminating pipeline bubbles. Leveraging this strategy, an automatic scheduling algorithm is developed and benchmarked against existing approaches, such as 1F1B and interleaved 1F1B. \n\noriginality: fair. The paper separated the backward loss and weight gradient calculation, and achieved zero bubble pipeline parallelism under the assumption that forward, backward loss, and weight gradient calculation time are identical.\nquality: good. The paper has theoretical calculation and also experimental data to support the zero bubble pipeline parallelism algorithm proposed in this paper.\nclarity: good. \nsignificance: medium. 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.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "X16nB2rALc", "reviewer": "Reviewer_b5Up", "summary": "Pipeline parallelism is a widely used technique in distributed training. However, the efficiency of pipeline parallelism suffers from pipeline bubbles. In view of this, this paper designs a new strategy to reduce the bubble rates. The idea behind the new pipeline strategy is to split the backward computation into two parts, which could reduce the bubble in the 1F1B strategy.", "questions": "The problem studied in this paper is significant in distributed training and the authors give an excellent solution to improve the existing pipeline strategy. \n\nHowever, the main issue of this paper is that the solution mentioned in this paper is likely to increase the memory usage of pipeline parallelism training. Moreover, the experimental results introduce bubble rate which is not important in the pipeline parallelism training. It is known to all that the efficiency of pipeline parallelism suffers from pipeline bubbles. But the bubble rate is not important. We only care about the throughput and memory usage of the pipeline parallelism strategy. Since the pipeline strategy developed in this paper will increase the memory usage of pipeline parallelism training, large memory usage is likely to decrease the size of the micro-batch. A small micro batch will decrease the efficiency of the pipeline parallelism strategy. The experimental results show the developed algorithm has limited improvement.\n\nThe writing of this paper could be improved. For example, we can use different colors to represent B in Figure 3 and Backward in Figure 2. Moreover, the title of this paper is too large. I do not think the method in this paper is zero bubble.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The problem studied in this paper is very important. With the increase of the model parameters, pipeline parallelism is a popular strategy for training large models. However, the bubble in the 1F1B strategy affects the efficiency of the pipeline strategy.\n2. The idea behind this paper is very easy. It is a good signal, I think the paper develops an easy solution with great performance.\n3. This paper contains a mathematical analysis.", "weaknesses": "1. The authors miss an important point in the experiment section.\n2. The writing of this paper can be slightly improved.\n3. The performance of the method developed in this paper is limited.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9j2AzBWAmS", "reviewer": "Reviewer_xbKE", "summary": "* The author proposed a new schedule in Pipeline Parallelism (PP) to reduce pipeline bubbles to close to zero\n* The key idea is instead of the 1F1B schedule of earlier work, splitting B to two stages (aka activation gradient and weight gradient) and a schedule like 1F1B1W can further reduce bubbles of 1F1B\n* The intuition behind it is only Forward calcluation and activation gradient calculation have a pipeline/stage dependency, while weight gradient calculation are not, so activation gradient calcluation should be eagerly scheduled while weight gradient calculation can be rescheduled to fill/balance pipeline bubbles\n* The authors have proposed two schedule algorithms to adapt for real-world workloads when F/B/W can have different runtime.", "questions": "1. 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?\n\n2. what's a \"local search\" in 3.1?\n\n3. what's the runtime of ILP, how much does it improve over HEURISTIC? can you do an ablation with HEURISTIC algo only?\n\n4. My understad is section 4 BYPASSING OPTIMIZER SYNCHRONIZATIONS only works for INF/NAN, how does it work for grad-norm when each stage need the global reduced norm? If the optimizer depends on partially reduced state how can it provide same numerical result of baselines like 1F1B?\n\n5. 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?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The author made a key observation that the backprop stage in PP can be split into two parts and scheduled in a finer granularity, i.e., `any pipeline (or producer/consumer) independent workload can be rescheduled to balance/fill pipeline bubbles`, which is not only benifitical to train NN with PP, but should hold generally true for any PP workloads\n\n* The author didn't limit themselves to just a simple TF = TB  = TW assumption but also proposed two optimization implementations to optimize the schedule for potentially unbalanced TF/TB/TW\n\n* The authors have conducted representative experiements and showed that their methods can outperform earlier methods by up to 20% and more importantly close to the upperbound (zero bubble) of PP\n\n* The presentation/feagures are easy to follow and illustrative", "weaknesses": "* While I agree the authors have discovered a valuable new angle optimizing PP and indeed bring the bubble rate close to zero, I think it is still a bit overselling to claim the method achives zero bubble PP, given:\n  1. The method still require a large enough number of microbatches, e.g., m >= 2p to achive an amortized zero bubble PP, how does it perform under m < p is rarely discussed. Experiments (at least a roofline analysis) with m < p can make the methods more sound\n  2. There are other methods that can achieve zero bubble PP like PipeDream and PipeMare. These are more strictly flushless (thus bubble-free) methods though they have altered the math of backprop and poses potential accuracy loss or incurs more weight memory. So it is a bit overclaim that this work is the first to achieve pippeline bubbles. Discussions on these previous related works are also encouraged.\n  3. Therefore \"near zero bubble\" is probably more accurate IMHO, emphasis on \"zero bubble\" actually overshadows the solid intuition mentioned as 1st bullet point in Strengths\n\n* Even though authors proposed two algos to schedule unbalanced TF/TB/TW, the experimetns however are not designed to show how their advantages, e.g., 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 sifinicantly different T) could be helpful to evaluate the efficacy of these algos", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)