Pipeline Parallelism with Controllable Memory
Authors: Penghui Qi
Xinyi Wan
Nyamdavaa Amar
Min Lin
Venue: NeurIPS 2024 poster (2024)
Keywords: LLM Training
Pipeline Parallelism
Balanced Memory
Zero Bubble
Building Block
Lifespan
Primary area: infrastructure
TL;DR: A unified framework for Pipeline Parallelism with balanced and controllable memory

Abstract: Pipeline parallelism has been widely explored, but most existing schedules lack a systematic methodology. In this paper, we propose a framework to decompose pipeline schedules as repeating a building block, and show that the lifespan of the building block decides the peak activation memory of the pipeline schedule. Guided by the observations, we find that almost all existing pipeline schedules, to the best of our knowledge, are memory inefficient. To address this, we introduce a family of memory efficient building blocks with controllable activation memory, which can reduce the peak activation memory to 1/2 of 1F1B without sacrificing efficiency, and even to 1/3 with comparable throughput. We can also achieve almost zero pipeline bubbles while maintaining the same activation memory as 1F1B. Our evaluations demonstrate that in pure pipeline parallelism settings, our methods outperform 1F1B by from 7\% to 55\% in terms of throughput. When employing a grid search over hybrid parallelism hyperparameters in practical scenarios, our methods demonstrate a 16\% throughput improvement over the 1F1B baseline for large language models. The implementation is open-sourced at https://github.com/sail-sg/zero-bubble-pipeline-parallelism.

Pipeline Parallelism with Controllable Memory
Penghui Qi " Xinyi " Nyamdavaa Amart?, "eart?,,inn Min Lin' 'Sea Lab ?National University of Singapore {qiph, wanxy, linmin)n@sea. com amara@u. nus edu
Abstract
Pipeline parallelism has been widely explored, but most existing schedules lack a systematic methodology. In this paper, we propose a framework to decompose pipeline schedules as repeating building block, and show that the lifespan of the building block decides the peak activation memory of the pipeline schedule. find that almost all existing pipeline schedules,
Guided by the observations, the best of our knowledge, are memory inefficient. To address this, we introduce a family of memory efficient building blocks with controllable activation memory, which can reduce the peak activation memory to 1/2 of IFIB without sacrificing efficiency, and even 1/3 with comparable throughput, We can also achieve almost zero pipeline bubbles while maintaining the same activation memory as IFIB. Our evaluations demonstrate that in pure pipeline parallelism settings, our methods outperform IFI IB by from 7% to 55% in terms of throughput. When employing a grid search over hybrid parallelism hyperparameters practical scenarios, our 16% throughput improvement over the IFIB baseline for large language models. The implementation open-sourced at this url
methods demonstrate a
Introduction
Distributed model training has attracted lot of attention in recent years, especially after the boom of large language models [Brown et al., 2020]. As the model size becomes larger and larger, data parallelism (DP) [Goyal et al., 2017] no longer capable hold all the parameters Under this background, model parallelism [Harlap et al., 2018 8, Huang et al., 2019, Shoeybi et al., 2019, Zheng et al., 2022] proposed to partition parameters into a set of devices to address the memory constraint. Tensor parallelism (TP) [Shoeybi et al., 2019] is commonly used model parallel strategy, which partitions weight parameters into several devices and performs matrix multiplication separately. A well-known shortcoming of TP is that, it requires a lot of communication volume, which makes it inefficient when bandwidth becomes the bottleneck Narayanan et al. [2021]. In such situations, pipeline parallelism [Harlap et al., 2018, Huang et al., 2019], which is another model parallel strategy, shows its advantage in low communication cost. The core idea of pipeline parallelism is split the entire model into several stages, which can be processed by several devices in streaming way. In typical large-scale training scenarios such as Narayanan et al. [2021], TP is generally used within compute node, and PP is used to scale up across nodes.
single device,
Although PP has been widely adopted and developed, suffers from two prominent disadvantages: pipeline bubbles and large activation memory. To eliminate pipeline bubbles, one line of work focuses on asynchronous PP [Gaunt et al., 2017, Yang et al., 2021], which is theoretically bubble free. However, it sacrifices the exact optimization semantics and may result in lower convergence performance [Lian et al., 2018, Tang et al., 2020J. A parallel line of works revolve around synchronous PP, focusing reducing pipeline bubbles and/or activation memory, GPipe [Huang et al., 2019] is
'Equal Contributors.
Work was done during an internship at Sea Al Lab.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

an early work to reduce the bubble rate by increasing the number of microbatches, at the cost of more activation memory. IFIB [Fan et al., 2021] avoids the activation memory growth with respect the number of microbatches by staggering forward pass and backward pass, keeping the same bubble rate with GPipe. Another notable work GEMS [Jain al., 2020], which stores activation memory of only one forward pass by scheduling microbatches one after another among two model replicas, thus with significantly large bubble rate. Chimera [Li and Hoefler, 2021] extends the ideas GEMS by combining two pipelines in different directions together, which reduces pipeline bubbles when the number microbatches is small, but with doubled parameter memory. Hanayo [Liu et al., 2023] is introduced to attain the same scheduling efficiency with Chimera without replicated models, but still suffering from scaling to more microbatches. Although its wave-like scheme is kind of similar to our V-shape building blocks, it is not motivated for memory balance, thus resulting in totally different pipeline schedules. In Megatron-LM Naraaannan et al., 2021], an interleaved strategy is proposed to further reduce the bubble rate, the cost of more communication cost and a portion of extra activation memory. BPipe [Kim et al., 2023] focuses on reducing the activation memory of 1F1B from another perspective, transferring activations across devices based on the memory imbalance of IFIB, However, it introduces lot of extra communication and increases the complexity of the system, which makes inefficient especially in settings with limited bandwidth. Zero Bubble [Qi al., 2023] splits the backward into activation gradient computation and weight gradient computation, which can either reduce the pipeline bubbles without changing the maximum peak activation memory, or achieve zero bubble at the cost of doubled activation memory compared to IFIB.
In this paper, first demonstrate all existing pipelines can be seen as repeating basic building block in time. We then identify a direct link between the activation memory and the lifespan of each building block, which reveals the core insight of this paper: lifespan decides the activation memory. Based on this insight, we present a family of novel and memory-efficient building blocks and their pipelines. Compared
IFIB, we reduce the activation memory to 1/2 asymptotically with even higher throughput, and to 1/3 asymptotically with comparable throughput. We can also achieve zero bubble under the same activation memory with IFIB. Notably, our strategy almost pure gain to the existing methods, only at the cost of doubled communication cost between pipeline stages, which relatively small and can be neglected.
2 How to Build a Pipeline
We propose a four-step framework to design pipeline schedules.
Building Block: It starts by laying out the passes for single microbatch, which we call a building block. For example, the building block of IFIB is made of a sequence of forward passes followed by backward passes in the reverse order. We highlight the building block of IFIB in color in Figure la.
More microbatches are then introduced. The building blocks are repeated and woven together form pipeline. In Figure 1 (top), the repeating building blocks are shown in different shades of gray. Notably, legit building blocks are required to repeat without a collision, namely, the passes from two building blocks should not overlap with each other.
Repeating:
Squeezing: Depending on the building block, there may be redundant bubbles in the pipeline, which can be simply removed by squeezing without changing the order of the passes. For example, Figure lb shows a case where squeezing produces a more efficient pipeline.
Litespon
Lifryynn
Inserval
Squesse
(a) IFIB
(b) Eager IFIB
Figure 1: A pipeline can be built by repeating a building block, and then squeezing redundant bubbles.

---

Reordering (optional): We can reorder the passes in the warm-up and cool-down phase to further improve the computation throughput. Intuitively, the peak of memory happens the stable phase of the pipeline, while in the warm-up and cool-down phases the RAM is under utilized, leaving some space for improving the computation throughput without changing peak memory. We leave the details in Appendix C.
Most of existing pipeline schedules can be explained under this framework. Besides the IFIB and eager IFIB shown in Figure I, we show the interleaved IFIB IShoeybi et al., 2019], ZB-HI [Qi et al., 2023] and a series of well-known pipelines in more extensive gallery (see Appendix I).
2.1 Building Blocks
The computation and memory efficiency various pipelines can be attributed to their building blocks. The diversity of the building blocks primarily comes from three factors, model partitioning, device placement, and offsets between passes. We follow the idea and notations in zero bubble PP [Qi et al., 2023], using F to denote forward pass, B to denote "backward for the activations", and W to denote "backward for the weights". Note that such finer granularity can be generalized to previous methods like IFIB, by always grouping B and W together.
Model partitioning deals with how the model is divided into pipeline stages. The most common pattern is to equally divide the model to match the number of devices. A prominent example is the IFIB schedule (Figure la). This is extended in interleaved IFIB where the number stages can be an integer multiple of the number of devices.
Device placement another key factor in the design of building blocks. While conventionally each pipeline stage is sequentially placed on a different device, it is not uncommon to place multiple stages on the same device like in interleaved 1FIB (Figure 18h). Another example of unconventional device placement is Chimera, where two pipelines are placed in reversed device order.
Last but not least, the offsets between F,B.W W. passes play major role in the computation and memory efficiencies of the pipeline. By simply enlarging the offsets between subsequent passes in the building block of 1F1B, obtain the eager 1F1B [Zhuang al., 2023] (Figure 1b) where more F passes are eagerly scheduled, resulting in higher memory consumption (but better communication overlapping). GPipe can be seen adding large offset between the last F and the first B in the IFIB building block. One more example on the effect of the offset is the comparison of ZB-HI (Figure 18c) and ZB-H2 (Figure 18d) schedules, schedules like 7B-H2. In this work, we assume that every F, B or W pass takes equally one unit of computation time, and only consider integer unit of offsets. Although this may limit the number of feasible building blocks, it greatly improves the simplicity of analysis.
that properly chosen offsets result in zero bubble
2.2 Calculating the Peak Memory
Not every pipeline is born equal, researchers are constantly looking for pipelines that are more efficient in computation and/or memory. While efficient pipelines could be discovered by enumerating every possible building block, it is nonetheless prohibitively expensive. We discover that the peak memory consumption of a pipeline can be calculated from its building block via a simple formula. This enables us to design pipelines with a controllable peak memory.
Two quantities crucial for the calculation of peak memory, the lifespan of stage, and the repeating interval of the building blocks, both of which are illustrated in Figure 1. The lifespan of a stage is detined as the amount of time between the starting of the pass and the ending of or pass. A piece of activation memory is allocated at the starting of F, and retained in the RAM throughout the lifespan until it consumed by both B and W. The peak memory consumption can be calculated by finding the maximum number of microbatches whose lifespans overlap with that of every other microbatch. Using denote lifespan, T to denote the repeating interval and m the size of activation memory for single microbatch, we have the relation.
1
peak memory < - 1m
When there are multiple stages on one device, e.g. interleaved IFIB, their contributions to the peak memory are independent, using St to denote all the stages allocated to device i, we sum the
3

---

contributions from every stage.
peak memory of device i 1
Σ sE8;
m
(I)
Another key insight is that the repeating interval T is readily determined from the building block. In an efficient pipeline, T should be equal to the number units of computation in each stage of the building block. Any T larger than that would cause pipeline bubbles in the stable phase, and T smaller than that would lead to collisions. A subtle exception is the interleaved IFIB whose repeating interval is not uniform. We leave the discussion to Appendix G.
2.3 Repeating without Collision
One constraint keep in mind when designing the building blocks is that legit building block is required to repeat without any collision. It may seem unintuitive how to design building blocks with this constraint. In practice, we design the building block first and perform post-hoc verification. Another useful observation is that legit building block usually produces a stable phase in the middle of the pipeline, which contains repeating x T rectangle, where d is the number of devices and T is the repeating interval. This offers an alternative to constrain the building blocks. We can start by ordering passes within this rectangle and convert it back to a building block.
3 Memory Efficient Building Blocks
With the above framework, we can conveniently analyze the memory consumption pattern of existing pipelines. To our knowledge, all existing pipelines are memory inefficient due two primary reasons: redundant dependency chain, and imbalanced memory usage. Before Zero Bubble [Qi et al., 2023], the backward often regarded as a single pass, resulting unnecessarily longer lifespan thus more memory footprint. In this paper, we leverage the backward splitting strategy to remove these redundant lifespan. The imbalanced memory comes from the innate heterogeneity of the lifespans across stages. From Figure la, we can easily see that the lifespan of the stages differs greatly from each other, with the first stage having the longest lifespan. Consequently, causes a memory bottleneck on the first device and under utilization of memory on all other devices. To resolve this problem, we introduce a family of novel building blocks, which we refer to as V-Shape building blocks. The core insight comes from Equation which says that the peak memory depends on the sum of the lifespans. Therefore, when we place multiple stages on the same device, we should always collocate stages of long lifespans with those of short lifespans. When the total sum of lifespans is fixed, balanced placement always means higher memory efficiency. This can be demonstrated by Figure 2, the parallel schedule (used in interleaved IFIB) is imbalanced and has memory bottleneck proportional to 1 + 4,, while in the V-Shape schedule is l + 16-
The V-Shape schedule requests the model to be partitioned into stages twice the number of devices and the device placement of the second half of stages to be in reverse order as the first half. As the offsets directly determine the lifespan of each stage and therefore the peak memory by Equation 1, we can then further control the offsets between passes to generate building blocks with diverse memory.
15
12
4
13
1/
Parallel
V-Shape
Figure 2: The V-Shape building block ensures balanced peak memory across all devices, whereas the parallel building block has a memory bottleneck in the first device.
3.1 Controllable Balanced Memory
We assume the model uniformly partitioned, namely, both the computation and memory of each stage are identical. For single microbatch, we denote the activation memory of each stage m, and

---

R
2d-1
2d-1
(a) IFIB
(b) V-Half
4d-3
2d-1
4d-3 2d-1
(c) V-Min
(d) V-ZB
Figure 3: V-Shape building blocks with 4 devices (d 4), where white text colors represent the first half of model stages and black text colors represent the second half. F, and W represent the forward, backward (for activation gradients) and backward for weight gradients, respectively.
5
7
6
4
3
6
5
3
4
6
6
8
(a) IFIB
1.2
3
(b) V-Min
2
8 8
8
8
(c) V-Half
2 413345550
(d) V-ZB
Figure 4: V-Shape schedules compared to IFIB, under the setting of 4 devices and 8 microbatches. The stable phases adhere the pattern their building blocks.
the total activation memory of the entire model as M. Note that M = 2dm, where d is the number of devices. To make simple and tractable, we use uniform offsets within each half of F and B passes to control the peak memory. Specifically, we apply the same offset 56 between two adjacent passes within the first d stages (e.g., 5%
Figure 3c and 8% 4 in Figure 3d). Similar constraints are applied to the other half of the F passes and both halves of the passes, denoted as 67. 60. 81. respectively. To guarantee balanced peak memory across devices, we add another two constraints, 59. 0%1 8 and 8). 0% 8¹, where we use notations 8 and
Figure 3b, 82
5

---

5' for simplicity. For example, in Figure 3d, we set 8 = 4 and 8' = 2. Note that we only control the offsets across different devices. For those adjacent passes within the same device (e.g.. F and of the last stage, two F and two B in the last device), we use brute force to find optimal solutions, ensuring their offsets are small (less than the repeating interval). Note that W can always be placed greedily after settling all F' and B passes, so we don't t need to search their offsets during brute force. According to Equation 1, we can analyze the asymptotic peak memory with respect to d,
peak memory of device 2d(6" + 6²) + 0(1) m
8 +
M
6
6
By ignoring the small constant, directly control the peak memory by the value of 5 and 8*.
Table 1:
Small constant values ignored for bubbles and peak memory of V-Min and V-Half, For 1FIB, /8' are redefined as the offsets between adjacent forward/backwart passes. M represents the total activation memory of the entire model, and d the number of devices.
Building Block
8"
8'
Peak Memory
Bubbles
1FIB
2
M
= 6d
V-Min
R
M/3
~ 4d
V-Half
M/2
3d
V-ZB
4
2
M
0
3.2 V-Shape Pipeline Schedules
By varying the values of 8 and 8 we come up with 3 novel V-Shape building blocks (Figure 3). and present their final schedules based on our framework in Figure 4. The building block of V-Min (Figure 3c) has the minimum offsets, namely 5 8² 1, thus the minimum memory consumption. With 5 4 and 5 2 as in Figure 3d, V-ZB eliminates the bubble to almost 0 (Figure 4d), pushing to extreme throughput. The building block of V-Half (Figure 3b), which uses 8 and 8 1, sits between the two extremes and consumes about half of the activation memory required by IFIB. Although both V-Min and V-Half have lower memory footprint than IFIB, V-Min contains about 2/3 and V-Half contains about 1/2 of FFIB's bubbles, assuming F, B, W have equal run time. We show the comparison between our proposed V-Shape schedules and IFIB in Table 1. Notably, the exact peak memory d+21 for V-Min, and [ [ M for V-Half. To avoid collisions the building blocks of V-Min and V-Half, the offsets (within the same device) are slightly different for different values of d. The details are Appendix F.
3.3 Repeating Bubbles in V-Min
In real-world scenarios where F. B and W have different run times, V-Min suffers from a repeating bubble. As shown in Figure 5, there exists bubbles for every repeating interval T. Consequently, the bubble grows as the number of microbatches increases. Although V-Half may encounter the same issue (when the times of F, B and W differ significantly), generates patterns that tessellate well in most empirical cases due to its loose dependencies. As illustrated in Figure 5b, the throughput of V-Half is robust to the variation of run times. Additionally, the bubbles of V-ZB will never grow when increasing the number of microbatches. We leave the related discussions in Appendix E.
12345
LI 1
44.5738
A
P136:4W-22
b
(a) V-Min
(b) V-Half
Figure 5: We take a repeating d x T grid from V-Min and V-Half schedules, and assign F/B/W with different values. The result shows V-Min has bubbles for every repeating grid, while V-Half does not.

---

3.4 Other Building Blocks
Besides V-Shape building blocks, we also propose some other interesting building blocks in Appendix H. to show the generalization ability of our framework. Some useful examples include a) IFIB-V V achieving 2/3 of IFIB's activation memory without doing B-W split; b) a schedule consumes less memory than interleaved 1FIB but with the same bubble rate (Figure 17c). Additionally, we design an adaptive scheduler control the memory at finer granularity in Appendix A.
4 Experiments
We construct our experiments to show three conclusions: a) The throughput and memory of V-Min, V-Half and V-ZB aligns with the theoretical analysis in Section 3.2; b) Memory-saving methods in- cluding V-Min and V-Half can bring accelcrations; c) Our methods still perform best when combining with other state-of-the-an techniques.
4.1 Setup
We evaluate our methods using a series of models detailed in Table 2 analogous to GPT-3 [Brown et al., 2020]. Our implementation is based on the open- open-source Megatron-LM, project [Narayanan et al., 2021] and experimented on up to 40 NVIDIA A100 SXM 80G GPUs distributed across nodes interconnected by a RoCE RDMA network. The running time of cach iteration is recorded after several warm-up iterations. Similar to the settings in [Qi et al., 2023], we deduct one transformer layer from both the initial and final pipeline stage compensate for the embedding and output layer in LM, which can otherwise become the bottleneck of the pipeline and interfere to the efficiency.
Table 2: Models used in experiments.
Model
Layers
Attention Heads
Hidden Size
GPUs
9.6B
30
40
5120
16
21B
46
48
6144
24
38.5B
64
7168
32
98.5B
78
80
10240
40
Our experiments majorly focuses on the following pipeline parallel schedules: a) V-Min, V-Half and V-ZB: schedules introduced in Section 3.2; b) IF1B and Interleaved 1FIB: methods implemented in Megatron- e) IFIB-R: IFIB with full activation rematerialization [Chen al., 2016]; d) ZB-1P and ZB- 2P: the adaptive zero-bo zero-bubble methods introduced in [Qi et al., 2023] with activation memory limit set to the 1x/2x times of 1FIB.
9.68 Model Pipelines
Model Pipelines
38.58 Model Pipelines
80-
VHaff
VHalf
Vhoif
30
WMN
30
Vilin
Viiin
IF1BR
1F1B-
FF3B-R
12a
96
142
32
12a
256
Number of Micccbbtcchs
Numbur of Nicrobatches
Number of Mirrottcche
9.68 Model Pipelines
Model Pipelines
38 Model
Pipelines
Actrvkionn Hemery (08)
50
(08)
Activation (C1)
40-
1F1B
Half
30
V-Min
20
Wehalt
Activatinn
50
WMin
UNin IF10-R
120
142
256
512
Number of ofMiccbbaccess
Number of Microbatches
Number of Mirroootches
Figure 6: Throughput and activation memory using the same microbatch size.
7

---

4.2 Comparing Pipeline Schedules
In Figure 6, present comparisons of the throughput measured in FLOPS utilization (MFU) and activation memory consumption across different pipeline schedules under various settings. From the results, V-ZB outperforms all other methods in terms of throughput, which aligns with Figure 4. When comparing the activation memory consumption, V-Min and V-Half stand out by significantly reducing activation memory approximately 1/3 and 1/2, while other methods' memory is similar except for 1FIB-R. More details of our experiments and definition of metrics can be found in Appendix D.
Notably V-Min has comparable throughput against IFIB, but its throughput falls behind IFIB at larger number of microbatches due to the aforementioned repeating bubble in Figure Sa, as discussed in Section 3.3. However, it still outperforms IFIB with full activation rematerialization, providing strong alternative for saving memory.
Model: 9. 68; d=16; n-64; mbs=2
Model: 218: d=24; n-96, mbs=1
Model: 36.58; d=32: n- mbs=1
28 2P
eV 26
55-
2B-1 gintelloweed 1FBB
55
2B IP enterioaved 1F18
2B-1P
Interieaved 1P1B
50
45
Min
45-
50-
-Min
FLB
NeU 080
45
50
-Min
F
F1B8
35
111-R
35
4P3
B-R
10
20
20
Activation Mamory (GB)
Activation Hemory (GS)
Activation (GB)
28 2P COM
Figure 7: Pareto frontier of MFU and memory for various setups.
We also plot both memory and MFU for the various methods in Figure in typical, but slightly different setting in which we reduced the microbatch size of 9.6B and 21B model to allow ZB-2P and Interleaved 1FIB to run which would otherwise run out of memory (OOM). It shows that the V-Shape pipeline schedules lie at the Pareto frontier.
4.3
When to Save Memory
While V-ZB provides optimal throughput, V-Half and V-Min methods mainly used when memory budget is tight. Conventionally, rematerialization is used when it runs out of memory (OOM). However, rematerialization leads repeated computation and consequently decrease the throughput. V-Ilalf and V-Min significantly outperforms rematerialization (IFIB-R) as we show in Table 7.
9.68 Model 16 Pipelines
218 Model on 24 Pipelines
38.58 Model on Pipelines
60
50-
55-
550
V.28
V.28 mbsil mosil
2B-1P mbsi4
mas
60
35
35
mbsi6
35
mbs33
30
32
128
256
45
96
282
384
64
125
258
522
Number of Nicrobatchees
Number of Nicrobatches
Number of Microbatches
9.68 Model 16 Pipelines
21B Model on Pipelines
38.58 Model Pipelines
60
celtottion Mermory (C10)
40
nettoon Mermory
60
Art Nution Mermury
V-28 mbs14
V-28 mbs12
V-28 mbsi]
20-
23-1P mosec
23-1P mose2
23-1P moscl
V4biif m
V44iif m
V45iif mos
1F18 mbs-4
1F1B mbs:2
1F1B mbs:1
Vi8h mos 12
vi8 mo0.6
Vi8 mo0.3
120
256
3M-
64
12n
258
512
Namber Nicrubalchhss
Number Nicrobatches
Namber Nicrobttchees
Figure& 8: Throughput and activation memory under similar memory limit.
Another benefit of saving memory is that we can potentially use the extra memory for an increased microbatch size, which leads to a higher arithmetic intensity. We present the results in Figure
8

---

8. On bigger models, where memory pressure is higher and hence microbatch size is smaller, V- Half schedule can surpass V-ZB and other baselines because of its arithmetic intensity gain. This observation does not apply for V-Min, implying its arithmetic intensity gain can not compensate for the increased bubble. Doubling/Tripling the microbatch size for V-Half/V-Min results in a slightly higher activation memory than the other methods. This reflects the constant factors we ignored in Section 3.2. The increase is less significant as the number of devices grows.
4.4 Combining with Existing Techniques
We present our methods in the context of LLM training together with various other techniques. The following techniques considered: a) Flash Attention [Dao al., 2022, Dao, 2023]; b) Tensor Parallelism [Narayanan et al., 2021] ] and Sequence Parallelism [Korthikanti et al., 2023]; c) Distributed Optimizer provided in Megatron-I.M. The implementations are all from Megatron-L.M [Narayanan et al., 2021], Both our methods and the baseline methods are combined with the above techniques. Similar to the evaluation method in Kim al. [2023], we perform grid search on the following parameters: the size of PP; the size of TP: the size of DP; the microbatch size (mbs). We use 40 GPUs in this experiment. For each method, the best result from the grid search is reported.
We present the best result for each pipeline parallel schedule in Table and the corresponding parameters. We find that when sequence length is smaller and hence the memory budget is more abundant, V-ZB performs the hest due to the elimination of bubbles. When we increase the memory pressure by increasing the sequence length, V-Half performs the best because of its memory efficiency. The detailed data and analysis of grid search can be found in the Appendix D.3 D.4.
Table 3: V-Shape schedules combined with other memory saving methods.
Common Setup
PP
Best MFU
Best Parameters
Method
(%)
DP
TP
PP
mbs
1F1B
54.77
5
98.5B Model Seq Length 1024
IFIB-R
40.84
2
10
ZB-1P
59.95
40
Batch Size 160
V-Half
57.83
20
V-Min
52.8
2
10
V-ZB
63.31
40
1FIB
62.95
2
4
5
98.5B Model
IFIB-R
50.37
20
Seq Length 3072
ZB-1P
62.18
4
10
Batch Size 640
V-Half
66.34
20
V-Min
61.04
40
V-ZB
62.56
10
IFIB
OOM
98.5B Model
FIBBR
42.05
10
Seq Length 16384 Batch Size 160
ZB-1P
0OM
V-Half
57.85
8
5
V-Min
48.58
5
V-ZB
OOM
Conclusion And Future Work
In this work, we present a framework that constructs pipeline schedules by focusing on their repeating building blocks. This framework enables direct computation of peak memory from the lifespan the building block. Based on this capability, we design family of memory-efficient building blocks. We discuss three representative methods from this family, namely V-Min, V-Half and V-ZB, and demonstrate with experiments that our methods advance the Pareto frontier of throughput and memory large model training. Furthermore, our methodology of designing pipeline schedules through building blocks may inspire the research community to explore more novel pipeline schedules. Notice that repeating building block the only way of building pipeline, other methods like greedy search could generate pipeline that has repeating patterns.

---

In the future, we plan to further explore more memory efficient pipeline schedules based on our framework. A major limitation of V-Min is that, it suffers from growing bubbles when increasing the number of microbatches. Although V-Half mitigates this issue, there is still a space to further reduce the memory consumption. Using continuous offsets or finer-granularity discretization a possible way to solve it. We leave it in our future work.
References
Tom Brown, Benjamin Mann, Nick Ryder, Mclanie 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, 33:1877-1901, 2020.
Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint
2016.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307 08691, 2023.
Tri Dao. Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 15:16344-16359, 2022.
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 SIGPLAN Symposium on Principles and Practice of Parullel Progrumming, pages 431-445, 2021.
Alexander L. Gaunt, Matthew A Johnson, Maik Riechert, Daniel Tarlow, Ryota Tomioka, Dimitrios Vytiniotis, and Sam Webster. Ampnet: Asynchronous model-parallel training for dynamic neural networks. arXiv preprint 1ath::1700.99900 2017.
Priya Goyal, Piotr Dollár, Ross Girshick. Picter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in I hour. arXiv preprint arXiv:1706.088
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: 1xXiv:8006.000 03377, 2018.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le. Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Arpan Jain, Ammar Ahmad Awan, Asmaa M Aljuhani, Jahanzeb Maqbool Hashmi, Quentin G Anthony, Hari Subramoni, Dhableswar K Panda, Raghu Machiraju, and Anil Parwani. Gems: Gpu-enabled memory-aware model parellpprallellsm system for distributed dnn training. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15. IEEE, 2020.
Taebum Kim, Hyoungjoo Kim, Gyeong-In Yu, and Byung-Gon Chun. Bpipe: Memory-balancee pipeline parallelism for training large language models. In International Conference on Machine Learning, pages 16639-16653. PMLR, 2023.
Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. Proceedings of Machine Learning and Systems, 5, 2023.
Shigang Li and Torsten Hoefler. Chimera: efficiently training large-scale neural networks with bidirectional pipelines. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-14, 2021.
Xiangru Lian, Wei Zhang, Ce Zhang, and Ji Liu. Asynchronous decentralized parallel stochastic gradient descent, In International Conference on Machine Learning, pages 3043-3052, PMLR, 2018.
10

---

Ziming Liu, Shenggan Cheng, Hao Zhou, and Yang You. Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency. The International Conference for High Performance Computing, Networking, Storage, and Analysis, pages 1-13, 2023. URL https: //api. www.etiismooollaaa... org :I
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, al. Efficient large-scale language model training on gpu clusters using megatron-Im. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-15, 2021.
Penghui Qi, Xinyi Wan, Guangxing Huang, and Min Lin. Zero bubble pipeline parallelism. In The Twelfth International Conference Learning Representations, 2023.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro. Megatron-Ing Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: (tt:::9990000555 2019.
Zhenheng Tang, Shaohuai Shi, Wei Wang, Bo Li, and Xiaowen Chu. Communication-efficientt distributed deep learning: A comprehensive survey. arXiv preprint arXiv:2003.06307 2020.
Bowen Yang, Jian Zhang. Jonathan Li, Christopher Ré, Christopher Aberger, and Christopher De Sa. Pipemare: Asynchronous pipeline parallel dan training. Proceedings of Machine Learning and Systems, 3:269-2999 296, 2021.
Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing, et al. Alpa: Automating inter-and {Intra- Operator} parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 559-578, 2022.
Yonghao Zhuang, Lianmin Zheng, Zhuohan Li, Eric Xing, Qirong Ho, Joseph Gonzalez, Ion Stoica, Hao Zhang, and Hexu Zhao. On optimizing the communication of model parallelism. Proceedings of Machine Learning and Systems, 5, 2023.

---

A Adaptive Scheduler Based on Search
Now we consider more general scenarios, where we want to minimize the bubbles given an activation memory limit. A straightforward approach should be simply searching over all possible offsets and picking the one with minimal bubbles. However, this naive method cannot work well due to there are exponentially many possible offsets, which makes intractable iterate thoroughly. In this section, we propose more practical searching method to solve this general problem.
We superscript E {0,1) }} to denote which stage in device, and use subscript i {1, .,,,,... d) to denote the index of the device. For example, FY represent the forward pass of stage cd + - in the i-th device. We define the offset from u to as 6(u, v) t(e) t(u), where t(e) represent the cell index along time horizon of pass v. To simplify the notations, we define S1.7 = 8(F P 5 = 6 F.11). b = 8( 1++1) and 6 B = 8B 1 to denote the offset from a pass to its next pass.
Instead of all possible offsets, we limit our search space to uniform offsets across devices. We also try to ensure each device has a balanced peak memory. Note that for the uniform offsets introduced in Section 3.1, the peak memory only falls into small discrete set (1 & M) where k is an integer). To make it work for a finer granularity of memory controlling, we split the repeating module into two parts, containing the first K rows and the last d - K rows respectively. More formally, we use the constraint as follows.
PP = =B, = 62 Sisk K
8F! 8B = k < kisk K 8F. SB. = 5% AK < <isc d
OF.- = SB; = 52 AK <i<d
Note that the above constraints have good properties that the peak memory is balanced across devices. As we can always greedily fill W and W when repeating, we only need to search over the permutation the first device, the values of 8 < 02 k; x2x S and K. The computational complexity is (d) if we regard repeating interval as a constant.
For each building block searched, we repeat the building block, check collision, do squeezing and reordering mentioned in 2.1. After searching over all possible building blocks, we pick the schedule with minimal bubbles. Note that we can use the true run times of F, B and W to calculate the bubbles, which will lead to more efficient schedule in real cases.
B Evaluation of Adaptive Scheduler
In Figure 9 we plot the bubble rate of adaptive V-Shape schedulers introduced in A under various settings and different memory limit. The run times of F, B and W are from profiled real cases, as in Table 5. We observe that the bubble rate drops as memory limit increases. Notably, there's sudden drop in bubble rate when the memory limit just goes above approximately 1/2 of 1F1B, at which point the repeating bubble mentioned in Figure 5a disappears.
9.68 Model, d=16
218 Model, d-24
36.58 Model, dm32 32
Microbatches=:32 32
Micrsbatches-488 68
Microbatchess 64
0.3
Merobatcheee112
Micrubetchee hee
Rate
Mirobaithw-192
Dusste Rate
Microbatchwss-56 Microbatches-511
0.2
Microbatchesee chatt 256
0.2
Mirobatche 384
0.2
0.1
0.0
0.50
0.75
1.00
1.25
1.50
1755
2.00
0.50
0.75
1.00
1.25
1.50
1.75
2.00
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Memory Limk/M
Mamory Limit/M
Mumory Link/M
Figure 9: Bubble rate of V scheduler under various settings.
We also compare scheduler with the adaptive zero bubble scheduler proposed in [Qi et al., 2023] in Figure 10. We find that V scheduler has a boarder range of memory configurations and smaller bubble rate compared to zero bubble scheduler. We also draw the bubble rate of 1F1 as a reference, though IFIB does not support a configurable memory.
12

---

9.68 Model, dm16, nm
218 Model, d-24, nm48
38.58 Model, dm32. n-64
Schhouukrr
0.3
V.Shhduulrr
S:2b1a Pate
0.3
Schwouuler
28 Schoduler
25 38-Scheullar
2B Schaduler 1F18
1F18
1F10
0.0
0.0
0.0
1.50 1.75
Nemmyy Limit M
Hemory
Memory
9.68 Madel, d-16, n-64
218 Model, d-24, n-96
38. 5B Model, d-32, n-128
Schoouler Scheddller
Buskll Rule
0.2
Buse Patto
0.2
Schcduer
1F18
1F18
1F18
Nemory M
Mamory Limit/
Memory /M
9.68 Model, d=16, n- 128
21B Model, d=24, n 192
38. 5B Model. d d=32, n 256
Rate
Scheduler 28 Schoduleer
02
Subbi hate
Schoollor
Schoduler
1F18
bubba
1F18
0.1
1F18
1.25 1.50 1.75 Memory
Hemory
Hamory
9.68 Model, d=16, 256
21B Model, d=24, 384
38. Model, d-32, n=11
BuDeR Rate
Bubble Rana
Scheedlerr
Bubbl hate
Scheduler
01
1718
1F18
6.1
1F18
1.50
1.00
1.2
1.50
1.75
Mamory
Figure 10: Comparison of scheduler and zero bubble scheduler.
C Reordering
In our framework, we may need to reorder the warm-up and cool-down phase after squeezing. Basically, employ simple greedy approaches to handle the reordering for warm-up and cool-down, and illustrate how zero bubble schedule is reordered in Figure 11.
6
S
B
Reordering
E
2
2
2
3
4
3
2
3
3
Figure 11: Top: the schedule after Squeezing. Bottom: the schedule after Reordering.
Warm-up In warm-up phase, bubbles mainly happen before the first B. We iterate all the cells from left to right. If a vacant cell (which means a bubble) is encountered, we try find computation pass to fill this bubble. We iterate all the following computation passes in the same device, and check whether it is possible to move if we keep all other passes unchanged. the check succeeds, we move it to the vacant cell, and the bubble is filled.
In cool-down phase, W can be scheduled at any time after its corresponding B. So we utilize a heuristic way to handle the bubbles. Firstly, we delete all the W passes in cool-down phase. Next, we squeeze the schedule to remove the bubbles caused by deleting W. After that, use W to
Cool-down
13

---

fill the remaining bubbles, ensuring each W is after its corresponding B. Finally, we schedule all the remaining W passes at the end.
Despite its simplicity, the above heuristics general and effective. However, it may not achieve the best performance in cases. We also design other greedy constructive methods complement for some building blocks. We will release all the related code in our repository.
D Detailed Experiment Data
Table 4: Comparing Pipeline Schedules
Model
9.6B
21B
38.5B
Setup
#GPU
16
24
32
Microhatch
4
2
1
#Microbatch
16
32
64
128
256
24
48
96
192
384
32
64
128
256
512
Samples
V-ZB
2.59 2.86
2.98
3.04
3.15
1.19
1.32
1.38
1.44
0.59
0.70
0.72 0.74 0.76
per
ZB-IP
2.24 2.58
2.84 2.99
3.08
1.02
1.18
1.30
1.39
1.43
0.53
0.62 0.68 0.72 0.75
second
V-Half
1.97 2.43 2.73 2.92 3.08
0.90
1.12
1.26
1.34
1.40
0.47
0.58 0.66 0.71 0.74
per
IFIB
1.69 2.18 2.55 2.80 2.94
0.76
0.99
1.16
1.29
1.35
0.40
0.52 0.61 0.67 0.71
GPU
V-Min
1.73 2.10 2.34 2.49 2.61
0.81
0.98
1.10
1.18
1.22
0.42
0.51 0.58 0.62 0.65
IFIB-R
1.26 1.62 1,90 2.08 2.18
0.57
0.74
0.87
0.96
1.01
0.29
0.38 0.45 0,50 0.53
V-ZB
50.4 55.7 57.9 59.1 61.3
50.5
56.1 58.4
59.8 61.2
45.6
53.9 56.0 57.1 58.5
ZB-1P
43.5 50.2 55.2 58.2 60.0
43.5
50.3
55.2
58.9 60.5
40.9
47.8 52.6 55.7 57.7
MFU (%)
V-Half
38.3 47.3 53.1 56.8 59.8
38.1 47.4
53.4
57.0
59.5
36.2
45.2 51.1 54.6 57.0
32.9 42.4 49.5 54.4 57.1
41.8
49.2
57.4
39.9
47.0
51.8 54.9
V-Min
33.6 40.8
45.6
48.5
50.7
34.4
41.8
46.9
50.0
52.0
32.4
39.7
44.9
48.1 50.1
IFIB-R
24.5
31.5
36.9
40.5
42.4
24.2
31.3
36.8
40.8
42.9
22.8
29.7
35.0
38.6
40.9
V-ZB
58
59
59
59
59
57
57
57
57
57
55
55
55
S5
55
Peak
ZB-IP
58
57
57
57
57
57
56
56
56
56
55
55
55
55
55
memory
V-Half
38
38
38
38
38
38
38
38
38
38
38
38
38
38
38
(GB)
IFIB
56
56
56
56
56
56
56
56
56
56
54
54
54
54
54
V-Min
29
29
29
29
29
31
31
31
31
33
33
33
33
33
IFIB-R
14
14
18
18
24
24
V-ZB
48
48
43
43
43
36
36
Activation
ZB-IP
48
47
47
47
47
43
42
42
42
42
36
36
36
36
36
memory
V-Half
28
28
28
28
28
24
24
24
24
24
19
19
19
19
19
(GB)
IFIB
46
46
46
46
46
42
42
42
42
42
35
35
35
35
35
V-Min
19
19
19
19
19
17
17
17
17
17
14
14
14
14
14
IFIB-R
1
1
1
V-ZB
18.7
8.88
4.57
2.32
1.16
18.9
8.51
4.39
2.19
1.09
18.7
8.61
4.3
2.24
1.14
Bubble
ZB-1P
31.7
20.1
11.2
5.92
3.06
31.5
19.8
11.0
5.82
3.0
32.1
20.1
11.1
5.89
3.04
rate
V-Half
40.5
24.2
7.41
3.84
41.3
24.6
3.89
24.9
14.3
3.97
(%)
IFIB
50.1 34.3
21.7
13.5
8.76
50.5
34.6
21.9
13.5
8.67
50.7
34.6
21.8
13.4
8.44
V-Min
48.4
36.5
28.0
23.1
20.3
47.8
35.8
27.4
21.9
19.1
48.6
36.5
27.8
22.5
19.5
IFIB-R
63.0 51.2 41.9
35.8
32.2
63.3
51.3
41.9
35.7
32.0
63.5
51.6
42.0
35.8
32.2
D.1
Comparing Pipeline Schedules
For Section 4.2, we present our detailed experiment data in Table 4. Specifically, the metrics are defined as:
* MFU: The FLOPS utilization of the training system. The calculation of FLOPS of model is following [Narayanan al., 2021].
Peak Memory: The maximum peak memory cross all devices.
14

---

Estimated as deducting the iteration-star start memory from peak memory on each device. The number presented the maximum activation memory cross all devices.
Activation Memory:
Bubble Rate: The theoretical bubble rate reported by scheduler, using profiled times of F, and W at starting iterations.
D.2 Single-pass MFU Gain When Increasing Microbatch Size
To evaluate how increasing microbatch size increases the arithmetic intensity, profile the run time of each single F/B/W pass and calculate their single-pass MFU. We list the results in Table 5. shows that whether there are significant MFU gain depends on both the model and the microbatch size.
Table 5: Single-pase MFU gain when increasing microbatch size
Model
9.6B
2
38.5B
Microbatch Size
4
6
2
Pass (ms)
12.96 26.30
39.45
9.30
18.11
26.81
6.72
12.27
18.05
B Pass (ms)
13.22
26.66
39.85
9.47
18.55
27.09
6.89
12.84
18.73
W Pass (ms)
9.76 19.62 28.93
7.19
14.03
21.82
5.06
9.63
15.46
FBW Average MFU (%)
72.13 71,45 71.86
71.11
72.83
73.14
66.86
71.87
71.69
D.3
More Details on Grid Search
Table MFU of grid search, with Senuthce 1024 and BatchSize 160
Parallelization
MicroBS
IFIB
IFIB-R
ZB-IP
V-Half
V-Min
V-ZB
DP=1
51.66
38.18
52.59
52.32
47.36
53.42
TP=4
54.0
40.32
56.37
55.92
50.24
58.25
PP=10
4
52.44
38.81
55,49
49.79
DP=2
54.17
40.84
57.26
56.85
52.8
59.03
TP=2
40.35
57.03
52.3
PP=10
4
35.14
DP=T
54.03
40.55
57.65
56.88
50.61
59.59
TP=2
2
53.3
39.83
57.55
56.9
50.4
60.23
PP=20
34.62
45.52
DP=2
40.34
57.83
52.78
TP=l
2
35.89
48.83
PP=20
27.78
DP=1
53.47
40.07
59.95
57.68
52.54
63.31
TP=1
35.69
54.08
48.27
PP-40
4
27.53
DP-I
44.71
33.21
44.6
44.05
36.88
45.32
TP=8
2
51.75
38.67
52.67
51.94
45.37
53.27
PP=5
52.0
38.49
53.84
52.98
46.67
53.98
8
50.2
36.84
51.87
46.83
DP-2
51.56
38.73
53.38
52.54
48.7
53.6
TP=4
54.77
40.69
57.48
56.32
52.24
58.05
PP=5
4
39.61
56.53
52.16
We show the MFU of every setup of our grid search in Table 6, 7 and 8 for three groups of experiments: one with Sequence = 1024 and BatchSize = 160, one with RequenceLength = 3072 and BatchSSize 640 and the other with Lenuth For the first experiment group, the best setup is V-ZB under pure PP because of its bubble elimination. For the second setup, the best setup is V-Half because its memory efficiency enables lower TP
16384 and BatchSize 160.
15

---

degree, which is otherwise impossible for V-ZB/ZB- For the last setup, due to high memory pressure only V-Min and V-Half can run without checkpointing. comparison of TP and PP can be found at D.4.
Table MFU of grid search, with SequenceLength 3072 and BatchSize 640
Parallelization
MicroBS
1FIB
1FIB-R
ZB--PP
V-Half
V-Min
V-Z.B
DP=1
62.06
45.85
62.18
62.17
56.29
62.56
TP=4
2
45.52
45,86
55.54
PP=10
45.34
DP=2
49.03
60.59
TP=2
47.39
PP=10
4
45.35
DP=1
48.63
66.34
57.92
TP=2
2
47.18
PP=20
4
45.3
DP=2
50.37
TP=1
2
46.01
PP=20
DP=1
49.7
61.04
TP=1
45.26
PP=40
4
42.73
DP=1
55,93
41.31
54.13
54,22
49.75
54.38
TP=8
57.64
42.37
55.25
55.64
51.88
55.76
PP=5
4
43,3
14.59
DP=2
62.95
46.49
61.95
57.74
TP=4
2
45.92
PP=5
44.44
Table 8: MFU of grid search, with Length 16384 and BatchSize = 160
Parallelization
MicroBS
IFIB
1FIB-R
ZB-1P
V-Half
V-Min
V-ZB
10
42.05
DP=2:1P-2:PP=100
DP=1:1P--:PP-20
1
41.52
DP-2:TP=1:PP-20
DP=1;1P=1;PP-400
DP=1,TP-8;PP-5
39.48
57.85
48.58
DP=1,TP-8:PP-5
2
DP-2:TP-4:PP-5
35.13
D.4
Model Parallelism: More PP or More TP?
Our grid search results in Appendix D.3 show strong favor of Pipeline Parallel (PP) over Tensor Parallel (TP), which might contradict with some existing industry experience where more degree of TP usually accelerates training. To understand the reason, we briefly compare TP and PP in this section.
Though PP also equally partition the model into PP shards. it usually needs to cache the activations for (pp) microbatches, resulting in the total activation memory demand same as the unpartitioned. On the other hand, TP, when used with sequence parallelism [Korthikanti et al., 2023], partitions most activation memory equally
TP shards, which one of the most significant benefit of TP over PP. However, this comes at the cost of collective communication and reducing the size of hidden dimension to 1. which can significantly decrease the single-pass (F/B/W) MFU. Though one can argue that the saved memory can be used to increase the microbatch size, our experiment measuring the MFU under different TP setups (Figure 12) demonstrates that a higher-degree of TP even with larger microbatch size still suffers from lower single-pass MFU.
16

---

Single-pass MFU under different TP
Single under different degree
0.85
mose2
mox
mox
mbs=6
0.65
0.50
0.55
0.53
0.50
9.50
0.45
0.45
C-2
Gequenneiingthon 102/4 102/4:95--100
95-1660
BS--640
Figure 12: Average single-past MFU (over FBW) for grid search under different TP degrees.
The throughput of PP also decreases PP degree increases for two major reasons: for PP schedules with a fixed global batch size, a higher PP degree usually results in higher pipeline bubbles. For example, IFIB has a bubble rate of p-1 p+n-11 which would increase p grows and n keeps unchanged; b) the first and last pipeline stage for a language model usually have an embedding or output layer which has innegligible additional computations, making the pipeline unbalanced and hurting the efficiency. With higher PP degree, the unbalance would be aggravated. However, theses two issues can be mitigated. For the first issue, larger number of microbatches using V-ZB can significantly reduce the pipeline bubble. For the second issue, we can use another standalone stage to handle the embedding or output layer instead avoid bottleneck. In our experiments, we simply deduct layer from both the first and last pipeline stage to mitigate this problem. As result, our methods essentially push the preference from TP PP.
E Bubble Analysis
Although we mainly focus on the memory efficiency in this paper, we also need to take care of the bubbles in pipeline parallelism. Typically, there trade-off between memory and bubble, where allowing more memory usage can reduce the bubble of a pipeline schedule. In this section, we will discuss how identify whether the bubble of schedule will grow with respect the number of microbatches. Additionally, we illustrate an empirical lower bound of bubble in our methods. We follow the notation in Section 3.1 and Appendix A this section.
E.1 Existence of Repeating Bubbles
In pipeline parallelism, usually expect there is repeating bubbles introduced in Section 3.3, namely, the bubble is only related to the number of pipeline devices (d), and won't grow when increasing the number of microbatches (π). This property guarantees good scalability with respect to n. We say a pipeline schedule with (d) bubble satisfies this property, otherwise with ((n) bubble. Note that the conclusion is based on the values of Tp. Ty and Tw, which are the times of F, B and W respectively. For example, the schedule of V-Min (Figure 4b) is with (d) bubble when TF Ta Tw. but is with (nn) bubble when Tw is significantly smaller than Tp and Tn.
If there are no repeating bubbles in schedule for any values of TF, TB. Tw, the minimal memory is 2dm. In pipeline schedule, there are two types of dependencies, streaming dependency within the same device and computational dependency across devices. We define dependency path as a sequence of passes (either streaming dependency or computational dependency). We define the time cost a dependency path as T, >i T., where To, is the time cost of Then the runtime of the whole pipeline should be T = max, T,,
(v1. U2, .... U/r where for any < < |r]. is dependent on U;-1
17

---

Obviously, there is a trivial dependency path that all the passes are from the same device, and Tr FT + Tn + T. Note that is the lower bound runtime, and any extra runtime would be considered as bubbles.
Let's consider another dependency path 72 containing only forward passes. To be simple, denote F3 as the forward sequence for the j-th microbatch. Then 72 concalenate (Fo, Fk...., F 1 )k
thus T₁ 2d| n+k-1 [TF, where 66 8(FP, F() 6(k 1) (greedily include many forward passes as possible). Note that if we want to guarantee (d)) bubble for any values of TF, Tn, Tw, we should choose k 2 d to make 2d] x = I 2n, otherwise 2d n |Tr 2n(TF + TB + Tw) E
((n) if we set Tв = Tw = 0. Then we can get ( F)) > 6d
Then we consider similar dependency path containing only backward passes, and we can get 8BB,, B)) > 6d 6. According to Equation I, we can get the peak memory in the first device is at least 6(F.. . 5(B) )m 2 2dm.
For most real cases, V-Half is enough to guarantee there are no repeating bubbles. Although the above proof shows that 2dm memory is required guarantee (dd) bubble for any values of Tr, Tn, Tw, we don't need that much memory in practice because the values of fr Tn and Tw are well constrained. As in Qi al. [2023], B and W have similar total FLOPS, and TF, B, Tw don't differ too much in real cases, Based on this insight, we can check the conditions where our methods are with (d) bubble.
Because both warm-up phase and cool-down phase have (dd) passes, we only need to consider the stable phase to identify whether schedule is with (d) bubble with O(n) bubble. Obviously, streaming dependency won't It block the device from executing the next pass. So bubble caused only by the computational dependency. Formally, schedule is with O(n) bubble if and only if there exist two passes u and o (within the same device) and there are two dependency paths and 7' between them, where r only contains streaming dependencies, 7' contains at least two computational dependencies, and T, Ty,, Based on our V-shape building blocks, we only need to check and with small distance (< and r' within two adjacent devices. In this way, we can conclude that the schedule in Figure 4c is with (dd bubble when Tww + 21' 2 27'F & Tw + 2TF 2 27' which is satisfied in most real cases.
E.2 Lower Bound
Deveed
Devis4
Ben
Figure 13: A dependency path of pipeline schedule.
The runtime of schedule can be bounded by any dependency path. In Figure 13, present dependency path which a non-trivial lower bound for most schedules after squeezing and reordering. Assuming the actual peak memory is km (k 2d) and Tp TB Tw 1, then the runtime of this dependency path is at least 6n + 6d 3k 1. Empirically, we find that we can always get close to this lower bound in adaptive scheduler.
F Building blocks of V-Min and V-Half for all values of d.
To avoid collisions when repeating building blocks, we need a slightly different offsets for different number of devices d. We list the details in Table 9 and 10 while continue using the notation defined in We do not list the offsets related to passes because they always fill the empty grids left by F and B. Some samples can also be found at Figure 14.
18

---

V-Min
modd
mod
S mom
V-Halt
mod 2
2
Figure 14: Building blocks of V-Min and V-Half on different settings of d
Table 9: Offsets for V-Min
1 <
O 8F.
8
OB.
6 P2)
8(F), B()
8(Bg, BJ)
d=0 mod 3 d 0=p 0
3
mod 3
1
1
Table 10: Offsets for V-Half
<<i<d1<js
OF
8F)
oB;
8(F2
F.)
8(F1
8(B By)
de
d 0 mode 2
2
2
4
d=1 1
mod 2
2
2
G
Non-uniform Repeating Interval of Interleaved 1F1B
While most existing schedules (Figure 18) are repeated with uniform interval, interleaved IFIB [Narayanan et al., 2021] is slightly different the repeating pattern. The official interleaved IFIB has repeating pattern shown in Figure 15a. If we also employ uniform repeating interval as in Figure 15b, we can obtain another schedule with the same memory footprint and bubble rate as the official interleaved 1F1B, shown in the gallery (Figure 18i).
(a) The non-uniform repeating interval of interleaved IFIB. The number in highlighted grids shows the microbatch with that index starts at this cell.
variation of interleaved IFIB with uniform repeat.
Figure 15: Different repeat pattern of the same building block of interleaved FFIB.
H Other Memory Efficient Building Blocks
Under the guidance of lifespan defined in Section 2.2, we also find some other building blocks besides the V-Shape building block family. We show the building blocks in Figure 16 and their full schedules
19

---

aa IFIB-V
W
W
W
B
W
W
(b) Zero bubble schedule with 2/3 IFIB memory
(c) variation interleaved 1FIB with the same bubble rate but lower memory
Figure 16: Other memory-efficient building blocks.
1
5
6 2
(a) IFIB-V V
8
8
(b) Zero bubble schedule with 2/3 IF1 memory
2
3
(c) A variation of interleaved IFIB with the same bubble rate but lower memory
Figure 17: Other memory-efficient schedules.
in Figure 17. All these schedules have lower bubble rate than IFIB. Specifically, 1FIB-V applies V-Shape to the building block of IFIB but without B- split, which can reduce the peak memory to asymptotically 2/3 of IFIB. We also find that utilizing B- W split, the zero bubble pipeline schedules proposed in [Qi et al., 2023] with configurable memory limit can support minimum of 2/3 activation memory of 1FBB, using the building block shown in Figure 16b. Note that two microbatches included a single building block to avoid collision. Using the building block defined in Figure 16c, can make schedule with the same bubble rate as interleaved FFIB but lower memory, shown in Figure 17c.
I A Gallery of Pipeline Parallel Schedules and Their Building Blocks
We show the building blocks and full schedules of some well-known existing methods in Figure 18.
20

---

(a) IFIB
(b) Eager IFIB
(c) ZB-HI
(d) ZB-H2
(c) GPipe
(f) GEMS
g Chimera
123312
(h) Interleaved 1FIB
125
(i) Interleaved IFIB with uniform interval (with reordering)
Figure 18: gallery pipeline schedules and their building blocks. The upper row of each schedule shows the building block and how it repeats. The lower row shows the final schedule after squeezing.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our paper focus on the memory efficiency of pipeline parallelism. Our major contribution contains a framework
design pipeline schedules, which can control the activation memory by adjusting the lifespan, and family of memory-efficient schedules. These are highlighted in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
It is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The major limitation of our work is that the performance V-Min method degrades when B and W times are different. We discussed the limitation in Section 3.2 and also reflected them in the evaluation Section 4
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
22

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: We don't have many theories in our work. The most theory part is the relationship between lifespan and memory, which is quite intuitive.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: Our implementation is based on Megatron-LM,, which is accessible from everyone and we also disclosed the detailed setup of each schedule. We also plan to open-source our implementation when camera-ready.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
23

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We'll opensource the implementation code on camera-ready version.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm
wwtssssinnooiiiyll. for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / /ttps.cccpubllt.../ wwtppcccccppceeeccccecccccccccce gggiiees//////e/eeoo http///ooSSSuiissseeeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
The experiment details are shared in Section 4 and Appendix D.I, D.3
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our stats in experiments are mostly profiled memory and runtime, which are quite stable in our environment.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
24

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
As mentioned in Section 4, we need at most 40 GPUs to run our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: atttps:/neerr/ /wwrrps....pp//
Answer: [Yes]
Justification:
We conforms to the NeurIPS code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our methods are focusing on the training systems only, not on the algorithms. They should not have any societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
25

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our methods has no risk of being misused.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Our implementation is based on Megatron-LMM which has lisence that allow other works to freely work on them.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
26

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasttss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We'll opensource code with documentation usage camera-ready
version.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: crowd sourcing involved
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: No human subjects
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
27

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "YnDPgapOBZ", "reviewer": "Reviewer_vShK", "summary": "Authors propose a way to identify the repeated building block a pipeline schedule is built from. By relating the peak activation memory of the schedule to the lifespan of the building block the authors show that existing schedules do not optimally use memory, and design higher-throughput schedules that use the same activation memory as the baseline. The methods demonstrate up to 55% improvement in throughput for pure pipeline parallelism settings and 16% improvement for large language models over the 1F1B baseline. The paper also provides a systematic methodology for designing pipeline schedules, emphasizing memory usage balance and reduced pipeline bubbles.", "questions": "Why is 1F1B the baseline used for the paper? It feels like there could be some more discussion here that is missing.\n\nIs the lack of bubble in V-Half (figure 5) purely empirical or is there some reason it should be expected? There doesn’t seem to be anything special about V-Half from the point of view of the discussion in the paper (only the activation memory is relevant to V-half; it’s not clear why this specifically should not have bubbles). Also, V-ZB has no bubbles in the case of differing runtimes, if I understand correctly; so maybe it is more accurate to say that the number of bubbles decreases to 0 as activation memory budget increases?\n\nHow is the number of stages in the pipeline determined? While the authors mentioned that it depends on the number of devices, what is the number of stages used in this work? How does the number of stages influence performance?\n\nGiven the constraints of memory or throughput, do the authors use grid search to find an optimal schedule? Is it possible to use a theoretical or approximation method to quickly find a good pipeline schedule?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper introduces a novel framework for decomposing pipeline parallelism schedules into memory-efficient building blocks. This approach addresses inefficiencies observed in existing schedules. The systematic methodology for designing pipeline schedules emphasizes memory usage balance and reduced pipeline bubbles.\n\nDiscussion of peak memory and new memory efficient building blocks is the core of the paper and the part which is most clear. The authors provide technical details, including the proposed building blocks, their implementation, and experimental results.", "weaknesses": "The paper suffers from quite a few clarity issues and could use some copyediting (there are lots of small grammatical errors, etc.). Main clarity issues in the paper are, for example, a lack of useful captions in figures, a lack of substantive discussion in some of the appendices (i.e., “we leave the discussion to appendix X” but there is not much discussion in appendix X).\n\nA lot could be done to improve clarity of the discussion in section 3. For example, although the paper explicitly details asymptotic behavior (d -> infinity?) this is really only mentioned in lines 166-167, which somewhat confuses the issue. Some brief discussion of what the effect is in low-d situations would be nice (not new experiments – just a qualitative idea). Fig. 3 and Table 1 contradict each other due to the presence/absence of the asymptotic limit, which is mentioned in the title of Table 1 and briefly in the text but is not very clear, especially since Fig. 3 and Table 1 are supposed to be read together (?).\n\nThe figure captions should at minimum restate what is shown in the figures and what relevant terms (d, l, etc.) mean – this makes it much easier to refer to the figure without searching through the text for definitions and explanations. This should be doable in 1 or 2 sentences per caption at most and should not take a lot of space.\n\nSome plots are hard to read. For example, Figure 4 and Figure 5 shows detailed pipelines using various colors of blocks and fonts without proper definition or explanation. It is also not clear the shown pipeline is the actual setting or a high-level demonstration of the design. \n\nV-Half seems to be a heuristic method based on V-Min and V-ZB. It may not be an optimal solution for the pipeline. Other configurations regarding the trade-off of memory and throughput are not evaluated.\n\nSection 4.4 is confusing. Table 3 is not referenced anywhere and not explained in this section.  The results mentioned in this section (Table 6) do not clearly show the combined performance of the proposed approach and existing techniques.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and insightful suggestions for improvement. We have revised the manuscript to address most of the mentioned issues. We respond to individual points from your review below.\n\n>grammatical errors, a lack of useful captions in figures, Some plots are hard to read.\n\n* We have used the grammar tool to check and correct the grammar errors.\n* We add captions for most of the figures to summarize the key points or conclusions, explain the notations in Figure 3 and Table 1.\n* We re-organize the layout of figures to make it more friendly to readers.\n\n>a lack of substantive discussion in some of the appendices.\n\nCould you provide specific pointers?\n\n>A lot could be done to improve clarity of the discussion in section 3... the effect in low-d situations... \n\nTo make it more clear, we rewrote the Section 3.1 to include more technical details in the main paper (refer to global response). Specifically,\n* We use uniform offsets between adjacent F/B passes across different devices, and leave those within the same devices flexible. This is where the asymptotic behavior comes from.\n* The exact peak memory is $\\lceil \\frac{d+2}{3} \\rceil \\frac{M}{d}$ for V-Min, and $\\lceil \\frac{d+1}{2} \\rceil \\frac{M}{d}$ for V-Half. Actually we intended to mention it in the paper, somehow we forgot to do so, thanks for pointing this out.\n\n>not clear the shown pipeline is the actual setting or a high-level demonstration of the design.\n\nThe pipeline shown in Figure 4 is the actual setting with 4 devices and 8 microbatches. (Refer to the PDF in global response)\n\n>V-Half seems to be a heuristic method... It may not be an optimal... Other configurations... are not evaluated.\n\nV-Half is not heuristically designed, all schedules are searched systematically with designated memory limit. There are indeed some other configurations from the search, e.g. a pipeline with 2/3 memory and 1/3 bubbles of 1F1B. As V-Min and V-ZB are optimal in memory and throughput, V-Half is to represent something in the middle. For the trade-off of memory and throughput, please refer to Figure 9 and 10 in the Appendix.\n\n>Section 4.4 is confusing. Table 3 is not referenced anywhere... (Table 6) do not clearly show...\n\nIt's a mistake due to duplicated label name in latex, it is fixed now, thanks for spotting it.\n\n>Why is 1F1B the baseline used for the paper? It feels like there could be some more discussion here that is missing.\n\nInstead of saying using 1F1B as a baseline, we are more of using it as a reference point. The memory consumption of other methods are normalized to a proportion of 1F1B, simply because 1F1B is the most well known schedule that readers should be familiar with.\n\n>Is the lack of bubble in V-Half purely empirical or ... expected? ... Also, V-ZB has no bubbles ... differing runtimes, ... so maybe it is more accurate to say that the number of bubbles decreases to 0 as activation memory budget increases?\n\nV-Half is free of **repeating bubbles** (in most cases), while it is not free of **warmup/cooldown bubbles**. As discussed in Appendix E.1, repeating bubbles occur in V-Half only when $T_W < 2\\max(T_F, T_B) - 2 \\min(T_F, T_B)$.\n\nGenerally, longer lifespan means looser dependency chain, making the schedule more robust to the variation of running times, and more friendly to overlap communication and computation. That's the intuitive reason why V-Min is vulnerable (minimal lifespan means tightest dependencies) and V-Half is more resilient. While at the memory budget of V-ZB, it is free of both types of bubbles.\n\n>How is the number of stages in the pipeline determined? ... what is the number of stages used in this work? How does the number of stages influence performance?\n\nFor a **balanced computation load**, each device is expected to host equal number of stages. That's why strategies like interleaved-1F1B use a stage number which is a multiple of the device number.\n\nFor a **balanced peak memory**, we need the total lifespan to be equal on each device. As shown in our paper, it can be achieved using the \"V-Shape\" schedule, where the number of stages is twice that of the device. Repeating the \"V-Shape\" into a \"W-Shape\" would double this ratio, and would be memory balanced as well, but it does not further decrease the memory, and on the throughput it improves at the rate of $\\frac{2+V}{V}2d$ (for V-Min) where $V$ is the number of stages per device, which is not very significant. Therefore we didn't bring this extra complexity to readers in the paper.\n\n>Given the constraints of memory ..., do the authors use grid search to find an optimal schedule? Is it possible to ... quickly find a good pipeline schedule?\n\nWe elaborate a simple search approach to quickly find the optimal schedule given the memory constraint (the number of possible building blocks is $O(1)$) in Section 3.1, and introduce an adaptive scheduler with finer control granularity (the number of possible building blocks is $O(d)$). We also evaluate its trade-off between memory and bubble in Figure 9 and 10 (Appendix B).", "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": "VfeQZIjYwL", "reviewer": "Reviewer_w4r8", "summary": "This work proposes systematic methodology for designing pipeline parallelism schedules and analyzing their performance (e.g. peak memory and pipeline bubbles).\nThe major observation is that a pipeline schedule can be viewed as repeating a building block, and the peak memory critically depends on the lifespan of a block.\nBased on these insights, the authors design a family of novel V-shape building blocks and pipeline schedules that are memory efficient and achieve higher throughput than strong baselines like 1F1B (though at the cost of more communication), which is validated both theoretically and empirically.", "questions": "A high-level question that might be worth a brief discussion in the manuscript is: \nis it true that a pipeline schedule *must be* repetition of a building block? \nActually the answer is no.\nA naive example I can think of is a variant of GPipe, where the order of backward passes can be arbitrary, and thus the building block for one microbatch can be different from that of another microbatch.\nThis extension of GPipe is designed just out of theoretical curiosity though.\n\nSo the next question is whether there could be any real advantage (in terms of peak memory, throughput, etc.) in a pipeline schedule that cannot be expressed as repetition of a building block.\nIt would be great if there is a clear answer with a formal proof.\nBut even if the authors do not know the answer yet (and neither do I), it might be worth mentioning, so that anyone who has read this work will keep this in mind, rather than blindly constrain themselves to repeating a building block when designing new schedules.", "limitations": "Limitations are discussed throughout the manuscript.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This work proposes a systematic and unified perspective for pipeline parallelism schedule,\nwith some novel insights that might foster the design of better schedules in the future.\n- Analysis and experiments are extensive and well executed, confirming the efficacy of the proposed methodology.\n- Overall, writing and presentation are clear.", "weaknesses": "There is still room for improvement in writing. For example:\n- There are quite some typos, some of which are listed here: Line 22 \"tensor\" -> \"Tensor\"; Line 189, \"V-Min\" -> \"V-ZB\" (?); Line 193, \"serious\" -> \"series\".\n- Line 210, \"while other methods' memory is similar\": one exception is 1F1B-R, whose activation memory is much smaller.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the valuable feedback and suggestions for improvement. We respond to individual points from your review below.\n\n>There are quite some typos, some of which are listed here: Line 22 \"tensor\" -> \"Tensor\"; Line 189, \"V-Min\" -> \"V-ZB\" (?); Line 193, \"serious\" -> \"series\".\nLine 210, \"while other methods' memory is similar\": one exception is 1F1B-R, whose activation memory is much smaller.\n\nThanks for the careful reading. We have corrected them in the revised version. We also used a grammar tool to check the whole paper and correct the grammar errors.\n\n>A high-level question that might be worth a brief discussion in the manuscript is: is it true that a pipeline schedule must be repetition of a building block? Actually the answer is no. A naive example I can think of is a variant of GPipe, where the order of backward passes can be arbitrary, and thus the building block for one microbatch can be different from that of another microbatch. This extension of GPipe is designed just out of theoretical curiosity though.\nSo the next question is whether there could be any real advantage (in terms of peak memory, throughput, etc.) in a pipeline schedule that cannot be expressed as repetition of a building block. It would be great if there is a clear answer with a formal proof. But even if the authors do not know the answer yet (and neither do I), it might be worth mentioning, so that anyone who has read this work will keep this in mind, rather than blindly constrain themselves to repeating a building block when designing new schedules.\n\nThis is very true, thank you for bringing up this insightful point. We add a sentence \"Notice that repeating a building block is not the only way of building a pipeline, other methods like greedy search could generate a pipeline that has no repeating patterns.\" in Conclusion.\n\nSome more discussions on the points that may favor repeating building blocks.\n1. Repeating the building blocks means uniform lifespans, and intuitively it is better than unbalanced lifespans where a single point could cause a larger peak memory.\n2. Mentally more tractable, the complexity of designing a \"good\" pipeline is reduced to designing a \"good\" building block.\n3. Additionally, the scalability with respect to the number of microbatches is guaranteed.", "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": "Z7RTd8AmfU", "reviewer": "Reviewer_b8L7", "summary": "The paper addresses the inefficiencies in current pipeline parallelism schedules used for training large-scale models, particularly focusing on high activation memory and pipeline bubbles. The paper proposes a new framework to decompose pipeline schedules into repeating building blocks and introduces V-shape building blocks (V-Min, V-Half, V-ZB) with controllable activation memory. Such building blocks reduce peak activation memory compared to the traditional 1F1B approach and even achieve higher throughput with reduced pipeline bubbles. The evaluations demonstrate substantial improvements in throughput and memory efficiency over 1F1B and 1F1B-I methods.", "questions": "- One of the problems with pipeline parallelism is that it increases implementation complexity because of its multi-program-multi-data (MPMD) nature. Do V-shaped abstractions bring any convenience in terms of implementation perspective?", "limitations": "No societal impact exists.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors decompose pipeline parallelism into repeating building blocks, providing a systematic methodology for designing more efficient schedules. This approach helps in understanding and optimizing pipeline parallelism in a structured manner.\n- The V-Shape building blocks demonstrate the structured understanding of pipeline parallelism.\n- The authors provide experimental evaluations demonstrating the practical benefits of the proposed methods.\n- By reducing pipeline bubbles, the paper demonstrates that PP can become a much more preferred option over TP at practical large-scale training.", "weaknesses": "- It would be helpful to have experimental results for long sequence lengths. 1024 and 3072 sequence lengths are too short compared to what SOTA LLMs can handle.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and suggestions for improvement. We respond to individual points from your review below.\n\n>It would be helpful to have experimental results for long sequence lengths. 1024 and 3072 sequence lengths are too short compared to what SOTA LLMs can handle.\n\nThanks for pointing out this. Theoretically, increasing the sequence lengths should not significantly affect our conclusions, as our methods are flexible with different F/B/W timings. On the contrary, increasing the sequence length will increase memory pressure, highlighting the advantages of our memory-saving schedules.\n\nWe added another set of grid-search experiment using almost the same setup but with sequence length=16384. Here is a summary of experiment result. In this setup, 1F1B, ZBV and ZB1P run OOM for all combination of distributed training parameters (dp/tp/pp degrees). The experiment shows V-Half is 38% faster than 1F1B-R, which is the only baseline method that can run.\n\n| Common Setup | PP Method | Best MFU (%) | Best Parameters |\n| ------ | ------ |------ | ------ |\n|  | 1F1B | OOM | - |\n| 98.5B Model | 1F1B -R | 42.05% | dp=1;tp=4;pp=10;microbs=1 |\n| Seq Length 16384 | ZB1P | OOM | - |\n| Batch Size 160 | V-Half | 57.85% | dp=1;tp=8;pp=5;microbs=1 |\n| | V-Min | 48.58% | dp=1;tp=8;pp=5;microbs=1 |\n| | V-ZB | OOM | - |\n\n\n>One of the problems with pipeline parallelism is that it increases implementation complexity because of its multi-program-multi-data (MPMD) nature. Do V-shaped abstractions bring any convenience in terms of implementation perspective?\n\nThanks for discussing this question from the implementation perspective. We believe that implementing pipeline parallelism primarily involves two components: a) a scheduler that determines the order of passes, and b) a runtime that executes the schedule.\n\nRegarding the scheduler, we believe the V-Shaped abstractions simplify the implementation. Developers only need to implement the building blocks, and we can use a framework to handle automated repeating, squeezing, and reordering. This approach is exactly how we implemented the schedules in our experiments.\n\nFor the runtime, implementing a pipeline parallelism (PP) runtime is indeed more complex. It requires careful insertion of communications between stages to avoid performance degradation (e.g., stalled send/recv operations) or deadlocks. Note that this complication is inherent to PP and not introduced by V-shape schedules. To address this, we built a uniform runtime that automatically inserts and fuses (or refrains from fusing when better) communication operations according to the schedule produced by the scheduler. This allows the scheduler to focus solely on the ordering of computation passes. We will also open-source our PP runtime in hopes that it will benefit the community.", "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:
Thanks for all reviewers for the valuable feedback and insightful suggestions. We updated our PDF accordingly. The changes mainly include:
- We change the title from "Efficient Pipeline Parallelism with Controllable Memory" to "Pipeline Parallelism with Controllable Memory.
- We've corrected the MFU numbers in the experiment tables and graphs. We initially calculated the peak FLOPs of A100 as 312 * 1024^4, which should be 312 * 1000^4 instead. This adjustment increases all MFU numbers by about 10%. This adjustment doesn't affect the conclusions or acceleration ratios, as it is applied equally to both the baseline and our methods. In the grid search experiment, our methods can reach 66.34% MFU (~206 TFLOPS/s) on 40 A100s.
- We rewrite Section 3.1 to include more technical details on the control of peak memory for V-Shape schedules. Specifically, we use uniform offsets between adjacent F/B passes across different devices (refer to the PDF), and leave those within the same devices flexible and use brute force to find optimal solutions. Notably, all of V-Min, V-Half and V-ZB result from systematic searches. Not designed by heuristics.
- We improve the presentation by reorganizing the placement of the figures, changing pipeline schedules from 5 devices 10 microbatches to 4 devices 8 microbatches, and adding more details in the captions. (refer to the PDF)


---------------------------------------------------------------------------
# Section 3.1

We assume the model is uniformly partitioned, namely, both the computation and memory of each stage are identical. For a single microbatch, we denote the activation memory of each stage as $m$, and the total activation memory of the entire model as $M$. Note that $M=2dm$, where $d$ is the number of devices. To make it simple and tractable, we use **uniform offsets within each half of F and B passes** to control the peak memory. Specifically, we apply the same offset $\delta_F^0$ between two adjacent F passes within the first $d$ stages (e.g., $\delta_F^0=2$ in Figure 3b, $\delta_F^0=1$ in Figure 3c and $\delta_F^0=4$ in Figure 3d). Similar constraints are applied to the other half of the F passes and both halves of the B passes, denoted as $\delta_F^1, \delta_B^0, \delta_B^1$, respectively. To guarantee balanced peak memory across devices, we add another two constraints, $\delta_F^0=\delta_B^1=\delta^0$ and $\delta_F^1=\delta_B^0=\delta^1$, where we use notations $\delta^0$ and $\delta^1$ for simplicity. For example, in Figure 3d, we set $\delta^0=4$ and $\delta^1=2$. Note that we only control the offsets across different devices. For those adjacent passes within the same device (e.g., F and B of the last stage, two F and two B in the last device), we use brute force to find optimal solutions, ensuring their offsets are small (less than the repeating interval). Note that W can always be placed greedily after settling all F and B passes, so we don't need to search their offsets during brute force.
According to Equation 1, we can analyze the asymptotic peak memory with respect to $d$,
$$\text{peak memory of device }i \leq \frac{2d(\delta^0 + \delta^1) + O(1)}{6} m  \approx \frac{\delta^0 + \delta^1}{6} M $$

By ignoring the small constant, we can directly control the peak memory by the value of $\delta^0$ and $\delta^1$.

By varying the values of $\delta^0$ and $\delta^1$, we come up with 3 novel V-Shape building blocks (Figure 3), and present their final schedules based on our framework in Figure 4.

Meta-review: All reviewers have given favorable ratings for this paper, ranging from 5 to 7.
Average Rating: 6.33 (Min: 5, Max: 7)
The following strengths of the paper have been highlighted. However, as there are few strong recommendations for this paper, an acceptance as a poster presentation is considered appropriate.

1: This paper proposes a novel framework that decomposes pipeline parallelism schedules into memory-efficient building blocks. This approach addresses inefficiencies in existing schedules, designing new schedules that reduce memory usage while improving throughput.

2: The proposed method demonstrates a 7% to 55% improvement in throughput in pure pipeline parallelism settings compared to 1F1B. Additionally, it achieves a 16% throughput improvement over the 1F1B baseline for large language models.

3: The paper presents a systematic methodology focused on balancing memory usage and reducing pipeline bubbles. The evaluations confirm the effectiveness of the proposed method through both theoretical and experimental validation.

4: The paper introduces new V-Shape building blocks, such as V-Min, V-Half, and V-ZB, which reduce memory usage by more than half compared to traditional methods while maintaining comparable throughput.
Decision: Accept (poster)