AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness
Authors: Dacheng Li
Hongyi Wang
Eric Xing
Hao Zhang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Machine Learning Systems
Model Parallelism
Automation
Heterogeneity

Abstract: Scaling up model sizes can lead to fundamentally new capabilities in many machine learning (ML) tasks. However, training big models requires strong distributed system expertise to carefully design model-parallel execution strategies that suit the model architectures and cluster setups. In this paper, we develop AMP, a framework that automatically derives such strategies. AMP identifies a valid space of model parallelism strategies and efficiently searches the space for high-performed strategies, by leveraging a cost model designed to capture the heterogeneity of the model and cluster specifications. Unlike existing methods, AMP is specifically tailored to support complex models composed of uneven layers and cluster setups with more heterogeneous accelerators and bandwidth. We evaluate AMP on popular models
and cluster setups from public clouds and show that AMP returns parallel strategies that match the expert-tuned strategies on typical cluster setups. On heterogeneous clusters or models with heterogeneous architectures, AMP finds strategies with 1.54$\times$ and 1.77$\times$ higher throughput than state-of-the-art model-parallel systems, respectively.

AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness
Dacheng Li, Hongyi Wang", Eric Xing"oo Hao Zhang' Carnegie Mellon University Petuum Inc. University of California, Berkeley
Mohamed Bin Zayed University of Artificial Intelligence
Abstract
Scaling up model sizes lead fundamentally new capabilities in many ma- chine learning (ML) tasks. However, training big models requires strong distributed system expertise carefully design model-parallel execution strategies that suit the model architectures and cluster setups. In this paper. we develop AMP, framework that automatically derives such strategies. AMP identifies valid space of model parallelism strategies and efficiently searches the space for nigh-performed strate- gies, by leveraging cost model designed capture the heterogeneity of the model and cluster specifications. Unlike existing methods, AMP is specifically tailored to support complex models composed of uneven layers and cluster setups with more heterogeneous accelerators and bandwidth. We evaluate AMP popular models and cluster setups from public clouds and show that AMP returns parallel strategies that match the expert-tuned strategies typical cluster setups. On heterogeneous clusters or models with heterogeneous architectures, AMP finds strategies with 1.54x and 1.77x higher throughput than state-of the-art model-parallel systems, respectively.
Introduction
Recent progress in language understanding [6, 4], multimodal learning [1], and many other ML applications can largely credit to the use of extremely big models. For example, the largest language models [4, 14] with billions of parameters are found to have fundamentally new capabilities. Due to the increased model size, training such models, however, calls for model parallelism execution strategies in order to place the gigantic model computation across multiple, heterogeneous accelerator devices ge. TPUs and GPUs), and ensure efficient distributed execution.
Different from typical data parallelism strategies [28, 5, 15), model parallelism considers a much larger space of parallelization techniques. In particular, when a model is too large fit in a single device memory, we can either place the computation of different layers across different devices (a.k.a. pipeline parallelism), or partition the computation of particular layers and dispatch parts onto parallel devices (i.e., tensor model parallelism (TMP) [23]). Due to their distinct communication and memory requirements, these model parallelism dimensions have different applicable scenarios, which may change with model architectures and cluster setups. For example, TMP results in high communication volumes, hence is favored on devices connected with high bandwidth, but may become less effective with limited communication bandwidth. Designing performant execution strategies requires combin- ing these parallelism dimensions to trade-off among memory, communication, and computation, then tuning the configurations in each parallelism dimension align with the heterogeneity of the model and the training cluster.
'Codes and experiment logs available httpe://ggtaab..cm///cceeeee for reproducibility.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Existing model model-parallel training strategies are either manually designed by domain experts for one or two specific models, or automatically generated with strong assumptions on the model architecture or cluster topologies. For example, the 3D parallelism strategy adopted in Megatron-LMM [23] and DeepSpeed [21] manually specialized for transformerrbased. language models, assuming fixed cluster setup and the model having the same layer repeated [23, 21, 19]. Therefore, they can hardly generalize to models with diverse layer compositions (i.e., heterogeneous models), or clusters with mixed types of accelerators and network switches (i.e., heterogeneous clusters). This motivates the core questions in this paper: Is it possible to automatically find the best model-parallel strategies for more heterogeneous models and clusters?
Given an arbitrary cluster and model setup, finding highly performed model-parallle strategies is difficult for the following reasons. First, when the model has diverse types of layers, existing heuristics on deciding the assignment of layers to pipeline stages (and balance stage workloads) can be invalidated, because different layers exhibit distinct execution and communication costs. Second, evaluating model-parallel strategy on large cluster can be expensive, and prevents search- based methods [11] that depend on extensive real evaluations. However, existing cost models do not capture the heterogeneity the model or cluster (e.g. [24] assumes uniform bandwidth the cost model), In this paper, we present solutions to these challenges. To address the layer-stage assignment problem, we develop new dynamic programming algorithm running in polynomial time. To evaluate large number of strategies within limited budget, we develop new cost model which serves a cheaper proxy to evaluate model-p model-parallet strategies. We name our method AMP. In summary, AMP makes the following contributions:
We identify factors that influence model-paralle performance in heterogeneous settings, and define strategies based on these factors.
We present a new cost model that handles heterogeneity in both the model and the cluster to evaluate the quality of model parallel strategies with minimal amount of real trials.
We develop dynamic programming approach to handle the load imbalance issue from heterogeneous models in the pipeline layer assignment problem.
We empirically show that existing systems output sub-optimal strategies in heterogeneous settings. In contrast, AMP achieves 1.54x speedup with heterogeneous clusters, and 1.77x speedup on heterogeneous model architectures compared to sfute-of-the-at heuristics.
2 Background and Related Work
In this section, we discuss the trade-offs of existing model-paralle, strategies, and the need for composing them with the awareness of model and cluster heterogeneity.
Pata-paralled (DP) strategies. DP partitions the input data batch evenly among workers, and each worker holds an entire model replica [15, 22, 9]. each iteration, each worker computes gradients over its assigned batch; the gradients are then synchronized among workers before the next iteration. DP requires each worker to hold an entire model replica, hence cannot be directly used to train models with massive parameters.
Tensor model-parallel (TMP) strategies. TMP, proposed by Megatron-LM [23]. is popular model-paralle: approach for large transformer models. In TMP, the layer weights of each two consecutive layers are partitioned row-wise (i.e., input dimension) first, then column (i.e., output dimension) [23]. TMP removes the need for synchronizing the intermediate output of the very first layer, but requires heavy cross-device communications afterward. TMP is normally combined with data parallelism to increase the training throughput [18].
Pipeline-paralle (PP) strategies. In PP. layers are placed across GPUs, the training mini-batch is split into smaller micro-batches. The forward and backward computation is then pipelined across micro-batches. PP requires less communication than DP and TMP, but suffers from device idle (i.e., pipeline bubbles) [10, 18). Despite synchronous pipelining schedules such as GPipe [10], PipeDream [18] proposes an IFIB asynchronous schedule to reduce the pipeline bubble. TeraPipe de- velops level pipelining schedules in single training sequence for auto-regressive models [16].

---

Automatic partition. Prior work, e.g. FlexFlow, PipeDream, and DAPPLE explore automatic computation graph partitioning over few typical cluster device setups [11, 18, 71, such as NVIDIA DGX workstations connected with Infiband switches. For partitioning the model over multiple devices, the aforementioned work leverages carefully designed cost models to select the desired strategy over the search space.
3D parallelism. To trade-off the scalability, memory footprint, and device utilization, DP, MP, and PP have be used together (which is also known 3D parallelism). DeepSpeed carefully designs 3D parallelism strategy to train a massive-scale language model with 17 billion parameters [21]. [19] manually tunes the degrees of DP. MP, PP to scale 3072 GPUs PrUs-bolt both of them are not automatic. The closest to our work is Piper [24], which leverages a cost model and a dynamic programming algorithm to search over the 3D strategy space. However, Piper works only with raanformer-basee models, and assumes the same type of devices in the cluster are homogeneously connected with equal bandwidth, ignoring heterogeneous setups. In contrast, AMP designed capture the heterogeneity of the given ML model and the deployed cluster, reflected in both the search space design and the searching algorithm ($ 3).
Heterogeneity in models and clusters. Newer generations of GPU devices and Ethernet switches emerge rapidly, thus it is common nowadays that multiple types of devices are deployed together in training cluster. Conducting distributed training over such types of heterogeneous clusters introduces extra challenges. For instance, poor load balancing strategy can cause straggler effects easily; heterogeneous memory capacity among GPUs introduces new constraints model and data partitioning. Apart from hardware heterogeneity, DL models also introduce another dimension of heterogeneity. Many today's 's popular models are composed with diverse layers with different sizes and computing characteristics, e.g. typical convolution networks have their layer width growing with depth [8], which would require distinct partitioning and placement strategies for each layer during model parallelism.
3 Methods
3.1 Problem Formulation
Formally, the inputs of AMP are (i) the model, (ii) the cluster, and (iii) the global batch size (gbs). A deep learning model W is represented as graph, i.e., W = (Cw, Vw). Cw is a set of vertices where each vertex ow,, E Cw, E {i..... , [Cwl} represents a layer, and each edge tw. (ww,i E Vw represents the tensor (activation) between layer ew.i and layer cw.++1- Additionally, we denote the execution cost of a layer Cy,, as C2, and the communication volume of an edge tw 24 A cluster is represented C BB where D {d < i S [D]} denotes set of devices. B RPP x|P| is a symmetric matrix where each matrix element bij represents bandwidth between device dj and dy. The goal of AMP is to output a high throughput parallel training strategy s given the hardware and model configuration. S denotes the entire space of candidate strategies. We formulate this as an optimization problem:
arg min f(s; W, C, gbs) SES
(1)
where f(-) denotes per iteration running time function of the strategy s, conditioned the given ML model, cluster, and user-specified global batch size.
3.2 Method Overview
To solve the optimization problem in Equation 1, we identify factors that influence the system performance ($3.3). However, there are two major difficulty in solving this optimization with these identified factors. First is how to evaluate f(-). One possible solution is to launch real trials, but this is prohibitively expensive because the entire optimization procedure may require evaluating large number of strategies. Rather we develop a cost model ($ 3.4) that can estimate f(-) quickly. Second, the discrete nature of the optimization problem disallows us to leverage gradient-based optimization methods. However, evaluating all strategies in the search space slow even with the cost model. One way to efficiently optimize the objective is to exploit structure in the strategy
3

---

space S, and prune areas with no promising candidates. Using this principle, we develop a dynamic programming approach ($3.5), which prunes out candidates with worse performance before invoking the cost model.
3.3 Representation of Strategy
At a high level, a strategy s defines how to train model [Cw, Vw3 on given cluster (D, B3 with the batch size gbs. Concretely, a strategy s that composes DP, TMP, and PP strategies can be represented by a few key aspects, e.g. how many model replicas should be created (i.e., how many pieces should a mini- batch be evenly split) in DP; how to partition model layers (for PP) or layer parameters (for TMP); and how shall the global batch size decompose into micro-batches in PP semantics. In AMP, we represent strategy with the following four aspects:
1. Parallelism degrees. Since PP, DP, and TMP have different characteristics, different degrees of each affect the system performance. We denote the degrees to be pp, dp, tmp. The tuple (pp, dp, tmp) defines a tiree-dimensiona grid, which satisfies the constraint pp x dp x tmp = [D]: pp, dp, tmp E {I..... [DI}. This tuple generates pp x dp x tmp virtual ranks that carry out corresponding tasks. For instance, ranks the same data parallelism group communicate with each other to finish the gradient update. These ranks are bijectively mapped to available devices D for execution.
2. Device placement. In a heterogeneous environment, physical devices can have different connection bandwidth to each other, thus placing physical devices in different virtual ranks leads to different performance. For instance, ranks in a communication-heay group are all allocated with physical devices with high bandwidth, the per iteration time will be shorter. Formally, We denote a device place- ment a bijective function from virtual ranks to devices p: (p. x {1.....dp dpp x {1,...,mmp} + D.
PP
whar rene.
your
DP
Tue
3. Micco-batch size. Micro-batch size (mbs) is the batch size a device proceeds at time in the pipeline parallelism execution ($2). When màs varies, devices usually have dif- ferent utilization [18, 26]. For instance, a larger micro-batel size usually leads to better GPU utilization because the GPU can better parallel the computation. However, a larger micro- batch size also leads to a larger pipeline bubble, which makes selecting màs hard.
Figure I: An example strategy with parallelism degrees = (2, 2).. Note that there are 8! possible GPU as- signments to these 8 virtual ranks. Assuming uniform bandwidth, uniform-layer assignment cuts at the 5th layer, while a more balanced as- signment cuts at the 3rd layer.
Pipeline layer assignment. Different assignments of layers to each device can also affect the system throughput. On the one hand, when activation volume is large in the boundary, the next device needs to wait longer, resulting in a longer per iteration time. On the other hand, when a device is assigned a much higher load than other devices, it blocks all other devices and thus incurs a longer iteration time. We denote a feasible pipeline layer assignment as a function a : Cw + R that maps layers to the stage index (1, (1.....p)) in pipeline.
In summary, a strategy can be defined by a 6-dimensional tuple:
S a (pp, dp, tmp, p, mbs, a)
3.4
Cost Model
Since launching real trials to obtain the actual f()) is expensive, we develop a cost model that serves as a fast estimation of f-)). In this section, we introduce the proposed cost model.
Communication primitive. We consider the ring topology for all-reduce collective communication in both DP and TMP [22]. Assume that the whole message size is M. In a communication group of workers with uniform bandwidth B, the communication time of single ring all-reduce operation can be modeled as [13]:
2(n 105-1)4 1)M
T Telireduce
(3)
nB
In a group with non-uniform bandwidth, we approximate the bandwidth with the lowest bandwidth B'. For point-to-point (P2P) communication operations between consecutive pipeline stages, we

---

Variable
Meaning
D
set of devices in the cluster
M
message size in all-reduce operation
n
number of workers in an all-reduce operation
B
bandwidth between workers in an all-reduce operation
activation volume for P2P operation
b
bandwidth for P2P operation
k
number stages a pipeline
communication time for the P2P operation between the 1th and (i + 1)th stage
4 T.p layer
execution time of the 1th stage in pipeline
exccution time of the ith pipeline execution time of single layer j
gas T.
number of micro-hatehes in pipeline number of layers in W
Table Major variables and their meanings in the cost model.
assume the communication operation use the full bandwidth: volume and b the bandwidth between these two GPUs.
[11]. Here e is the activation
Top-down run time decomposition. We now present our cost model designed from lop-down view. First, DP is at the top gramularity, where there are dp pipelines running in parallel. synchronizes gradients when all pipelines finish their schedules. It does not consider how each pipeline is executed. At the DP level, the system aun-time can be decomposed into two terms:
T = max (7 S S dp} + max Tripeym dED
T."h is the time taken to complete a single pipeline i. Typsyme denotes the time for a device d to synchronize gradients with an all-reduce operation. We model the latter term by using Equation with the message size equal the number of model parameters held device d. Since different pipelines may finish with different speeds, and different devices may finish their data parallelism synchronization with different speeds, we take maximum over these two terms respectively.
At a lower granularity, each pipeline executes training schedule that includes activation//radiiets exchange at layer level. It does not need to consider how each layer is executed, which is carried out by MP workers. Assume the number of stages a pipeline k. and the communication time between the ith and (i + t) is ej, and the execution time for the 2th stage is ti. We model t; the
sum execution time foyer in the stage for layer in the stage: 1, Ea(j)=i dayer t We model the execution time for a single pipeline as
TPP = (gas - 1) max(t,]1 s pp} + Σ Σ t.
(5)
Intuitively, the first term captures the straggler effect, which gets amplified by the number of micro-
batches gas. If we assume gas = 1, Equation 5 reduces to TPP -> e: + > 4. This aligns with the actual schedule: the mini-batch simply performs single forward pass and single backward pass, and sends or receives gradients or activations. If we assume all layers have the same execution cost and the cluster is equipped with the same bandwidth, the optimal solution simply balance the number of layers. This heuristic is often adopted as the default layer assignment solution in current systems[23).
At the bottom granularity, MP workers consider how a single layer is executed. The run-time of single layer r Jayer is composed of the layer computation time (e.g., Matrix multiplication) and all-reduce communication time. One way to estimate / is using the number floating point operations (FLOPs) for computation time, and Equation for communication time. However, we find this inaccurate in our setting. Concretely, when mp grows, the arithmetic intensity of each matrix multiplication becomes lower. In such scenarios. a matrix multiplication with only few FLOPs
can take a long time execute. Thus, simply obtain playe by profiling several tensor model parallelism settings and store results in a lookup table. This profiling is feasible because it only needs be done once the entire optimization.
We describe how consider overlapping in the P2P communication and computation in $6.1.
5

---

Cluster Size
2x2
4x4
8x8
16x 16
Optimization time (seconds) 58.0
325.6
890.3
1627.4
Table AMP optimization time with respect to the cluster size. Results are obtained with 24-layer GPT-2 model and global batch size 32.
3.5 Pipeline Layer Assignment
A key observation in the layer assignment problem is that only consecutive layers be assigned to the same stage. In particular, computation up to layer can be reused for that of layer + 1. This pattern fits neatly a dynamic programming solution. To handle the max term in objective an auxiliary variable tolerance m. m takes value from all possible execution time for single stage. For instance, can be the sum of the execution time of the first and the third layer because they can form a valid pipeline stage. However, it can not be that of the first and the third layer because they are not consecutive. Using m, we rewrite an objective for the dynamic programming algorithm:
use
Pp
Pp
(gas 1) m0 ma S S pp} m
6
(6)
Note that when m 0, Equation 6 reduces the desired objective 5. Formally, store the optimal assignments and the corresponding cost up to layer i with j stages and variable m in table dp [[]]] Let L be the number of layers, then dp [R]]]]]] stores the solution. We induct on the number of stages in the pipeline.
Base case (j = 1). The only possible assignment assigns all layers into single stage. Denote the sum of execution time for these layers f1. The associated optimal cost is thus
(gas 1) max mm + 4
We store the table entry | [m] with (i, g()) for all possible and m.
Recursive step. To compute dp when j 1, enumerate all possible positions of the last cut i' for all i' < i. Denote the execution time from layer i' to layer as 2,,', let the value stored dpp -1] []][(x(tt, m)] be (ap, cv). Denote the longest stage time in ap t1, the between the (j 1)th and jth device at layer i' as ev The cost associated with the last cut at i' is:"
communication
g(i') cv + (gas 1) max (0, t2 m) t2, + egy
(8)
Select the optimal last cut position l'on = argmin, (,, (f)). The dynamic programming table entry [m] updated as: (a.; Uii ). g((pp)t)
Running time. There are ELn i = O(L2) possible values for m. Thus, our problem size is O(L x L?), each sub- bebb-roblem takes (L)) time. Thus the total running time for the dynamic programming algorithm is OkL*)), while a brute force solution takes O ((k-))
L-11
3.6 Optimization Procedure
In this section, introduce the overall optimization procedure to find the best strategies defined at 3.3. First, we enumerate all possible parallelism degrees and micro-batch sizes. For each parallelism degree and micro-batch size, we deterministically solve for the optimal pipeline layer assignment using the algorithm at 5 3.5. We adopt the heuristics at [19] to solve the device placement problem: prioritize placing model parallelism workers in the same node. If there is more available mode bandwidth, then place data parallelism workers and pipeline parallelism workers. The procedure is written in Algorithm 1. More details be found in 6.2
prooff can be found in $ 6.4

---

Algorithm 1: Optimization procedure
Input: C. W. gbs, budget
degrees = enumerate_ enumerate_degrees(((
record set(
3 for d in degrees do
/* Possible micro-bard sasro-bacch size given data parallel degree for mbs in enumerate (d.dd) do
p placement(C,
Device placement heuristics
pipe_ast d.pp, p) // Optimal pipeline assignment 8 (a, mbss, d, p) // current strategy
cost estimate(s) // our cost model record.add(s, cost)
/* run top predicted strategies
best run(record, budget)
11 return best_s
Megatron
AMP
SA
Min
1.32
1.20
1.57
Mean
2.41
1.29
1.43 + 0.18
1.82 0.17
Table 3: Best and average strategies under homogeneous setup over Top 10 candidates (scale: seconds)
Complexity of the overall optimization. The number of factors of an integer N can be loosely upper bounded by (/N1/2). Thus, the number of possible degrees can be upper bounded by ((ID)1/2) x ([1D)1/2 = ((D)), and the number of possible micro-batch sizes can be upper bounded by O(gbs//2). The pipeline assignment algorithm run-time is bounded by O(KL*) = ((I x L*) (section 3.5). Thus, the total run time can be upper bounded by Ogbbs 1/2 x |D] x [D] x L*) = (gbs 1/2 x |D] L*). Empirically, we verify the optimization time with respect to the cluster size in Table 2. We find that the actual optimization time is near linear with respect to the cluster size.
4 Experiments
In this section, we evaluate AMP on different types of clusters and models. Besides homoge- neous setup, we study two heterogeneous setups, one with heterogeneous cluster and one with heterogeneous model. We use the state-of-the-art Megatron heuristics as our baseline [19].
4.1 Baseline
Megatron [23] is runtime engine that supports 3D parallelism training DNN models. It suggests several heuristics on how to train models with 3D parallelism faster. We summarize these heuristics in the following procedure.
Mgetraon explores all possible micro-batch sizes as in AMP. For each micro-batch size, forces (mp to be smaller than the smallest number of devices in node. Then it selects parallelism degrees with the smallest tmp X pp. It solves the device placement problem by prioritizing tensor model parallelism workers in the same node. If there are additional available intra-node connections, it places data parallelism workers next. It solves the pipeline layer assignment problem by either balancing the number of parameters or the number of layers in each stage.
Megatron proposed based on observations on homogeneous cluster and model setup. The model is composed of simple stacks of identical layers. The cluster is equipped with the same type of devices with the same bandwidths. Thus, We hypothesize under such as homogeneous setup. is (near) optimal. To validate this hypothesis, we conduct experiments on homogeneous setup (5 4.3). Moreover, we are interested in how Megatron heuristics perform when the homogeneous assumption does not hold. Concretely, ask two questions. What is optimal when the cluster is equipped with
7

---

spearann
Spearman
Spcaran
74
Extimated Per Ber. Time [sec.]
Extimated Per Bee. Time [uec.]
Extimated Per Ber. Time (sec.)
(a) homogeneous
(b) heterogeneous cluster
heterogeneous model
Figure 2: Predicted strategy cost versus actual of top candidates. The Spearman correlation ranges from -1 to 1, and 0.5 is generally considered a reasonably strong positive correlation.
different types of devices and different bandwidths to each other? What is the effect of heterogeneity models? To answer these, we design two sets of experiments under heterogeneous cluster and heterogeneous model in $ 4.3.
4.2 Alternative Device Placement Algorithm
As described in Section 3.6, adopt the heuristic in Megatron solve the device placement algorithm in our main method. During the development time, we also explore an alternative device placement optimization method. Concretely, we adopt a Domino-tiling style device placement representation. Since this representation introduces much more possible strategies, enumerating them is costly. We instead aspect [25]. We name this method SA. Details can found in 3 6.3.
Simulated Annealing(SA) algorithm to optimize the device placement
4.3 Experiment Setup
Homogeneous. We conduct experiments using GPT-2 (L 24, = 1024) [20] on 4 AWS EC2 g4dn. 12xlarge nodes with a global batch size 32, Each instance is equipped with T4 GPUs with 50 Gbps PCIc connection intra-node bandwidth, and 50 Gbps inter-node bandwidth, In all experiments, we place instances in the same AWS placement group to make sure they have the full bandwidth.
Heterogeneous cluster. To simulate heterogeneity in both different types of GPUs and different bandwidths, ensemble cluster using 3 AWS EC2 p3.8xlarge instances (V100 GPUs) and I 12xlaree (T4 GPUs) instance. Each p3.8x 33l8gea instance is equipped with 4 V100 GPUs with approximately 170Gbps NVLink intra-node connection and 10 Gbps inter-node bandwidth. We evaluate the performance on GPT-2 (L = 24, = 1024) with a global batch size 32.
Heterogeneous model. We use a modified version of Transgan generator [12], which stacks transformer layers in the way that later layers have smaller hidden sizes. Specifically, we use 12 layers with dimensions 1024, and 12 layers with dimensions 16. We use p3.8xlarge instances and global batch size 64.
4.4 Evaluation Metric
We evaluate strategies by the running time per iteration. Specifically, we run each strategy by 40 iterations and take the average iteration time for the latest 20 iterations [27]. The underlying system Deepspeed (built on top of the Megatron engine) [21] with fp16 optimization enabled.
4.5 Results and Discussion
Homogeneous setup. In this experiment, we compare three methods: (1) Megatron: Megatron Heuristics, (2) AMP: method described in 3.6, and (3) SA: optimize the device placement aspect using the Simulated Annealing. Under homogeneous setup in & 4.3, we hypothesize that Megatron heuristics (near) optimal. We also expect that method can find such optimal strategies.
In our experiment setting, Megatron proposes feasible strategies. To understand the difference in strategies between proposal and the Megatron proposal, evaluate the Top 10 strategies from

---

Megatron-LM
RS
AMP
SA
Min (sec.)
1.97
1.71
1.28
1.46
Mean (sec.)
2.62+0.37
3.04
1.73+0.43
2.29+0.68
Table 4: Best and average strategies under heterogeneous cluster setup (scale: seconds).
Megatron-LM
ours w/ parameter
ours w/ uniform
AMP
Min. (sec.)
1.89
2.19
1.25
1.07
Mean (sec.) 1.94
2.19+0.0
1.25
1.66
Table 5: Best and average strategy under heterogeneous model setting (scale: seconds)
our method. The results are presented in Table 3. We find that Megatron is indeed near-optimal and AMP can find such strategies (in fact slightly better). We analyze the selection procedure below.
Since the model can be fit into a single GPU, Megatron selects an extreme configuration - dp- = 16. This may be sub-optimal because the communication in the data parallelism group is high. Suppose the micro-batch size is small, so that the bubble overhead in the pipeline parallelism may hidden, may be better assign some workers the pipeline parallelism. In fact, AMP finds that dp 4, pp with the smallest micro-batch size 1 is a better choice. We also find that exploring more device placement strategies with SA finds worse results. This is because considering has limited influence in the homogeneous setup. For instance, whether the tensor model parallelism communication happens within node or cross node shall not change the performance. However, introduces much more strategies that make the optimization hard.
device placement
Heterogeneity in the cluster. In the heterogeneous cluster setup, devices have different connectivity to each other, and different execution speeds. We hypothesize that some heuristics from Megatron can lead to bad performance, but our method should be able to automatically address the heterogeneity in (especially) bandwidth since we explicitly model it (3.4). We are also interested in how much heterogeneity can influence the performance of strategies. Concretely, we examine how much Megatron and AMP improve compared to random strategies in the space. In this experiment, we have four methods to compare: Megatron, AMP, SA, and
RS: Randomly selecting strategies without heuristics and a cost model.
We run RS for 50 iterations. Results are presented in Table 4. We find that the mean of random strategies is 1.16x worse than that of Megatron, and 1.76x worse than that of ours. By using optimization procedure, Specifically, the heuristic Megatron that minimize mp x pp is sub-optimal. When using micro- batch size 1, the best strategy shall be dp 2,pp performance. However, since the model can be in a single GPU, Megatron selects dp=16, which results in 2.72 second/iteration performance. This strategy is much slower because some GPUs have a low bandwidth (10Gbps) to others. On the contrary, by modeling each GPU's bandwidth in Equation 3, our cost model penalizes this strategy heavily. We let the cost model to estimate f()) for all strategies and find that AMP ranks this strategy in the 21 place.
find strategies that achieve 1.54x speedup compared to Megatron.
8. which results in a 1.28 second/iteration
Heterogeneity in the model. When the layers in the model have different execution costs, the pipeline assignment methods in Megatron, which balances the number of layers or parameters, can not balance the actual workload. In such setup, expect these two. In this experiment, We compare four methods: (1) Megatron: Megatron with the uniform layer assignment, (2) AMP, (3) ours w/ uniform: the best strategy found by AMP but replaced with the uniform layer assignment, and (4) ours w/ parameter: the best strategy found by AMP but replaced with the uniform parameter assignment.
pipeline assignment can outperform
Results are presented Table 5. AMP finds a 1.77 faster strategy than Megatron ones. We analyze the speedup in two dimensions. First, if we only consider the micho size and the parallelism

---

degrees aspect, the strategy found by AMP is 1.51x x faster than the one by Megatron. The reason similar to the one in the heterogeneous cluster setting. Second, if we additionally consider the pipeline layer assignment aspect, the strategy gains an additional 1.17x speedup.
To study the relationship between heterogeneity in the model and the system performance, we fix the parallelism degrees and vary only the pipeline layer assignment. Results are presented in 3. Since we only use transformer layers with two different hidden sizes, we simply define heterogeneity as:
h=16 layers time
heterogeneity h=1024 layers time
Cost model accuracy To correctly select the top strategies, AMP must rely on an accurate cost model. We consider the cost model to be accurate if it can correctly rank different strategies, where the absolute value of its prediction is of less importance. To study the accuracy, we obtain more ground truth f(-) (more than 10 for the main result) and compare them with the estimated f(.) by the cost model. The results are presented in Figure 2. We find that generally a lower estimate f(-) corresponds
lower actual f()) More importantly, strategies with low actual f(-) are ranked among the top. Concretely, the best strategy is ranked 3rd, 2nd and 1st in the three settings. As a result, the best strategy can be selected because AMP runs several real trials at the end of the optimization procedure (Algorithm 1). While AMP outperforms the baseline as it needs to launch 10 real trials, we would like to improve our cost model in the future so that fewer real trials are needed (each trial takes around 1-2 minutes in our setting).
Optimization cost. In the heterogeneous model exper- iment, AMP takes 140.7 seconds for profiling and 235.6 seconds for 65 iterations of the optimization procedure. In particular, 230.8 seconds are spent in the dynamic pro-
bs-64, vs dparams.
bs=16.ss apparms.
bs-64, uriffrm
bs-16,
Baseline Strategies
2.00
Speedup over
gramming algorithm. We also launch several real trials for top predicted candidates to seleet the final best one. In
1.75
total, the cost of our optimization is around 10 minutes
1.50
with several real trials (around 1 - 2 minutes per trial).
1.25
Limitations. AMP does not model memory footprint. Consequently, it can suggest strategies that do not fit in the device's 's memory, and launch unnccessary real trials. In our experiments, small number of top strategies are out of memory (i.e., 7th and 9th of the top 10 strategies in the TransGAN experiment). Additionally, the solution for the pipeline assignment problem is not efficient enough. In our experiments, it takes the major time of the whole optimization procedure. When the number of layers L goes up, the running time is increasing in the ((L.). In the future, we aim to either improve the dynamic programming solution itself, or combine the solution with techniques such operator clustering to keep I. low [2, 29].
1.00
1.0
1.5
2.0
2.5
Hetergeneity
Figure Speedup against heterogeneity versus two Megatron layer assignment methods. Numbers obtained when a top parallelism strategy: mp=1, pp=4, dp=4.
5 Conclusion
We present AMP, an automatic algorithm that can find model parallel training strategies with high system throughput. AMP is equipped with an expert-design cost model that considers cluster and model configurations. We show that when heterogeneity exists in cluster or model setup, current heuristics are sub-optimal. Our automatic algorithm can find strategies with 1.51x x and 1.76x speedup within similar budgets.
Acknowledgements
We would like to thank Sangkeun Choe and Aurick Qiao for their comments in the early stage of the paper. We thank Tailin Zhang and Tianyi Yu for their help in deriving the pipeline objective and the dynamic programming solution. We thank all reviewers for their invaluable comments and feedback. This research was supported by NSF IIS I1S1563887, NSF CCF1629559, NSF IIS1617583, NGA HM04762010002, NSF IIS 19555322, NSF CNS2008248, NSF IIS2 and NSF BCS2040381.
10


Reviews:
[{"review_id": "jRxk1W8FQjw", "reviewer": "Reviewer_9vj5", "summary": "This paper presents a strategy in model-parallel setting, where the model is being trained in a cluster. The strategy considers heterogeneity of computational resources, diverse characteristics of different layers, communication costs in networks, and uses a dynamic programming algorithm to intelligently assign layer-stages.\n", "questions": "1) The reviewer is interested in more study on how accurate the cost model is, for example, what if we replace the analytical approximation with an oracle cost model?\n2) Compared with previous work like Alpa which uses homogeneous setting, how far does the proposed method go besides adding more terms in the dynamic programming equaltion?", "limitations": "As the author stated, memory footprint is right now not under consideration. However, the reviewer does not consider this as a significant drawback.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n1) A dynamic programming-based algorithm is prposed to find out the optimial placement strategy that is significantly time-efficient that bruteforce exhausive search;\n2) The paper proposes an analytical cost model that saves the cost of estimating the run time of each layer on actual hardware;\n3) Experiments show that the strategy is highly effective and provides concrete speedup for end-to-end workloads.\n\nWeakness:\n1) The paper doesn't present comprehensive ablation study of the effectiveness of the cost model.\n2) The paper seems to overlook some latest advances in distributed training and model parallelism, for example, Alpa [1] that considers homogeneous setting but still formulates the problem similarly, while having heterogeneity seems to be a plain extension. Would love to hear from the reviewer about this particular comparison.\n\n[1] Zheng L, Li Z, Zhang H, Zhuang Y, Chen Z, Huang Y, Wang Y, Xu Y, Zhuo D, Gonzalez JE, Stoica I. Alpa: Automating Inter-and Intra-Operator Parallelism for Distributed Deep Learning. arXiv preprint arXiv:2201.12023. 2022 Jan 28.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "POoVYfgLvZ0", "reviewer": "Reviewer_tyEZ", "summary": "This paper presents a method to automatically find the optimal parallelization strategy for large model training. Unlike previous frameworks, the proposed model considered the heterogeneity of the model and hardware. The parallelization problem is defined as the parallelism degrees, micro batch size, device assignment, and the arrangement of the layers in the pipeline. The authors built a cost model to estimate the cost of the candidate strategy. The parallelism degrees and micro batch size are selected through enumeration. The device assignment is solved by applying a heuristic while the pipeline arrangement is found through dynamic programming.", "questions": "1. Could you provide extra information on the accuracy of the proposed cost model? For example, the correlation coefficient against the real latency. Could you also indicate how will the accuracy of the cost model affect the result?\n\n2. Could you comment on whether it is possible to perform joint optimization and how does it compare with the discrete optimization of each subproblem?", "limitations": "The author discussed the limitation in terms of the lack of memory model.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strength\n\n1. This paper present a detailed cost model for parallelization strategy on heterogeneous models and hardware. The difference in hardware is considered as different bandwidths and the profiling result of the layer execution time.\n\n2. The pipeline arrangement is modeled as a dynamic programming problem and given a determined solution. \n\n3. The ablation study shows the effectiveness of each component optimized by the proposed method.\n\n\n## Weakness\n\n1. **The accuracy of the cost model.** The evaluation focuses on the iteration latency. But as a major contribution of this paper, it is also important to understand the accuracy of the cost model and how will it affect the strategies found by the method. Based on what is shown in Figure 2, it seems that the estimation for most strategies is optimistic, especially for the optimal one indicated by the estimation. \n\n2. **Discrete optimization.** The authors decompose the strategies into four different subproblems and separately optimize for each problem. Will this lead to a suboptimal solution compared with a joint optimization method?", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "DM-1B-IQZAv", "reviewer": "Reviewer_5Naa", "summary": "This paper proposes a method to automatically determine the parallel strategy that includes data parallelism (DP), pipeline parallelism (PP), and tensor model parallelism (TMP) over a distributed runtime with heterogeneous computation and communication resources.", "questions": "- In line 116, does the tensor refer to output activations from layer i? \n\n- Would it be possible to be more concrete about the optimization procedure? For example, can the dynamic programming approach explore all the space? What is the complexity of the optimization? \n\n- Would it be possible to provide the execution time of the optimization procedure in the experiments? ", "limitations": "It is helpful to note the limitation on the lack of memory constraints in the cost model. ", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Automatically determining the combined parallel strategy is an interesting and important problem for distributed machine learning considering the advanced of giant foundation models. \n\n- The formalization of the problem and the proposed cost model is reasonable.\n\nWeaknesses:\n- The presentation of the methodology part can be further polished; some of the details are difficult to follow for the first time. For example, a table defining all the variables demanded in the cost model will make Section 3 easier to follow; and there should be a clear discussion about how communication and computation are considered as overlapped execution; perhaps a Gantt chart will be helpful to illustrate this. \n\n- The optimization procedure is not clearly stated in Section 3.6. In fact, Section 3.6 is just a sketch instead of a concrete discussion of the technique. Many details are missing for optimization. \n\n- The experiments are conducted on a relatively small GPT-2 model with a small batch-size; the cluster scale is also small with 4X4 GPUs in total. In fact, the practical challenge for automatic layout is for very large clusters with hundreds or even thousands of GPUs. In the current scope of experimental results, it is difficult to verify the effectiveness of the proposed method.   ", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "BN_-IBHIWCL", "reviewer": "Reviewer_Nnic", "summary": "This paper provides a model to optimize large DL models. It takes into account TMP, DMP, and PMP. Moreover, the proposed model accounts for heterogeneity that exists across devices. ", "questions": "Please address the weaknesses mentioned above.", "limitations": "* Major heterogeneity that exists today in clusters are in terms of GPU types, e.g., use of A100, V100 etc. Paper falls short of its claims\n\n* It is unclear as to why only select few device characteristics such as bandwidth was used. \n\n* Overall it is bunch of heuristics rather than a global and generic optimization framework. Perhaps the paper is more suited for a system’s conference such as OSDI/Eurosys/ATC", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n—-\n\n* models device heterogeneity \n* performs better than state of the art\n\n\nWeakness\n—-\n\n* model is not generic. Lot of heuristics that have been integrated; thereby, exploring few possible optimizations\n* For example, why was latency differences among devices not modelled. It is well known that both bandwidth and latency play a role\n* Similarly, why not consider different generations of GPUs. Is not that more realistic?\n* Does not really account for virtualisation and multi-tenent environment. \n* moreover, the pruning strategies are not really clear. Overall it seems to me that most of the optimization can be specified by a rule-based system", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept