#  GraphPipe: Improving Performance and Scalability of DNN Training with Graph
Pipeline Parallelism

Byungsoo Jeon<EMAIL_ADDRESS>OctoAISeattleWashingtonUSA ,  Mengdi Wu
<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Shiyi
Cao<EMAIL_ADDRESS>UC BerkeleyBerkeleyCAUSA ,  Sunghyun Kim
<EMAIL_ADDRESS>MITCambridgeMAUSA ,  Sunghyun Park<EMAIL_ADDRESS>OctoAISeattleWashingtonUSA ,  Neeraj Aggarwal<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Colin Unger<EMAIL_ADDRESS>Stanford UniversityPalo AltoCAUSA ,  Daiyaan Arfeen<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Peiyuan Liao
<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Xupeng
Miao<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Mohammad
Alizadeh<EMAIL_ADDRESS>MITCambridgeMAUSA ,  Gregory R. Ganger
<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA ,  Tianqi
Chen<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA and  Zhihao
Jia<EMAIL_ADDRESS>Carnegie Mellon UniverisityPittsburghPAUSA

###### Abstract.

Deep neural networks (DNNs) continue to grow rapidly in size, making them
infeasible to train on a single device. Pipeline parallelism is commonly used
in existing DNN systems to support large-scale DNN training by partitioning a
DNN into multiple stages, which concurrently perform DNN training for
different micro-batches in a pipeline fashion. However, existing pipeline-
parallel approaches only consider sequential pipeline stages and thus ignore
the topology of a DNN, resulting in missed model-parallel opportunities.

This paper presents graph pipeline parallelism (GPP), a new pipeline-parallel
scheme that partitions a DNN into pipeline stages whose dependencies are
identified by a directed acyclic graph. GPP generalizes existing sequential
pipeline parallelism and preserves the inherent topology of a DNN to enable
concurrent execution of computationally-independent operators, resulting in
reduced memory requirement and improved GPU performance. In addition, we
develop GraphPipe, a distributed system that exploits GPP strategies to enable
performant and scalable DNN training. GraphPipe partitions a DNN into a graph
of stages, optimizes micro-batch schedules for these stages, and parallelizes
DNN training using the discovered GPP strategies. Evaluation on a variety of
DNNs shows that GraphPipe outperforms existing pipeline-parallel systems such
as PipeDream and Piper by up to 1.6$\times$. GraphPipe also reduces the search
time by 9-21$\times$ compared to PipeDream and Piper.

##  1\. Introduction

Deep neural networks (DNNs) grow more rapidly in size against hardware
developments, making them computationally costly to train
(openai2018computecost, ; patterson2021carbon, ). A recent language model
GPT-4 (openai2023gpt4, ) supposedly uses a much larger number of parameters
(mittechreview2023gpt4size, ) compared to the previous model GPT-3 with 175
billion parameters (brown2020language, ). As a result, training modern DNNs
requires distributing the model architecture across multiple devices.

Figure 1. Pipeline parallelism for DNN training with basic terms used in this
paper.

To address this challenge, existing DNN systems apply model parallelism
(dean2012large, ; shazeer2018mesh, ; shoeybi2019megatron, ;
lepikhin2020gshard, ; narayanan2021efficient, ; zheng2022alpa, ) where a DNN
is partitioned into smaller pieces, each of which fits into the memory of a
single device. Pipeline parallelism (huang2019gpipe, ; fan2021dapple, ;
narayanan2021memory, ; narayanan2019pipedream, ; tarnawski2021piper, ) is a
particular form of model parallelism that further improves device utilization
and throughput. As shown in Figure 1, a key idea of pipeline parallelism is to
split both a DNN and a mini-batch of samples into smaller pieces. First, the
DNN is partitioned into multiple disjoint stages, each of which is a sub-model
and links to other stages to form a pipeline. Second, a mini-batch of samples
is further divided into multiple micro-batches, which are executed on
different stages in a pipeline fashion. This approach reduces device idle time
in training iterations, during each of which a single data mini-batch is
processed, and thus improves throughput.

#### Shortcomings of existing sequential pipeline parallelism.

Existing schemes of applying pipeline parallelism form a _sequential_ pipeline
from partitioned stages, which we refer to as _sequential pipeline
parallelism_ (SPP).

Figure 1 illustrates a DNN training scheme that employs it. A micro-batch
traverses the pipeline’s stages ($S_{1}$ to $S_{4}$) in sequence to perform
the computations ($O_{1}$ to $O_{4}$) dictated by the DNN (forward pass:
$F_{1}$’s), and traverses in reverse for all stages to update their assigned
model weights (backward pass: $B_{1}$’s). Each stage needs to store the
intermediate activations of a forward pass until its corresponding backward
pass is completed. For a given stage, a micro-batch is _in-flight_ until its
backward pass finishes. As micro-batches are continuously injected into the
pipeline, there is a warm-up of in-flight micro-batches. The earlier the stage
in the pipeline, the longer the warm-up. As described, SPP is simple to
construct and operate, but has three key limitations.

First, opportunities to exploit the inherent parallel structures of a DNN are
left unseized. DNN applications such as healthcare (krumholz2016data, ;
tan2020multimodal, ; suresh2020deep, ), chatbot (openai2023gpt4, ), and
recommendation (naumov2019deep, ) jointly process heterogeneous data types
(e.g., text, images, and tabular data). DNNs employed therein are designed to
feature multiple branches, which are computationally independent and thus can
be executed concurrently. But existing DNN systems with SPP first linearize
the computation graph of a DNN to construct the stages of a sequential
pipeline and process these stages sequentially, falling short in harnessing
the opportunity to blend such branch-level parallelism with pipeline
parallelism.

Second, pipeline depth (i.e., number of sequential stages in SPP) is unduly
increased by missing parallelism opportunities that arise from inherent DNN
structures (e.g., parallel branches). Under an alternative arrangement in
which some pipeline stages are parallelized by exploiting such structures, the
number of sequential stages a micro-batch traverses in a forward (or backward)
pass can be smaller. That is, the elongated pipeline formed by SPP unduly
increases pipeline depth, which in turn increases the number of in-flight
micro-batches to manage. This imposes a higher burden of managing memory,
especially for early stages in the pipeline. Recall that the tight memory
constraint in training large DNNs is a primary reason to apply pipeline
parallelism. Thus, it is critical to curb the heightened memory requirement.

Third, today’s devices employed for DNN training (e.g., GPUs) have high
parallel-computing capabilities, requiring a large amount of training samples
to be fetched to achieve peak performance. The increased memory consumption
that results from applying SPP impedes doing so. As a consequence, devices
perform computations at an operational intensity lower than their desired
capacity, resulting in suboptimal training performance.

Figure 2. A high-level comparison between existing (SPP) and our (GPP)
approaches. SPP (top) produces sequential pipeline stages that miss the
opportunity of parallelizing the branches in the DNN. In contrast, GPP
(bottom) generates graphical pipeline stages that enable parallel execution of
the branches. This leads to lower training iteration time (i.e., higher
training throughput) and smaller memory footprint in pipeline-parallel DNN
training.

#### Our approach.

To address the above challenges, we introduce graph pipeline parallelism
(GPP), that enables performant and scalable DNN training. Figure 2 highlights
the key difference between GPP and SPP. Instead of enforcing a strictly
sequential execution order of pipeline stages, GPP allows partitioning a DNN
into stages whose dependencies are identified by a directed acyclic graph. GPP
includes SPP as a special case and can preserve the inherent topology of the
DNN during stage partitioning. As a result, GPP enables concurrent execution
of computationally-independent components, resulting in reduced memory
requirement and improved GPU performance compared to SPP.

GPP involves a significantly larger and more complicated search space of
parallelization strategies compared to the SPP strategies considered by
existing DNN systems. Discovering GPP strategies with superior performance
over existing SPP baselines requires weighing subtle trade-offs between
pipeline depth, memory consumption, and micro-batch schedule. To unleash the
power of GPP, we develop GraphPipe, a system that automatically discovers
efficient GPP strategies to enable performant and scalable DNN training.
GraphPipe includes three key components. First, a pipeline stage partitioner
automatically determines how to partition the operators of a DNN into a graph
of stages, while balancing the computational load among these stages and
minimizing inter-stage communication. Second, a static micro-batch scheduler
schedules the forward and backward passes of different micro-batches within a
mini-batch to minimize the peak GPU memory requirement of a GPP strategy. The
stage partitioner and micro-batch scheduler jointly partition a DNN into
stages and determine the micro-batch schedules for each stage. Finally, a
distributed runtime uses the discovered GPP strategy to enable performant and
scalable DNN training.

Through experiments on three multi-branch DNNs (e.g., Multi-Modal Transformer
(vaswani2017attention, ; wang2023large, ; radford2021learning, ;
openai2023gpt4, ; ramesh2021zero, ; jia2021scaling, ), DLRM (naumov2019deep,
), and CANDLE-Uno (2018candleuno, )), we show that GraphPipe can achieve up to
1.6$\times$ training throughput improvements over existing pipeline-parallel
systems such as PipeDream (narayanan2019pipedream, ) and Piper
(tarnawski2021piper, ). GraphPipe also reduces the search time by 9-21$\times$
compared to PipeDream and Piper.

To summarize, we make the following contributions:

  * •

We introduce graph pipeline parallelism, a new parallelization scheme that
promotes concurrent stage execution, reduces memory requirement, and improves
GPU utilization compared to existing SPP schemes.

  * •

We design algorithms to partition a DNN into a graph of stages and schedule
micro-batches for these stages, which jointly discover efficient GPP
strategies.

  * •

We develop GraphPipe, a distributed runtime that enables fast and scalable DNN
training with GPP.

##  2\. Graph Pipeline Parallelism

Figure 2 describe the key differences between sequential pipeline parallelism
(SPP) employed by existing DNN systems (tarnawski2021piper, ;
narayanan2021memory, ) and graph pipeline parallelism (GPP). Given a DNN and a
set of devices, SPP and GPP produce strategies with different partitioning of
stages and pipeline schedules.

Concurrent execution of stages. SPP _linearizes_ all operators of a DNN while
preserving data dependencies between these operators, and then partitions the
linearized DNN into a sequence of pipeline stages. As a result, each stage has
at most one predecessor and one successor. The execution order of these stages
is thus strictly sequential.

In contrast, GPP preserves the topology of a DNN when partitioning it into
pipeline stages. To avoid circular dependencies between pipeline stages, the
relationships between these stages form a directed acyclic graph. The
execution order of the stages can be thus more general compared to SPP. This
topology-aware partitioning and pipeline stage execution provides GPP a clear
advantage: (potentially) concurrent execution of stages that are
computationally-independent.

For the GPP strategy in Figure 2, the three stages $S_{1}$, $S_{2}$, and
$S_{3}$ are computationally-independent. Accordingly, the forward and backward
passes of the three stages can be executed concurrently. However, in the SPP
strategy, the two stages $S_{2}$ and $S_{3}$ are partitioned such that they
have a sequential data dependency (due to the dependency between operator
$o_{6}$ in $S_{2}$ and operator $o_{7}$ in $S_{3}$) since the SPP partitioner
does not consider the topology of the DNN and fails to exploit it. Moreover,
while the two stages $S_{1}$ and $S_{2}$ in the SPP strategy should be
computationally-independent according to the original DNN, the SPP scheduler
executes the forward and backward passes of these two stages sequentially.
This is because new data dependencies are imposed between them when
linearizing the operators of a DNN to construct a sequential pipeline.

This distinction directly leads to a performance gap. Specifically, both SPP
and GPP involve a warm-up phase during which micro-batches are injected into
the pipeline until all stages can perform work concurrently. However, as shown
in Figure 2, the warm-up phase of GPP (i.e., 2) is shorter than that of SPP
(i.e., 4). This performance improvement also applies to the cool-down phase
during which in-flight micro-batches are resolved. As a result, GPP achieves a
shorter per-iteration training time (hence, a higher throughput) than SPP. The
topology-aware stage partitioning and scheduling of GPP address the first
shortcoming of SPP (§1).

Reduced memory requirement. There is a close relationship between the memory
requirement of a pipeline-parallel strategy and its pipeline depth, which is
defined as the diameter of its stage graph. As mentioned earlier, the depth of
the pipeline becomes excessively extended due to overlooked opportunities for
parallelism within DNN architectures, such as parallel branches. This results
in a higher memory footprint compared to that of pipeline stages that are
parallelized by taking advantage of these structures.

In Figure 2, GPP and SPP have a pipeline depth of 2 and 4, respectively. As a
result, the first stage with the highest activation memory pressure needs to
store the forward pass results for 2 micro-batches in GPP and those for 4
micro-batches in SPP. All else being equal (i.e., an identical model partition
by both), GPP has a lower total memory footprint than SPP. Note that memory
saving is likely to grow as model size grows since a bigger model with deeper
pipeline depth requires larger number of in-flight micro-batches (especially
for early stages). The activation memory saving by GPP addresses the second
shortcoming of SPP in §1.

Improved GPU utilization. Devices employed in DNN training (e.g., GPUs) are
designed to parallelize DNN computation of a micro-batch efficiently. Thus,
larger micro-batches (i.e., more training samples within a micro-batch) can
improve the operational intensity, thus GPU utilization of DNN operators. Note
that larger micro-batches lead to reduced numbers of micro-batches, which in
turn increases the warm-up and cool-down time of pipeline that GPP can
significantly reduce. For simplicity of presentation, Figure 2 assumes that
the same micro-batch size is used by GPP and SPP. However, a lower device
memory requirement of GPP over SPP allows integrating more training samples in
a micro-batch, which increases the operational intensity and overall GPU
utilization, and therefore further reduces the per-iteration training time. We
evaluate this aspect in more detail in §7.

##  3\. Problem Formulation

In this section, we formulate the problem of devising a GPP strategy for
distributed DNN training. Compared to existing works (tarnawski2021piper, ;
narayanan2021memory, ; zheng2022alpa, ), we further generalize the formulation
to support graphical pipeline stages with fine-grained per-stage micro-batch
size and schedules. As input, we are given (a) a computation graph
$\mathcal{G}_{C}=(\mathcal{V}_{C},\mathcal{E}_{C})$ that represents the neural
architecture of a DNN model, (b) a mini-batch size $B$, and (c) a device
topology graph $\mathcal{D}=(\mathcal{V}_{D},\mathcal{E}_{D})$ where each node
$v\in\mathcal{V}_{D}$ represents a device with memory budget $M_{v}$ and each
edge $e\in\mathcal{E}_{D}$ represents a communication link with bandwidth
$C_{e}$ between the two adjacent devices.

As output, we generate a pipeline stage graph
$\mathcal{G}_{S}=(\mathcal{V}_{S},\mathcal{E}_{S})$ that optimizes the
performance metric of interest. In this work, we limit the scope to strategies
that combine pipeline-parallel and data-parallel techniques, and aim to
minimize the Time-Per-Sample (TPS) of the bottleneck pipeline stage since the
pipeline throughput performance hinges upon the straggler stage. The stage
graph $\mathcal{G}_{S}=(\mathcal{V}_{S},\mathcal{E}_{S})$ is a directed
acyclic graph (DAG), where each node $S_{i}\in\mathcal{V}_{S}$ specifies a
pipeline stage and each directed edge $(S_{i},S_{j})\in\mathcal{E}_{S}$
indicates that stage $S_{i}$ must precede $S_{j}$ for forward passes and that
$S_{j}$ must precede $S_{i}$ for backward passes.

The goal is to solve the min-max optimization problem:

(1) $\displaystyle\min\quad$
$\displaystyle\max_{S_{i}\in\mathcal{V}_{S}}\mathsf{TPS}(S_{i};\mathcal{G}_{C},B,\mathcal{D})$
(2) s.t.
$\displaystyle\max_{S_{i}\in\mathcal{V}_{S}}\mathsf{DeviceMemoryUsage}(S_{i};\mathcal{G}_{C},B,\mathcal{D})\leq
M.$

Formally, GPP devises a strategy $\mathcal{G}_{S}$ as follows. We define
$S_{i}\in\mathcal{V}_{S}$ in further detail as a four-element tuple:
$S_{i}=\langle\mathcal{G}_{i},b_{i},\mathcal{D}_{i},\Pi_{i}\rangle$:

  1. (1)

$\mathcal{G}_{i}$ represents a subgraph of $\mathcal{G}_{C}$,

  2. (2)

$b_{i}$ is the micro-batch size of $S_{i}$ (i.e., there are $B/b_{i}$ micro-
batches for each mini-batch),

  3. (3)

$\mathcal{D}_{i}$ is a set of devices allocated to process the forward and
backward passes of $S_{i}$ (we apply data parallelism within $S_{i}$ if
$|\mathcal{D}_{i}|>1$), and

  4. (4)

$\Pi_{i}$ is a micro-batch schedule that specifies the order in which the
$B/b_{i}$ forward and $B/b_{i}$ backward passes are processed. We use
$\mathsf{fw}^{i}_{j}$ (or $\mathsf{bw}^{i}_{j}$) to denote the forward (or
backward) pass of the $j$-th micro-batch for $S_{i}$.

$\mathcal{G}_{S}$ is a valid GPP strategy if and only if the memory constraint
(Equation 2) and all following conditions are met:

  * C1.

$\mathcal{G}_{i}$ is a convex subgraph of $\mathcal{G}_{C}$ , and
$\mathcal{G}_{1},\dots,\mathcal{G}_{|\mathcal{V}_{s}|}$ form a partition of
$\mathcal{G}_{C}$.

  * C2.

If there exists $(u,v)\in\mathcal{E}_{C}$ such that $u\in\mathcal{G}_{i}$ and
$v\in\mathcal{G}_{j}$, then $(S_{i},S_{j})\in\mathcal{E}_{S}$.

  * C3.

$\mathcal{D}_{1},\dots,\mathcal{D}_{|\mathcal{V}_{s}|}$ form a partition of
$\mathcal{D}$.

  * C4.

For each micro-batch schedule $\Pi_{i}$, $\mathsf{fw}_{k}^{i}$ precedes
$\mathsf{fw}_{k+1}^{i}$, $\mathsf{bw}_{k}^{i}$ precedes
$\mathsf{bw}_{k+1}^{i}$, and $\mathsf{fw}_{k}^{i}$ precedes
$\mathsf{bw}_{k}^{i}$.

In words, C1 mandates that all operators be covered by stages that do not
overlap with each other, and C2 mandates that a strict sequential execution
order between two stages be established if according to the computation graph
there exists a data dependency between two operators each in either of the
stages. C3 ensures that at least one device is allocated to every stage. C4
dictates the orderings of forward and backward passes.

##  4\. System Overview

Figure 3. Overview of GraphPipe. It consists of a pipeline stage partitioner
and a micro-batch scheduler. Given a DNN computation graph, mini-batch size,
and device configuration, they interact with each other to produce an
optimized GPP training strategy as output. The output can be launched on the
distributed runtime framework we also develop to execute it and evaluate its
real-world performance.

Figure 3 illustrates an overview of GraphPipe, a system that accelerates
distributed DNN training at scale using GPP. Taking as input (a) the
computation graph of a DNN, (b) mini-batch size, and (c) the topology of
assigned GPUs, GraphPipe produces an optimized GPP strategy for parallel DNN
training. GraphPipe includes three key components: a pipeline stage
partitioner, a static micro-batch scheduler, and a distributed runtime. The
first two components jointly discover a high-performance GPP strategy for a
given DNN model, mini-batch size, and assigned devices, which will be executed
by the distributed runtime.

Pipeline stage partitioner. The partitioner performs three tasks. First, it
partitions a DNN, aimed at achieving an effective distribution of workloads
across stages. It examines the amount of computation and communication needs
associated with the operators in each stage. Importantly, it _leverages the
inherent topology_ of the DNN at hand in order to exploit _concurrent_
execution opportunities. To this end, it performs a sequence of series-
parallel decompositions of the given DNN. Second, it adjusts the micro-batch
size for each stage. This fine-grained adjustment aims to exploit
heterogeneous compute efficiencies of different types of operators. Finally,
it determines how many devices to assign to each stage to achieve an effective
allocation of resources. Note that all three functions are jointly performed,
as no one function is independent of the others. We provide further details in
§5.

Static micro-batch scheduler. The scheduler performs two tasks. First, it
optimizes micro-batch schedules for forward and backward passes while ensuring
the integrity of distributed DNN training. This involves examining both intra-
and inter-stage data dependencies between the passes (see C4 in §3). Next, it
checks if the memory usage that results from the schedule is within the given
device memory constraint (see Equation 2). Memory usage is closely related to
the numbers of in-flight micro-batches of a stage, which can be computed based
on the schedule of the forward and backward passes of the stage. §6 provides
further details.

Distributed runtime framework. We develop a distributed DNN runtime system
that executes GPP training strategies generated by the optimizer of GraphPipe.
Using the distributed runtime as the testbed, we compare the performance of
the generated GPP strategies against existing SPP strategies for various DNNs.
We provide details in §7.

##  5\. Pipeline Stage Partitioner

Algorithm 1 Pipeline stage partitioner.

Input: Computation graph $\mathcal{G}_{C}$, number of devices
$|\mathcal{V}_{D}|$
Output: Optimized stage graph $\mathcal{G}_{S}$

1:// $\mathsf{MAXTPS}$: safe upper-bound for TPS of bottleneck stage.

2:$t_{l}=0$, $t_{r}=\mathsf{MAXTPS}$, $\mathcal{G}_{S}=\varnothing$

3:while $t_{r}-t_{l}>\epsilon$ do

4: $t_{m}=(t_{l}+t_{r})/2$

5: $\mathcal{G}_{S}^{best}=$
SearchStageGraph($\mathcal{G}_{C},|\mathcal{V}_{D}|,t_{m},B$)

6: if $\mathcal{G}_{S}^{best}==\varnothing$ then

7: $t_{l}=t_{m}$

8: else

9: $t_{r}=t_{m}$

10: $\mathcal{G}_{S}=\mathcal{G}_{S}^{best}$

11:return $\mathcal{G}_{S}$

12:

13:function SearchStageGraph($\mathcal{G}_{C},|\mathcal{V}_{D}|,t_{m},B$)

14: // $C$ is a set of candidate schedule configurations ($c$)

15: for $c\in C$ do

16: // $c_{0}$: dummy schedule configuration

17: $\mathcal{G}_{S}^{new}$ =
DP($\mathcal{G}_{C},c_{0},c,|\mathcal{V}_{D}|,t_{m}$)

18: // PickBetter($\cdot$) picks one with less memory

19: $\mathcal{G}_{S}^{best}=$
PickBetter($\mathcal{G}_{S}^{best},\mathcal{G}_{S}^{new}$)

20: return $\mathcal{G}_{S}^{best}$

21:

22:function DP($\mathcal{G},c_{f},c_{b},d,t_{max}$)

23: if this DP state has been visited then

24: return corresponding $\mathcal{G}_{S}^{best}$ to this DP state

25: // Consider a given DP state as a SINGLE stage

26: $\mathcal{G}_{S}^{best}=\varnothing$

27: if EstimateTPS($\mathcal{G},c_{f},c_{b},d$)$\leq t_{max}$ then

28: // Optimize schedule via Algorithm 2

29: $\Pi_{opt}$ = ScheduleStage($\mathcal{G},c_{f},c_{b},d$)

30: $\mathcal{G}_{S}^{best}=$ StageGraph($\mathcal{G},\Pi_{opt},d$)

31: // Decompose a given DP state into two stages

32: if $\mathcal{G}$ can be decomposed in series then

33: for $(\mathcal{G}_{1},\mathcal{G}_{2})\in$ SeriesDecompose($\mathcal{G}$)
do

34: for $d_{2}\leftarrow 1$ to $d-1$ do

35: $d_{1}=d-d_{2}$

36: for $c_{m}\in C$ do

37: $\mathcal{G}_{S_{2}}^{new}=$
DP($\mathcal{G}_{2},c_{m},c_{b},d_{2},t_{max}$)

38: Update $i_{m}$ based on $\mathcal{G}_{S_{2}}^{new}$

39:
$\mathcal{G}_{S_{1}}^{new}=$DP($\mathcal{G}_{1},c_{f},c_{m},d_{1},t_{max}$)

40: else if $\mathcal{G}$ can be decomposed in parallel then

41: for $(\mathcal{G}_{1},\mathcal{G}_{2})\in$
ParallelDecompose($\mathcal{G}$) do

42: for $d_{1}\leftarrow 1$ to $d-1$ do

43: $d_{2}=d-d_{1}$

44: $\mathcal{G}_{S_{1}}^{new}=$
DP($\mathcal{G}_{1},c_{f},c_{b},d_{1},t_{max}$)

45: $\mathcal{G}_{S_{2}}^{new}=$
DP($\mathcal{G}_{2},c_{f},c_{b},d_{2},t_{max}$)

46: $\mathcal{G}_{S}^{best}=$
PickBetter($\mathcal{G}_{S}^{best},\mathcal{G}_{S_{1}}^{new}\cup\mathcal{G}_{S_{2}}^{new}$)

47: return $\mathcal{G}_{S}^{best}$

The pipeline stage partitioner of GraphPipe aims to minimize Time-Per-Sample
(TPS) of the bottleneck pipeline stage as described in §3. It takes as input a
DNN computation graph $\mathcal{G}_{C}$, a mini-batch size $B$, and a device
topology graph $\mathcal{G}_{D}$, and generates an optimized stage graph
$\mathcal{G}_{S}$ by searching over different model partitions, device
assignments, and micro-batch sizes simultaneously. A key challenge we must
address is the large and complex search space of potential GPP strategies. To
reduce the complexity of the search task, we employ a binary search method
combined with series-parallel decomposition and dynamic programming. We next
describes these three components.

Binary search. Given the large search space of potential solutions, GraphPipe
does not attempt to directly find an optimal solution. Instead, GraphPipe
employs binary search to iteratively narrow down the target performance range
and examines whether there exist valid solutions within the range. By
iteratively reducing the range, GraphPipe discovers solutions arbitrarily
close to an optimal one, and thus there is little difference in performance
for practical purposes. Lines 2–11 of Algorithm 1 shows GraphPipe’s binary
search process.

Series-parallel decomposition. Since most DNNs structurally reflect series-
parallel graphs (series-parallel, ; unger2022unity, ), GraphPipe applies
series-parallel decomposition to an input graph $\mathcal{G}_{C}$ in order to
decompose it into smaller, manageable subgraphs, and perform model
partitioning, device allocation, and task scheduling for each subgraph. In the
unusual cases where a DNN does not possess such a structural property,
GraphPipe bypasses this issue by converting the DNN to an arithmetically
identical one whose structure is a series-parallel graph.

Dynamic programming (DP). GraphPipe adopts a dynamic programming algorithm
where the value of each DP state indicates the existence of a strategy
achieving a throughput within a target range (Lines 13–20 of Algorithm 1). At
each DP level, GraphPipe applies series-parallel decompositions to split an
input graph (say $\mathcal{G}$) into two new subgraphs (say
$\mathcal{G}_{1},\mathcal{G}_{2}$), each of which serves as the input
computation graph of a new DP subproblem at one DP level below. GraphPipe
recursively solves the DP subproblems to construct a solution of the original
problem where the input computation graph is $\mathcal{G}_{C}$ (Lines 22–47 of
Algorithm 1).

DP subproblem. We ensure that each DP subproblem maintains a certain structure
(i.e., having a unique pair of source and sink nodes and a subgraph
$\mathcal{G}$ comprised of them). The input to a DP subproblem includes a
computation graph $\mathcal{G}\subseteq\mathcal{G}_{C}$, the number of devices
$d$, and some schedule-related information for its predecessor and successor
stages, which we furnish by enumeration if not available.

The solution of a DP subproblem involves devising a training strategy such
that (1) the number of in-flight micro-batches for the source stage (i.e., the
pipeline stage that includes the source node) is minimized; and (2) the Time-
Per-Sample (TPSes) for all stages do not exceed the target TPS range. These
results are returned back to the parent DP subproblem at one DP level above
where the results are gathered for the parent DP subproblem to produce its
own.

We consider three cases in a DP subproblem:

  * •

Base case: We consider the entire subgraph $\mathcal{G}$ as a single stage and
apply data parallelism with a data-parallel degree of $d$ (Line 27 in
Algorithm 1). We check if the target TPS range is achievable with the memory
constraint, and compute the number of in-flight micro-batches according to
Algorithm 2 (see §6).

  * •

Series decomposition: We perform a series decomposition to create two
subgraphs $\mathcal{G}_{1}$ and $\mathcal{G}_{2}$, where the sink node of
$\mathcal{G}_{1}$ coincides with the source node of $\mathcal{G}_{2}$ (Line 32
in Algorithm 1). We first solve the subproblem associated with
$\mathcal{G}_{2}$. To do so, we enumerate all feasible schedules for the
source node of $\mathcal{G}_{2}$. We then solve the subproblem associated with
$\mathcal{G}_{1}$.

  * •

Parallel decomposition: We perform a parallel decomposition to create
$\mathcal{G}_{1}$ and $\mathcal{G}_{2}$, where $\mathcal{G}_{1}$ and
$\mathcal{G}_{2}$ share the same source and sink nodes (Line 40 in Algorithm
1). As there is no data dependency between these subgraphs, the pipelines can
be executed in parallel. The subproblems associated with $\mathcal{G}_{1}$ and
$\mathcal{G}_{2}$ may produce different optimal numbers of in-flight micro-
batches for the shared source node. To ensure continuous pipelining, we take
the larger number of in-flight micro-batches as the solution.

Overall process. Figure 4 visualizes the overall process. At the top, a DP
subproblem is provided with its initial conditions: computation graph
$\mathcal{G}$, the number of available devices $d$, and the target TPS range
$[0,t_{max}]$. Suppose the number of in-flight micro-batches for the sink node
is $i_{b}$, the micro-batch sizes for the source and sink nodes are $b_{f}$
and $b_{b}$, the stage containing the source node (i.e., source stage) uses
the $k_{f}$F$k_{f}$B schedule, and the stage containing the sink node (i.e.,
sink stage) uses the $k_{b}$F$k_{b}$B schedule (we introduces GraphPipe’s
micro-batch schedules in §6). These supposed conditions comprise a schedule
configuration denoted by $c:=(i,b,k)$ in Algorithm 2. They are either
available as the results of some other DP subproblems solved previously, or
furnished by enumeration. The solution of this DP subproblem computes the
smallest possible number of in-flight micro-batches for the source stage
(i.e., $i_{f}$ in Figure 4) that meets the target TPS range $[0,t_{max}]$.

Figure 4. Pipeline stage partitioner performing series-parallel
decompositions. Black arrows indicate subproblem formulations. Red arrows
indicate solutions of subproblems.

Time complexity. We analyze the time complexity of the stage partitioner to
gauge the impacts of design parameters. Let $N$ be the number of series-
parallel subgraphs of $\mathcal{G}_{C}$, $\mathcal{B}$ be the set of possible
micro-batch sizes, $\mathcal{D}$ be the set of possible data-parallel degrees.
The maximal element of $\mathcal{B}$ is upper-bounded by $B$. We consider
powers of 2 for micro-batch sizes (i.e., $|\mathcal{B}|<\log_{2}B$). Likewise,
the maximal element of $\mathcal{D}$ is upper-bounded by $|\mathcal{V}_{D}|$
and $|\mathcal{D}|<\log_{2}|\mathcal{V}_{D}|$ holds.

The number of candidates for $\mathcal{G}$ is $O(N)$, that for
$c_{f}=(b_{f},k_{f})$ is $O(|\mathcal{B}|^{2})$, that for
$c_{b}=(i_{b},b_{b},k_{b})$ is $O(B|\mathcal{B}|^{2})$, and that for $d$ is
$O(|\mathcal{D}|)$ in each DP subproblem. To compute a DP value, it takes
$O(|\mathcal{D}||\mathcal{B}|^{2})$ time for series decompositions and
$O(|\mathcal{D}|)$ time for parallel decompositions. Therefore, the time
complexity for a single DP run is $O(NB|\mathcal{B}|^{6}|\mathcal{D}|^{2})$
and the overall time complexity is
$O((\log\mathsf{MAXTPS})NB|\mathcal{B}|^{6}|\mathcal{D}|^{2})=O((\log\mathsf{MAXTPS})NB(\log_{2}B)^{6}(\log_{2}|\mathcal{V}_{D}|)^{2})$.

Algorithm 2 Static micro-batch scheduler.

Input: Model partition $\mathcal{G}$, initial current and next stage schedule
configurations $c_{f},c_{b}$, number of devices $d$
Output: Optimized schedule $\Pi_{opt}$

1:function ScheduleStage($\mathcal{G},c_{f},c_{b},d$)

2: // Optimize schedule by minimizing number of

3: // in-flight micro-batches

4: // while respecting data dependencies

5: $i_{f}=$ ComputeInFlight($k_{f},b_{f},k_{b},b_{b},i_{b}$)

6: $c_{opt}=(i_{f},k_{f},b_{f})$

7: if $c_{opt}$ violates device memory constraint then

8: $c_{opt}=\varnothing$ // Invalidate schedule $c_{opt}$

9: $\Pi_{opt}\leftarrow$ ScheduleTask($c_{opt}$)

10: return $\Pi_{opt}$

##  6\. Static Micro-Batch Scheduler

The static micro-batch scheduler of GraphPipe optimizes micro-batch schedules
to minimize training time and memory footprint. Specifically, we design our
scheduler to address the unique challenges presented by graph-like data
dependencies in GPP pipeline stages. These dependencies complicate both
scheduling and memory footprint optimization. Our scheduler aims to optimize
schedules while ensuring their validity by meeting the data dependency
requirements between all forward and backward passes.

We further generalize our scheduler so that it can support different micro-
batch sizes and schedules over pipeline stages. This can be effective when
running heterogeneous models (e.g., multi-modal models with different ideal
micro-batch size over pipeline stages across different modalities). We also
find it can save extra memory footprint with fine-grained per-stage micro-
batch scheduling (i.e., Figure 5).

The scheduler (Algorithm 2) takes as input (1) a configuration of model
partition $\mathcal{G}$, (2) current and next stage schedule configurations
$c_{f},c_{b}$, and (3) the number of devices $d$ from the pipeline stage
partitioner, and produces an optimized micro-batch schedule $\Pi_{opt}$ for a
given stage configuration. As in Figure 3, the input is fed by the stage
partitioner, and the output is returned back to the stage partitioner to form
a stage graph with an optimized micro-batch schedule.

GraphPipe’s pipeline stage partitioner (Algorithm 1) first calls Algorithm 2
to discover an optimized micro-batch schedule for the last stage. It then
traces back all directed edges $(S_{i},S_{j})\in\mathcal{E}_{S}$ of the stage
graph $\mathcal{G}_{S}$ in the reverse direction and determines a schedule for
each stage $S_{i}$ until a schedule for the first stage is determined. The
reason for backward traversal is that computing the activation memory usage,
and thus the total usage, for a stage $S_{i}$ requires complete schedule
information of its subsequent stages $S_{j}$.

Figure 5. A comparison between (all-stage) $k$F$k$B and per-stage $k$F$k$B
schedules with different micro-batch sizes over stages. F$\\{i,j\\}$,
B$\\{i,j\\}$ indicate forward and backward passes for a micro-batch including
samples $i$ and $j$. It showcases how per-stage $k$F$k$B scheduling can save
memory footprint.

$\textsc{ComputeInFlight}(\cdot)$ and $\textsc{ScheduleTask}(\cdot)$ in
Algorithm 2 are two key functions. First, $\textsc{ComputeInFlight}(\cdot)$ is
a key subroutine to optimize schedule by effectively minimizing the number of
in-flight micro-batches for a given stage without increasing per-iteration
training time. In existing SPP methods, it is trivial to implement
$\textsc{ComputeInFlight}(\cdot)$ since it adopts simple $k$F$k$B schedule
(i.e., forward pass for $k$ micro-batches followed by backward pass for $k$
micro-batches)111A schedule $\Pi_{i}$ is said to be $k$F$k$B when there exist
$\ell$ and $k$ such that $\Pi_{i}$ starts with $\ell$ forward passes (for
warm-up), alternates between $k$ backward and $k$ forward passes, and ends
with $\ell$ backward passes (for cool-down). (wang2023ada, ; huang2019gpipe, )
or its subset (e.g., 1F1B) (narayanan2019pipedream, ; tarnawski2021piper, )
with sequential dependency. However, in GPP, _graph-like dependencies_ between
pipeline stages complicate memory footprint optimization, as they require
consideration of all parallel successor stages of a current stage.

In addition, we generalize $\textsc{ComputeInFlight}(\cdot)$ to support 1)
different micro-batch sizes over pipeline stages and 2) fine-grained $k$F$k$B
schedules _per stage_ (i.e., varying $k$ over stages). Figure 5 illustrates
how fine-grained $k$F$k$B schedules can save activation memory footprint. In
Figure 5, $S_{2}$ has a different micro-batch size (=1) from that of
$S_{1},S_{3}$ (=2). $S_{2}$ can employ either 1F1B or 2F2B without degrading
training iteration time. If $S_{2}$ employs 2F2B, it can save the activation
memory footprint of $S_{1}$ by reducing the number of in-flight micro-batches
from 4 to 3 in comparison to 1F1B.

(a) Multi-Modal Transformer (MMT)

(b) DLRM

(c) CANDLE-Uno

Figure 6. End-to-end performance evaluation. GraphPipe outperforms both
PipeDream (narayanan2021memory, ) and Piper (tarnawski2021piper, ) in three
different models: Multi-modal Transformer-based model (radford2021learning, ),
DLRM (naumov2019deep, ), and CANDLE-Uno (2018candleuno, ) at all but one GPU
count configurations tested. Missing data points indicate that no training
strategy can be found within reasonable timeframes.

Second, $\textsc{ScheduleTask}(\cdot)$ produces an optimized schedule of
forward and backward passes for a given stage with optimized schedule
configuration ($c_{opt}$). It adopts greedy scheduling that schedules backward
pass as early as possible. It reduces both memory consumption and training
iteration time since it quickly resolves the corresponding in-flight forward
pass. However, graphical pipeline stages and different micro-batch sizes over
stages make the data dependencies of forward and backward passes convoluted.
Still, we find a greedy scheduling handling these factors is _optimal_ in
terms of iteration time and memory footprint.

GraphPipe uses the following default schedule configurations: 1) synchronous
1F1B schedule (narayanan2021memory, ) adjusted to support graph-like
dependencies and 2) the same micro-batch size across stages. The synchronous
1F1B avoids gradient staleness with the same pipeline latency and lower
activation memory footprint in comparison to alternatives (e.g., GPipe
(huang2019gpipe, )). Furthermore, except for some corner cases, we observe
that performance improvements from per-stage micro-batch sizes and $k$F$k$B
schedules are incremental to justify the increased search times for models and
device clusters we explored. Still, with GraphPipe, users can choose to search
over per-stage micro-batch sizes and $k$F$k$B schedules for more heterogeneous
models and larger device clusters.

##  7\. Evaluation

We develop GraphPipe on top of FlexFlow (jia2019beyond, ), a distributed
multi-GPU runtime for DNN training. Major modifications are to replace
FlexFlow’s partitioner and scheduler with ours described in §5 and §6
respectively. We evaluate GraphPipe on the Summit supercomputer (summit, ).
For each compute node of Summit, we use 2 IBM POWER9 CPUs and 4 NVIDIA V100
GPUs with 512GB of main memory. GPUs within a node are interconnected via
NVLink while nodes are connected via Mellanox EDR 100Gb InfiniBand. We use the
default schedule configurations of GraphPipe mentioned in Section 6. Note that
we omit error bars for our plots, as we observe marginal standard deviations
(less than 3%) for all results.

DNNs. We explore three multi-branch DNNs: Multi-Modal Transformer-based model
(MMT) (vaswani2017attention, ; radford2021learning, ), DLRM (naumov2019deep,
), and CANDLE-Uno (2018candleuno, ). Multi-Modal Transformer (MMT) is a
backbone of most state-of-the-art multi-modal models (wang2023large, ;
radford2021learning, ; openai2023gpt4, ; ramesh2021zero, ; jia2021scaling, ).
DLRM is a popular deep learning recommendation model for personalization and
ads recommendation. CANDLE-Uno is a specialized model in the medical domain
(i.e., precision medicine). We describe the detailed model configurations in
the supplemental material. Despite different applications, all these models
feature parallel branches, each processing a different type of data.

###  7.1. End-to-End Evaluation

We compare the training throughput of GraphPipe with existing pipeline-
parallel systems: PipeDream (narayanan2021memory, ) and Piper
(tarnawski2021piper, ). We choose these two baselines since their combined
search space encompasses all possible model partitions covered by other
approaches including DAPPLE (fan2021dapple, ), Alpa (zheng2022alpa, ), etc
(narayanan2019pipedream, ; unger2022unity, ). We implemented their stage
partitioning strategies in GraphPipe to conduct fair comparisons; note that it
is also necessary for Piper given that it is a partitioning algorithm and
doesn’t have a runtime.

# GPUs | MMT | DLRM | CANDLE-Uno
---|---|---|---
Piper | PipeDream | Ours | Piper | PipeDream | Ours | Piper | PipeDream | Ours
4 | 52.9 (440.5$\times$) | 2.57 (21.4$\times$) | 0.12 | ✗ | 6.39 (19.3$\times$) | 0.33 | ✗ | 3.84 (20.2$\times$) | 0.19
8 | 126 (165.7$\times$) | 11.9 (15.6$\times$) | 0.76 | ✗ | 31.3 (11.4$\times$) | 2.73 | ✗ | 17.0 (11.8$\times$) | 1.43
16 | 304 (101.3$\times$) | 44.3 (14.7$\times$) | 3.00 | ✗ | 131 (9.9$\times$) | 13.28 | ✗ | 66.10 (10.7$\times$) | 6.14
32 | 745 (73.7$\times$) | 151 (15.0$\times$) | 10.11 | ✗ | 505 (9.2$\times$) | 54.6 | ✗ | 234 (10.4$\times$) | 22.37
Table 1. Solution search times (in seconds) for Piper, PipeDream, and Ours
(GraphPipe) on the Apple M1 Max; ✗ indicates search cannot be completed.
Numbers in parentheses indicate the search time ratio of the algorithm to that
of GraphPipe.

We compare the training throughput of GraphPipe with existing pipeline-
parallel systems such as PipeDream (narayanan2021memory, ) and Piper
(tarnawski2021piper, ). We choose these two baselines since their combined
search space encompasses all possible model partitions covered by other
approaches (fan2021dapple, ; zheng2022alpa, ; narayanan2019pipedream, ).

Figure 6 showcase the results. We measure the training throughput (i.e.,
number of samples processed per second) as we increase the number of GPUs and
mini-batch sizes. Note that Piper cannot generate training strategies for DLRM
and CANDLE-Uno since its time and space complexity increases exponentially
with respect to the number of parallel branches. GraphPipe outperforms
PipeDream and Piper at all but one GPU configuration. Moreover, the
performance gap widens as the number of GPUs increases.

Our analysis reveals that we can attribute the widening performance gap to the
pipeline depths greatly reduced by GraphPipe compared to PipeDream and Piper
for the multi-branch models. As we use more devices, the number of sequential
pipeline stages tends to increase to achieve a higher throughput, particularly
when the model size is too large to apply data parallelism at the cost of
weight memory footprint and weight synchronization. With a larger number of
stages, sequential pipeline schemes by generated by PipeDream or Piper suffer
from extended warm-up and cool-down phases. Directly, these extended pipeline
bubbles negatively affect training throughput. Indirectly, these bubbles
increase activation memory footprints, which in turn impede effective model
partitioning. We visualize this analysis in detail via a case study (see
§7.4).

###  7.2. Search Time

Table 1 presents the search times by the three optimizers (GraphPipe,
PipeDream, and Piper) for the three models (Multi-Modal Transformer, DLRM, and
CANDLE-Uno). The Multi-Modal Transformer-based model has two branches and the
DLRM and CANDLE-Uno models have eight branches.

GraphPipe is at least $9\times$ faster than the baselines irrespective of the
models or GPU configurations. In addition, GraphPipe’s efficient partitioner
produces a strategy within a minute for all configurations. The SPP baselines
are much slower by comparison, and this search time discrepancy can be
attributed in large part to the fact that the baselines rarely leverage DNN
topology in expediting search. Note that Piper does not produce strategies for
the DLRM and CANDLE-Uno models for the aforementioned reasons.

To see the large search space of each SPP baseline, it is helpful to
approximate their time complexities. Let us consider a simple multi-branch
model with each branch having $k>n$ operators, where $n$ is the number of
branches. Recall that Piper considers model partitions in which cross-branch
stages exist. This level of granularity of model partitions significantly
increases the number of model partitions to examine. Piper’s optimizer runs in
$O(|\mathcal{D}|^{2})$ time (Appendix D in (tarnawski2021piper, )), where
$\mathcal{D}$ is the set of downsets (Definition 4.1 in (tarnawski2021piper,
)). According to the definition, model partitions in which one stage spans
multiple branches and all other stages are formed within a branch are valid
candidates. Since we can choose one operator out of $k$ from each branch to
form a cross-branch stage, the number of such model partitions is at least
$|\mathcal{D}|\geq\prod_{i=1}^{n}k=k^{n}$. Thus, Piper’s time complexity is
lower-bounded by $O(k^{2n})$. This time complexity implies that unless we
employ a set of clever heuristics, Piper’s time complexity can be
significantly high for multi-branch DNNs.

On the other hand, PipeDream considers a converted DNN that linearizes all
branches and the operators within. Thus, it deals with a single chain of
operators, where the number of model partitions to consider is much smaller
than Piper.

Still, GraphPipe considers significantly fewer model partitions than PipeDream
(and hence Piper) particularly when a given DNN features multiple branches.
Instead of solving a single long chain of $nk$ operators as in PipeDream,
GraphPipe solves $n$ short chains of $k$ operators separately. As empirically
shown in Figure 6, GraphPipe barely demonstrates throughput degradation, which
could have resulted from examining much fewer model partitions. Explicitly
leveraging DNN topology in examining model partitions in search for a training
strategy turns out to be critical to reducing the search space and time
complexity.

###  7.3. Different Numbers of Branches and
Micro-Batch Sizes

Figure 7 shows the results of two experiments in which we change the number of
parallel branches for the CANDLE-Uno model (left) and change the number of
micro-batch sizes for the two-branch multi-modal Transformer-based model
(right). The purpose of the experiments is to investigate the effects of main
parameters on the performances of GraphPipe and the SPP baselines (i.e.,
PipeDream and Piper).

The left sub-figure depicts the throughputs of different systems normalized by
that of PipeDream with respect to the number of parallel branches for the
CANDLE-Uno model.222Piper was not able to produce a strategy for the CANDLE-
UNO model. We see that the performance gap achieved by GraphPipe scales with
the number of branches, reaching up to 2$\times$ at 16 branches. Intuitively,
the performance gain mostly stems from the fact that GraphPipe is able to
reduce the pipeline depths at all configurations allowing concurrent execution
of parallel branches, reducing the inefficient pipeline warm-up and cool-down
phases significantly. The gain scales because the larger the number of
branches, the larger the differentials of the phases between GraphPipe and
SPP. This experiment result demonstrates that (1) reducing pipeline depth is
critical to training performance; and (2) GraphPipe is better at it than SPP
especially when multiple branches of non-negligible workload are present. The
larger the number of branches in a given DNN to train, the more opportunities
for GraphPipe to exploit and reduce pipeline depths.

Figure 7. Throughput vs. different numbers of branches using 4, 8, 16 GPUs
respectively (left). Throughput vs. different micro-batch sizes using 8 GPUs
(right).

The right sub-figure depicts the throughput performances for the multi-modal
Transformer-based model with four branches. We use a mini-batch size of 128
and eight GPUs. We intentionally fix a micro-batch size (instead of using the
best ones chosen by the optimizers) in comparing the performances, for the
purpose of examining the benefits (or harms) of using large micro-batch sizes.
If increasing micro-batch size turns out to be beneficial, then it is worth
reducing pipeline depth so as to reduce activation memory footprints, and in
turn create room for using a larger micro-batch size.

We can observe the key role of reduced pipeline depth by GraphPipe in
improving throughput. For each micro-batch size, GraphPipe always outperforms
SPP. Since there is no difference in operational intensity with the same
micro-batch size used for both GraphPipe and SPP, the performance gap can be
solely attributed to the difference in pipeline depth. The reduced pipeline
depth by GraphPipe leads to a shorter execution time for the warm-up and cool-
down phases, hence a higher throughput.

Figure 8. Pipeline schemes devised by SPP (top) and GraphPipe (bottom). They
produce an identical model partition. The selected micro-batch sizes are
different: 2 (SPP) v.s. 4 (GraphPipe), which results in a better compute
efficiency for GraphPipe. Both methods deem it unnecessary to employ data
parallelism primarily because doing so would have split a smaller micro-batch
size even further, which would have harmed compute efficiencies. The pipeline
depths are also different: 8 (SPP) v.s. 4 (GraphPipe), which results in a
smaller pipeline depth for GraphPipe. This improvement comes purely from the
fact that GraphPipe can produce a pipeline scheme that allows for concurrent
execution of parallel branches.

###  7.4. Case Study

Figure 9. A synthetic Transformer-based two-branch DNN for case study. A
sequence of one multi-head attention and two linear layers is repeated four
times to compose a single branch. One concatenation layer combines two
branches.

It is instructive to take a close look at the strategies produced by GraphPipe
and SPP. We run both GraphPipe and SPP optimizers for a synthetic model,
execute the strategies, and observe a 20% throughput improvement by GraphPipe
over SPP. Our analysis finds that the aggregate gain comes from two sources,
and the contributions are nearly equal.

Figure 9 depicts the two-branch Transformer-based model synthesized for the
experiment. Each branch consists of four repeated sequences of one multi-head
attention and two linear (dense) layers. The branches are merged by a
concatenation operator.

Both GraphPipe and SPP produce the identical model partition on a budget of
eight devices. Each stage contains one multi-head attention and two linear
layers. There are eight such stages, four per branch, except that one stage
necessarily contains the concatenation operator. A key difference between the
two strategies, however, is the way the stages are pipelined. Figure 8 depicts
the pipeline schedules. Note that the pipeline depth for SPP is eight since
all eight stages form a sequential pipeline. In contrast, the pipeline depth
for GraphPipe is four. The two branches are computationally-independent, hence
stage $1+i$ and $5+i$ for $0\leq i\leq 3$ can be executed in parallel, and
this is precisely what the training strategy produced by GraphPipe suggests.
This concurrent execution reduces the warm-up phase by half in terms of number
of micro-batches from eight to four. This warm-up phase reduction leads to 10%
performance improvement.

There is another subtle, yet key difference. Since GraphPipe reduces the
pipeline depth by half, the activation memory footprints for early stages are
smaller for the GraphPipe strategy. As a result, GraphPipe can choose a micro-
batch size from a wider range of candidates, and indeed selects a size of 4.
The compute efficiency improvement from choosing a larger micro-batch size
over SPP (which chooses a size of 2 due to larger activation memory
footprints) leads to a larger number of samples processed per unit time. This
means that when the pipeline operates at full capacity, it processes training
samples at a faster rate for GraphPipe than for SPP. Our measurements show
that the gain from this compute efficiency improvement is 10%. The two gain
sources combined, GraphPipe achieves 20% higher throughput over SPP.

##  8\. Related Work

Pipeline parallelism. Existing DNN frameworks (TensorFlow, ; PyTorch, ;
jia2019beyond, ; rajbhandari2020zero, ; shazeer2018mesh, ) employ sequential
pipeline parallelism (SPP) where pipeline stages are strictly sequential. As
we discuss in Section 2, SPP hinders parallel execution of computationally-
independent components of a DNN and memory savings from reduced pipeline
depth. While this limitation still exists as long as SPP is adopted, there are
a variety of pipeline parallelism approaches to improve pipeline performance
in other ways. These approaches fall into one of two paradigms: synchronous
and asynchronous pipeline parallelism.

_Synchronous pipeline parallelism_ (huang2019gpipe, ; narayanan2021memory, ;
fan2021dapple, ; zheng2022alpa, ; unger2022unity, ) refers to a set of
techniques in which the model parameters spread across devices are updated
synchronously after every training iteration. The DNN training semantics is
preserved, thus statistical convergence issues do not arise. But the
synchronous updates fill and drain the pipeline periodically over iterations,
hurting throughput. Our graph pipeline parallelism mitigates this issue by
reducing pipeline bubbles better than sequential pipeline parallelism.

_Asynchronous pipeline parallelism_ (narayanan2019pipedream, ;
narayanan2021memory, ; tarnawski2021piper, ; wpipe, ) refers to a set of
techniques in which the model parameters spread across devices are updated
asynchronously. Although this mode may suffer from statistical convergence
issues as devices execute their stages using out-of-sync model parameters, it
keeps the pipeline full at nearly all times. Graph pipeline parallelism helps
us reduce total device memory usage, thus use a larger micro-batch size to
execute operators at a higher operational intensity compared to sequential
pipeline parallelism. This enables us to process training data faster while
the pipeline is full.

Multiple pipeline stages per device. In the pipeline parallel techniques
above, each device contains only one pipeline stage. It has been shown that
assigning multiple non-contiguous stages to a device can reduce pipeline
bubbles (narayanan2021efficient, ; lamy2022breadth, ) and reduces memory
consumption imbalances across stages (li2021chimera, ; hanayo, ). Earlier work
GEMS (gems, ) has a similar idea but does not utilize the pipeline well —
devices are idle for most of the time and waiting for results from other
stages. These techniques are orthogonal to graph pipeline parallelism, and
thus can be applicable upon some modifications.

Data parallelism. Data parallelism (valiant1990bridging, ; krizhevsky2014one,
; li2014scaling, ; goyal2017accurate, ; mudigere2021high, ) is one of parallel
DNN training techniques in which every device has a local copy of a DNN to
train and a batch of training data is split across devices. Each device
updates its model parameters based on its share of training data and
synchronizes the parameters periodically with other devices. In our work, we
apply data parallelism within a pipeline stage to which we assign multiple
devices, in order to balance stage execution times in a more fine-grained
manner compared to applying pipeline parallelism only.

Automatic DNN parallelism. There are a number of automated approaches
(zheng2022alpa, ; unger2022unity, ; jia2019beyond, ; tarnawski2021piper, ;
narayanan2019pipedream, ; mirhoseini2017device, ; wang2019supporting, )
combining data, pipeline, and tensor parallelisms (shoeybi2019megatron, ).
Existing works first partition a DNN into sequential pipeline stages (SPP) and
then apply data and tensor parallelism to each stage. GraphPipe follows this
same high-level process as well. However, the key difference is that it
generalizes stage partitioning to produce graphical stages and exploit
concurrent execution opportunities from DNN structures (i.e., parallel
branches). Note that it is also feasible to combine our approach with tensor
parallelism by adding a subroutine of applying tensor parallelism (e.g.,
intra-op pass in Alpa (zheng2022alpa, )) in our partitioner while our
scheduler and runtime are already compatible.

##  9\. Conclusion

We have developed _graph pipeline parallelism_ where pipeline stages form a
directed acyclic graph whose edges indicate execution orders of forward and
backward passes in pipeline-parallel DNN training. This design encourages
_concurrent_ execution of parallel branches for superior performance. We have
also developed a distributed system GraphPipe, and through experiments using
three multi-branch models, showed that GraphPipe achieves up to 1.61$\times$
higher training throughputs and $>9\times$ faster solution search times over
existing baselines that operate in a strictly sequential manner.

## References

  * [1] https://github.com/ECP-CANDLE/Benchmarks/tree/master/Pilot1/Uno.  Accessed: 2023-05-15.
  * [2] Ai and compute.  https://openai.com/research/ai-and-compute.  Accessed: 2023-05-15.
  * [3] Summit supercomputer.  https://www.olcf.ornl.gov/summit/.  Accessed: 2023-09-06.
  * [4] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al.  Tensorflow: A system for large-scale machine learning.  In 12th $\\{$USENIX$\\}$ symposium on operating systems design and implementation ($\\{$OSDI$\\}$ 16), pages 265–283, 2016.
  * [5] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  Advances in neural information processing systems, 33:1877–1901, 2020.
  * [6] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc’aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al.  Large scale distributed deep networks.  Advances in neural information processing systems, 25, 2012.
  * [7] 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 Parallel Programming, pages 431–445, 2021.
  * [8] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  arXiv preprint arXiv:1706.02677, 2017.
  * [9] Will Douglas Heaven.  Gpt-4 is bigger and better than chatgpt—but openai won’t say why.  https://www.technologyreview.com/2023/03/14/1069823.  Accessed: 2023-05-15.
  * [10] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, et al.  Gpipe: Efficient training of giant neural networks using pipeline parallelism.  Advances in neural information processing systems, 32, 2019.
  * [11] Arpan Jain, Ammar Ahmad Awan, Asmaa M. Aljuhani, Jahanzeb Maqbool Hashmi, Quentin G. Anthony, Hari Subramoni, Dhabaleswar K. Panda, Raghu Machiraju, and Anil Parwani.  GEMS: gpu-enabled memory-aware model-parallelism system for distributed DNN training.  In International Conference for High Performance Computing, Networking, Storage and Analysis (SC), page 45. IEEE/ACM, 2020.
  * [12] Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parekh, Hieu Pham, Quoc Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig.  Scaling up visual and vision-language representation learning with noisy text supervision.  In International conference on machine learning, pages 4904–4916. PMLR, 2021.
  * [13] Zhihao Jia, Matei Zaharia, and Alex Aiken.  Beyond data and model parallelism for deep neural networks.  Proceedings of Machine Learning and Systems, 1:1–13, 2019.
  * [14] Alex Krizhevsky.  One weird trick for parallelizing convolutional neural networks.  arXiv preprint arXiv:1404.5997, 2014.
  * [15] Harlan M Krumholz, Sharon F Terry, and Joanne Waldstreicher.  Data acquisition, curation, and use for a continuously learning health system.  Jama, 316(16):1669–1670, 2016.
  * [16] Joel Lamy-Poirier.  Breadth-first pipeline parallelism.  arXiv preprint arXiv:2211.05953, 2022.
  * [17] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen.  Gshard: Scaling giant models with conditional computation and automatic sharding.  arXiv preprint arXiv:2006.16668, 2020.
  * [18] Mu Li, David G Andersen, Jun Woo Park, Alexander J Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J Shekita, and Bor-Yiing Su.  Scaling distributed machine learning with the parameter server.  In 11th USENIX Symposium on operating systems design and implementation (OSDI 14), pages 583–598, 2014.
  * [19] 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, SC ’21, New York, NY, USA, 2021. Association for Computing Machinery.
  * [20] Ziming Liu, Shenggan Cheng, Haotian Zhou, and Yang You.  Hanayo: Harnessing wave-like pipeline parallelism for enhanced large model training efficiency.  CoRR, abs/2308.15762, 2023.
  * [21] Azalia Mirhoseini, Hieu Pham, Quoc V Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean.  Device placement optimization with reinforcement learning.  In International Conference on Machine Learning, pages 2430–2439. PMLR, 2017.
  * [22] Dheevatsa Mudigere, Yuchen Hao, Jianyu Huang, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal, Jade Nie, Jongsoo Park, Liang Luo, et al.  High-performance, distributed training of large-scale deep learning recommendation models.  arXiv preprint arXiv:2104.05158, 2021.
  * [23] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia.  Pipedream: Generalized pipeline parallelism for dnn training.  In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1–15, 2019.
  * [24] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia.  Memory-efficient pipeline-parallel dnn training.  In International Conference on Machine Learning, pages 7937–7947. PMLR, 2021.
  * [25] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, et al.  Efficient large-scale language model training on gpu clusters using megatron-lm.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–15, 2021.
  * [26] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G Azzolini, et al.  Deep learning recommendation model for personalization and recommendation systems.  arXiv preprint arXiv:1906.00091, 2019.
  * [27] OpenAI.  Gpt-4 technical report, 2023.
  * [28] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  Advances in neural information processing systems, 32:8026–8037, 2019.
  * [29] David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean.  Carbon emissions and large neural network training, 2021.
  * [30] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.  Learning transferable visual models from natural language supervision, 2021.
  * [31] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.  Zero: Memory optimizations toward training trillion parameter models.  In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16. IEEE, 2020.
  * [32] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In International Conference on Machine Learning, pages 8821–8831. PMLR, 2021.
  * [33] Noam Shazeer, Youlong Cheng, Niki Parmar, Dustin Tran, Ashish Vaswani, Penporn Koanantakool, Peter Hawkins, HyoukJoong Lee, Mingsheng Hong, Cliff Young, et al.  Mesh-tensorflow: Deep learning for supercomputers.  Advances in neural information processing systems, 31, 2018.
  * [34] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism.  arXiv preprint arXiv:1909.08053, 2019.
  * [35] Annamalai Suresh, R Udendhran, and S Vimal.  Deep neural networks for multimodal imaging and biomedical applications.  IGI Global, 2020.
  * [36] Kazuhiko Takamizawa, Takao Nishizeki, and Nobuji Saito.  Linear-time computability of combinatorial problems on series-parallel graphs.  Journal of the ACM (JACM), 29(3):623–641, 1982.
  * [37] Wei Tan, Prayag Tiwari, Hari Mohan Pandey, Catarina Moreira, and Amit Kumar Jaiswal.  Multimodal medical image fusion algorithm in the era of big data.  Neural Computing and Applications, pages 1–21, 2020.
  * [38] Jakub M Tarnawski, Deepak Narayanan, and Amar Phanishayee.  Piper: Multidimensional planner for dnn parallelization.  Advances in Neural Information Processing Systems, 34:24829–24840, 2021.
  * [39] Colin Unger, Zhihao Jia, Wei Wu, Sina Lin, Mandeep Baines, Carlos Efrain Quintero Narvaez, Vinay Ramakrishnaiah, Nirmal Prajapati, Pat McCormick, Jamaludin Mohd-Yusof, et al.  Unity: Accelerating dnn training through joint optimization of algebraic transformations and parallelization.  In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 267–284, 2022.
  * [40] Leslie G Valiant.  A bridging model for parallel computation.  Communications of the ACM, 33(8):103–111, 1990.
  * [41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [42] Minjie Wang, Chien-chin Huang, and Jinyang Li.  Supporting very large models using automatic dataflow graph partitioning.  In Proceedings of the Fourteenth EuroSys Conference 2019, pages 1–17, 2019.
  * [43] Siyu Wang, Zongyan Cao, Chang Si, Lansong Diao, Jiamang Wang, and Wei Lin.  Ada-grouper: Accelerating pipeline parallelism in preempted network by adaptive group-scheduling for micro-batches.  arXiv preprint arXiv:2303.01675, 2023.
  * [44] Xiao Wang, Guangyao Chen, Guangwu Qian, Pengcheng Gao, Xiao-Yong Wei, Yaowei Wang, Yonghong Tian, and Wen Gao.  Large-scale multi-modal pre-trained models: A comprehensive survey.  Machine Intelligence Research, pages 1–36, 2023.
  * [45] Pengcheng Yang, Xiaoming Zhang, Wenpeng Zhang, Ming Yang, and Hong Wei.  Group-based interleaved pipeline parallelism for large-scale DNN training.  In International Conference on Learning Representations (ICLR), 2022\.
  * [46] 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.