Pfeife: Automatic Pipeline Parallelism for PyTorch
Authors: Ho Young Jhoo
Chung-Kil Hur
Nuno P. Lopes
Venue: ICML 2025 poster (2025)
Keywords: Parallel Training
Automatic Parallelism
Pipeline Parallelism
PyTorch
Primary area: general_machine_learning->hardware_and_software
TL;DR: Implementation of fast and general automatic pipeline parallelism for PyTorch using torch.compile.

Abstract: The memory requirements of machine learning (ML) models has been growing quickly. However, the memory capacity of GPUs has not kept pace. Despite significant research on reducing the memory usage of ML models, the larger models do not fit in a single device. A popular solution to the memory capacity issue is to use multiple devices in parallel. In this paper, we focus on a particular form of parallelism called pipelining, as it offers a good balance between cost and performance for many ML models. We present Pfeife, the first tool that integrates with PyTorch to provide automatic pipelining of ML models. Pfeife intercepts the execution of models and parallelizes them transparently, requiring no manual work. We show that Pfeife can execute large models that would otherwise not run due to not fitting in a single device. Moreover, Pfeife can pipeline non-sequential models such as Stable Diffusion, which are not supported by existing pipelining parallelism tools. Pfeife outperforms state-of-the-art tools by up to 22%.

Pfeife: Automatic Pipeline Parallelism for Py
Ho Young Jhoo' Chung- Kil Hur!
Nuno P. Lopes
Abstract
device. The maximum memory capacity of the top NVIDIA GPU released this year (GB200) is only 384 GiB, still unable to fit Llama 3.1.
The memory requirements of machine learning (ML) models has been growing quickly. How- ever, the memory capacity of GPUs has not kept pace. Despite significant research on reducing the memory usage of ML models, the larger models do not fit in a single device. A popular solution to the memory capacity issue is to use multiple devices in parallel. In this paper, we focus on particular form of parallelism called pipelin- ing, as it offers good balance between cost and performance for many ML models. We present Pfeife, the first tool that integrates with PyTorch to provide automatic pipelining of MI. models. Pfeife intercepts the execution of models and par- allelizes them transparently, requiring no manual work. We show that Pfeife can execute large mod- els that would otherwise not due to not fitting in a single device. Moreover, Pfeife can pipeline sequential models such as Stable Diffusion. which are not supported by existing pipelining par- allelism tools. Pfeife outperforms state-of-the af-t tools by up to 22%.
Several techniques have been developed to circumvent the hardware memory capacity limitations, especially using multiple devices on the same server and across servers. This encompasses multiple orthogonal techniques that can be combined, including data parallelism, which replicates the model across devices and splits each input batch between devices; weight sharding, which splits parameters across devices; and model and tensor parallelism that partition the execution of one instance of the model.
These approaches face two main challenges: scalability and usability. Data parallelism and weight sharding, for instance, eventually require full synchronization of the weights across all devices, thus require heavy communication. For model and tensor parallelism, most tools are notoriously hard to use. They often require writing models in particular way or to annotate models in a non-trivial way. Or even to pick the partitions or identify the parallelism manually. None of these tasks are easy for humans, especially for ML practi- tioners, who are usually not experts in distributed systems, parallel computing, or HPC.
In this paper, we focus on pipeline parallelism, which is a particular form of model parallelism. We built Pfeife, a tool that integrates with PyTorch and that pipelines models automatically with no user intervention. PyTorrhh's 's tracing JIT compiler constructs, at run time, a data-flow graph that represents the model being executed, and it is capable of looking through lot of Python's 's dynamism. This enables tools like Pfeife to grab the data-flow graph of ML models. Moreover, it allows us to intercept the execution of models and parallelize the execution in a transparent way to users.
1. Introduction
ML models have been growing in terms of memory re- quirements very quickly. The memory needed just for the weight matrices has gone from 0.4 GiB in 2018 (ELMo (Pe- ters et al., 2018)) to 754 GiB in 2024 (Llama 3.1 405B FP16 (Grattafiori et al., 2024)).
While memory requirements of ML models have increased by several orders of magnitude in just few years, hardware has not kept up the pace. For example, NVIDIA's high-end GPUs, the most used platform for training and deploying MI. models, cannot fit state-of-the-art MI models in a single
Pfeife also allows selectively parallelizing a model in a training loop, so that can ensemble multiple models like pre-trained frozen models. This functionality helps users train complex models like Stable Diffusion (Rombach et al., 2022) which are traditionally not supported by automatic model parallelism tools.
'Seoul National University, Republic of Korea PINESC-ID / Instituto Superior Técnico University of Lisbon, Portugal 2 Furiosaal, Republic of Korea, Correspondence to: Ho Young Thoo <hoyoung jhoo@ sf. sm..
The contributions of this paper are as follows:
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
specification language for describing synchronous

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
and asynchronous pipeline parallelism that can simu- late broad types of pipeline schedules.
2.2. Model and Tensor Parallelism
Any partition of model to run whole operations in separate devices is called model parallelism.
2. An efficient algorithm that finds good pipeline sched- ules automatically from data-flow graphs with thou- sands of operations.
Tensor parallelism (TP) splits individual operations and distributes the operations with sliced input or weights to multiple devices (Bian et al., 2021; Xu et al., 2021a; Wang et al., 2023a; Sousa et al., 2023; Singh et al., 2023). can distribute any kind of data but requires per-operation gather/scatter communication and it is hard automate.
3. Investigation into the optimal workload distribution conditions. We show that unbalanced schedules with prefetching are often the optimal solution, contrary to the popular belief that work must be distributed evenly across devices.
Mixture-of-experts (MoE) selectively activates subset of layers per input (Shazeer et al., 2017: Rajbhandari et al., 2022; Masoudnia & Ebrahimpour, 2014), allowing experts to be placed on separate devices and run in parallel. Large- scale frameworks increasingly combine such techniques to scale to thousands of GPUs (Jiang et al., 2024).
4. Pfeife: our tool that integrates with PyTorch to provide automated and transparent way to pipeline models across multiple devices. Pfeife offers three key func- tionalities: (1) fine-grained operation-level pipeline parallelism, (2) an implementation of the scheduling algorithm that generates schedules in our language, and (3) a runtime to execute the schedules with partial pipeline parallelism and data parallelism.
2.3. Pipeline Parallelism
Pipeline parallelism (this paper) is a particular form of model parallelism. It slices the data-flow graph of a model and executes each sub-graph in different device.
We show that Pfeife is capable of rumning large ML models that do not fit in a single device due to memory capacity constraints. Moreover, we show that Pfeife outperforms state-of-the-art pipelining tools by up to 21% on large mod- els, while being fully automatic.
Pipeline parallelism is an effective technique to run large models by distributing the weights and activations across devices. However, it suffers from having devices idle at times, the so-called pipeline bubbles.
After PipeDream (Narayanan et al., 2019) and GPipe (Huang et al., 2019) introduced the concept of pipeline parallelism for ML models, number of approaches have appeared to overcome some of the drawbacks of pipeline parallelism, including more effective communication/coompuaalie overlap (Chen et al., 2024a).
2. Background and Related Work
There are many techniques to parallelize large ML models. We give a brief overview, with special focus on pipelining.
2.1. Data Parallelism and Weight Sharding
DAPPLE (Fan et al., 2021) uses IFIB schedules to reduce peak memory, despite noting that uneven schedules can yield better performance. AutoPipe (Liu et al., 2022) applies critical path analysis like Pfeife, but only for 1FIB without prefetching and assumes balanced partitioning is optimal. It also partitions at sub-layer granularity, whereas we support per-operation slicing.
Data parallelism (DP) replicates the model and splits input batch into multiple mini-batches. Each device then processes a different slice of the input. Since only the size of the activations depends on the size of the mini-batch, DP reduces the size of the activations but not of the weights.
Sometimes single weight does not fit in a single device/- core's memory. For that case, weight sharding spreads the weights across devices and gathers them on demand (Rajb- handari et al., 2020: Wang et al., 2023b; Xu et al., 202 1b:b: Zhao et al., 2023). Sharding reduces the number of dupli- cate copies of weights and enables sharing parts of weights across parallel instances, but it incurs in additional commu- nication.
There are other schedules that reduce pipeline bubbles. For example, the looped pipeline from TeraPipe (Li et al., 2021), BFS (Lamy-Poine) Poirier, 2023), and AIAO (Li et al., 2023a) all reduce the initial bubble and can hide gradient synchroniza- tion latency in data parallelism. The kFkB schedule (Wang et al., 2023c) aims at hiding communication latency with a cluster of k forward passes, and Chimera (Li & Hoefler, 2021) uses bi-directional pipelines to increase utilization.
Data parallelism is the simplest and easiest way to par- allelize a model, but it requires full synchronization (i.e., all-reduce) of gradients between devices since it replicates the weights. Its communication workload is linear in the size of weights, and thus requires a high-bandwidth network between devices.
In terms of implementations, PiPPy (Reed et al., 2022) uses the TorchFX tracer to generate graph, but users must man- ually specify cutpoints. Slapo (Chen et al., 2024b) provides more flexible scheduling language. Varuna (Athlur et al.,
2

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
PyTorch Module
Compile Training Loop
new, sodd tarch, complia (nod,
pipe curnez = Pipultunoe:
relu(e) self. 1in2
inpet dataset:
return
exec_ [Amput]
Ieww and Lngutt
Replace the Original Module
Python Function Frame
pape
Benchmark
Complled Module
CALL METROD
Time & Memory
Pipeline Graph
Pipeline Schedule
Workers
Static FX Graph
G GPU
Pfeife
TorchDynamo
Compiler
Figure 1: Overview of the architecture of Pfeife.
2022) runs on vanilla PyTorch but requires users to manually inject no-op cutpoint layers.
flow graph (DFG) of a PyTorch model from its function frame through symbolic execution at rum time. It also allows intercept the execution of models and replace the original model with parallel version of them.
Hybrid Parallelism It is possible to combine multiple kinds of parallelism. 3D parallelism combines DP, TP, and PP for the same model.
Pfeife is implemented as a backend for torch. comp le. After getting the DFG, Pfeife distributes the operations in the graph through the available devices. It then builds a pipeline schedule so that maximizes throughput while still fitting the memory capacity of each device.
Megatron-LM (Narayanan et al., 2021) trained transformer- like models using 3D-parallelism with the looped 1F1B pipeline. DeepSpeed (Rajbhandari et al., 2020) supports 3D-parallelism with ZeRO-DP. However, their support for TP and PP requires manual work.
Once schedule is generated, Pfeife replaces the original model. Since Pfeife captures and parallelizes the actual execution of PyTorch APIs, the tool works with any user- defined model. Execution is parallelized transparently to the user, except that will execute a large model that would otherwise go out of memory.
Colossal-AT (Li et al., 2023b) and Merak (Lai et al., 2023) support automatic 31D-paraleeissm but they are limited to a few selected transformer models of HuggingFace (Wolf
al., 2020).
The main training loop is expressed a single closure exec_fn. From the closure, the user can mix pipeline module with any kind of model, such as pretrained frozen models before or after the main module. P PipeRunner will call the closure with a set of matchebatccees, catch the in- put entered into the pipelined model new.mod, and execute pipeline schedule through the workers. This design al- lows pipeline parallelism for general models and use cases, whereas existing frameworks usually require a fully lin- earized forward pass from input to loss calculation. The appendix contains more details about the implementation.
Alpa (Zheng al., 2022) and Sagemaker (Karakus al., 2021) use several search strategies to optimize the perfor- mance of 3D- 3D-parallelism, but users have to rewrite the model using their libraries, or compile the Py PoToch module with TorchSeript (DeVito, 2022).
Recently, PyTorch revealed TorchTitan (Liang al., 2024), which enables 3D parallelism in PyTorch, but limited to train LLMs, especially Transformer models and requires manual parallelization.
There are also other parallelism tools (Jia et al., 2019; Xu et al., 2021b) that do not perfectly fit into this classification.
3.2. Problem Setting
Graph Construction A core part of pipeline parallelism is how we distribute the set of operations of a model through multiple devices. Since a device can only execute one func- tion at a time, we must know the order in which the functions have to be run on each device, and what results have to be sent the following device. That order and the data de- pendencies can be expressed computation graph. A computation graph is expressed as TorchFX graph generated by TorchDynamo.
3. Pfeife
3.1. Architecture Overview
Figure shows the architecture of Pfeife. ML models are written in plain PyTorch. They are then compiled using PyTrrch 2's torch. compi le (Ansel et al., 2024), as it is now common.
graph, which call
TorchDynamo is PyTorch 2's new tracing JIT compiler in torch. compilollleooooo.... le. TorchDynamo captures static data-
3

---

Pfcife: Automatic Pipeline Parallelism for Py PyTorch
10
10
20
20
Y2
V/
1
30
1
20
F122
912
C122
15
15
15
30
Out
Figure 2: Example dependency graph with two mini-batches and two nodes. The computation time of on and d2 is, respectively. 15 and 10, communication time is 1, and backward nodes take twice the forward time. The red line is the critical path, and blue circles indicate the longest distance from the source node (In).
Computation graphs of large ML models can have thousands of nodes. To reduce the search space, we group several vertices together and treat them single node in reduced graph, referred to the fused graph G = (V, Producing good fused graph is key to balance scheduling time and model running time. We call this process slicing graph, and we delve into this in Section 3.4.
graph of all the Ops by interpreting the orderings as edges. Figure 2 shows an example of that graph, which we refer to as the operation dependency graph, or dependency graph.
For simplicity, we add dependencies based on the mini- batch index without loss of generality. Also, we assume that communication of output values is executed right after the computation is finished. Existence of a dependency between two communication nodes (2),))) is governed by the number of the communication channel. We assume that there is single channel (e.g., PCIe NVLink) between two devices. The order between e and e2 is decided by the scheduling algorithm, given the network topology.
Scheduling of Pipeline Parallelism Pipeline parallelism involves executing multiple forward and backward passes across devices according to a predefined schedule. We as- sume that each device can only run a single function (i.e., a vertex in a fused graph) on a single mini-batch at a time. We now formalize the notions of a schedule and the associated operations.
Some implementations of pipeline parallelism such as torchgpipe (Kim et al., 2020) strictly divide the computation and communication phases. By assigning more dependen- cies, we can simulate synchronization operations between two devices. Considering the graph of Figure 2, if we explic- itly set cross dependencies like e u and u + e then these work like a synchronization barrier which forces communication to be executed only after all the computation is finished.
Definition 3.1 (Schedule). A Schedule S = R) defines the order of execution of a fused graph G = E), Let device operations be P: D + P(Op), where is a set of devices and P(Op) is the set of operations to run on each device. Operation dependency
P(Op x Op) is the set of dependencies (partial order) between operations.
P denotes the set of computation and communication oper- ations assigned to each device and mini-batch. represents the set of dependencies between these operations. These dependencies are either automatically derived from the struc- ture of the fused graph, or manually specified by the user for, e.g., synchronization or reducing memory usage.
3.3. Cost Model
To achieve our goal of reducing the total execution time of models, we need a cost model to estimate the running time and the peak memory usage of a given graph slicing and schedule, and use it as the objective for our optimization algorithm.
The set of operations, Op, consists of two primitive types: (1) forward or backward computation of the i-th vertex v, with the n-th mini-batch, denoted as oşt = n, F/B), where oz represents forward computation and v, repre- sents backward computation; and (2) communication be- tween two nodes vith and - connected by an edge e,, +n for the n-th mini-batch, denoted as cit (ci,), n, P/B).
Following the strategy of DAPPLE (Fan et al., 2021), we use profiler to measure the computation time and memory consumption of each operation in a computation graph be- fore pipelining the model. We then use a planner to search for the optimal graph slicing and execution order using the results of the profiler.
By combining fused graph and schedule, we simulate most scheduling strategies for pipeline parallelism described in the literature. We give some examples in Section 3.4.
Unlike previous pipeline plamners, cost model is based on critical path analysis and thus is not dependent on any specific pipeline schedule nor needs an in in-depth analysis of the schedule to build a heuristic. It is also agnostic to the structure of the model or its graph. Because of this generality, our planner found several interesting pipelines
Operation Dependency Graph Since R is a partial order of operations in schedule, we can derive a dependency
4

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
7
S
6
11 m4
2
| Figure 3: A looped IFIB schedule with parameters = B, 8, B -, device. The x-axis is time. Light-colored (resp. shaded) boxes represent the forward (resp. backward) pass. Red/green colors indicate each of the halves of each pass.
00: (2, 2. Each row represents
schedules that are faster than ffte-of-the-ar frameworks.
DAG.
For the communication time, we run set of micro- benchmarks offline to measure throughput and latency of sending data between all pairs of devices. This information then used by the planner to estimate the communication time based on source/destination devices and data size.
Selecting a particular topological order when the compu- tation graph contains parallel branches is obviously an ap- proximation and may not yield the optimal running time. However, most current models, such as Transformer-ban models or CNNs, have striet total order between their op- erations or only few negligible parallel operations. There- fore, opted to use the order of operations generated by the TorchDynamo compiler. We show later that this approx- imation yields good results in practice.
Running Time Estimation We can interpret dependency graphs flow graphs with source and sink nodes. We want to know the earliest possible starting time of the sink node. To achieve this, we have to compute the earliest possible starting time of each of the individual nodes.
The remaining job consists in slicing the linearized sequence of nodes, i.e., selecting split points. The detailed algorithm is described in the next section.
Nodes can have multiple parents. For example, in Figure 2, us has two parents: 1/2 and +3. The child node's possi- ble starting time (31) is the maximum of the parent's ending time (start execution times). This be viewed as vertex weight (starting time, blue circles) + edge weight (execution time). The running time of operations is obtained by running each operation at a time on the target device and measuring the time it takes.
3.4. Pipeline Parameters and Optimization Algorithms
We now describe how Pfeife generates fused graphs and schedules.
Looped Training Schedule Any ML model has many valid pipelining schedules. Trying to find the optimal sched- ule run time without any constraint is not viable current solving techniques do not scale. Inspired by previous work, we reduce the search space to finding the following parame- ters for the looped I FFB schedule (Narayanan et al., 2021):
Since the dependency graph DAG, we can use, e.g., Dijkstra's shortest path algorithm with negated weights to find the longest path, which corresponds to the estimated run time of the model. By finding the longest path from the source to the sink node, we can find the critical path and the estimated running time of a single iteration of the schedule.
(B) Total batch count: Number of mini batches
(N: Loop count: How many times the forward loop is executed.
Linearization and Graph Slicing by Split Points Slicing a computation graph is one of the most important tasks when generating a parallel schedule. The issue is that slicing an the min-cut simplify
(Bi) Loop batch count: How many mini-batches go through the forward pass single stage.
arbitrary computation graph is equivalent problem, and thus it is NP-hard. However, the decision of where to split the model by linearizing the graph and finding a set of split points in the linearized graph.
(B,) Prefetch batch count: A list with the number of forward passes each device runs in addition to B, before it runs its first backward pass. [Byl [D].
Linearizing a graph is equal to selecting one of the partial orders of the computation graph. can be easily found through a topological sort since the computation graph is
Obviously, each set of parameters offers a different tradeoff in terms of performance and memory consumption. With
5

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
this setting, we can express the schedules proposed in previ- ous work including BFS (Lamy-Poirier, 2023) and looped- IFIB (Narayanan et al., 2021). See Appendix B) for the tradeoff's and examples.
Algorithm Graph-schedule co-optimization.
1: Input: Computation graph G, batch count B.
2: Output: Fused graph mmins dependency graph Smin and pipeline latency Tmin-
Figure 3 shows an example schedule with 8 mini-batches and devices. We sliced the graph into Ni-[D) nodes and let the forward pass loop the sequence of devices. This makes the length of single forward pass equal 1/M; stages, so it reduces the pipeline bubble in the beginning (i.e., the top-left idle zone in the figure).
3: Tú. Tmin + 00, Xc
4: for Paet + posssible.param(P
B) do
1 (p1,'.. ,P,,',,PP) = init_split_pon popiit_po((]) G)
T, Mp & cale_time timemmemmG, P Pecet: P)
7: K {(T,p)}
8: while K is changing do
9: for (-, (P1, .. ((Pp,,...PPP..... .Pr)) E K do
One of our major findings is discovering the role of prefetch- ing forward passes before the first backward pass specified as By, as we explain next.
10: for 1 S S n1 0 < |k| < 10 do
11:
Psingle + (p1, ... -Ps + K
Pn))
12:
Perop t (pi+ h1, k., .ps 1 k, Pn kn)
13:
for p t {Poingle: Pprop) do
Co-optimization of Slicing and Scheduling Both graph slicing and scheduling matter to reduce latency and peak memory consumption. However, existing planners choose only one part as their optimization target. They usually fix schedule and then find the optimal graph for that, or equally slice the graph first and then find the best schedule for that slicing.
14:
TV, M₂ + cale_ime timemmemmC.... (G, sstt: P) I' TỤ and Mp < mmax (D) then K +- K U ((T",
15:
16:
17:
end if
18:
K > L then
19:
K + select select_K_min K.min Lx)
20:
TB + max((T) (T, ) K)}
Our algorithm optimizes both slicing and scheduling at the same time so that it can find the best pair of scheduling and slicing. We use superset of the schedules proposed
21:
end if
22:
end for
23:
end for
in previous work and find the best graph slicing for each schedule.
24:
end for
25:
end while
26:
T', Pmm + 1--mm((T p) E K})
Algorithm 1 shows the pseudo-code. We first generate a set of schedules and then find split points for each sched- ule based on our cost model (line 4). The set of sched- ules is given as a set of scheduling parameters Pset (B, N2, B,, By)-
27:
if T' < mmi then
28:
T Tmin T
29:
gmin: Smin build graphs(G, Poet (minn)
30: end if
31:
end for
We find the split points for each schedule using beam search. The initial set of split points p has the sum of the weights and activations evenly distributed across devices (line 5). Then, we make a beam set K that has a maximum size Lt, and insert a pair of latency and split points (T, calculated from the computation graph. Next, we try to improve each set of split points using the following steps:
7. Repeat (2)-(6) until the set of split points does not change (line 25).
The number of movements in (3) is round (Pe-Pi- 1),,++)))))
Po if i < and rownd( (pee) PO(() otherwise. For all pairs of schedule and split points, select the pair that has the lowest running time (lines 26-30).
[V -
1. Take out a set of split points in the beam set (line 9).
2. Move one of the points by 1 to 10 nodes (line 10).
The purpose of steps (2) and (3) is to fill the pipeline bubble below the schedule by assigning more operations at the beginning of the pipeline and fewer at the last devices.
3. For each point from (2), move the other points propor- tional to the movement of (p., k) (lines 11-12).
4. Calculate the latency and memory consumption for the split points from (2) and (3) (line 14).
For example. Figure shows the overall steps with the actual split points and structure. From the first to the last schedule, the beam searcher moved the first and then the second points and finally found the optimal split points (6, 11).
5. If the latency is improved and the peak memory does not exceed the memory capacity, insert the new split points in the beam set (lines 15-17).
The running time can be further improved with prefetching as shown in the last schedule. If we focus in the middle of the schedule, we see that it is 33 that is increasing the
6. Select the best LK valid split points in the beam set (lines 18-21).
6

---

Pfeife: Automatic Pipeline Parallelism for Py
15
250
Equal Blicing (5: 5::5 5)
12'
75
800
2229, 787.5m))
Uneven Sticing (7: 4)
Latency (ms)
750
(225, 89.88ms)
73
(236, 639)
800
Uneven Slicing (6: 5 4)
102 Y2?
(238, 629.4ms)
72
(245, 499.7ms)
Prefetching Forward node (6: 5: 4)
(238, 408.7ms)
175
200
225
250
275
Split point (VIT-g/14, batch size 4, prefetch 1)
#batch -expect
#batchh actuall
#batch expect
Figure 4: Latency comparison of four IFIB schedules with 15 ops and devices. The number of operations for each device is shown on the left for each schedule; the latency is shown on the right. Critical paths are shown in red.
#batch 3 actual
.. #bstch 4 expect #batch 4- actual
Figure 5: Expected vs actual latency of the ViT-g/14 model (407 nodes) with devices. Bottom to top: 2-4 mini-batches. Circles and stars indicate the minimum expected and actual running time, respectively.
critical path. Moving it to earlier reduces the critical path.
This schedule shows that uneven slicing changes the critical path and reduces the length of the pipeline bubbles between the forward and backward passes. This is in contrast with the popular belief that even slices are optimal.
and GPUs, so we have a baseline to compare against. We excluded models that do not support training by default, and models with quantized integer weights or that are too small.
We run the models with and without Pfeife and compared the results. As expected, 37 out of 49 models run successfully and produce the exact same output with and without pipelin- ing, including complex models such as pytorch unet, hf.clip, and hf. hff.Ber..large
In our benchmark suite, beam search with 4 batches, 11 split points, and 406 vertices (ViT-g/14 with N₁- 3) completes within 10 seconds. For the largest model (StableDiffuusin- XL, 1645 vertices, 8 devices), with N₁- = 2 (15 split points) it takes less than 5 minutes, and with N₂- = 3 (23 split points) it takes about 10 minutes. In most cases, the whole optimization process finishes within few seconds.
Of the remaining models, 11 cannot be compiled with torch. comp le due issues within PyTorch itself, and one model (t titimmeff www.efffiiii.. wwwfffiienteeeeee sults. Appendix C provides further details. These results show that Pfeife can successfully run a broad range of state- of-the-art ML models.
produces incorrect re-
4. Evaluation
We evaluate Pfeife in three ways: (1) applicability of the approach, (2) accuracy of cost estimations, and (3) end-to- end performance comparison with existing frameworks.
4.2. Cost Estimations and Effect of Scheduling Parameters and Split Points
We used two servers for the experiments. For coverage and correctness, we used small server with 8x NVIDIA RTX 3090 24 GiB GPUs with 4 NVLink connections. For the end-to-end experiments, we used larger server with 8x A100 40GB GPUs with NVSwitch. The main reason for splitting the work across two machines was because we did not have sufficient compute time available for validation in the A 100 cluster.
We now study how the scheduling parameters affect the per- formance and peak memory consumption. See Appendix for more details.
Table 5 summarizes the results. We used two models: Vi- sion Transformer (ViT-g/14) (Zhai et al., 2022) and GPT2- large (Radford et al., 2019) with 512 tokens per batch. We observe that Pfeife safely overestimates the peak memory usage for every device (to avoid unexpected OOMs). Also, the estimated value is at most 10% higher than the actual.
4.1. ML Model Coverage
To check how applicable Pfeife is, we used Torch- Bench (Hao et al., 2023), which is the official PyTorch benchmark suite. It includes wide range of models.
Figure 5 presents a study on how the choice of split points affects the running time. Two observations are noteworthy: (1) the predicted and optimal split points yield similar ex- ecution times, demonstrating the effectiveness of Pfeife's 's heuristics, and (2) the optimal split is not at the midpoint. The model has 407 nodes, while the best split occurs be- tween nodes 229-245, indicating that the first device handles
We selected 49 trainable models that allow customization of the mini-batch size and that can be run both on CPUs
Available at https: : //github./ com/ cwm/ htt.////SSSSSSfffffeee
7

---

Pieife:
Automatic Pipeline Parallelism for PyTorch
Expected and Actual Latencies (VIT-g/14, devices, 4 mini-batches)
Table 1: Throughput comparison of pipeline parallelism
1150
Distribute weight evenly
(item/s). [D]: number of devices. B: number of micro-
1100
Distribute time evenly Pfeife
batches. C-Al: Colossal-AT. Higher is better. (* 7eRO
Latency (ms)
3)
1050
1000
Model
[D], B
DeepSpeed
C-AI
Pfeife
VIT
2, 2
72.0
70.9
84.8 (+18%)
950
VIT
2, 4
86.5
86.1
92.8 (+7.3%%
900
VIT
4, 4
124.5
122.6
151.4 (+22%)
1/4/0 1/4/1
2/2/0 2/2/1
2/4/0 2/4/1
Scheduling parameters (loop count loop batch prefetch)
VIT
4, 8
157.7
139.7
161.5 (+2.4%)
VIT
8,
230.9
197.0
265.7 (+15%)
VIT
8,
297.7
216.7
282.2 (-5.5%)
Expected and Actual Latendies (MT-g//4,, devices, mini batches)
LLama
8
0OM(*)
6.80
7.46 (+10%)
weight evenly Distribute time evenly
LLama
8, 16
0OM(*) 8.69
8.81 (+1.4%)
Latency (mg)
1158-
1120-
Table 2:
Throughput comparison of SDXL. (item/s). Batch count is equal to the number of devices (except FP16 x2, which uses twice more batches). DP, PP: number of devices in the groups of DP and PP.
1058
450
1ib1s 21470 Scheduling parameters (loap count batch prefetch)
Method
DP
PP
FP16
FP16 x2
FP32
Figure 6: Expected (red points) actual (bars) latency of ViT-g/14 4 with 2 and 4 devices by changing the scheduling parameters and split points.
ZeRO2
3.78
OOM
OOM
ZeRO3
2
0.81
0.94
0OM
Pfeife
1
3.72
4.74
2.38
ZeRO2
4
9.12
10.8
5.32
ZeRO3
4
2.01
1.84
2.01
larger portion the computation.
Pfeife
1
7.12
9.32
3.98
Pfeife
2
7.06
8.67
4.54
Figure shows the per-device running time of sched- ules with varying loop count (L,), loop batch count (Bi), prefetch batch count (By), and split points. For all schedul- ing parameters, the actual latency follows the expected la- tency with a small (30-40ms) difference at most (I/O over- head). Therefore, Pfeife can select the optimal scheduling parameters by simply choosing the parameters which show the minimal expected latency.
ZeRO2
8
19.8
19.2
11.2
ZeRO3
8
3.82
3.35
3.22
Pfeife
1
12.1
15.4
6.09
Pfeife
2
13.9
15.9
7.32
Pfeife
13.8
17.2
9.02
None of the state-of-the-art tools offers fully automatic pipeline parallelism, unlike Pfeife. Colossal-Al supports just few selected models and DeepSpeed requires hand- written, linearized sequence of layers, which makes it hard to use with models that reuse a layer (e.g., transformers).
We also compare three slicing strategies: distribute weights uniformly (DeepSpeed), distribute operations by running time uniformly (Alpa), and Pfeife's uneven slicing, For all cases with prefetching, Pfeife's uneven slicing outperforms the other strategies. We find that prefetching necessary to take advantage of uneven slicing.
Sequential Models Table 1 shows the results. Pfeife out- performs the state-of-the-att tools in all but one case. The odd lines VIT have higher bubble/compptation ratio, and therefore there are more opportunities for optimizing the schedules and thus the performance gap is wider.
4.3. End-to-End Comparison
For the end-to-end experiments, we compared Pfeife with two semi-automatic PyTorch parallelization frameworks: DeepSpeed (Rajbhandari al., 2020), and Colossal-AI (Li et al., 2023b). We used HuggingFace models for bench- marking, namely a vision model (ViT-g/14) (Zhai et al., 2022), an LLM (Llama2-7B) (Touvron et al., 2023), and a diffusion model (Stable Diffusion-XL -XL) (Podell et al., 2023) as an example of sequentill model. Note that the weights of Llama and SDXL. do not fit in single GPU.
Pfeife outperforms state-of-the-art tools mainly for two rea- sons: prefetching and looped schedules. Figure shows the profiles of 4 GPUs training the ViT-g/14 model with Colossal-Al and Pfeife with and without loops and prefetch. When we use a naive IFIB schedule (top two profiles), there is limited opportunity to exploit uneven slicing due to the fixed critical path. By prefetching some forward passes,
8

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
Colossal-Al
Pfeife
Loog:
Pfelfe
Logp:
Prefetch:
Pfeife
Loop:
Figure 7: Profiling of ViT training with 4 GPUs and B Blue, red, and green blocks: forward. backward, and optimizer steps.
additional idle time (blank spaces) appears between the backward passes on the first device. Pfeife extends the for- ward pass on the first device to fill these idle slots, resulting in an 8% speedup. When combined with looped schedule, the idle time at the beginning further reduced (by 9%), bringing the total latency reduction 16%.
significantly larger models that cannot be trained with other frameworks while maintaining comparable performance.
5. Conclusion
In this paper we presented Pfeife, the first tool that inte- grates transparently with PyTorch to provide fully automatic pipelining. It works by (1) capturing the whole data-flow graph of MI models through Py PyTorch's new JIT compiler, (2) distribute operations across devices and schedule the input data into mini-batches to establish the order in which they will travel the pipeline, and (3) intercept the execution of models and run their pipelined version instead.
Since Transformer-basee models share embedding weights at both input and output layers, building pipeline by hand with DeepSpeed is not straightforward. Our attempt to use ZeRO-DP resulted in out-of-memory errors, showing the limitations of data parallelism for models of this size. Pfeife, in contrast, successfully trained L.Lama and achieved up to 10% higher throughput than Colossal-AI. This demonstrates Pfeife's ability to handle large models more efficiently than existing parallelism approaches.
We show that Pfeife runs large models that do not fit in single GPU transparently, thus saving significant amount of manual work. Our fine-grain ine-grained, general pipeline opti- mizer enables Pfeife to train models up to 22% faster than state-of-the-art tools.
The only case in which Pfeife underperforms relative to an- other framework with the ViT model using 8 devices and a batch size of 16, where DeepSpeed faster. As the batch size increases and pipeline bubbles consequently decrease, performance increasingly reflects DeepSpeed's maturity and extensive engineering optimizations. However, the number of mini-batches per iteration (i.e., gradient accumulation steps) is inherently limited by model accuracy considera- tions and interactions with other parallelism strategies.
Acknowledgements
This work supported in part by national funds through FCT, Fundação para a Ciência Tecnologia, under project UIDB/50021/2020 (DOI: 10.54499/UIDB/5002 1/2020), and a National Research Foundation of Korea (NRF) grant funded by the Korean Ministry of Science and ICT (MSIT) (Grant No. RS- 2S-2024-00355455)
Non-sequential models & mixing with Data Parallelism Pfeife can train complex models and be combined with other parallelism methods such data parallelism. For example, SDXL uses pre-trained autoencoder in front of the U-Net model. Due to its complex training pass and multiple skip- connections between its front and latter parts, U-Net models are not typically a target for pipeline parallelism.
We acknowledge EuroHPC Joint Undertaking for awarding access to Vega IZUM, Slovenia.
Impact Statement
This work aims to advance the field of machine learning, specifically in the area of large-scale model training. We are not aware of any specific societal or ethical concerns arising from this work, beyond the standard considerations that apply to the development and deployment of machine learning technologies.
We parallelized the trainable U-Net submodule SDXL; Table 2 shows the results. We compared Pfeife with Deep- Speed's ZeRO-DP. The results show that Pfeife is compara- ble to ZeRO2. With more batches, Pfeife is sometimes faster than ZeRO2 since the proportion of the pipeline bubble is reduced. Also, Pfeife is 3-4 times faster than ZeRO3 and does not raise OOM in FP32.
These results demonstrate that Pfeife enables training of
9

---

Pfcife: Automatic Pipeline Parallelism for PyTorch
References
Jiang, Z., Lin, H., Zhong, Y., Huang, Q. Chen, Y., Zhang,
Z., Peng, Y., Li, X., Xie, C.,, Nong, S., Jia, Y., He, S., Chen, H., Bai, Z., Hou, Q. Yan, S., Zhou, D., Sheng, Y., Jiang, Z., Xu, H., Wei, H., Zhang. Z., Nie. P., Zou, I.,, Zhao, S., Xiang, I.., Liu, Z.,, Li, Z., Jia, X., Ye, J., Jin, X., and Liu, X. MegaScale: scaling large language model training to more than 10,000 GPUs. In NSDI, 2024. URL ht /www usenix wwg/ Files/ #.les/nsdi22 24 24 wwwaannnn-iiggggeee pdf.
Ansel, J. et al. PyTorch 2: Faster machine learning through dynamic Python bytecode transformation and graph com- pilation. In ASPLOS, 2024. 3640366.
doi: 10.1 10.1145/360065.
Athlur, S., Saran, N., Sivathanu, M., Ramjee, R., and Kwa- tra, N. Varuna: Scalable, low-cost training of mas- sive deep learning models. In EuroSys, 2022. doi: 10.1 (0.1145349232233331953
Karakus, C., Huilgol, R., Wu, F., Subramanian, A., Daniel, C., Cavdar, D., Xu, T., Chen, H., Rahnama, A., and Quin- tela, L. Amazon SageMaker model parallelism: general and flexible framework for large model training, 2021. URL https: / arxiv. orgg /abs/ / /21//22111111 05972.
Bian, Z., Xu, Q. Wang, B., and You, Maximiz- ing parallelism in distributed training for huge neural networks, 2021. URL https: /tt//axiivvv org/ abs/ 2105. 14450.
Chen, C., Li, X., Zhu, Q., Duan, J., Sun, P., Zhang, X., and Yang, C. Centauri: Enabling efficient scheduling overlap in large model
Kim, C., Lee, H., Jeong, M., Back, W. Yoon, B., Kim, I., Lim, S., and Kim, S. torchgpipe: On-the-fly pipeline parallelism for training giant models, 2020. URLhttps / //arxiv.. org/ abs / 2004. 09910.
for communication-compieeee training via communication partitioning. In ASPLOS, 2024a. doi: 10.1 145/36200666336117799
Lai, Z., Li, S., Tang, X., Ge, K., Liu, W., Duan, Y., Qiao, L., and Li, D. Merak: An efficient distributed DNN training framework with automated 3D parallelism for giant foundation models. IEEE Transactions on Parallel and Distributed Systems, 34(5); 34(5)) 1466-17886 2023. doi: 10.1
Chen, H., Yu, C. H., Zheng, S., Zhang, Z., Zhang, Z., and Wang, Y. Slapo: A schedule language for progressive optimization of large deep learning model training. In ASPLOS, 2024b. doi: 10.1 00.114//620065536660333
DeVito, Z. TorchScript: Optimized execution of Py-
Torch programs, 2022. URL https://program
Lamy-Poirier, J. Breadth- first pipeline parallelism, 2023, URL https: www.a.i.oooo...... abs 2211.111222111 05953.
transformations wwt.gh.th..oooo htb.io.....i....... pytorch_neeee wwwpcchhneuripp.po
Li, F., Zhao, S., Qing, Y., Chen, X., Guan, X., Wang, S., Zhang, G., and Cui, H. Fold3D: Rethinking and par- allelizing computational and communicational tasks in the training of large DNN models. IEFE Transactions on Parallel and Distributed Systems, 34(5):1432-1449, 2023a. doi: N0.1109/TDDS220233222222
Fan, S., Rong, Y., Meng, C., Cao, Z., Wang, S., Zheng, Z Wu, C., Long, G., Yang, J., Xia, L., Diao, L., Liu, X., and Lin, W. Dapple: A pipelined data parallel approach for training large models. In PPoPP, 2021. doi: 10.1/45/ 3437801.3441595.
Grattafiori, A. et al. The Llama 3 herd of models, 2024.
Li, S. and Hoefler, T. Chimera: Efficiently training large- scale neural networks with bidirectional pipelines. In SC, 2021. doi: 111444//458877334761644
URL https: / //axivv... org/ abs /407 21783.
Hao, Y., Zhao, X., Bao, B., Berard, D., Constable, W., Aziz, A., and Liu, X. TorchBench: Benchmarking PyTorch with high API surface coverage, 2023. URL https: //arxiv iisiviooo www.vii.orr////ooo abs/ 2344 14226.
Li, S., Liu, H., Bian, Z., Fang. J., Huang. H., Liu, Y., Wang, B., and You, Y. Colossal-AI: A unified deep learning system for large-scale parallel training. In ICPP, 2023b. doi: 10.1 141444055533.66655555
Huang, Y., Cheng, Y. Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H. Ngiam, J., Le. Q. V., Wu, Y., and Chen, Efficient training of giant neural networks using pipeline parallelism. In NeurlPS, 2019. doi: 10. 5555/3454287.3444999
Z. Gpipe:
Li, Z., Zhuang, S., Guo. S., Zhuo D., Zhang, H., Song, D., and Stoica, Terapipe: Token-level pipeline paral- lelism for training large-scale language models. ICML, 2021. URL httpt: http::/:s//pppeeeipp//ee wwwiinggsssrr....mm mlr. mir.press/ v139 v139/1121y/11211 pdd.
Jaa, Z., Zaharia, M. and Aiken, A. Beyond data and model parallelism for deep neural networks. In MrSSs, 2019. URL
.mlys......mmm.....m...m..mm.... D422600E1a00099487811182122 Paper.pdf.
Liang, W., Liu, T., Wright, L., Constable, W., Gu, A., Huang, C.-C., Zhang, I., Feng, W., Huang, H., Wang, J., Purandare, S., Nadathur, G., and Idreos, S. TorchTi- tan: One-stop PyTorch native solution for production
10

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
ready LLM pre-training, 2024. URL httpt:ttpp///aee org/abs 2410.. 06511.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and
Ommer, B. High-resolution image synthesis with la- tent diffusion models. In CVPR, 2022. doi: 10.1109/
Liu, W., Lai, Z., Li, S., Duan, Y., Ge, K., and Li, D. Au- topipe: A fast pipeline parallelism approach with bal- anced partitioning and micro-batch slicing. In CLUSTER, 2022. doi: 10.110//LUSTTRR1411.20220..02
CVPRS2688.2022.01000
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q. Hinton, G., and Dean, J. Outrageously large neural net- works: The sparsely- gated mixture-of-experts layer, 2017. URL htw.tttp///// /arxivv.aiiivvv//aa///i
Masoudnia, S. and Ebrahimpour, R. Mixture of experts: literature survey. Artificial Intelligence Review, 42: 275-293, 2014. doi: 10.1007/11446--01--93333
httpppps//aa....oooo/aab//////// http://wwwrrr///////eee/ 06538.........
Singh. S., Ruwase, O., Awan, A. A., Rajbhandari, S., He, Y., and Bhatele, A. A hybrid tensor-expert-datt parallelism approach to optimize mixture-of-experts training, In ICS, 2023. doi: 10. 15/1314799333359933
Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, P. B., and Zaharia, M. PipeDream: Generalized pipeline parallelism for DNN training. In SOSP, 2019, doi: 14.145/3341101.335566
Sousa, R., Pereira, M., Kwon, Y., Kim, T., Jung, N., Kim, C. S., Frank, M., and Araujo, G. Tensor slicing and optimization for multicore NPUs. J. Parallel Distrib. Comput., 175(C):66-79, may 2023. doi: 10.1016/j.jpde. 2022.12.008.
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P. Pat- wary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., Phanishayee, A., and Za- haria, M. Efficient lcrge-scale language model training on GPU clusters using Megatron-L.M. In SC, 2021. doi: 10.1 00.11435558887.3476000
Touvron, H. al. Llama 2: Open foundation and fine- tuned chat models, 2023. URL https://arxiv.
org/ abs 2307 09288.
Wang. B., Xu, Q. Bian, Z., and You, Y. Tesseract: Paral- lelize the tensor parallelism efficiently. In ICPP, 2023a. doi: 00.11455355008.355 35440077.
Peters, M. E., Neumann, M., lyyer, M., Gardner, M., Clark, C., Lee, K., and Zettlemoyer, L. Deep contextualized word representations, 2018. URL https: org/abs/1802.00 053655
Wang, G., Qin, H., Jacobs, S. A., Holmes, C., Rajbhandari, S., Ruwase, O., Yan, F., Yang, L., and He, Y. ZeRO++: Extremely efficient collective communication for giant model training, 2023b. URL URLhttps://exxivvorr abs 0223066.1009
Podell, D., English, Z., Lacey, K.,, Blattmann, A., Dockhorn, T., Müller, J., Penna, J., and Rombach, R. SDXL: Im- proving latent diffusion models for high-resolution image synthesis, 2023. URL httpt:tt///// /wwpiv..ivvvvvv////a//..///aaa org/ / 2307. 01952.
Wang, S., Cao, Z., Si, C., Diao, L., Wang, J., and Lin, W. Ada-Grouper: Accelerating pipeline parallelism in pre- empted network by adaptive group-scheduling for micro- batches, 2023c. URL https: /ttpss///axiv..///////a org/ aba/ 2303. 01655.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., and Sutskever, I. Language models are unsupervised multitask learners, 2019. URL
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C.. Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. I.., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. HuggingFace's transformers: ofte--of-the-an natural language processing, 2020. URL https: /tas:///axiiii www.wxiwwwv......ommmmoo
Languageenoodelss/nngggaa.e.gn saaplaaeeeeemees__a_______ mnsupervisee__mlliittsss..
Rajbhandari, S., Rasley, J., Ruwase, O., and He. Y. ZeRO:
Memory optimizations toward training trillion parameter models. In SC, 2020. doi: 10.1 1109/SC4144..2220000022
1910. 03771.
Xu, Q., Li, S., Gong, C., and You, Y. An efficient 2D method for training super-large deep learning models, 2021a. URL https: /awxvviivvv. wrrg /21 044 05343.
Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He. Y. DeepSpeed- MoE: Advancing mnixurre-of-experr inference and train- ing to power next-generation AI scale, 2022. URL. org/ wwg/abs/////oooooooo..oooeeesrbe abs//2011111 220122211112 05596.
Xu, Y., Lee, H., Chen, D., Hechtman, B., Huang. Y., Joshi, R., Krikun, M., Lepikhin, D.,, Ly, A., Maggioni, M., Pang, R., Shazeer, N., Wang, S., Wang, T., Wu, Y., and Chen, Z. GSPMD: General and scalable parallelization for ML computation graphs, 2021b. URL https: //arx/ar...v org/abs/2105.044 04663......
Reed, J., Belevich, P., and Wen, K. PiPPy: Pipeline par- allelism for PyTorch, 2022. URL https: / //github. com/pytorch/Pp. pPippy..ccppp...

---

Pieife: Automatic Pipeline Parallelism for Py
Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers. In CVPR, 2022. doi: 10.1 1L1109CVVPR2268.222220011
Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., Xu, M., Wright, L., Shojanazeri, H., Ott, M., Shleifer, S., Desmai- son, A., Balioglu, C., Damania, P., Nguyen, B., Chauhan, G., Hao, Y., Mathews, A., and Li, S. PyTorch FSDP: Experiences on scaling fully sharded data parallel, 2023. URL https:////// / /arxiv... org/ abs/2300 2304. 11277.
Zheng, L., Li, Z., Zhang, H. Zhuang, Y., Chen, Z., Huang, Y., Wang. Y., Xu, Y.. Zhuo, D., Xing, E. P., Gonzalez, J. E., and Stoica, I. Alpa: Automating inter- and intra-operator paral- lelism for distributed deep learning. 2022. htttt:s:: //ww. wwwuus.nn...oooo org/system/ siles/osdi222225 2hensi222-heeggiaaaeea pdf.
OSDI,
12

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
A. Implementation Details
A.1. Implementation of Partial Pipeline Parallelism
import terch
from some Some_Library import PreProcessot, KainModel, set_dataset from pfeife import initialize pfeite, seiite_copplery Pipelptions, FipelineRunner
initialize_peeife))
preproc_model PreProcessor () frezen preprocessing module
main_moces MainModel
10 optimizer terch. ortch..ptm.... Adam (main seain_mmdell. pww.ddeeeaaammtersseeo (), 1r=1e-5)
I1 criterion = torch.nn. nn, crossEntropyLoss ()
12 dataset get_dataset
14 options peOptions. from_arge
18 runner PipelineRunne inRunnner (options, optirizer)
16 ma hainnne main_model
corch.compi 1e (ma m backend-pfei fe_compiles ler)
for inputs, labels in dataset:
19
н** Original training loop:
20
with torch.no_gradl,
inputs preproc_model (inputs) outputs medel (inputs)
loss criterion (outputs, labels)
24
loss. (asswaackwrrd()
25
optimizer.stepll
26
print (f"loss: (Loss.item())
Do the same with a closure with inputs and labela 28
29
30
def fer_fniinputs, labels): with torch. norch.no_grad():: inputs preproc model (inputs)
33
34
The execution of the model is automatically pipelined outputs = main_model (inputs) loss criterion (outputs, labels) loss backwarwd))
35
37
38
print (f"loss: (loss.item() F") allows side-effects
39
return loss # returned to runner.step ()
Set the target of pipelining runner, runner.set_ene www.
42
43
44
Execute the pipeline by slicing 'inputs' and 'labels' into N micro-bbtcchs losses runner. step (inputs, labels) loss sum ([loss item for loss losses])
46
Listing 1: Example training loop using Pfeife.
Listing shows an example of the full code required to train a model with Pfeife. It demonstrates Pfeife's strength: easy integration and ability to mix multiple modules, some of which may be frozen. As we can see from the example, pipelining a training loop using Pfeife requires changing less than dozen lines.
Before training, compile the main module to be pipelined with Pfeife (line 16). The Pfeife runtime inserts synchroniza- tion point in the forward function of the compiled module and attaches a backward hook to the result of the forward function. When input tensors or gradients flow to the pipelined module, Pfeife automatically sends the tensors to the pipelined devices and workers.
The core part Pfeife involves enclosing forward to backward passes in a training loop within closure (lines 29-39). We must enclose the training loop with closure for two reasons: first, we want to execute the reeprrcessing part with sliced
13

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
(B,, By) 1001.8ms (opt.)
3 00
3
18
3
6
O
15
0
1
0
O
4
B
2
B
R
5
3
6
0
1
2
3
0
1
2
2
H
5
6
D
1
3
4
(B,, B-) (8. 2) 975 (no opt.
0
1
2
3 4
8
3
4
R
2
6
8
4
3
4
(B,, By) 2) 924.3ms (opt.)
O
2 4
3
0
2 3
5 6
0
0 3 5
3 4
Figure 8: Three schedules of ViT-g/14 training with different Bd and By. The second schedule uses uniform slicing in terms of the forward pass running time, while the other two are optimized by beam search.
micro-batch rather than full mini-batch. Second, we want to control the timing of the execution of the pre-processor so that we can run the pre- pre-processing part on the main device concurrently with the pipelined parts on the other devices.
B. Scheduling Parameters
B.1. Tradeoffs about Scheduling Parameters
Pfeife restricts the search space of our scheduling algorithm with combination of scheduling parameters that simulates looped- IFIB schedule. The tradeoff of each parameter is as follows:
(B) Total batch count: Enhances throughput per training iteration if increased, but can harm the performance and convergence if the batch size is too large.
(N₂) Loop count: Reduces the pipeline bubble the beginning, but requires higher activation memory in the last device.
(B,) Loop batch count: Communication latency can be hidden if increased, but uses significantly higher activation memory.
(By) Prefetch batch count: Reduces the pipeline bubble between forward and backward passes, but requires higher activation memory in the first device.
For example, the schedule of Figure 3 suffers from high memory consumption due to storing many activations. The device at the bottom accumulates all the activations of 4 mini-batches in the first part the loop. To mitigate this problem, we can reduce the number of mini-batches B, which are fed a single loop. To fully utilize all the devices, B, should be divisor of B and Bi 2/1 (D].
We now emphasize the role of prefetching the forward passes before the first backward pass. Each device must run (N, - 1) B, + forward passes before executing the first backward pass where i is the index of the device (indicated with the red vertical line). This configuration has the lowest peak memory usage for the looped IFIB schedule.
By specifies that each device should execute a few more forward passes before the first backward pass. There are two positive effects of prefetching: (1) hide communication latency, and (2) fill in otherwise empty slots (the pipeline bubble). For example, observe the nodes of and oJ" (the two light red "7" blocks of the two bottom devices). If we execute vo" before %0- (the bold green "1" block at the bottom), can hide the latency between ut? and o. which passes the result of the first device to the second device.
Figure& 8 shows a concrete example with the actual timings of ViT-g/14, Through repeated local optimization with prefetching by beam search, we effectively fill the pipeline bubbles between forward and backward passes with a looped schedule. This level of optimization is achievable only with a fine- fine-grainn aoneggraied, operation-level pipeline. The subtle differences between forward passes in the second and third schedules demonstrate this precision.
B.2. Examples of Pipeline Schedules
To show the expressiveness of this setting, we summarize the schedules proposed in previous work with parameters (B, Nr, Bi, By) and assuming 4 devices:
14

---

Pieife: Automatic Pipeline Parallelism for PyTorch
GPipe (Huang al., 2019), Figure 2: (8, 8, (4, 6, 7))
Pripeiipe (Narayanan et al., 2019), Figure (7, 1, 8, 1.
IFIB-DAPLLE (Fan al., 2021), Figure 3: (6, 1, 6, (0, 0, 0))
TeraPipe (Li et al., 2021), Figure 2: (8, 2, 4, (0, 1, 2))
Mcgatron-LM2 (Looped-1FIl) I (Narayanan et al., 2021), Figure 4: 28, 4, 1,
BFS Pipeline (laamy-Porrier, 2023), Figure 4: (8, 4, 4, (4, 5, 6. 7))
We note that there are some scheduling strategies that cannot be expressed in this framework, namely kFkB (Wang et al., 2023c) and Chimera (Li Hoefler, 2021). However, these strategies double number of activations or weights, and thus they require lot of memory, meaning they are not useful for the current memory- heavy models.
Nevertheless, optimization algorithm is specifie any pipelining strategy. These strategies be expressed with scheduling notation. If user wants consider those strategies, they just need to specify the strategy in our notation and give it to the optimizer.
C. Coverage Test Details
Table 3: Program coverage of TorchBench.
Total
Success
Multiple graphs
Model sharing error
Dynamo error
Incorrect result
49
37
5
2
Our coverage test compares the result (i.e., loss) between pipelined model and the vanilla model. Also, we execute 5 training loops and compare their results after the parameters are trained.
Table summarizes results of the coverage test. As we allow graph breaks in torch, torch.compil one of the model (hf_Whi (hf.Whisper) is trained correctly even it returns subgraphs while its compilation.
However, 5 models result in tens of graphs so they raise an error from their side-effect handlers in the forward passes of the compiled module. Note that side-effect handlers connects the sub-graphs when there are graph breaks.
Four models raise error while sharding sliced subgraphs to worker processes. Since we support multi-node training, we need to send a sliced Py Torch GraphModule through process pipes. However, their subgraphs cannot currently be marshaled by Python, preventing multi-process training.
Two models fail to compile with torch. compile. Since Pfeife depends on torch. compil le to extract a static TorchFX graph, those models cannot be target of pipeline parallelism until the Py Torcch development team add supports for those models.
The last column refers to the tim.effficinttnee benchmark. The cause of the error is that some parameters not captured in the generated TorchFX GraphModule. This seems to be an engineering problem, and fix is expected soon.
The complete list of the coverage test is as follows:
Success: alexnet, BERT pytorch, dcgan, densenet 21, functorch dp ci cifarl.. 0, hf Albert, hf_Bert, hf.Bert.large hficlip, hf.DistilBert, hf.Roberta.base, hf_GPT2, hf.GPT2.large, hf Whisper, LearningToPPannt, dennard_jones, 1lama, mnasnet1.0, mobilenet. v2, mobilenet.v3.large, phlippe.densenet, phlippe-resnet, pytorch.unet, resnet152, resnet18, resnet50, resnext55.32x4d. shufflenet.v2.x1100 squeezenet1.1, timm-nfnet. timm.regnet, timm.resnest, cim.vissio..rrassfomom. timm vovnet, torch multinodal clip, vggl6.
Multiple graphs: hf BigBi rd, hf Longformer, hf Reformer, hf T5 generate, opacus ci fir1000,
Model sharing emor: hf.75, hf_r5_base,, hf_55_large, d1r,,
Dynamo error: cm3leon generate, hf Bart
Incorrect result: timm efficientnet.
15

---

Pfeife: Automatic Pipeline Parallelism for PyTorch
D. End-to-end Model Details
Table 4: Models used for evaluation. Train memory (GiB) = weights + gradients optimizer state. Total Estimated total
memory usage when training with a single mini-batch.
Model
Batch size
Weights
Activations
Train
Total
ViT-g/14
24
3.1
21.1
12.3
33.4
LLama2-7B
(1024 tokens)
25.7
15.3
103
118
StableDiffusion-XL XL (FP16)
(1024px)
5.1
8.4
41.1
49.5
StableDDiffusion-Xi. (FP32)
(1024px)
10.3
13.2
41.1
54.3
Table 4 shows the breakdown of memory usage in our end-to-end benchmarks. We reused the initial mini-batch through the whole training iterations to measure the performance of parallelization only.
ViT-g/14 (Zhai al., 2022) is vision model that we use as an example of computation-intet -nnttnio-- workload. Llama2 (Tou- al., 2023), large language model, represents a memory- ceinory-intensivee workload. Stable Diffusion XI. (Podell et al., 2023) is in between the two, exhibiting both significant computational demands and complex computation graph.
We exclude the GPTe-Large (Radford et al., 2019) model from our end-to-end training experiments due to its relatively small size, which results in suboptimal utilization the A100 GPU.
E. Estimated vs Actual Run Time and Memory Consumption
Table 5: Accuracy of the run time and memory consumption estimations with mini-batches. The columns are: [D]: number of devices (and of the first and last devices (GiB); expected/actual running time (ms). (*): Rejected by the scheduler.
Lg: size of mini-batch; Nr, B₁, By: scheduling parameters; expected/actual peak memory
Model
D
LB
Ni
B
Bj
Est. peak mem.
Act. peak
Est. time
Act. time
VIT
2
1
20.5 / 11.4
20.2 / 11.2
1899
1960
VIT
2
2
4
19.1 14.8
18.8 / 14.6
1808
1891
ViT
2
4
19.0 16.6
18.3 15.9
1776
1856
ViT
17.4 6.21
17.1 5.96
1121
1149
VIT
4
2
8
18.5 14.6
17.6 13.2
966
1002
VIT
4
3
8
20.6 18.7
19.2 / 16.8
936
974
VIT
8
1
11.6//333
10.8 / 3.05
760
768
VIT
8
2
8
13.5/ 8.27
11.7 7.37
621
683
VIT
8
3
8
13,8 10.3
12.0/9.10
591
697 (*)
GPT2
2
1
16.4/114
16.9 11.6
1230
1283
GPT2
2
2
19.6 13.9
20.1 13.8
1142
1190
GPT2
19.4 5.25
19.9 5.45
699
725
GPT2
4
8
20.5 15.3
20.4 14.6
622
726
GPT2
8
1
13.0 4.07
12.5 / 3.96
486
661
GPT2
8
2 8
12.2
9.57 5.32
420
744 (*)
Table 5 summarizes the results of the cost model accuracy for two models, ViT-g/14 14 and GPT2-large, with 512 tokens batch. We run these models with a varying number devices and batch sizes. For each setting, we let Pfeife find the best scheduling parameters, and record actual peak memory and running time.
Regarding loop count (Nj), we can see in Table that the performance increases with more loops, as reduces the pipeline bubble. However, more loops also mean increased communication latency and also smaller kernels and thus higher overhead.
The estimated running time is very accurate, modulo a constant 30-60 ms warm-up time to distribute the mini-batches. However, when the slicing is too fine-grained (less than 10 ms forward time per node), each node finishes too quickly, and thus the Python threads cannot queue CUDA kermels fast enough. In this case, the overhead between the threads and CUDA dominates the running time, hence we made the optimizer reject schedules with an estimated running time the forward pass less than ms (marked with *).
16


Reviews:
[{"review_id": "QLrK5PkXrL", "reviewer": "Reviewer_aSA1", "summary": "This paper presents Pfeife, a tool that automatically performs pipeline parallelization of PyTorch models. Compared to prior methods, the main innovation is that the pipelining is performed in a manner that is completely transparent to the developer, requiring no manual annotations. Specifically, Pfeife is implemented as a backend to $torch.compile$. Starting from a computation graph traced by the PyTorch compiler, Pfeife uses an estimated cost model tuned to the specific hardware, and performs a co-optimizing search that alternates between fusion (\"slicing\") and pipeline scheduling, aiming to minimize the critical path. Experiments with an 8xA100 (40GB) cluster shows that Pfeife exhibit improvements over DeepSpeed and ColossalAI tools on Llama2-7b and is comparable on SDXL, despite requiring no manual effort.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "This paper makes two main claims\n1) Pfeife performs pipelining with no manual effort. This is generally true: the schedules are generated without any manual efforts, and the schedules are generally (though not always) correct from the coverage tests.\n2) Pfeife discovers pipeline schedules that are as good as semi-automated pipeline tools. This is also supported by the experimental results, albeit for relatively \"small\" models (Llama2-7b and SDXL)", "methods_and_evaluation_criteria": "The proposed methods and evaluation both make sense.\n\nIt would have been nice to evaluate performance on even larger models / diverse hardware topologies (where pipelining is even more necessary and tricky to get right), but I recognize that this is out of the reach for the majority of researchers.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The way the experiments are properly designed (both the correctness and performance) and the analysis all make sense.\n\nHowever, I would have liked to see at least one hand-tuned pipeline as a baseline for comparison.", "relation_to_broader_scientific_literature": "This paper focuses specifically on pipeline parallelism for multi-GPU workloads. Compared to existing pipeline tools, the main novelty is that the pipelining is done in a completely automatic fashion, without any developer input. The cost model is inspired by DAPPLE (Fan et al., 2021). The search is unique in that, for every pipeline schedule in the scheduling search space, a beam search is conducted for an optimal graph slice; the search space subsumes many (but not all) prior pipeline schedules (1F1B).", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "I think the contribution of completely automatic pipelining is strong. I can see the ideas here being incorporated into future works. It's also interesting that this work allows SDXL to be pipelined, though I'm unclear about the practical implications\n\nThe main weaknesses are\n1) the co-optimization algorithm is quite complicated and the design decisions receive no justification, either theoretical or practical\n2) the experiments are conducted in a relatively limited setting (8x A100s, llama 7b)\n3) the correctness is somewhat low. Even being generous in what is counted as a bona fide failure, approximately 10% of the evaluation benchmark set (which consists of relatively small models) yields failures", "other_comments_or_suggestions": "For Table 4 in Appendix C.3, it would be good to have additional columns that display the errors.", "questions_for_authors": "\"Also, the A100 cluster is too powerful to benchmark multiple sets of pipeline parameters (i.e., forward time of each\ncomputation node is too short to correctly estimate total latency as we increase loop count or number of devices). Therefore,\nwe need to limit the performance of a device itself.\" Can you clarify this statement?", "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": "g7luW92uKq", "reviewer": "Reviewer_wSiC", "summary": "The paper introduces Pfeife, a system integrating with PyTorch's `torch.compile` to automate pipeline parallelism without user intervention. Pfeife partitions models at an operation-level granularity across multiple GPUs, employing a cost model combined with beam search to optimize pipeline scheduling. Key claimed contributions include a novel algorithmic framework allowing simultaneous optimization of pipeline slicing and scheduling, demonstrating that uneven workload distribution with prefetching can improve performance significantly. Evaluations demonstrate Pfeife achieving throughput improvements of up to 22% over DeepSpeed and Colossal-AI across various large models, including ViT-g/14, Llama2-7B, and StableDiffusion-XL.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Claims about performance improvement and automation are generally supported through rigorous experiments comparing Pfeife to DeepSpeed and Colossal-AI.", "methods_and_evaluation_criteria": "The evaluation methods are well designed and carried out for demonstrating Pfeife’s practicality and efficiency within PyTorch. However, the limited scope in evaluation criteria—lacking comprehensive comparisons with established frameworks like GPipe, PyTorch FSDP and PipeDream—undercuts the robustness of its methodological validation. Moreover, Pfeife’s reliance on PyTorch's compiler limits the scope of its application, especially considering the reported failure rates on several TorchBench models.", "theoretical_claims": "The manuscript does not contain substantial new theoretical results but correctly refers to established concepts such as cost models of pipeline parallelism. No formal proofs requiring validation are presented; theoretical claims referenced from literature are appropriate and correctly applied.", "experimental_designs_or_analyses": "Experimental setups using ViT-g/14, Llama2-7B, and StableDiffusion-XL are adequate for evaluating memory management and throughput in large-scale models: they cover a wide range of recent architectures and are of adequate size.", "relation_to_broader_scientific_literature": "Pfeife situates itself within the context of automatic pipeline parallelism, effectively highlighting distinctions from manual annotation-required systems like PiPPy and Varuna. However, the literature review fails to deeply engage with critical predecessors such as GPipe, PipeDream, GraphPipe, and ZeRO-3. Particularly, the omission of GraphPipe (which introduced DAG-based scheduling with significant performance gains) severely undermines the perceived novelty of Pfeife's algorithmic contributions.", "essential_references_not_discussed": "Not necessarily essential, but GraphPipe is of interest to this work:\n\nJeon, Byungsoo, et al. \"Graphpipe: Improving performance and scalability of dnn training with graph pipeline parallelism.\" arXiv preprint arXiv:2406.17145 (2024).", "other_strengths_and_weaknesses": "Strengths:\n\nPfeife's integration with PyTorch’s TorchDynamo presents practical usability.\n\nThe idea of automatic processing in lieu of manual annotation is fresh and interesting.\n\nWeaknesses:\n\nPfeife lacks fundamental novelty - there is a plethora of existing pipeline parallel systems, and in the absence of strong, comprehensive benchmarking results backed by available artifacts, it is not a strong submission to this venue.", "other_comments_or_suggestions": "The author may consider making the figures more b/w friendly.", "questions_for_authors": "I do wonder if the author has a reason for omitting other well-known pipeline parallel systems like GPipe and PipeDream in the work. Since the manuscript is built on PyTorch, a comparison with FSDP may also be warranted. There is also more recent works on more sophisticated flavors of pipeline parallelism, such as GraphPipe.", "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": "VaQN73NmkH", "reviewer": "Reviewer_NzCV", "summary": "The paper introduces Pfeife, a new tool that integrates with PyTorch to provide automatic pipelining of machine learning models.\nPfeife aims to address the memory limitations of GPUs when training large models by parallelizing the execution of these models across multiple devices.\nIt leverages PyTorch's tracing JIT compiler (TorchDynamo) to capture a static data-flow graph of a PyTorch model and then automatically generates a pipeline schedule to distribute the operations across devices.\nThe authors claim that Pfeife can run large models that would otherwise not fit on a single device and that it outperforms state-of-the-art pipelining tools.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Claim 1. Pfeife can run large models that would otherwise not run due to not fitting in a single device.\nThe paper provides empirical results demonstrating Pfeife's ability to run large models. This advantage is shared by all parallelism techniques, which is not an exclusive benefit of Pfeife.\n\nClaim 2. Pfeife outperforms state-of-the-art tools by up to 22% in terms of training throughput.  \nThe paper presents comparative evaluations against other pipeline tools. This number is from Table 1.\n\nClaim 3. Pfeife can pipeline non-sequential models such as Stable Diffusion that are not supported by previous pipeline parallelism tools. The paper mentions this capability, and the evaluation includes Stable Diffusion. The authors should provide a more detailed explanation of why previous tools cannot support such models and how Pfeife overcomes these limitations.", "methods_and_evaluation_criteria": "Pfeife leverages PyTorch 2's TorchDynamo to capture the data-flow graph of ML models. It employs a pipeline scheduling algorithm that considers both graph slicing and scheduling parameters to optimize performance.\nThe authors use a cost model based on critical path analysis to estimate running time and memory usage.\n\nThe evaluation includes experiments with large models, and the primary performance metric appears to be running time or throughput.  It would be beneficial to include memory usage as a key evaluation criterion, given that the motivation is to address memory limitations.  The choice of benchmark models seems appropriate for demonstrating the capabilities of the system.", "theoretical_claims": "The paper describes synchronous and asynchronous pipeline parallelism and an algorithm to find pipeline schedules.  The authors also investigate optimal workload distribution conditions.\n\nThe paper uses a cost model based on critical path analysis.  This cost model and the scheduling algorithm are central to the paper's contributions. The correctness of the algorithm and the accuracy of the cost model in predicting performance should be rigorously established, possibly with proofs or theoretical analysis in the supplementary material.\n\nI suggest to avoid using the term \"optimal\" unless there are solid proof for that claim.", "experimental_designs_or_analyses": "The experimental design involves evaluating Pfeife on large ML models and comparing its performance with existing pipelining tools.    \n\nThe paper mentions using a profiler to measure computation time and memory consumption.  More details on the profiling methodology, including the granularity of the profiling (e.g., per-operation), the profiling overhead, and the accuracy of the profiler, would strengthen the analysis.    \n\nThe authors acknowledge that linearizing the computation graph is an approximation.  A more detailed analysis of the impact of this approximation on the generated schedules would be valuable.", "relation_to_broader_scientific_literature": "The paper effectively situates Pfeife within the context of existing parallelization techniques for machine learning models, including data parallelism, weight sharding, model parallelism, and tensor parallelism. It provides a good overview of pipeline parallelism and discusses relevant prior work such as PipeDream, GPipe, DAPPLE, AutoPipe, TeraPipe, and others. The authors clearly articulate how Pfeife builds upon and differs from previous approaches, such as by using a more fine-grained partitioning (per operation) and co-optimizing slicing and scheduling.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "Strengths:\n* The paper addresses an important problem in deep learning: memory limitations in training large models.    \n* Pfeife offers an automated and transparent way to pipeline models, which can improve usability.    \n* The approach of co-optimizing slicing and scheduling is novel and promising.    \n\nWeaknesses:\n* The theoretical claims and the cost model would benefit from more rigorous analysis.    \n* The limitations of the approach, such as those arising from the linearization of the computation graph, should be discussed in more detail.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "1. What are the limitations of Pfeife, and how do you plan to address them?\n2. How difficult can we apply Pfeife to other machine learning frameworks and compilers, such as JAX/XLA?", "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": "J3R7N7cbkO", "reviewer": "Reviewer_8hKh", "summary": "This paper proposes Pfeife, an automated tool that integrates with PyTorch to transparently partition and pipeline large machine learning models across multiple GPUs. It leverages PyTorch's JIT tracing to construct a data-flow graph of the model and then optimizes the pipeline schedule. Experimental evaluations show that Pfeife can outperform existing pipelining tools by up to 22% while handling complex, non-sequential models.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would to thank the reviewers' time and feedback.\n\n\nReviewer aSA1\n - Unfortunately, we only have access to the 2 machines we used in the experiments.\nWe don't have budget for more.\n\n - Regarding correctness, first we note that we run an order of magnitude more\nmodels than most academic papers. It required substantial engineering effort to\nensure the results were correct and that we supported most features of PyTorch's\ntorch.compile/TorchFX. These are still under development and thus a moving target.\nAs such, we believe we deserve some credit for being able to run so many models\nand being open that not all models work (we could have simply left these out).\nAs shown in table 3, only 1 out of 49 models produces incorrect results,\nwhich is actually due to a limitation in the model capturing of torch.compile.\nIt's not a bug in our algorithm or implementation.\n\n - \"Also, the A100 cluster is too powerful to benchmark multiple sets of pipeline\nparameters (i.e., forward time of each computation node is too short to correctly\nestimate total latency as we increase loop count or number of devices). Therefore,\nwe need to limit the performance of a device itself.\"\n\nApologies for the broken English; we'll fix that.\nWhat we mean is that some configuration parameters generate forward passes that\nare too small. At some point, the overhead of spinning a new GPU kernel and\nthe communication starts to dominate, and then the results are not interesting\nanymore. As we mention in the paper, we prevent the scheduler from generating\nthese configurations since they are unlikely to yield good performance.\nThe slower the device the more interesting configurations we can test.\nAnyway, this is nothing fundamental, we just wanted to try all configurations\nwe could run.\n\n\nReviewer wSiC\n - We run an order of magnitude more models than most academic\npapers. We believe we deserve some credits for that.\nWe have also run a comprehensive study of the optimal schedules to show\nthat unbalanced schedules are better. We understand we didn't run experiments\nwith a thousands GPUs due to not having them, but our experiments meet or\neven exceed the bar set by previous pipelining papers. And we do offer new\ninsights into the problem and a new algorithm.\n\n - We compared with the systems that are available. GPipe and PipeDream are\nnot. FSDP requires manual work and it overlaps significantly with\nDeepSpeed. We believe the experiments we did are sound and provide good insights.\n\n\nReviewer NzCV\n - Profiler: this is done only once when the model starts and is done per\noperation. Hence, the overhead is negligible for long-running training\nsessions. The time it takes is about the same as running the model\nsequentially.\nProfiling of inter-device communication is done only once on installation.\n\n - Limitations: Pfeife's current implementation only supports natively\nlooped schedules. The algorithm itself is obliviously of the schedule,\nbut we use a template to produce looped schedules. We would like to\nhave a language to specify templates for schedules. That would be a great\nextension to this work.\n\n - Using Pfeife from other frameworks: Pfeife has a frontend part that\nreads TorchFX graphs. That part is specific to PyTorch, but it's thin;\nwe only need to create a graph with nodes.\nThen, we have the execution engine that has some bits tied to PyTorch\nsince we implement a backend for torch.compile. But it can as well be\nreplaced.\nThe core, the scheduling algorithm, etc, they are framework independent\nsince they operate over model graphs annotated with costs.\n\n\nReviewer 8hKh\n - We don't have results for multi-server deployments. We only have 2\nservers that are shared.\n\n - PiPPy is a work in progress. It is not ready for comparisons (we tried).\n\n - We could compare with TP, but this requires more manual work than we\ncould afford.\n\n - Pfeife's scheduling algorithm can handle DualPipe. Our current implementation\nis hardcoded with a scheduling template that spawns looped schedules only.\nBut that's not a fundamental limitation, we just didn't implement other\ntemplates. The scheduling algorithm itself is oblivious since it operates over\ngraphs. Implementing other templates can be done with a dozen lines of code.\n \n - Pfeife can work with powerful GPUs. We are sorry for the broken English.\nWhat we meant is that we should not slice graphs too thinly for faster devices\nsince then the overhead of launching GPU kernels and communication starts to\ndominate.\n\n - We don't include results for other models mostly because we didn't have\nenough compute credits. Running a couple of times to test correctness is one\nthing, but taking stable performance results for multiple configurations\nand attempting to compare against other frameworks requires a lot more\ncompute. Another reason is that the other frameworks require manual work\nfor each model, and our team is small. We cannot spend months trying to\nrun a dozen models in multiple frameworks. Plus most tools are fragile\nand crash often.", "claims_and_evidence": "- Any specific reasons are provided for TorchTitan's limitation to LLMs only, yet the described techniques could be applied to arbitrary models. Please explicitly state what TorchTitan cannot do that Pfeife can.\n- Distributed training is not solely due to limited memory capacity but also driven by computation requirements. Techniques like quantization and activation checkpointing reduce memory usage, so using memory constraints as the main motivation for pipeline parallelism is insufficient.\n- Please specify what kind of data parallelism you are talking about. Traditional data parallelism requires minimal communication (just gradient averaging). When you refer to \"it requires full synchronization of the weights across all devices\", I suppose you are referring to ZeRO3/FSDP. Please cite them explicitly.\n- The usage of \"model parallelism\", \"tensor parallelism\", and \"pipeline parallelism\" is inconsistent. For example, \"For model and tensor parallelism\" implies exclusivity, yet you also say \"we focus on pipeline parallelism, which is a particular form of model parallelism.\" Please ensure consistent usage: when referring to model parallelism, it should include tensor parallelism and pipeline parallelism.\n- Alpa does not require users to rewrite the model, and it is based on JAX instead of PyTorch.", "methods_and_evaluation_criteria": "Only three models (ViT, LLaMA, and Stable Diffusion) are evaluated. Beyond Stable Diffusion, can Pfeife handle other model types, such as MoE models (e.g., DeepSeek-V3) or state-space models (e.g., Mamba)?", "theoretical_claims": "- How do you decide which parts of a model can be parallelized? For instance, in a Stable Diffusion pipeline, do users have to manually define which parts go into different pipeline stages, as shown in Listing 1?\n- Many HuggingFace models cannot be fully traced by TorchDynamo. How do you handle multiple partial graphs? Do you use a custom torch.fx tracer to deal with such models?", "experimental_designs_or_analyses": "The experimental section is weak, as it includes only a small set of models and configurations.\n- Pipeline parallelism is typically used in multi-node environments, where tensor parallelism may be used within a node, and pipeline parallelism spans multiple nodes. How is communication modeled in a multi-node environment? Do you have any results on the multi-node setting?\n- The comparison with automatic parallel frameworks like DeepSpeed and Colossal-AI is insufficient. PiPPy also has an automatic mode for pipeline parallelism (https://pytorch.org/docs/main/distributed.pipelining.html#option-2-splitting-a-model-automatically) —please compare with that as well. Additionally, comparisons to manually optimized systems like TorchTitan and Megatron-LM would show how Pfeife fares against specialized manual implementations.\n- 3D parallelism is prevalent in nowadays large model training. The experiments only contain DP+PP results. What about including TP? Can your PP method incorporate tensor parallelism?\n- Only one specific configuration (ViT, 4, 4) shows significant improvements (over 20%). Why is this configuration particularly favorable? Why should Pfeife generalize well to other scenarios?\n- What is the cost of profiling? For example, if users want to run a job on an 8-node cluster, do they need to firstly launch the profiling job for this 8-node cluster? This would be very time-consuming and costly.\n- See the below question for C.1, how do you handle faster, newer-generation GPUs? If latency estimations are prone to fail under extreme speed, does that mean Pfeife cannot be generalized to other devices?", "relation_to_broader_scientific_literature": "This paper mainly focuses on automating the pipeline parallelism process in distributed training, which can potentially reduce programming efforts for deploying models in a distributed environment.", "essential_references_not_discussed": "Most cited works in the related section are at least two years old. Please incorporate more recent references, such as [B], [C], and [D], and clarify how Pfeife compares or improves upon these newer approaches. They all analyze or consider different kinds of parallelism in a distributed setting. Also, as mentioned in the above question, it would be good to discuss the applicability of Pfeife to DualPipe proposed in [A].\n* [A] DeepSeek-AI, \"DeepSeek-V3 Technical Report\", https://arxiv.org/pdf/2412.19437v1\n* [B] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, Chao Yang, \"Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning\", ASPLOS, 2024.\n* [C] Hongzheng Chen, Cody Hao Yu, Shuai Zheng, Zhen Zhang, Zhiru Zhang, Yida Wang, \"Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training\", ASPLOS, 2024.\n* [D] Ziheng Jiang et al., \"MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs\", NSDI, 2024.", "other_strengths_and_weaknesses": "See the above sections.", "other_comments_or_suggestions": "See the above sections.", "questions_for_authors": "See the above sections.", "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:

Meta-review: Reviewers consistently point to the utility of an automated pipeliner.  The rebuttal responds well to the criticisms: as the reviewers note, the experiments are on sufficiently large (and sufficiently many) models to give confidence that the technique works.  As the authors note, they are to be commended for listing the models on which it does not work.

The reviewers' numerical scores are not consistent with the review text, so the review text takes priority in decision making.
Decision: Accept (poster)