Piper: Multidimensional Planner for DNN Parallelization
Authors: Jakub Tarnawski
Deepak Narayanan
Amar Phanishayee
Venue: NeurIPS 2021 Poster (2021)
Keywords: pipeline model parallelism
DNN partitioning
tensor parallelism

Abstract: The rapid increase in sizes of state-of-the-art DNN models, and consequently the increase in the compute and memory requirements of model training, has led to the development of many execution schemes such as data parallelism, pipeline model parallelism, tensor (intra-layer) model parallelism, and various memory-saving optimizations. However, no prior work has tackled the highly complex problem of optimally partitioning the DNN computation graph across many accelerators while combining all these parallelism modes and optimizations.

In this work, we introduce Piper, an efficient optimization algorithm for this problem that is based on a two-level dynamic programming approach. Our two-level approach is driven by the insight that being given tensor-parallelization techniques for individual layers (e.g., Megatron-LM's splits for transformer layers) significantly reduces the search space and makes the global problem tractable, compared to considering tensor-parallel configurations for the entire DNN operator graph.

Piper: Multidimensional Planner for DNN Parallelization
Jakub Tarnawski
Deepak Narayanan Microsoft Research
Microsoft Research
j jakub. tarnawski 0mi wwwsoft. com
dharayaaaanemicroosfftcccm
Amar Phanishayee
Microsoft Research amar amar@micoozoft. com
Abstract
The rapid increase in sizes of state-of-the thrt DNN models, and consequently the increase in the compute and memory requirements of model training. has led to the development of many execution schemes such data parallelism, pipeline model parallelism, tensor (intra-layer) model parallelism, and various memory-saving optimizations. However, no prior work has tackled the highly complex problem of optimally partitioning the DNN computation graph across many accelerators while combining all these parallelism modes and optimizations. In this work, we introduce Piper, an efficient optimization algorithm for this problem that is based two-level dynamic programming approach. Our two-level approach is driven by the insight that being given tensor- parsllapaallllization techniques for individual layers (e.g., Megatron-LM's splits for transformer layers) significantly reduces the search space and makes the global problem tractable, compared to considering tensor-parallel configurations for the entire DNN operator graph.
Introduction
Deep Neural Network (DNN) models have grown exponentially in size over the past two decades; consequently, modern DNNs are extremely computationally expensive. For example, state-of-ther art language model, GPT-3 (2) has an astounding 175 billion parameters, requiring 314 ZettaFL.OPS (3.14 x 28 FLOPS) of compute.
This trend of computationally expensive models with working-set sizes for training and inference that far exceed the memory capacities of individual DNN accelerators (e.g., GPUs, TPUs, FPGAs, ASICs) has reinvigorated the study of parallel training techniques that split DNN model state across different devices. In the "modern era", such model-parallel training techniques trace their roots back to AlexNet [14] and early influential systems such as DistBelief (6 and Project Adam (3] Recent advances in model-parallel training include tensor (intra-layer) model parallelism [22]. where individual operators of the model are partitioned over multiple workers, and pipeline model parallelism that combines inter-layer model parallelism with pipelining [17] [T]. Various hybrid
parallelism techniques that combine tensor, pipeline, and data parallelism across operators [13 19
have also been developed to more efficiently train wider range of models. Additionally, memory- saving optimizations like activation recomputation, which trade off throughput for memory, also being used 18 to scale training to larger models such as GPT-3.
Combining these dimensions, however, is non-trivial 19 since each dimension has trade-offs with respect to computational efficiency, amount of communication, and memory footprint. Given the importance of efficient model-parallel training (and inference 7 4)), partitioning model across
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

devices has thus received recent interest, evolving from a manual process driven by human experts to using algorithms for automated device placement of DNN operators. Focusing on algorithms. some
take black-box approach using learning- techniques [6
15
8
26
21
25 ; others build
23]. However, each of the prior algorithmic approaches tackle slightly different subset of parallelization techniques, system constraints, and DNN optimizations; for example, FlexFlow [3] does not consider pipelining. PipeDream 17 does not consider memory-saving optimizations that are useful for training large both do not consider tensor parallelism.
cost model and algorithmically solve an optimization problem [2
17
18
models, and PipeDream and PipeDream-2B4 [8]
Our contribution. In first, this paper covers a broad search space of parallelization and optimiza- tion techniques while determining placement of DNN layers on the provided hardware deployment (a given number of accelerators of specific type with certain performance characteristics, memory ca- pacity, and interconnect bandwidth). Specifically, our optimization algorithm, called Piper, considers three different dimensions of parallelism (data, tensor model, and pipeline model parallelism), and a general framework of memory-saving optimizations that trade off throughput for lower memory footprint. In this work, we consider pipeline parallelism schedules without pipeline flushes [17 We leave extensions to pipeline parallelism schedules with flushes [1] 19 to future work.
The focus of our work is on algorithms for DNN partitioning across devices to maximize throughput (for training and large-batcl inference of DNN models). Piper uses non-trivial combinations of data, tensor, and pipeline parallelism, as well recomputation, carefully chosen for each stage of the partition. Piper supports general DAG (Directed Acyelic Graph) model topologies and is agnostic of the type of DNN accelerator used. Piper is given a target "safe" batch size (e.g.. 1920 for large GPT-style models (20]) and returns a training configuration with this batch size. It is efficient both in theory and in practice.
memory-saving optimizations such activation
In Section
we evaluate Piper on real-world DNN profiles, and study the effects of combining the various parallelism modes and memory-saving optimizations on performance. We find that Piper's s large search space combined with its principled algorithmic approach allows it to find non-tri non-trivial high-quality parallelization configurations for each given problem instance. Piper compares favorably to planners from prior work (PipeDream, PipeDream-2BW in many settings, pipeline model parallelism combined with per-shard data parallelism is enough to obtain high throughput; however, tensor parallelism is also required when one needs a batch size smaller than the number of devices, as well as in highly nemory-constraingg settings. Piper's 's ability to form heterogeneous stages also gives it an advantage over prior work.
For example, we find (as expected) that
2 Background
In this section, we briefly describe methods to train models scale.
2.1 Pipeline model parallelism
Pipeline model parallelism is a recent technique [17 | used to accelerate the distributed training of both small and large models. The operators of model are partitioned over the available devices (i.e., each device hosts a subset of the operators of the model; typically these subsets are contiguous). A batch then split into smaller microbatches, and execution is pipelined across these microbatches: while the second device starts executing the first microbatch, the first device starts executing the second microbatch, and so on. Pipeline model parallelism features cheap point-to-point communication.
A key problem with deploying pipeline parallelism is determining how to schedule forward and backward microbatches on the available devices, and how to stash state to ensure a final high-quality model. Different schedules have different throughput, memory footprint, and semantics trade-offs. For example, GPipe [ maintains just single weight version, but introduces periodic pipeline flushes where microbatches in the pipeline are allowed to drain. This leads to idle resources and lower throughput, but also lower memory footprint. PipeDream [77 maintains a weight version for every in flight input microbatch, but does not use pipeline flushes, leading to both high throughput and memory footprint. PipeDream-225
[8 tries to find a middle ground, where gradients are coalesced across microbatches, reducing the number of weight versions, while still keeping throughput high. In this work, we assume that a non- non-flushing pipelining scheme used when considering pipelines with multiple stages (i.c., we do not use inter-layer model parallelism without pipelining).
2

---

2.2 Tensor (intra-layer) model parallelism
model can partitioned in other ways as well. Megatron-LM [22] slices transformer blocks in NLP models like BERT and GPT across available devices, in effect performing distributed matrix multiplications. All-to-all communication the form of scatters, gathers and all-reductions needs to be performed to ensure that these distributed operators retain their semantics. Megatron-LM uses a specific, hand-tuned strategy to distribute transformer layers across multiple GPUs. Tensor parallelism reduces per-device memory usage, model parameters are split across devices, and shows good scaling performance when using bighbbandwidth interconnects like NVLink [19.
Trying to find an optimal tensor parallelization strategy for the entire model by deciding how to split every operator a highly challenging problem with a large search space, and requires understanding the semantics of each operator. Modern models consist of a large number of operators, often with considerable branching. Moreover, each operator can be split in many different ways. For example, matrix multiplication can be split over the rows columns of the left (input) matrix or right (weight) matrix, or a combination of these. Some splits will require inserting an all-reduction operator to recover the result, and some will require an all-gather. Some will divide weights across devices, which helps with memory usage. The optimal split in terms of latency depends on the matrix dimensions and also downstream operators; matrix multiplication split over some dimension would accept input tensors split over the same dimension, but a dropout operator requires the entire input tensor.
2.3 Data parallelism
Data parallelism is the de-facto way of parallel training today. With data parallelism, every worker has copy of the full model. Inputs are sharded, and each worker computes weight gradients independently. Periodically, these weight gradients are aggregated using all-to-all communication or parameter server. Vanilla data parallelism assumes that model parameters fit on single device. For large models, data parallelism can be used on model shards; gradients need to only be exchanged among devices responsible for the relevant model shard. Data parallelism can be used with both pipeline (as shown in PipeDream [17) and tensor model parallelism (as shown in Megatron-LMM [22)
2.4 Memory-saving optimizations
When training large models, is common to use memory-saving optimizations that trade off a reduction in throughput for lower memory footprint. One example is activation recomputation, where the intermediate activations produced during the forward pass are not stored. Instead, the forward pass is performed a second time, just before the backward pass, from a stashed inpur activation, which is much smaller than the full set of intermediate activations. This drastically reduces the footprint of pipelined training, at the cost of executing the forward pass twice.
2.5 Related work on algorithms for partitioning and placement of DNN graphs
A number of systems have looked at the problem of partitioning DNN models across devices. However, each of these systems considers a subset of the parallelism dimensions considered this paper. FlexFlow [3] considers many dimensions (data, tensor, and inter-layer model parallelism), but not pipelining: moreover, it uses heuristic MCMC-based approach that requires domain knowledge to be encoded into the "parallel plan" generator. PipeDrream [1 and Tarnawski et al. (23) use dynamic programming and consider a combination of pipeline model parallelism and data parallelism; however, they support different model topology classes. PipeDream supports only linear (path) graphs while (23] supports general DAGs. While PipeDream oblivious to memory usage, its enhancement, PipeDream-2B1 [18), targets large models that do not necessarily fit on a single accelerator. Exploiting the repetitive structure of some of these large models, such as transformer- based language models, FipeDream-2BW3 planner only considers configurations where every stage in the pipeline is replicated an equal number of times (equi-replicated stages), thus simplifying the search for parallelization configurations. PipeDream-2B1 per-device memory capacity constraints and includes activation recomputation in its configuration search space. None of these three approaches consider tensor parallelism.
planner also explicitly considers
Learning-basee approaches (Mirhoseini et al. 16 15), Spotlight R Placeto 0. GDP 26 RE-
GAL (21) treat the objective (lateney or throughput) black box to be optimized using RL;
3

---

for runtimes are measured in an online fashion on the target system, which is very accurate but computa- tionally expensive. Moreover, some of them optimize for latency or peak memory usage rather than throughput, and they do not explicitly consider splitting individual operators or pipelining.
All existing algorithms, except PipeeDream-2BW also do not consider memory-saving optimizations, that trade off throughput for reduction in memory footprint, in combination with different modes of parallelism; activation recomputation is one such important optimization. These optimizations are commonly needed for training models that exceed the memory capacities of modern accelerators.
3 Problem Setup
In this section, we describe the search space of parallelization configurations that Piper explores, as well as the objective function (cost model) that we minimize in order to find the best solution. These give rise to the description of the optimization problem that Piper solves.
Piper is given as input a DNN workload (model) which is represented as a Directed Acyelic Graph (DAG), G = (V, E.) Each node corresponds to a single DNN layer such as LSTT or transformer (each layer is composed
lower-level operators, e.g., matrix multiplications or additions). Each edge (u, v) corresponds to a data transfer: node/layer requires the output of u as its input. The input graph is annotated with a number of attributes related to compute times, communication costs, and memory usage; we introduce and motivate these attributes in this section, and then formally list them in Section 3.6 To obtain an input to be supplied to Piper, these attributes should be profiled (or estimated). Furthermore. Piper is given: the number of devices (K), available memory per device (M), the network bandwidth (B), and the target number of microbatches in a batch (N). N is the ratio of the chosen batch size (usually the maximum that is safe for convergence, e.g., large transformer-bast LMs) to the provided microbatch size.
3.1 Space of configurations
Piper's 's search space consists of partitions of G into consecutive contiguous subgraphs, also called stages, which are executed using disjoint sets of devices. Each node/layer is assigned to exactly one stage. Each stage is executed using some number d I of devices, combining data parallelism of degree d with tensor parallelism of degree t. That is, we want to find partition V = S] U... U Se and, for each stage parallelism, (c) further configuration: how (for each layer in the subgraph) the tensor parallelism is carried out, as well as what other optimizations are used.
1,... &: (a) the degree d, of data parallelism, (b) the degree t, of tensor
Thus the i-th stage will be run on de -, devices, and so we must have Ei d. S K. Furthermore, each stage should be feasible in terms of memory usage, the sum of data aata-parallel degrees Ei= d)) must be at most N (in order to saturate the pipeline), and we need to divide N (to be able divide microbatches evenly among the data matta-paallel replicas). We call this output solution. We use the definition of contiguity from (23 (see Fig. in the Appendix for an illustration):
Definition 3.1 We say that a set S C V is contiguous if there do not exist nodes u E S. V.S. and w such that e reachable from u and w is reachable from v.
3.2 Peme-Per-Sample
Our objective is to maximize throughput; equivalently, minimize its inverse, the Time-Per- Sample (TPS). For non fo--flusiin pipeline execution schemes such as PipeDream or PipeDream-299 it has been argued 23 as well shown in practice [17 that the Time-Per-Sample of a solution equals the maximum Time-Per-e Per-Spee of a device. The problem finding solution with minimal TPS (absent data tensor parallelism and other optimizations) that model, which we extend, the TPS of a device/subgraph is given as the sum of compute times of all nodes in the subgraph, plus costs of communication along edges coming into leaving the subgraph. We now discuss how to incorporate the modes of parallelism from Section into the model.
addressed by Tarnawski et al. (3); in

---

3.3 Tensor Model Parallelism Configurations (TMPCss
As we remarked in [Section 2.5 the task of making global, joint decision about how tensor- parallelize every operator in a DNN computation graph is very complex. In this work we therefore reduce this task the problem of coming up with good tensor parallelization techniques for individual layers. Solving the latter problem is not a goal of Piper; instead, we assume that we already know (i.e., Piper given as input) some number of good tensor parallelism configurations (for some layers v and some degrees t of tensor parallelism). In practice, such configurations can be devised by hand (as done e.g. for transformer layers in Megatron-LM [22]; we use their technique in our experimental evaluation), or the problem can be attempted using automated methods similar to e.g. FlexFlow [I3 or Tofu [24] or RL.-based approaches. This problem has to only be solved for each layer type, e.g.. even if the input DNN has many transformer layers, one just needs good technique for a transformer layer. It is also not necessary to provide configurations for all layers v and possible degrees t. On the other hand, there could be many reasonable tensor parallelism configurations for the same and t; for example, there might be a trade-off between memory usage and runtime.
For each available configuration, Piper gets as input certain attributes related to compute, communica- tion and memory costs for the target model (listed formally in in[Section 3.6). tuple containing these attributes is called a Tensor Model Purallelism Configuration (TMPC). Piper receives a (possibly empty) list T (v, 1) TMPCs for each node v and degree of tensor parallelism (including 1). Memory-saving optimizations (e.g., activation recomputation) are modeled using TMPCs as well
3.4 Data parallelism
we model the effect of this on TPS follows. The compute load of the stage is spread evenly across devices (with a factor 1/d). However, data-parallel execution necessitates periodic synchronization of weight gradients, which has total communication cost of 4 - 1) w. where w is the size of parameters in all layers in the stage. This gives cost of 4 ==1
Consider a stage that is data- aata-parallelzzed across d > 1 devices. Following PipeDream 17
bytes per device (per batch, i.e., N microbatches).
3.5 Memory usage in pipelined configurations
With pipelining, every stage needs to store some number of weight and activation versions; these depend on the number of in-fiight microbatches (per data- mata-parallel replica). Specifically, the memory usage of a stage can be modeled an affine function of this number. That is, for every stage with data-parallel degree d and suffix sum of data-parallel degrees 8, memory usage would be a- [x/d] + for appropriate coefficients and b. See Apeendiddxx A for explanation. E.g., in PreasDrem--ttyle execution, a would he the size of weights plus the size of all activations in the subgraph, and b would be the size of optimizer state and temporary buffers. PipeDream-2B8 2BW uses only two stashed weight versions instead of [s/d] many, so the size of weights would contribute only to b (multiplied by 2). not a. If activation recomputation is used, then a accounts for only the sizes of input activations of every layer, and b accounts for all other activations. This accounting is more precise than PipeDream's ss which ignores memory usage, or Tarnawski al.'s 233, which uses a constant function.
3.6
Input and output
Now can formally (re-)state Piper's input:
a Directed Acyelic Graph G (V, E): layer graph of DNN,
for every edge (u, v) E E, an associated communication cost c(u, v) (in bytes),
number K of devices, memory MM per device, network bandwidth B, [maximum] number N of microbatches in a batch,
for every node/layer v E V. and every degree / 1= 1,..,, K of tensor parallelism: a (possibly empty) list T(v, 1) of TMPCs Section 3.3
The description of a TMPC in T(v, f) must correspond to tensor parallelism configuration where, for both the forward and the backward pass, each of the devices start execution already with full
1 As an example, node e might have the following TMPCs as input: for 1. one TMPC corresponding to regular execution, and another with activation recomputation; and for t = 8, one TMPC corresponding to tensor parallelization in dimension I and another using dimension 2, with and without activation recomputation (four in total).
5

---

copy of each input activation, and must end execution having in its memory a full copy of every output activation (see Appendixx B for discussion). Each TMPC description X E T(,, contains:
a Peme-Per-Sample X.p that encompasses the compute time together with the communica- tion costs of tensor parallelism (data transfers between the devices) |
for every incoming edge (u, u) E E, a communication cost X.c (u, v) (in bytes) that it would take to synchronize the tensor incoming over that edge between t devices (going from a state where one device has the tensor to a state where all have it); for more explanation and motivation this, see Appendix
analogous quantity X.c"(v, w) for every outgoing edge (v, w) E, the size X.w (in bytes) of parameters/weights on each device,
memory usage: two amounts X.a, X.b of bytes such that if layer v is located on a stage with data-parallel degree d and data-parallel degrees of this and later stages sum up then the memory consumption on each device is X.a [s/d] + X.b (see Section 3.5
Piper outputs the following low-TPS solution: collection of contiguous subgraphs/stages Si, Sp such that V S₁ U... U St, and for each stage i: (a) the degree de of data parallelism, (b) the degree of tensor parallelism, (c) for each node/layer E Sy: the index of the TMPC selected in T(v 14).
All nodes in stage use the same degrees of data and tensor parallelism, but different nodes (even all correspond to e.g. transformer layers) can use different TMPCs. For example, some but not all layers in a stage might employ activation recomputation.
Algorithm
Our algorithm is based on dynamic programming downsets, extending recent work 23
Definition 4.1 We call set D V of nodes downset if (u, u) E and D implies D. As proved there, going from downset to downset can obtain every possible contiguous subgraph:
Fact 4.2 A set S C V of nodes is contiguous (see Definition 3.7 and only if it is the difference of Iwo downsets: S = D D' where D' C D.
We rely on the domain insight that layer-granularity computation graphs of modern DNNs do not have many downsets, making tractable the approach of considering all downsets and all contiguous sets using dynamic programming. We define the following dynamic programming table: dp [k][s] is the minimum TPS of solution that partitions the downset D among & devices (recall that the TPS of solution is the maximum TPS over all & devices) such that the sum of data-parallel degrees of all stages is s. key novel point in our approach is that using the single variable s, which ranges from to N, we can precisely control both memory usage (see feetion 3.5) and batch size (Appendix A).
for every is and The answer (minimum TPS) when using any number & of devices and sum s of data-parallel degrees will be held in [<][][ and the final answer is mink==1..... min,==1....N dpl [A]ss).
We initialize the table follows: dy 0J [
We proceed to the main part: the recursion. For every possible sub-downset D', we consider partitioning the contiguous subgraph S D D' with data-parallel degree d and tensor-paralle degree i, so that d devices are used for S and the rest for D'. For 0 Ç C V, we have
[k/d]
dpl [p]|][
min
min
min max(dp[D'\kk d' t][s d). TPS(D) D',t,
downset D'C CDD d=1. 1....
d|N
where we define TPS(S, d, s) as the minimum Time-Per-Sampi when optimally partitioning $ with degrees d and t of data and tensor parallelism, respectively, and so that the sum of data-parallel degrees of this and all further stages is s. We now describe how we compute this quantity.
The knapsack subproblem. How compute TPS(S, t, d, s)? Given parallelism degrees d and f, we need to choose one TMPC X E t(u,t) for each E S. Some TMPCs may have better
Inncludes the runtime of both the forward and the backward passes.

---

TPS. but others may use less memory. More concretely, we want a combination that minimizes the total TPS of the stage while staying under the memory limit M. Both the total TPS and the total memory usage of any combination can be expressed as sums of the TPSes and the memory usages of TMPCs of individual nodes v E S (more details below). However, there are still exponentially many (II (T(u, f)]) possibilities, and in fact this subproblem is NP-hard, as it generalizes the (-1 knapsack problem. At the same time, we require a very fast solution as it needs to be solved for all S, t, d, We use simple bang-per-buck heuristic:
start by picking the lowest-TPP (fastest) TMPC for each U E S. as long the memory usage is too high, pick the TMPC in UPes T(e, 1) whose ratio of decrease memory usage (compared to the currently picked TMPC for v) to increase in TPS is the highest, and pick it for (instead of the currently picked one).
If no combination of TMPCs is memory-feaiblle we return +00. This is the only non-exact component of the algorithm: however, we find that this heuristic almost always finds the optimal solution in practice. See Appendixx for more discussion.
TPS of a fixed configuration. Finally, how do we compute the TPS of a stage for given S, t, d, s and selection of an TMPC ((v) T(u, t) for every E S? We need take compute, communication and memory into account. Compute is the simplest: compute a ees (()).p. For communication, we have costs stemming from sending activations into S, sending activations out of 5, and costs of all-reductions from using data parallelism (all in bytes):
Σ
(c(u, + X(u)).e"uu (u)
comm
(u,D)s (8)
d d-1 d. VES
Σ
2 (c(v, w) + ((t))cc(u,,
(1,)))&+($)
where we used the notation 8 (S) and 8 (5) for the sets of incoming and outgoing edges of S. Memory usage is given by: mem = Eves (X(v).a [s/d] + (().b). The final TPS for the selected S, f, d, and (X()))wes is compute + comm if mem S M; otherwise we return +00.
B
Running time. Under natural assumptions, Piper's running time is Ö(V K d(N))), where d(N) is the number of divisors of N. A detailed analysis can be found in Appendix D
Evaluation
We now evaluate the Piper algorithm on real-world DNN models. Besides the efficiency and scalability of Piper, we are also interested in the advantages of the configurations produced by Piper over best possible configurations that forgo least one of the modes of parallelism memory-saving optimizations, as well as over best possible partitioot configurations such as those found by PipeDream-22BW planner algorithm. We stress that we focus on evaluating the Piper partitioning algorithm in this work, not any particular pipelined DNN training system. Moreover, with fixed batch size and pipelining scheme, minimizing TPS is equivalent to minimizing time-to-convergence.
Inputs. For our comparisons, we use BERT-32 32 model, which consists of an embedding layer, 32 transformer layers and a pooling layer. This model has 5.9 billion parameters, with similar dimensions to the GPT-3 model with 6.7B parameters (2] and does not fit on a single GPU. We provide TMPCs for non-tensor-paras parnor-paaallllizeed (t 1) and censor-parallelizer executions of transformer layers [22 (t E (2, 4,8 8)), each with and without activation recomputation. These TMPCs are obtained by profiling models implemented in PyTorch on NVidia A GPUs interconnected with 300 GB/s bandwidth NVSwitch within server, and 25 GB/s across servers. In our evaluation, assume that all K < 2048 devices are connected by a flat network topology with a (full bisection) bandwidth of B = 25 GB/s (used for all non-tensor--pralleliis communication). Our TMPCs model the memory footprint assuming the PipeDream-2B3 pipelining scheme.
Baselines.
We compare Piper to:
7

---

3.0x
Piper
no TP
equi-no TP
throughput
2.0x
1,000 x
00.78x
1.000
10000
1.000 1.000
1.000 X
1.000
1.000 x
1,000 x
1,000 x
1.000
1.000 x
1,000 x
1,000 x
1,000
1.000 x
1,000
1.000
1,000
1.00 x
3880
no DP
C980x
1.000
no AR
1.000
0.987
1.000 x
equi
1.000 x
1.000 x
(0.434x
0.672
0.6559
00.428 x
0.419x
1.0x
0.123X
0.037 X
0.025
0.016 x
0.0X
(a) BERT-32, 80GB of memory on each & devices, N 512.
throughput
1.000
1.000x
1.000
0718%
D.718x
1.000x
1.000x
0.703x
0.719x
0.719x
1.000x
1.000x
0. 0.706x
0.722x
0.722
1.000x
1.000x
0.727x
0.735x
0.735x
1.000x
0.6766x
0.960x
0.834x
1.000x
1.000x
10.971x
D0.56
1.0x
U.144x
0.038x
0.028 x
0.
0.497
0.017
0.449x
0.316x
m
128
612
1024
2045
(b) BERT-32, 8GB of memory on each of devices, N = 512.
1.000x
0.900 x
1.000x
1.000x
1.000x
x.00xx
1.00xx
throughput
0.734x
0.735x
10 0,745x
1.0x
3.000
0.000
2.00
X
0.000
0..45 x
0.2113 13x
0.000
0.27xx
22226
0
0.234x
5.000x
10.271x x
10058
0.238
0. 0.000x
0.2766x
0.031 x
404x
0585xx
0.000
10 0.300x
0.023 x
0.293x
0.370x
0.000
00 x333
x.0x
64
128
512
1024
&
2048
(c) BERT-32. 32, 2GB of memory on each of k devices, N = 512.
1.000x
1.000
x.000x
1.000x
throughput
0.603x
1.0x
0.
0.000
x.000
0.000
0.000
0.151x
x.000x
0.000
0.000x
0,000x
0.076x
0.000
0.000
0.000
0.000x
0.038 x
0.000
0.000
1.000x
0.000
0.0
128
512
1024
2048
(d) BERT-32. 32, 1GB of memory cach of l devices, N = 512.
Figure 1: Comparison of Piper to baselines on the basis of throughput normalized to the best configuration. Throughput is the inverse of TPS (Time-Per-Sample); higher throughput is better. Piper always obtains the highest throughput (1.000x). We vary the numbers of devices (from 2048) and the memory per device, which impacts the space of feasible solutions. For 1GB- or 2GB-devices, baselines are often unable to find any -emorr-ffasibll configuration (0.000x), 1GB models scenario where TP and AR are both required fit even single layer on device.
no DP: a version of Piper that forgoes data parallelism (i.e.,
TP: version of Piper that forgoes tensor parallelism (i.c., 1),
no AR: version of Piper that forgoes activation recomputation,
equi: an algorithm that equi-partitions layers into w equal groups and uses d-data parallelism and I-tensor parallelism, choosing the lowest-TPS among all valid (w, triples: also considers performance with and without activation recomputation (for the entire model), and places the embedding and pooling layers either on separate stages with the first and last transformer layer respectively,
equi-no TP: a version equi that forgoes tensor parallelism (i.e., t = 1); we show this as a surrogate for PipeDream-22WW planning algorithm.
Both no and no AR are more sophisticated than PipeDream's partitioning algorithm, which forgoes tensor parallelism, activation recomputation, and also assumes M co.

---

TPS (logarithmic scale)
TPS (logarithmic scale)
10-1
10-1
10
10
Piper
equi
Piper
equi
no DP
equi-no no TP
no DP
equi-no no TP
no TP
no TP
10 3
10
10a
10a
10s
10a
10a
10a
Number of devices (logarithmic scale)
Number & of devices (logarithmic scale)
(a) BERT-32, 8 GB RAM per device, N = 512.
(b) BERT-32, 8 GB RAM per device, N = 960.
Figure A single run Piper produces table of optimal Penne-e (TPS) for every number of devices. Here, we show the best TPS for Piper and the baselines versus the number of devices, for two different batch size (N) values. Lower TPS better. Tensor parallelism needed use more than N devices, TPS than Piper for most
by the TP baseline plateauing after &s N. equi obtains 10-15% worse
BERT-48; N=480
BERT-96; N=480
400
BERT-48; N=512
BERT-96; N=512
Runtime [s]
BERT-48; N=1920
BERT-96; N=1920
BERT-48; N=2048
BERT-96: N=2048
200
200
400
600
800
1,000
1,200
1,400
1,600
1,800
2,000
Number devices
Figure 3: Piper's running time a function of the number of devices, for various DNN model sizes and batch sizes (N). Bars show standard deviation. Single-core implementation executed Intel Xeon E5-2673 v4 CPU. Runtime depends on the number divisors of N (over which is iterated).
5.1 Results
Results of our evaluation in terms of the quality (TPS) of the obtained configurations are given in The running times of Piper are shown in Fig. 3 We discuss the main takeaways.
Figs. and 2
Data parallelism (DP). Data parallelism is crucial, as expected. Without it, scaling to more devices is impossible; e.g., in Fig. the TPS of no DP stops decreasing beyond A = 22 devices.
Tensor model parallelism (TP). We found that, in a vacuum, the speedup obtained by TP much less than that of DP due frequent and costly synchronization. However, TP also shrinks the parameter size per device, which helps with memory feasibility, and does not increase s (in contrast DP), which enables one to stay below the specified number of microbatches in batch. In particular, TP is essential to use k > N devices (see the no TP baseline in e.g. Fig. 1a or Fig. 2b. TP also required in when even a single layer cannot fit onto device (Fg.. ld).
Activation recomputation (AR). Activation recomputation is required to obtain any memory- feasible solution in certain cases. Even when it is possible to do so without AR, foregoing it can force quite imbalanced (in terms TPS) splits in order fit memory. For example, Fig. Ic without AR, one is forced to use small values of stay within memory, which limits the number of stages and data-paran parallll degrees, and hence performance (using TP replacement yields worse throughput). In scenarios with enough memory relative to DNN size (Fgg. Ia), AR is not useful.

---

Predicted TPS
True TPS
Figure 4: Real vs. Piper-predicted Time-Per-Sample for 16 configurations of pipeline, data, and tensor parallelism for BERT-32 with K = 64 devices. Training times were measured on system with NVidia DGX A100 machines, each with GPUs. The Pearson correlation coefficient is 0.95.
Power of Piper's search space. Piper has the ability to return configurations with stages that different from each other (e.g., different numbers of layers, degrees d and 4, memory optimizations). We found that this leads a lower optimal TPS than the baseline equi even for extremely repetitive DNNs such as BERT (see Figs. I [nd This is useful for two reasons. First, in earlier stages, the parameter that controls memory usage (Section 3.5) is higher (as is "suffix sum" of DP degrees over stages), which means that it is optimal to have larger degrees d or t or smaller numbers of layers in the earlier stages. other words, in an optimal equi-partitioned configuration, if memory is bottleneck, then the memory limit will be attained almost tightly in the first stage, which needs to store activations for many in in-ftight microbatches, but will be loose in the last stage, which only stores activations for one microbatch. In contrast, Piper will effectively adjust all stages to ensure that almost all available memory used on every device. Second, Piper also able adjust different layer types; for instance, the embedding layer can use a different TP degree than transformer layers.
Running time Piper. Fig. shows the scalability of Piper's algorithm as a function of and [V] as well N, on larger BERT inputs created by expanding our BERT-32 profiles. We find that Piper terminates minutes, which a very small percentage of the actual training times of the DNNs (days to weeks). We also note that Piper's s algorithm is embarrassingly parallel, as for every I, computing the table TPS(I \ .) and updating dp can be done in parallel over all sub-downsets I' C I. With such implementation, runtimes for the algorithm would scale linearly a multi-core CPU server.
Quality of cost model. Fig. 4 shows that Piper's cost model very close to real performance. In this work, we assumed flat network topology. While this already yields precise estimates, Piper can be extended to handle hierarchical network topologies (see (eppendix E}.
6 Conclusions and Future Work
In this work, we presented a two-level dynamic programming approach that determines how to partition a provided model over given number of accelerators, while optimally composing tensor and pipeline model parallelism with data parallelism and memory-saving optimizations. evaluation, we found that Piper finds high-throughput parallelization strategies that in many cases improve upon planners from prior work, and that Piper's power of finding non-trivial combinations of parallelism modes is especially beneficial in large-scale
our
mnemory-constraae
scenarios.
While we consider pipeline parallelism schedules without pipeline flushes in this paper, one way of extending Piper to consider schemes with pipeline flushes is by adding the pipeline depth to Piper's dynamic programming state. However, this would worsen Piper's time complexity. We leave such extensions and improvements to the algorithm for future work. The main open question resulting from this work is automatically finding good lensor-parallelism schemes for each individual layer type (rather than the entire DNN operator graph). Piper can then determine how to best compose these per-layer tensor censor-parallelism schemes to obtain a high-performance training configuration for the entire model.
10

---

Acknowledgments and Disclosure of Funding
This work was done in the context of Project Fiddle MSR. During this work, the second author was at Stanford University, where he was supported in part by NSF Graduate Research Fellowship grant DGE- 165651... Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors alone.
References
[1] Ravichandra Addanki, Shaileshh Bojja Venkatakrishnan, Shreyan Gupta, Hongzi Mao, and Mohammad Alizadeh, Learning Generalizable Device Placement Algorithms for Distributed Machine Learning. In Advances Neural Information Processing Systems 32, pages 3981-3991
Curran Associates, Inc., 2019.
[2] Tom B Brown, Benjamin Mann, Nick Ryder, Mclanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language Models are Few-Shot Learners. arXiv preprint arXiv:20005.4415. 2020.
[3] Trishul M Chilimbi, Yutaka Suzue, Johnson Apacible, and Karthik Kalyanaraman, Project Adam: Building an Efficient and Scalable Deep Learning Training System. In Ithh USENIX Symposium on Operating Systems Design and Implementation (OSDI '14), volume 14, pages 571-582, 2014.
[4] Eric S. Chung, Jeremy Fowers, Kalin Ovtcharov, Michael Papamichael, Adrian M. Caulfield, Todd Massengill, Ming Liu, Daniel Lo, Shlomi Alkalay, Michael Haselman, Maleen Abeydeera, Logan Adams, Hari Angepat, Christian Boehn, Derek Chiou, Oren Firestein, Alessandro Forin, Kang Su Gatlin, Mahdi Ghandi, Stephen Heil, Kyle Holohan, Ahmad El Husseini, Tamás Juhász, Kara Kagi, Ratna Kovvuri, Sitaram Lanka, Friedel van Megen, Dima Mukhortov, Prerak Patel, Brandon Perez, Amanda Rapsang, Steven K. Reinhardt, Bita Rouhani, Adam Sapek, Raja Seera, Sangeetha Shekar, Balaji Sridharan, Gabriel Weisz, Lisa Woods, Phillip Yi Xiao, Dan Zhang, Ritchie Zhao, and Doug Burger. Serving DNNs in Real Time at Datacenter Scale with Project Brainwave. In JEEE Micro, 2018.
[5] J. Csirik, J. B G Frenk, M. Labbe, and S. Zhang. Heuristics for the 0-1 Min-Knapsack Problem. Acta Cyberneticu, 10(1-2):15-20, January 1991.
[6] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Andrew Senior, Paul Tucker, Ke Yang, Quoe V Le, et Large Scale Distributed Deep Networks. In Advances in Neural Information Processing Systems, pages
2012.
[7] Jeremy Fowers, Kalin Ovtcharov, Michael Papamichael, Todd Massengill, Ming Liu, Daniel Lo. Shlomi Alkalay, Michael Haselman, Logan Adams, Mahdi Ghandi, Stephen Heil, Prerak Patel, Adam Sapek, Gabriel Weisz, Lisa Woods, Sitaram Lanka, Steve Reinhardt, Adrian Caulfield, Eric Chung, and Doug Burger. A Configurable Cloud-Scale DNN Processor for Real-Time AI. In ACM/IEEE 45th Annual International Symposium on Computer Architecture (ISCA '18), 2018.
[8] Yuanxiang Gao, Li Chen, and Baochun Li. Spotlight: Optimizing Device Placement for Training Deep Neural Networks. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1676-1684, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018. PMLR.
[9] L.L. Gurobi Optimization. Gurobi Optimizer Reference Manual, 2019.
[10] G.H. Hardy, E.M. Wright, D.R. Heath-Brown R. Heath-Brown, J. Silverman, and A. Wiles. An Introduction the Theory Numbers. Oxford mathematics. OUP Oxford, 2008.
[11] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V. Le, Yonghui Wu, and Zhifeng Chen. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. Advances in Neural Informa- tion Processing Systems, 2019.

---

[12] Zhihao Jia, Sina Lin, Charles R Qi, and Alex Aiken. Exploring Hidden Dimensions in Paral- lelizing Convolutional Neural Networks. In Proceedings of the 28th International Conference on Machine Learning (ICMI. '8), 2018.
[13] Zhihao Jía, Matei Zaharia, and Alex Aiken. Beyond Data and Model Parallelism for Deep Neural Networks. In Proceedings of the 2nd SysML Conference, SysML '19, Palo Alto, CA, USA, 2019.
[14] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. ImageNet Classification with Deep Convolutional Neural Networks. In Advances in Neural Information Processing Systems, 2012.
[15] Azalia Mirhoseini, Anna Goldie, Hieu Pham, Benoit Steiner, Quoc V. Le, and Jeffrey Dean. A Hierarchical Model for Device Placement. ICLR, 2018.
[16] Azalia Mirhoseini, Hieu Pham, Quoc V Le, Benoit Steiner, Rasmus Larsen, Yuefeng Zhou, Naveen Kumar, Mohammad Norouzi, Samy Bengio, and Jeff Dean. Device Placement Opti- mization with Reinforeement Learning. In Proceedings of the 34th International Conference on Machine Learning- -aarning-Volume 70, pages 2430-2439. JMLR. org, 2017.
[17] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, Phil Gibbons, and Matei Zaharia. PipeDream: Generalized Pipeline Parallelism for DNN Training. In Proc. 27th ACM Symposium Operating Systems Principles (SOSP), Huntsville, ON, Canada, October 2019.
[18] Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory- Efficient Pipeline-Paralle DNN Training, 2020.
[19] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, al. Efficient Large-Scale Language Model Training on GPU Clusters. arXiv preprint arXiv:2104.04473 2021.
[20] NVIDIA and Microsoft. Using DeepSpeed and Megatron to Train Megatron- NLG 530B, the World's Largest and Most Powerful Generative Language Model, 2021.
[21] Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs. In International Conference on Learning Representations, 2020.
[22] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LNM Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020.
[23] Jakub Tarnawski, Amar Phanishayee, Nikhil R. Devanur, Divya Mahajan, and Fanny Nina Paravecino. Efficient Algorithms for Device Placement of DNN Graph Operators. In NeurlPS, 2020.
[24] 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, EuroSys 19, New York, NY, USA, 2019. Association for Computing Machinery.
[25] Yanqi Zhou, Sudip Roy, Amirali AbdoIrashidi, Daniel Wong, Peter Ma, Qiumin Xu, Hanxiao Liu, Phitchaya Phothilimtha. Shen Wang, Anna Goldie, Azalia Mirhoseini, and James Laudon. Transferable Graph Optimizers for ML Compilers. In H. Larochelle. M. Ranzato, R. Had- sell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18844-1855 Curran Associates, Inc., 2020.
[26) Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter C. Ma, Qiumin Xu, Ming Zhong, Hanxiao Liu, Anna Goldie, Azalia Mirhoseini, and James Laudon. GDP: Generalized Device Placement for Dataflow Graphs, 2019.
12


Reviews:
[{"review_id": "U0nU6o8IBOT", "reviewer": "Reviewer_bfSz", "summary": "Parallel/distributed training of machine learning models is a big and complex task. This paper proposed a dynamic programming algorithm in searching optimal parallelization configurations for machine learning models, considering tensor parallelization, data parallelization, and pipeline parallelization. Designing such algorithm is important and interesting, and potentially relieve machine learning app writers from manually designing parallelization configurations. To provided an optimal (or close to optimal) parallelization configuration without breaching memory limitation is already good contributions in real-world scenarios.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Finding parallelization configurations is not a new problem, but the problem is so complex that new contributions are always needed. This paper didn't bring up any new ways to parallelize machine learning models, but tried to provide a better (or optimal) combinations of known parallelization techniques, which is also very interesting. Finding good combinations of parallelization techniques isn't new idea either, but previous publications have missed opportunities. \n\nDescribing such dynamic programming algorithm is not easy. This paper did a good job in clearly presenting the algorithm with formal definitions. The main dynamic programming recursion is deciding where to cut the next pipeline stage from the difference of 2 downsets. This is a clearly good idea. Although the algorithm can probably be optimized to not consider very small pipeline stages or very large pipeline stages, since all the pipeline stages should still probably have similar latencies for optimal performance, this doesn't hurt the correctness of the algorithm.\n\nFirst question I have is: are the authors only considering PipeDream-style pipelining, and discarding GPipe-style pipelining? It is unclear to me if the algorithm is considering the extra cost of pipeline-filling and pipeline exhausting in GPipe. If the answer is yes, that is probably fine, since PipeDream-style pipelining is probably higher in throughput. If the answer is no, I would like to hear why.\n\nEach pipeline stage is handling a part of the machine learning model, including both its forward computations and backward computations. This is probably a reasonable limitation. Within each pipeline stage, the computation of TPS depends on a heuristics. This is also quite reasonable I think, since solving this problem exactly is almost impossible, not only because the problem is NP-hard, but also because the estimations of memory consumption is very unlikely to be accurate. Each device might reserve memories for other purposes, and some details of memory recycling can be chaotic. Using a heuristics and trying not to use to the memory limit is a wise choice.\n\nHowever, I think one significant problem of this paper is that the handling of activation recomputation depends on the TMPC. That is to say, recomputation can only happen in a layer. For instance, given a pipeline stage of A -> B -> C -> D -> E, the backward computation of the layers (A', B', C', D', and E') might choose to recompute something in each layer, but you can never recompute multiple layers together, such that the backward computation becomes B -> C -> D -> E' -> D' -> C' -> B' -> A'. In my example, the whole layers of B C and D are recomputed. This limitation might not be so significant in this paper, since the evaluations used models that have huge layers (such as transformer layers). But if the model has many small layers, then it is more likely that consecutive layers need to be recomputed to save memory. How difficult is it for this type of recomputation to be considered?\n\nWhen modeling the cost of parallelization, the cost of communication is for sure critical. That is why when the authors say that in tensor parallelization, each of the t devices has a full copy of the inputs and outputs, the reviewer is a bit doubtful. One likely reason for this to not be a problem, is that tensor parallelization is not heavily used or optimized at runtime, or most tensor parallelizations likely require full copy of the inputs and outputs to be on all t devices. Based on the supplemental data, it seems that the latter is the reason? This limitation is thus likely reasonable, since the modeling of cost doesn't have to be very very accurate.\n\nIn the algorithm, it seems that each pipeline stage has one d (data parallelization) and t (tensor parallelization) value. This is a limitation such that layers in the same pipeline stage must have the same d and t. The reviewer is not very sure how big this limitation is. For instance, if a pipeline stage has layer A -> B -> C -> D, and only layer C has a very efficient tensor parallelization of t = 4, does that mean that the algorithm cannot make use of it, since the other layers (A, B, and D) do not have a TMPC for t = 4? I guess the authors can create a (TMPC of t = 4) by duplicating a (TMPC of t = 1) 4 times (just like data parallelization)? but I am not sure if that works, since it breaks the conceptional separation of data parallelization and tensor parallelization. The authors can certainly argue that, the lack of TMPCs is not the problem for this paper. However, this limitation seems to be quite practical. Can the authors elaborate here?\n\nOf course, in the end, we care about if the new algorithm is providing better throughput than existing methods, and if we learn anything exciting from the dynamic searching algorithm. Evaluation seems to be the main reason of rejection from previous submissions. Indeed, it sticks out that most evaluations are just running the algorithm itself and presenting the estimated throughput. The Figure 1 showed interesting ablation tests, and demonstrated that removing some options in the search space caused worse throughput. It does increase the confidence that the algorithm is considering the right things, but it makes reviewers less confident, since the modeling/estimation itself could be incorrect. The added Figure 4 does boost that confidence.\n\nThe take aways are not so exciting but rather expected. The \"Power of Piper's search space\" is particularly interesting, since it provides a theoretical explanation of why Piper is better. \n\nIn summary, this paper addressed a very relevant problem of high importance. The solution is a small delta to what has been published, but it is nevertheless important. There are 2 significant limitations (about recomputation and fixed d/t for all layers in one stage), which the reviewer hope to get elaborations from the authors. The evaluations are not ideal, but the challenges of running all configurations in real machines is probably too high to request. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors didn't say that their work can potentially solve the climate change problem, but that is totally fine.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15 hours"}, "extra_text": {}}, {"review_id": "D3EchZJnbLH", "reviewer": "Reviewer_cpuf", "summary": "This paper proposes Piper, an efficient dynamic programming algorithm to combine data parallelism, tensor model parallelism, pipeline model parallelism, and activation rematerialization.\nThe paper evaluates the algorithm on BERT models with real per-layer profiling data and simulated end-to-end results.\nIt provides insightful analyses of the results.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality and clarity:  \nThis paper tackles an important problem. The DP algorithm and the idea of using TMPC are novel.\nThe paper is well written with the necessary background information. It precisely discusses its relations to prior pipeline-parallelism works.\n\nStrength:\n1. It is the first work that uses dynamic programming to search in a non-trivial search space that includes data parallelism, tensor model parallelism, pipeline model parallelism, and activation rematerialization.\n2. Clear analysis and explanation of the results.\n\nWeakness:\n1. The evaluation results in Figure.1 are not measured in real systems. They are obtained by simulation. The y-axis\n\"throughput\" is computed by formulas proposed by authors under various assumptions. It does not reflect the\nbehavior of real systems. As a paper targeting a system problem, only running evaluation in a simulator is not\nenough. The authors show real measured throughput in Figure.4, meaning that the authors have access to the machines.\nCould you please use real profiling data in Figure.1, at least for small-scale experiments? The insight drew from these simulation experiments can be misleading if we do not consider some important factors in real systems.\n2. In the evaluation, the strategy generated by this paper only provides significant advantages in unrealistic settings. For example, it only shows significant advantages over \"equi-no TP\" on settings with an extremely low memory budget or an extremely high number of devices.\nTraining large models under an extremely low memory budget distributedly is not realistic and nobody will do this. When the number of devices becomes large, the simulator model used in the paper is not accurate anymore.\n3. The paper does not consider the effect of batch size when measuring the throughput of TMPC. A small batch size typically cannot utilize the full compute power of a modern high-end GPU. Even with the same tensor model parallel configuration, different batch sizes can have different TPS.\n\n\nAdditional questions:\n1. In Figure.2, the blue line (Piper) is almost a straight line. Does this mean Piper can get perfect **strong** linear scaling? I feel like this is impossible in a real system. For example, even in this paper (https://arxiv.org/pdf/2104.04473.pdf), they can only get good **weak** scaling in a cluster with very high internode bandwidth.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The author adequately addressed the limitation of this paper in Section 6.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "7IQmJr6PCnh", "reviewer": "Reviewer_BM6L", "summary": "This paper proposes a novel algorithm to find a pipeline parallel configuration where each pipeline stage also includes data and tensor model parallelism. The method assumes that the input computational graph consists of multiple “layers” that form a DAG and each layer has a  set of predefined tensor model parallelism configurations. The algorithm is a two-stage dynamic programming algorithm and is mainly evaluated in a simulated environment.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I lean towards accepting this paper. Although there are limitations with the proposed method and the evaluation method, I believe the proposed two-step problem formulation and dynamic programming algorithm are novel and should be useful for the community. Detailed comments:\n\nStrengths:\n- The two-step formulation proposed for pipeline and data/tensor-model parallelism is novel and will greatly reduce the complexity of the overall search space. I believe this is a positive contribution and will be useful for future work working in space.\nA two-stage dynamic based on the two-step formulation solves the problem.\n- The paper is clear and easy to follow.\n\nWeaknesses:\n- The algorithm assumes the knowledge of predefined “layers” and for every layer, there is a list of predefined model-parallel methods. For a more complicated neural network, this is not an easy task. This limits the application of the method on more complicated neural networks.\n- There are non-negligible limitations on the experiments in the paper:\n  - The paper only evaluates the BERT model. In this case, the DAG formulation actually degenerates to a chain, making the algorithm proposed in the paper overkill. The method should be evaluated on more neural networks with more complex architecture (that actually forms a DAG).\n  - Most of the evaluations in the paper are simulated. Although Figure 4 partly proves the simulation is realistic, I’m still not convinced and I would like to see larger-scale experiments.\n\nMinor Issues:\n- Some key concepts will be easier to be understood with examples (e.g. Downsets).\nLine 26: Maybe also add \\emph for “pipelined model parallelism”?\n\nOther questions\n- Do you profile the Time-Per-Sample for every downset? How long does this take? Will this be a potential bottleneck when the neural network structure getting more complex?\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See weaknesses in the main review section.\n\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "5pFfkOS2r2K", "reviewer": "Reviewer_eU38", "summary": "This paper presents a search based method to plan DNN parallelism. The search space includes pipelining, recomputation choices and data partitioning. \n", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "First of all, this paper studies an important problem of automatic parallelism. Two of the primarily related works are FlexFlow and Pipedream. Piper's search space can be viewed as the combination of the above two, plus the choice of recomputation. The overall contribution over the prior work is slightly incremental.\n\nGiven the limited technical contributions, the paper would benefit from stronger empirical evaluations. On that end, it would be beneficial to show a wide variety of model settings and how Piper applies to those models. It would also be useful to bring baselines that reflect the search space used by FlexFlow and PipeDream. Finally, it would be helpful to evaluate Piper's plan on a real end-to-end setting and see the benefit brought by the additional configuration space.\n\nThe paper is borderline due to the above considerations.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NA", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}]

Author rebuttal:

Meta-review: After some discussion, there's general consensus toward accepting the work. As Reviewer cpuf notes: "It is the first work that uses dynamic programming to search in a non-trivial search space that includes data parallelism, tensor model parallelism, pipeline model parallelism, and activation rematerialization." There are a number of challenges in finding parallelizable configurations, and while the paper's specific ideas may have overlap with prior works or may be argued as incremental, it makes an important delta in solving this important problem area.

Flaw-wise, there are several experimental limitations that would be worth resolving as Reviewer BM6L and cpuf note. I highly recommend the authors address the rebuttal concerns as they further polish the paper.
Decision: Accept (Poster)