Original: AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness
Revised: AMP: Automatically Finding Model Parallel Strategies with Heterogeneity Awareness

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Scaling up model sizes can lead to fundamentally new capabilities in many machine learning (ML) tasks.
Revised: Scaling up model sizes can lead to fundamentally new capabilities in many machine learning (ML) tasks.

Original: However, training big models requires strong distributedsystem expertise to carefully design model-parallel execution strategies that suitwith the model architectures and cluster setups.
Revised: However, training big models requires strong distributed system expertise to carefully design model-parallel execution strategies that suit the model architectures and cluster setups.

Original: In this paper, we develop AMP, aframework to automatically derive such strategies.
Revised: In this paper, we develop AMP, a framework that automatically derives such strategies.

Original: AMP identifies a valid space ofmodel parallelism strategies, and efficiently searches the space for high-performedstrategies, by leveraging a cost model designed to capture the heterogeneity ofthe model and cluster specifications.
Revised: 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.

Original: Unlike existing methods, AMP is specifically tailored to support complex models composed of uneven layers and clustersetups with more heterogeneous accelerators and bandwidth.
Revised: Unlike existing methods, AMP is specifically tailored to support complex models composed of uneven layers and cluster setups with more heterogeneous accelerators and bandwidth.

Original: We evaluate AMPon popular models and cluster setups from public clouds and show that: AMPreturns parallel strategies that match the expert-tuned strategies on typical clustersetups.
Revised: 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.

Original: On heterogeneous clusters or models with heterogeneous architectures, AMP demonstrates up to 1 . 37 × and 1 .
Revised: On heterogeneous clusters or models with heterogeneous architectures, AMP finds strategies with 1 . 54 × and 1 .

Original: 76 × higher throughput than state-of-the-artmodel-parallel systems, respectively.
Revised: 77 × higher throughput than state-of-the-art model-parallel systems, respectively.

Original:  Introduction
Revised:  Introduction

Original: Recent progress in language understanding [5, 3], multimodal learning [1], and many other MLapplications can largely credit to the use of extremely big models.
Revised: 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.

Original: For example, the largest languagemodels [3, 13] with billions of parameters are found to have fundamentally new capabilities.
Revised: For example, the largest language models [4, 14] with billions of parameters are found to have fundamentally new capabilities.

Original: Dueto the increased model size, training such models, however, calls for model parallelism executionstrategies in order to place the gigantic model computation across multiple, heterogeneous acceleratordevices (e.g., TPUs and GPUs), and ensure efficient distributed execution.
Revised: 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 (e g., TPUs and GPUs) and ensure efficient distributed execution.

Original: Different from typical data parallelism strategies [24, 4, 14], model parallelism considers a muchlarger space of parallelization techniques.
Revised: Different from typical data parallelism strategies [28, 5, 15], model parallelism considers a much larger space of parallelization techniques.

Original: In particular, when a model is too large to fit in a singledevice 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 paralleldevices (i.e., tensor model parallelism (TMP) [22]).
Revised: In particular, when a model is too large to 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]).

Original: Due to their distinct communication and memoryrequirements, these model parallelism dimensions have different applicable scenarios, which maychange with model architectures and cluster setups.
Revised: 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.

Original: For example, TMP results in high communicationvolumes, hence is favored on devices connected with high bandwidth, but may become less effectivewith limited communication bandwidth.
Revised: 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.

Original: Designing performant execution strategies requires combining these parallelism dimensions to trade-off among memory, communication, and computation, thentuning the configurations in each parallelism dimension to align with the heterogeneity of the modeland the training cluster.
Revised: Designing performant execution strategies requires combining these parallelism dimensions to trade-off among memory, communication, and computation, then tuning the configurations in each parallelism dimension to align with the heterogeneity of the model and the training cluster.

Original: Existing model-parallel training strategies are either manually designed by domain experts for one ortwo specific models, or automatically generated with strong assumptions on the model architecture or cluster topologies.
Revised: Existing 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.

Original: For examples, the 3D parallelism strategy adopted in Megatron-LM [22] and
Revised: For example, the 3D parallelism strategy adopted in Megatron-LM [23] and

Original: DeepSpeed [20] is manually specialized for transformer-based language models, assuming a fixedcluster setup and the model having the same layer repeated [22, 20, 18].
Revised: DeepSpeed [21] is manually specialized for transformer-based language models, assuming a fixed cluster setup and the model having the same layer repeated [23, 21, 19].

Original: Therefore, they can hardlygeneralize to models with diverse layer compositions (i.e., heterogeneous models ), or clusters withmixed types of accelerators and network switches (i.e., heterogeneous clusters ).
Revised: 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 ).

Original: This motivates thecore questions in this paper:
Revised: This motivates the core questions in this paper:

Original: Is it possible to automatically find the best model-parallel strategies formore heterogeneous models and clusters?
Revised: Is it possible to automatically find the best model-parallel strategies for more heterogeneous models and clusters?

Original: Given an arbitrary cluster and model setup, finding highly performed model-parallel strategies isdifficult for the following reasons.
Revised: Given an arbitrary cluster and model setup, finding highly performed model-parallel strategies is difficult for the following reasons.

Original: First, when the model has diverse types of layers, existingheuristics on deciding the assignment of layers to pipeline stages (and balance stage workloads) canbe invalidated, because different layers exhibit distinct execution and communication costs.
Revised: 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.

Original: Second,evaluating a model-parallel strategy on a large cluster can be expensive, and prevents search-basedmethods [10] that depend on extensive real evaluations.
Revised: Second, evaluating a model-parallel strategy on a large cluster can be expensive, and prevents search-based methods [11] that depend on extensive real evaluations.

Original: However, existing cost models do not capturethe heterogeneity in the model or cluster (e.g. [23] assumes a uniform bandwidth in the cost model).
Revised: However, existing cost models do not capture the heterogeneity in the model or cluster (e.g. [24] assumes a uniform bandwidth in the cost model).

Original: In this paper, we present solutions to these challenges.
Revised: In this paper, we present solutions to these challenges.

Original: To address the layer-stage assignment problem,we develop a new dynamic programming algorithm running in polynomial time.
Revised: To address the layer-stage assignment problem, we develop a new dynamic programming algorithm running in polynomial time.

Original: To evaluate a largenumber of strategies within a limited budget, we develop a new cost model which serves as a cheaperproxy to evaluate model-parallel strategies.
Revised: To evaluate a large number of strategies within a limited budget, we develop a new cost model which serves as a cheaper proxy to evaluate model-parallel strategies.

Original: We name our method AMP .
Revised: We name our method AMP .

Original: In summary, AMP makesthe following contributions:
Revised: In summary, AMP makes the following contributions:

Original: • We identify factors that influence model-parallel performance in heterogeneous settings, and define strategies based on these factors.
Revised: • We identify factors that influence model-parallel performance in heterogeneous settings, and define strategies based on these factors.

Original: • 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 a minimal amount of real trials.
Revised: • 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 a minimal amount of real trials.

Original: • We develop a dynamic programming approach to handle the load imbalance issue from heterogeneous models in the pipeline layer assignment problem.
Revised: • We develop a dynamic programming approach to handle the load imbalance issue from heterogeneous models in the pipeline layer assignment problem.

Original: • We empirically show that existing systems output sub-optimal strategies in heterogeneoussettings.
Revised: We empirically show that existing systems output sub-optimal strategies in heterogeneous settings.

Original: In contrast, AMP achieves 1 . 37 × speedup with heterogeneous clusters, and 1 .
Revised: In contrast, AMP achieves 1 . 54 × speedup with heterogeneous clusters, and 1 .

Original: 76 ×speedup on heterogeneous model architectures compared to state-of-the-art heuristics.
Revised: 77 × speedup on heterogeneous model architectures compared to state-of-the-art heuristics.

Original:  Background and Related Work
Revised:  Background and Related Work

Original: In this section, we discuss the trade-offs of existing model-parallel strategies, and the need forcomposing them with the awareness of model and cluster heterogeneity.
Revised: In this section, we discuss the trade-offs of existing model-parallel strategies, and the need for composing them with the awareness of model and cluster heterogeneity.

Original: Data-parallel (DP) strategies.
Revised: Data-parallel (DP) strategies.

Original: DP partitions the input data batch evenly among workers, and eachworker holds an entire model replica [14, 21, 8].
Revised: DP partitions the input data batch evenly among workers, and each worker holds an entire model replica [15, 22, 9].

Original: At each iteration, each worker computes gradientsover its assigned batch; the gradients are then synchronized among workers before the next iteration.
Revised: At each iteration, each worker computes gradients over its assigned batch; the gradients are then synchronized among workers before the next iteration.

Original: DP requires each worker to hold an entire model replica, hence cannot be directly used to train modelswith massive parameters.
Revised: DP requires each worker to hold an entire model replica, hence cannot be directly used to train models with massive parameters.

Original: Tensor model-parallel (TMP) strategies.
Revised: Tensor model-parallel (TMP) strategies.

Original: TMP, proposed by Megatron-LM [22], is a popularmodel-parallel approach for large transformer models.
Revised: TMP, proposed by Megatron-LM [23], is a popular model-parallel approach for large transformer models.

Original: In TMP, the layer weights of each twoconsecutive layers are partitioned row-wise ( i.e., input dimension) first, then column-wise ( i.e., outputdimension) [22].
Revised: In TMP, the layer weights of each two consecutive layers are partitioned row-wise ( i.e., input dimension) first, then column-wise ( i.e., output dimension) [23].

Original: TMP removes the need for synchronizing the intermediate output of the very firstlayer, but requires heavy cross-device communications afterwards.
Revised: TMP removes the need for synchronizing the intermediate output of the very first layer, but requires heavy cross-device communications afterward.

Original: TMP is normally combined withdata parallelism to increase the training throughput [17].
Revised: TMP is normally combined with data parallelism to increase the training throughput [18].

Original: Pipeline-parallel (PP) strategies.
Revised: Pipeline-parallel (PP) strategies.

Original: In PP, layers are placed across GPUs, the training mini-batchis split into smaller micro-batches.
Revised: In PP, layers are placed across GPUs, the training mini-batch is split into smaller micro-batches.

Original: The forward and backward computation is then pipelinedacross micro-batches.
Revised: The forward and backward computation is then pipelined across micro-batches.

Original: PP requires less communication than DP and TMP, but suffers from deviceidle ( i.e., pipeline bubbles) [9, 17].
Revised: PP requires less communication than DP and TMP, but suffers from device idle ( i.e., pipeline bubbles) [10, 18].

Original: Despite synchronous pipelining schedules such as GPipe [9],
Revised: Despite synchronous pipelining schedules such as GPipe [10],

Original: PipeDream [17] proposes an 1F1B asynchronous schedule to reduce the pipeline bubble.
Revised: PipeDream [18] proposes an 1F1B asynchronous schedule to reduce the pipeline bubble.

Original: TeraPipedevelops token-level pipelining schedules in a single training sequence for auto-regressive models [15].
Revised: TeraPipe develops token-level pipelining schedules in a single training sequence for auto-regressive models [16].

Original: Automatic partition.
Revised: Automatic partition.

Original: Prior work, e.g., FlexFlow, PipeDream, and DAPPLE explore automaticcomputation graph partitioning over a few typical cluster device setups [10, 17, 6], such as NVIDIA
Revised: Prior work, e.g., FlexFlow, PipeDream, and DAPPLE explore automatic computation graph partitioning over a few typical cluster device setups [11, 18, 7], such as NVIDIA

Original: DGX workstations connected with Infiband switches.
Revised: DGX workstations connected with Infiband switches.

Original: For partitioning the model over multiple devices, the aforementioned work leverages carefully designed cost models to select the desiredstrategy over the search space.3D parallelism.
Revised: 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.

Original: To trade off among the scalability, memory footprint, and device utilization, 
Revised: To trade-off the scalability, memory footprint, and device utilization, DP, MP, and

Original: DP, MP, and PP have to be used together (which is also known as 3D parallelism ).
Revised: PP have to be used together (which is also known as 3D parallelism ).

Original: DeepSpeedcarefully designs a 3D parallelism strategy to train a massive scale language model with 17 billionparameters [20].
Revised: DeepSpeed carefully designs a 3D parallelism strategy to train a massive-scale language model with 17 billion parameters [21].

Original: [18] manually tunes the degrees of DP, MP, PP to scale to 3072 GPUs – bothof them are not automatic .
Revised: [19] manually tunes the degrees of DP, MP, PP to scale to 3072 GPUs – both of them are not automatic .

Original: The closest to our work is Piper [23], which leverages a simple costmodel and a dynamic programming algorithm to search over the 3D strategy space.
Revised: 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.

Original: However, Piperworks only with transformer-based models, and assumes that same type of devices in the cluster arehomogeneously connected with equal bandwidth, ignoring heterogeneous setups.
Revised: However, Piper works only with transformer-based models, and assumes the same type of devices in the cluster are homogeneously connected with equal bandwidth, ignoring heterogeneous setups.

Original: In contrast, AMPis designed to capture the heterogeneity of the given ML model and the deployed cluster, reflected inboth the search space design and the searching algorithm (see section 3).
Revised: In contrast, AMP is designed to capture the heterogeneity of the given ML model and the deployed cluster, reflected in both the search space design and the searching algorithm (§ 3).

Original: Heterogeneity in models and clusters.
Revised: Heterogeneity in models and clusters.

Original: Newer generations of GPU devices and Ethernet switchesemerge rapidly, thus it is common nowadays that multiple types of devices are deployed togetherin a training cluster.
Revised: Newer generations of GPU devices and Ethernet switches emerge rapidly, thus it is common nowadays that multiple types of devices are deployed together in a training cluster.

Original: Conducting distributed training over such types of heterogeneous clustersintroduces extra challenges.
Revised: Conducting distributed training over such types of heterogeneous clusters introduces extra challenges.

Original: For instance, a poor load balancing strategy can cause straggler effectseasily; heterogeneous memory capacity among GPUs introduces new constraints on model and datapartitioning.
Revised: For instance, a poor load balancing strategy can cause straggler effects easily; heterogeneous memory capacity among GPUs introduces new constraints on model and data partitioning.

Original: Apart from hardware heterogeneity, DL models also introduce another dimension ofheterogeneity.
Revised: Apart from hardware heterogeneity, DL models also introduce another dimension of heterogeneity.

Original: Many today’s popular models are composed with diverse layers with different sizesand computing characteristics, e.g., typical convolution networks have their layer width growing withdepth [7], which would require distinct partitioning and placement strategies for each layer duringmodel parallelism.
Revised: Many today’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.

Original:  Methods 
Revised:  Methods 3.1 Problem Formulation

Original: 3.1 Problem Formulation

Original: Formally, the input of AMP contains: (i) a given model, (ii) a given cluster, and (iii) the global batchsize ( gbs ).
Revised: Formally, the inputs of AMP are (i) the model, (ii) the cluster, and (iii) the global batch size ( gbs ).

Original: A deep learning model W is represented as a graph, i.e., W = ( C W , V W ) .
Revised: A deep learning model W is represented as a graph, i.e., W = ( C W , V W ) .

Original: C W is aset of vertices where each vertex c W,i ∈ C W , i ∈
Revised: C W is a set of vertices where each vertex c W,i ∈ C W , i ∈

Original: { i, · · · , | C W |} represents a layer, and each edgev W ,i = ( c W ,i ; c W ,i +1 )
Revised: { i, · · · , | C W |} represents a layer, and each edge v W ,i = ( c W ,i ; c W ,i +1 )

Original: ∈ V W represents the tensor (activation) between layer c W ,i and layer c W ,i +1 .
Revised: ∈ V W represents the tensor (activation) between layer c W ,i and layer c W ,i +1 .

Original: Additionally, we denote the execution cost of a layer c
Revised: Additionally, we denote the execution cost of a layer c

Original: W ,i as c i , and the communication volume of anedge v W ,i as v i .
Revised: W ,i as c i , and the communication volume of an edge v W ,i as v i .

Original: A cluster is represented as C =
Revised: A cluster is represented as C =

Original: {D , B } where D = { d i | 1 ≤ i ≤ |D|} denotes a setof devices.
Revised: {D , B } where D = { d i | 1 ≤ i ≤ |D|} denotes a set of devices.

Original: B ∈ R |D|×|D| is a symmetric matrix where each matrix element b ij represents bandwidthbetween device d i and d j .
Revised: B ∈ R |D|×|D| is a symmetric matrix where each matrix element b ij represents bandwidth between device d i and d j .

Original: The goal of AMP is to output a high throughput parallel training strategy sgiven the hardware and model configuration.
Revised: The goal of AMP is to output a high throughput parallel training strategy s given the hardware and model configuration.

Original: S denotes the entire space of candidate strategies.
Revised: S denotes the entire space of candidate strategies.

Original: Weformulate this as an optimization problem:
Revised: We formulate this as an optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where f ( ·; · ) denotes per iteration running time as a function of the strategy s , conditioned on thegiven ML model, cluster, and user-specified global batch size.
Revised: where f ( · ) denotes per iteration running time as a function of the strategy s , conditioned on the given ML model, cluster, and user-specified global batch size.

Original: 3.2 Method Overview
Revised: 3.2 Method Overview

Original: To solve the optimization problem in equation 1, we identify factors that influence the systemperformance (§3.3).
Revised: To solve the optimization problem in Equation 1, we identify factors that influence the system performance (§3.3).

Original: However, there are two major problems in solving this optimization problemwith these identified factors.
Revised: However, there are two major difficulty in solving this optimization with these identified factors.

Original: First is how to evaluate f ( · ) .
Revised: First is how to evaluate f ( · ) .

Original: One possible solution is to launch realtrials, but this is prohibitively expensive given that we possibly need to evaluate a large number ofstrategies.
Revised: One possible solution is to launch real trials, but this is prohibitively expensive because the entire optimization procedure may require evaluating a large number of strategies.

Original: Rather we develop a cost model (§ 3.4) that can estimate f ( · ) quickly.
Revised: Rather we develop a cost model (§ 3.4) that can estimate f ( · ) quickly.

Original: Second, the discretenature of the optimization problem disallows us to leverage gradient-based optimization methods.
Revised: Second, the discrete nature of the optimization problem disallows us to leverage gradient-based optimization methods.

Original: However, enumerating all strategies  is slow even with the cost model.
Revised: However, evaluating all strategies in the search space S is slow even with the cost model.

Original: One way to efficiently optimizethe objective is to exploit structure in the strategy space S , and prune areas with no promisingcandidates.
Revised: One way to efficiently optimize the objective is to exploit structure in the strategy space S , and prune areas with no promising candidates.

Original: Using this principle, we develop a dynamic programming approach (§3.5), which prunesout candidates with worse performance before invoking the cost model.
Revised: Using this principle, we develop a dynamic programming approach (§3.5), which prunes out candidates with worse performance before invoking the cost model.

Original: 3.3 Representation of a Strategy
Revised: 3.3 Representation of a Strategy

Original: Intuitively, a strategy s defines how to train model { C W , V W } on a given cluster {D , B } with thebatch size gbs .
Revised: At a high level, a strategy s defines how to train model { C W , V W } on a given cluster {D , B } with the batch size gbs .

Original: 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 amini-batch be evenly split) in DP; how to partition model layers (for PP) or layer parameters (for
Revised: 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

Original: TMP); and how shall the global batch size decompose into micro-batches in PP semantics.
Revised: TMP); and how shall the global batch size decompose into micro-batches in PP semantics.

Original: In AMP,we represent a strategy s with the following four aspects: 
Revised: In AMP, we represent a strategy s with the following four aspects: 1.

Original: Parallelism degrees
Revised: Parallelism degrees.

Original: Since PP, DP, and TMP have different characteristics, different degreesof each of them affect the system performance.
Revised: Since PP, DP, and TMP have different characteristics, different degrees of each affect the system performance.

Original: Denote the degrees to be pp , dp , tmp .
Revised: We denote the degrees to be pp , dp , tmp .

Original: The tuple( pp, dp, tmp ) defines a three-dimensional grid, which satisfies the constraint pp × dp × tmp =|D| ; pp, dp, tmp ∈ { 1 , . .
Revised: The tuple ( pp, dp, tmp ) defines a three-dimensional grid, which satisfies the constraint pp × dp × tmp = |D| ; pp, dp, tmp ∈ { 1 , . .

Original: . , |D|} .
Revised: . , |D|} .

Original: ( i.e., this tuple generates pp × dp × tmp virtual ranks that will bebijectively mapped to available devices). Two ranksare in the same DP(TMP/PP) communicationgroups if they have different values in the DP(TMP/PP) dimension and the same values in others ( e.g.,rank (0 , 0 , 0) and rank (0 , 0 , 1) forms a TMP group).
Revised: This tuple generates pp × dp × tmp virtual ranks that carry out corresponding tasks. For instance, ranks in 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.

Original: [Figure]
Revised: [Figure]

Original: Assuminga uniform bandwidth, a uniform-layer assignment cuts
Revised: Assuming uniform bandwidth, uniform layer assignment cuts

Original: balanced

Original: Device placement.
Revised: Device placement.

Original: Since physical devices may have different connectionspeeds, i.e., bandwidth to each other,  placing physical devices in different virtual ranks leads to dif- ferent performance.
Revised: 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.

Original: For instance, if ranks in a communication- heavy group are all allocated with physical devices with highbandwidth, the per iteration time will be shorter Formally,
Revised: For instance, if ranks in a communication-heavy group are all allocated with physical devices with high bandwidth, the per iteration time will be shorter.

Original:  We denote a device placement as a bijective function fromvirtual ranks to devices p : { 1 , . .
Revised: Formally, We denote a device placement as a bijective function from virtual ranks to devices p : { 1 , . .

Original: . , pp } × { 1 , . . .
Revised: , pp } × { 1 , . . .

Original: , dp } ×{ 1 , . .
Revised: , dp } × { 1 , . .

Original: . , tmp } (cid:55)→ D .3.
Revised: . , tmp } (cid:55)→ D . 3.

Original: Micro-batch size.
Revised: Micro-batch size.

Original: Micro-batch size ( mbs ) is the batchsize a device holds at a time.
Revised: Micro-batch size ( mbs ) is the batch size a device proceeds at a time in the pipeline parallelism execution (§2).

Original: When mbs varies, devicesgenerally have different utilization.
Revised: When mbs varies, devices usually have different utilization [18, 26].

Original: This effect is empiricallyverified in [18], where the observation is that different microbatch sizes lead to noticeably different in system throughput.4.
Revised: For instance, a larger micro-batch size usually leads to better GPU utilization because the GPU can better parallel the computation. However, a larger microbatch size also leads to a larger pipeline bubble, which makes selecting mbs hard.

Original: Pipeline layer assignment.
Revised: Pipeline layer assignment.

Original: How to assign layers to eachdevice can also affect the system throughput.
Revised: Different assignments of layers to each device can also affect the system throughput.

Original: On the one hand,when activation volume v is large in the boundary, the nextdevice needs to wait longer, resulting in a longer per iteration time.
Revised: On the one hand, when activation volume v is large in the boundary, the next device needs to wait longer, resulting in a longer per iteration time.

Original: On the other hand, when a deviceis assigned a much higher load than other devices, it blocks all other devices and thus incurs a longeriteration time.
Revised: 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.

Original: We denote a feasible pipeline layer assignment as a function a : C W (cid:55)→ R that mapslayers to the stage index (1 , . .
Revised: We denote a feasible pipeline layer assignment as a function a : C W (cid:55)→ R that maps layers to the stage index (1 , . .

Original: , pp ) in a pipeline.
Revised: , pp ) in a pipeline.

Original: In summary, a strategy s can be defined by a 6 -dimensional space:
Revised: In summary, a strategy s can be defined by a 6 -dimensional tuple:

Original: [Equation]
Revised: [Equation]

Original: 3.4 Cost Model
Revised: 3.4 Cost Model

Original: In this section, we introduce our cost model that serves as a fast estimation of f ( · ) . 
Revised: 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.

Original: Communication primitive.
Revised: Communication primitive.

Original: We consider  ring topology for all-reduce collective communicationin both DP and TMP [21].
Revised: We consider the ring topology for all-reduce collective communication in both DP and TMP [22].

Original: Assume that the whole message size is M .
Revised: Assume that the whole message size is M .

Original: In a cluster of n workerswith uniform bandwidth B , the communication timecost of a single ring all-reduce operation can bemodeled as [12]:
Revised: In a communication group of n workers with uniform bandwidth B , the communication time of a single ring all-reduce operation can be modeled as [13]:

Original: [Equation]
Revised: [Equation]

Original: In a cluster with non-uniform bandwidth,denote B ′ to be the slowest bandwidth in the all-reducecommunicative group, we approximate the bandwidth with  B ′ . 
Revised: 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

Original: For point-to-point (P2P) communication operations between consecutive pipeline stages, we assume the communication operationcan always enjoy the full bandwidth: e = vb where v is the activation volume and b is the bandwidthbetween these two GPUs.
Revised: [Table] assume the communication operation can use the full bandwidth: e = vb [11]. Here v is the activation volume and b is the bandwidth between these two GPUs.

Original: Top-down run time decomposition.
Revised: Top-down run time decomposition.

Original: We now present our cost model designed from a top-downview.
Revised: We now present our cost model designed from a top-down view.

Original: First, DP stands at the highest granularity, where there are dp pipelines running in parallel.
Revised: First, DP is at the top granularity, where there are dp pipelines running in parallel.

Original: It synchronizes gradients when all pipelines finish.
Revised: It synchronizes gradients when all pipelines finish their schedules.

Revised: It does not consider how each pipeline is executed.

Original: At the DP level, the system runtime can bedecomposed into two terms:
Revised: At the DP level, the system run-time can be decomposed into two terms:

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: T ppi is the time taken to complete a single pipeline i .
Revised: T ppi is the time taken to complete a single pipeline i .

Original: T dpsyncd denotes the time for a device d tocommunicate data parallelism gradient updates.
Revised: T dpsyncd denotes the time for a device d to synchronize gradients with an all-reduce operation.

Original: The formula is given at equation 3 with message sizeequal to the number of model parameters held at device d. Since different pipelines may finish withdifferent speeds, and different devices may finish its data parallelism synchronization with differentspeeds, we take maximum over these two terms respectively.
Revised: We model the latter term by using Equationwith the message size equal to the number of model parameters held at 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.

Original: At a lower granularity, each pipelineexecutes its own training schedule such as activation/gradients exchange at a layer level ( i.e., it doesnot need to consider how each layer performs computation, which is carried out by MP workers).
Revised: At a lower granularity, each pipeline executes a training schedule that includes activation/gradients exchange at a layer level. It does not need to consider how each layer is executed, which is carried out by MP workers.

Original: Assume that asingle pipeline consists of k stages. The communication time between the i th and( i + 1) th isdenoted as e i  .
Revised: Assume the number of stages in a pipeline is k , and the communication time between the i th and ( i + 1) th is e i , and the execution time for the i th stage is t i .

Original: Execution cost for the i th stage t i is the sum over all layer j executiontime t layerj in the  stage: t i = (cid:80) a ( j )=
Revised: We model t i as the sum over execution time t layerj in the stage for layer j in the stage: t i = (cid:80) a ( j )=

Original: i t layerj .
Revised: i t layerj .

Original: Assuming that the P2P communication and computation can be overlapped 1 , the runtime for a singlepipeline can be modeled as: [Equation]
Revised: We model the execution time for a single pipeline as 2 :

Original: [Equation]
Revised: [Equation]

Original: Intuitively, the first term models the straggler effect, which gets amplified by the number of gradient accumulation steps. When gas 
Revised: Intuitively, the first term captures the straggler effect, which gets amplified by the number of microbatches gas .

Original: = 1 , the schedule is straightforward to understand - themini-batch simply performs a single forward pass and a single backward pass, sends or receives gradients/activations at the boundary.

Original: In this case, T pp given at equation 5 reduces to  (cid:80) pp − 1 i =

Revised: If we assume gas = 1 , Equation 5 reduces to T pp = (cid:80) pp − 1 i =

Original: e i + (cid:80) ppi =
Revised: e i + (cid:80) pp i =

Original: 1 t i .
Revised: 1 t i .

Revised: This aligns with the actual schedule: the mini-batch simply performs a single forward pass and a single backward pass, and sends or receives gradients or activations.

Original: Note that when all layers have the same execution cost and the cluster is equipped with the samebandwidth, the optimal solution is simply to balance the number of layers.
Revised: If we assume all layers have the same execution cost and the cluster is equipped with the same bandwidth, the optimal solution is simply to balance the number of layers.

Original: This heuristics is oftenadopted as the default layer assignment  in current systems[22].
Revised: This heuristic is often adopted as the default layer assignment solution in current systems[23].

Original: At the lowest granularity, MP carries out results of individual layers: t layerj =
Revised: At the bottom granularity, MP workers consider how a single layer is executed.

Original: c j tmp + T mpsyncj . T mpsyncj is  the  all-reduce communication time, which can also be calculated by equation 3.
Revised: The run-time of a single layer t layerj is composed of the layer computation time ( e.g., Matrix multiplication) and all-reduce communication time.

Original: Oneway to estimate the execution cost for a single layer c j is by counting the number of floating-pointoperations (FLOPs).
Revised: One way to estimate t layerj is using the number of floating point operations (FLOPs) for computation time, and Equation 3 for communication time.

Original: However,  when matrix size is small, its arithmetic intensity is low, and GPUs canhardly parallel efficiently.
Revised: However, we find this inaccurate in our setting. Concretely, when mp grows, the arithmetic intensity of each matrix multiplication becomes lower.

Original: In such scenarios,  matrix multiplication with  few FLOPs can take a longtime to execute.
Revised: In such scenarios, a matrix multiplication with only a few FLOPs can take a long time to execute.

Original: To handle general deep learning models, we simply obtain t layerj by profiling withseveral tensor model parallelism settings.
Revised: Thus, we simply obtain t layerj by profiling several tensor model parallelism settings and store results in a lookup table.

Original: This profiling  only needs to run once, which is negligiblecompared to model training time.
Revised: This profiling is feasible because it only needs to be done once in the entire optimization. [Table]

Original: 3.5 Pipeline Layer Assignment
Revised: 3.5 Pipeline Layer Assignment

Original: A key observation in the layer assignment problem is that only consecutive layers can be assigned tothe same stage.
Revised: A key observation in the layer assignment problem is that only consecutive layers can be assigned to the same stage.

Original: Intuitively, computation up to layer i can be reused for that of layer i + 1 .
Revised: In particular, computation up to layer i can be reused for that of layer i + 1 .

Original:  To handlethe max term in objective 5,  an auxiliary variable tolerance m is introduced, where semantically mmeans a possible execution time for a single stage:
Revised: This pattern fits neatly in a dynamic programming solution. To handle the max term in objective 5, we use an auxiliary variable tolerance m . m takes value from all possible execution time for a single stage.

Original: [Equation]
Revised: For instance, it 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:

Original: [Table]
Revised: [Equation]

Original: Note that when m = 0 , objective 6 reduces to the desired objective 5.
Revised: Note that when m = 0 , Equation 6 reduces to the desired objective 5.

Original: Formally, we record theoptimal assignments and the corresponding cost up to layer i with j stages and variable m in tabledp[i][j][m].
Revised: Formally, we store the optimal assignments and the corresponding cost up to layer i with j stages and variable m in table dp [ i ][ j ][ m ] .

Original: Let L be the number of layers and k be the number of stages, then dp[L][k][0] stores thesolution.
Revised: Let L be the number of layers, then dp [ L ][ k ][0] stores the solution.

Original: We induct on the number of stages in the pipeline.
Revised: We induct on the number of stages in the pipeline.

Original: Base case  j = 1 .
Revised: Base case ( j = 1 ).

Original: The only possible assignment assigns all i layers into a single stage.
Revised: The only possible assignment assigns all i layers into a single stage.

Original: Denote thesum of execution time for these i layers as t 1 .
Revised: Denote the sum of execution time for these i layers as t 1 .

Original: The associated optimal cost is thus
Revised: The associated optimal cost is thus

Original: [Equation]
Revised: [Equation]

Original: We store the table entry dp [ i ][0][ m ] with ( i, g ( i )) for all possible i and m .
Revised: We store the table entry dp [ i ][0][ m ] with ( i, g ( i )) for all possible i and m .

Original: Recursive step.
Revised: Recursive step.

Original: To compute dp [ i ][ j ][ m ] when j > 1, we enumerate all possible position for lastcut i ′ for all i ′ <
Revised: To compute dp [ i ][ j ][ m ] when j > 1 , we enumerate all possible positions of the last cut i ′ for all i ′ <

Original: i .
Revised: i .

Original: Denote the execution time from layer i ′ to layer i as t 2 ,i ′ , let the valuestored at dp [ i ′ ]
Revised: Denote the execution time from layer i ′ to layer i as t 2 ,i ′ , let the value stored at dp [ i ′ ]

Original: [ j − 1][max( t 2 ,i ′ , m )] be ( a i ′ , c i ′ ).
Revised: [ j − 1][max( t 2 ,i ′ , m )] be ( a i ′ , c i ′ ).

Original: Denote the longest stage time in a i ′ as t 1 ,i ′ , thecommunication cost between the ( j − 1) th and j th device at layer i ′ as e i ′ ,j .
Revised: Denote the longest stage time in a i ′ as t 1 ,i ′ , the communication cost between the ( j − 1) th and j th device at layer i ′ as e i ′ ,j .

Original: The cost associatedwith the last cut at i ′ is:[-
-]
Revised: The cost associated with the last cut at i ′ is:[-
-]

Original: [Equation]
Revised: [Equation]

Original: Select the optimal last cut position i ′ opt = argmin i ′ <
Revised: Select the optimal last cut position i ′ opt = argmin i ′ <

Original: i g ( i ′ ) .
Revised: i g ( i ′ ) .

Original: The dynamic programming table entrydp [ i ][ j ][ m ] is updated as: ( a i ′ opt ∪ ( i − i ′ opt ) , g
Revised: The dynamic programming table entry dp [ i ][ j ][ m ] is updated as: ( a i ′ opt ∪ ( i − i ′ opt ) , g

Original: ( i ′ opt )) .
Revised: ( i ′ opt )) .

Original: Run time.
Revised: Running time.

Original: There are (cid:80) Li
Revised: There are (cid:80) Li

Original: =

Revised: =

Original: i = O
Revised: i = O

Original: ( L 2 ) possible values for m .
Revised: ( L 2 ) possible values for m .

Original: Thus, our problem size isO
Revised: Thus, our problem size is

Original:  ( L ×k × L 2 ) , each sub-problem takes O ( L ) time.
Revised: O ( L × k × L 2 ) , each sub-problem takes O ( L ) time.

Original: Thus the total run time  is O ( kL 4 ) , while a brute forcesolution takes O (cid:0) L − 1 K − 1 (cid:1) .
Revised: Thus the total running time for the dynamic programming algorithm is O ( kL 4 ) , while a brute force solution takes O (cid:0) L − 1 K − 1 (cid:1) .

Original: 3.6 Optimization Procedure
Revised: 3.6 Optimization Procedure

Original: To find  best strategies defined at § 3.3, we enumerate all possible parallelism degrees and microbatch sizes.
Revised: In this section, we 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.

Original: We deterministically solve for the optimal pipeline layer assignment using the dynamicprogramming approach developed at § 3.5.
Revised: For each parallelism degree and micro-batch size, we deterministically solve for the optimal pipeline layer assignment using the algorithm at § 3.5.

Original: We adopt the heuristics at [18] to solve the deviceplacement problem: prioritize placing model parallelism workers in the same node, over dataparallelism workers, and pipeline parallelism workers. 
Revised: 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 intra-node bandwidth, then place data parallelism workers and pipeline parallelism workers. The procedure is written in Algorithm 1. More details can be found in § 6 .

Revised: in § 6.

Revised: Input: C , W , gbs , budget 1 degrees =

Original: We also explore an alternative optimizationfor device placement at § 4.2.
Revised: enumerate_degrees( C ) 2 record = set() 3 for d in degrees do /* Possible micro-batch size given data parallel degree */ for mbs in enumerate_( d.dp ) do p = placement( C , d ) // Device placement heuristics a = pipe_ast( W , d.pp , p ) // Optimal pipeline assignment s = ( a , mbs , d , p ) // current strategy cost = estimate( s ) //

Revised: our cost model record.add( s , cost) /

Revised: * run top predicted strategies */ 10 best_s =

Revised: run(record, budget) 11 return best_s

Original:  The concrete algorithms and complexity analysis can be found inappendix.
Revised: [Table] Complexity of the overall optimization. The number of factors of an integer N can be loosely upper bounded by O ( N 1 / 2 ) . Thus, the number of possible degrees can be upper bounded by

Revised: O ( |D| 1 / 2 )) × O ( |D| 1 / 2 ) =

Revised: O ( |D| ) , and the number of possible micro-batch sizes can be upper bounded by O ( gbs 1 / 2 ) .

Revised: The pipeline assignment algorithm run-time is bounded by O ( kL 4 ) =

Revised: O ( |D| × L 4 ) (section 3.5).

Revised: Thus, the total run time can be upper bounded by O ( gbs 1 / 2 × |D| × |D| × L 4 ) =

Revised: ( gbs 1 / 2 × |D| 2 × L 4 ) .

Revised: Empirically, we verify the optimization time with respect to the cluster size in Table 2.

Revised: We find that the actual optimization time is near linear with respect to the cluster size.

Original:  Experiments
Revised:  Experiments

Original: In this section, we evaluate our algorithm and several of its variants on different types of clusters andmodels.
Revised: In this section, we evaluate AMP on different types of clusters and models. Besides a homogeneous setup, we study two heterogeneous setups, one with a heterogeneous cluster and one with a heterogeneous model.

Original: Specifically, we compare against the state-of-the-art Megatron heuristics [18].
Revised: We use the state-of-the-art Megatron heuristics as our baseline [19].

Original: 4.1 Baseline
Revised: 4.1 Baseline

Original: Megatron  is a runtime system that supports 3D parallelism training on DNN models.
Revised: Megatron [23] is a runtime engine that supports 3D parallelism training on DNN models.

Original: It suggestsseveral heuristics on how to train models with 3D parallelism faster.
Revised: It suggests several heuristics on how to train models with 3D parallelism faster.

Original: We summarize these heuristicsas the following procedure.
Revised: We summarize these heuristics in the following procedure.

Original: Megatron heuristics loops over all possible micro-batch sizes.
Revised: Mgetraon explores all possible micro-batch sizes as in AMP.

Original: For each micro-batch size, it selectsparallelism degrees with the smallest tmp × pp , whose tmp is smaller than the  number of devices inthe smallest node. 
Revised: For each micro-batch size, it forces tmp to be smaller than the smallest number of devices in a node. Then it selects parallelism degrees with the smallest tmp × pp .

Original: It solves the device placement problem by prioritizing tensor model parallelismworkers in the same node.
Revised: It solves the device placement problem by prioritizing tensor model parallelism workers in the same node.

Original: If there are additional available intra-node connections, it places data2 proof can be found at the supplementary materials.
Revised: If there are additional available intra-node connections, it places data parallelism workers next.

Revised: It solves the pipeline layer assignment problem by either balancing the number of parameters or the number of layers in each stage.

Revised: Megatron is proposed based on observations on a homogeneous cluster and model setup.

Revised: The model is composed of simple stacks of identical layers.

Revised: The cluster is equipped with the same type of devices with the same bandwidths.

Original: 1.5 2.0 2.5 3.0 3.5 4.0 4.5 Estimated Per Iter.
Revised: Thus, We hypothesize under such as homogeneous setup, it is (near) optimal. To validate this hypothesis, we conduct experiments on a homogeneous setup ( § 4.3). Moreover, we are interested in how Megatron heuristics perform when the homogeneous assumption does not hold. Concretely, we ask two questions.

Revised: What is optimal when the cluster is equipped with

Revised: 8 Estimated Per Iter.

Original: Time (sec.)
Revised: Time (sec.)

Original: P e r I t e r .
Revised: P e r I t e r .

Original: T i m e ( s e c . )
Revised: T i m e ( s e c . )

Revised: spearman correlation = 0.

Original: (a) homogeneous 1.
Revised: (a) homogeneous 1.

Original: 3.0 Estimated Per Iter.
Revised: 3.0 Estimated Per Iter.

Original: Time (sec.)
Revised: Time (sec.)

Original: P e r I t e r .
Revised: P e r I t e r .

Original: T i m e ( s e c . )
Revised: T i m e ( s e c . )

Revised: Spearman Correlation = 0.

Original: (b) heterogeneous cluster 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.55.0 Estimated 
Revised: (b) heterogeneous cluster 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 Estimated Per Iter.

Original: [Figure]

Original: Time (sec.)1.
Revised: Time (sec.)

Original: P e r I t e r .
Revised: P e r I t e r .

Original: T i m e ( s e c . )
Revised: T i m e ( s e c . )

Original: (c) heterogeneous model parallelism workers next. It solves the pipeline layer assignment problem by either balancing thenumber of parameters or the number of layers in each stage. Megatron heuristics assume a homogeneous cluster and model setup. The model is composedof simple stacks of identical layers.
Revised: Spearman Correlation = 0.

Original: The cluster is equipped with the same type of devices withsimilar bandwidths to each other.

Original: We hypothesize that under this setup, it is (near) optimal. Tovalidate this hypothesis, we conduct experiments on a homogeneous setup (section 4.3). Moreover,we are interested in how Megatron heuristics perform when the homogeneous assumption breaks. Specifically, we want to answer two questions.
Revised: [Figure]

Original: What will happen when the cluster is equipped withdifferent types of devices and different bandwidths to each other?
Revised: heterogeneous model different types of devices and different bandwidths to each other?

Original: What is the effect of heterogeneityin models?
Revised: What is the effect of heterogeneity in models?

Original: Thus, we design two sets of experiments under  heterogeneous cluster and heterogeneousmodel insection 4.3.
Revised: To answer these, we design two sets of experiments under a heterogeneous cluster and a heterogeneous in

Original: 4.2 Alternative Device Placement Algorithm
Revised: 4.2 Alternative Device Placement Algorithm

Original: Since the number of possible device placements is large and often difficult to optimize efficiently[16], we adopt  Megatron heuristics in our  method.
Revised: As described in Section 3.6, we adopt the heuristic in Megatron to solve the device placement algorithm in our main method.

Original: To study the optimality of this heuristic, we use a 
Revised: 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.

Original:  Simulated Annealing(SA) algorithm to explore candidates with different placements.
Revised: We instead use a Simulated Annealing(SA) algorithm to optimize the device placement aspect [25].

Original: We name thismethod SA .
Revised: We name this method SA .

Original: Details can be found in the supplementary materials.
Revised: Details can be found in § 6.3.

Original: 4.3 Experiment Setup
Revised: 4.3 Experiment Setup

Original: Homogeneous.
Revised: Homogeneous.

Original: We conduct experiments using GPT-2 ( L = 24 , H = 1024 ) [19] on 4 AWS EC2g4dn.12xlarge nodes with a global batch size 32.
Revised: We conduct experiments using GPT-2 ( L = 24 , H = 1024 ) [20] on 4 AWS EC g4dn.12xlarge nodes with a global batch size 32.

Original: Each instance is equipped with 4 T4 GPUs with[-
-]
Revised: Each instance is equipped with 4 T4 GPUs with[-
-]

Original: Gbps PCIe connection intra-node bandwidth, and 50 Gbps inter-node bandwidth.
Revised: Gbps PCIe connection intra-node bandwidth, and 50 Gbps inter-node bandwidth.

Revised: In all experiments, we place instances in the same AWS placement group to make sure they have the full bandwidth.

Original: Heterogeneous cluster.
Revised: Heterogeneous cluster.

Original: To simulate heterogeneity in both different types of GPUs and differentbandwidths, we ensemble a cluster using 3 AWS EC2 p3.8xlarge instances (V100 GPUs) and 1g4dn.12xlarge (T4 GPUs) instance.
Revised: To simulate heterogeneity in both different types of GPUs and different bandwidths, we ensemble a cluster using 3 AWS EC2 p3.8xlarge instances (V100 GPUs) andg4dn.12xlarge (T4 GPUs) instance.

Original: Each p3.8xlarge instance is equipped with 4 V100 GPUs withapproximately 170Gbps NVLink intra-node connection and 10 Gbps inter-node bandwidth.
Revised: Each p3.8xlarge instance is equipped with 4 V100 GPUs with approximately 170Gbps NVLink intra-node connection and 10 Gbps inter-node bandwidth.

Original: evaluateperformance on GPT-2 ( L = 24 , H = 1024 ) with a global batch size 32 .
Revised: We evaluate the performance on GPT-2 ( L = 24 , H = 1024 ) with a global batch size 32 .

Original: Heterogeneous model.
Revised: Heterogeneous model.

Original: We use a modified version of Transgan generator [11], which stackstransformer layers with gradually smaller hidden sizes.
Revised: We use a modified version of Transgan generator [12], which stacks transformer layers in the way that later layers have smaller hidden sizes.

Original: Specifically, we use 12 layers with dimensions1024, and 12 layers with dimensions 16.
Revised: Specifically, we uselayers with dimensions 1024, and 12 layers with dimensions 16.

Original: We use 4 p3.8xlarge instances and a global batch size 64.
Revised: We use 4 p3.8xlarge instances and a global batch size 64.

Original: 4.4 Evaluation Metric
Revised: 4.4 Evaluation Metric

Original: We evaluate strategies by  time per iteration.
Revised: We evaluate strategies by the running time per iteration.

Original: Specifically, we run each strategy by 40 iterations andtake the average iteration time for the latest 20 iterations.
Revised: Specifically, we run each strategy byiterations and take the average iteration time for the latest 20 iterations [27].

Original: The underlying system is Deepspeed (builton top of  Megatron system) [20] with fp16 optimization enabled.
Revised: The underlying system is Deepspeed (built on top of the Megatron engine) [21] with fp16 optimization enabled.

Original: 4.5 Results and Discussion
Revised: 4.5 Results and Discussion

Original: Homogeneous setup.
Revised: Homogeneous setup.

Revised: In this experiment, we compare three methods: (1) Megatron : Megatron

Revised: Heuristics, (2) AMP : our method described in § 3.6, and (3) SA : optimize the device placement aspect using the Simulated Annealing.

Original: Under homogeneous setupas specified in  4.3, we hypothesize that Megatronheuristics is (near) optimal.
Revised: Under homogeneous setup in § 4.3, we hypothesize that Megatron heuristics is (near) optimal.

Original: We also expect that our method can find such optimal strategies.
Revised: We also expect that our method can find such optimal strategies.

Revised: In our experiment setting, Megatron proposes 10 feasible strategies.

Revised: To understand the difference in strategies between our proposal and the Megatron proposal, we evaluate the Top 10 strategies from

Original: [Table]
Revised: [Table]

Original: Additionally, we conjecture that our device placement algorithm can not propose better strategiesthan the default one from Megatron. In this experiment, we compare three methods: • Megatron : Megatron Heuristics. • AMP : our cost model that enumerates parallelism degrees and micro-batch size, and uses
Revised: [Table]

Original: dynamic programming to solve the pipeline layer assignment. • SA : additionally optimize device placement using the simulated Annealing algorithm. In our experiment setting, Megatron proposes around 10 feasible strategies. To understand thedifference in strategies between our proposal and the Megatron proposal, we evaluate the top 10strategies from our method.
Revised: our method.

Original: The results are presented in Table 2.
Revised: The results are presented in Table 3.

Original:  AMP finds slightly better strategiesthan Megatron Heuristics.
Revised: We find that Megatron is indeed near-optimal and AMP can find such strategies (in fact slightly better).

Original:  Since the model can be fit into a single GPU, Megatron heuristics selectsan extreme configuration - dp = 16 .
Revised: We analyze the selection procedure below. Since the model can be fit into a single GPU, Megatron selects an extreme configuration - dp = 16 .

Revised: This may be sub-optimal because the communication in the data parallelism group is high.

Revised: Suppose the micro-batch size is small, so that the bubble overhead in the pipeline parallelism may be hidden, it may be better to assign some workers to the pipeline parallelism.

Original: Our method automates the strategies selection process andfinds that dp = 4  is a better choice.
Revised: In fact, AMP finds that dp = 4 , pp =with the smallest micro-batch size 1 is a better choice.

Original: By enabling our SA algorithm to explore device placement, SA findsslightly worse results.
Revised: We also find that exploring more device placement strategies with SA finds worse results.

Original: The reason is that in such a cluster setup where devices have thesame connection speeds, device placements have limited influence on the system performance.
Revised: This is because considering more device placement has limited influence in the homogeneous setup.

Original: Forinstance, whether the tensor model parallelism communication happens within a node or cross nodeshall not change the performance.
Revised: For instance, whether the tensor model parallelism communication happens within a node or cross node shall not change the performance.

Original: However, it introduces  more strategies that make optimizationhard.
Revised: However, it introduces much more strategies that make the optimization hard.

Original: This effect is illustrated in Figure 3.

Original: [Figure]

Original: top

Original: SA strategies are higher, leading to a higher actual cost.

Original: Heterogeneity in cluster setup.
Revised: Heterogeneity in the cluster.

Original: In the setup at section 4.3,devices have different connectivity to each other, and different execution speeds.
Revised: In the heterogeneous cluster setup, devices have different connectivity to each other, and different execution speeds.

Original: We hypothesize that some heuristicsfrom Megatron can lead to bad performance, but our methodshould be able to automatically address the heterogeneity in(especially) bandwidth since we explicitly model it (3.4).
Revised: 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).

Original: Withheterogeneity in the cluster, we are also interested in examining how much Megatronheuristics and our method improvecompared to average strategies in the space.
Revised: 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.

Original: In this experiment,we have four methods to compare: Megatron, AMP, SA, and• RS : Randomly selecting strategies without heuristicsand a cost model.
Revised: In this experiment, we have four methods to compare: Megatron, AMP, SA, and • RS : Randomly selecting strategies without heuristics and a cost model.

Original: We run RS for 50 iterations.
Revised: We run RS for 50 iterations.

Original: Results are presented in Table3.
Revised: Results are presented in Table 4.

Original: Wefound that average strategies have a high cost, butrandomly evaluating them can find better strategies  than Megatron.
Revised: We find that the mean of random strategies is 1 . 16 × worse than that of Megatron, and 1 .

Original: However, this comes with a muchhigher cost (50 vs 10 real trials), and cannot reliably produce good results.
Revised: 76 × worse than that of ours.

Original: By using our cost model,we can find strategies up to 1 . 37 × speedup compared to Megatron.
Revised: By using our optimization procedure, we find strategies that achieve 1 . 54 × speedup compared to Megatron.

Original: Specifically, the heuristics of Megatron that minimize mp × pp is sub-optimal.
Revised: Specifically, the heuristic of Megatron that minimize mp × pp is sub-optimal .

Original: When using micro-batch size 1, the best strategyshall be dp=4 and pp=4, which results in a 1.44 second/iteration performance.
Revised: When using microbatch size 1, the best strategy shall be dp = 2 , pp = 8 , which results in a 1.28 second/iteration performance.

Original: However, since themodel can be fit in a single GPU, Megatron selects dp=16, which results in 2.72 second/iterationperformance.
Revised: However, since the model can be fit in a single GPU, Megatron selects dp=16, which results in 2.72 second/iteration performance.

Original: This strategy is much slower becauseit requires all GPUs to perform a data-parallelsynchronization with large message size (the whole model parameter), while some GPUs have lowbandwidth (10 Gbps).
Revised: 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

Original: On the contrary, our cost model is aware of such a slow connection speed inthe data-parallel communication group and punishes such strategies heavily. 
Revised: 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 st place.

Original: Heterogeneity in  model.
Revised: Heterogeneity in the model.

Original: When model is no longer homogeneous, layers  have different executiontime. Simply balancing the number of layers or parameters can not balance the  workload.We expect Megatron-LM
Revised: 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.

Original: [Table] our pipeline assignment is better than those in Megatron.
Revised: In such a setup, we expect our pipeline assignment can outperform these two.

Original: In this experiment, We compare fourmethods: • Megatron : Megatronheuristics with  uniform layer assignment. • AMP  : our algorithm with dynamic programming layer assignment algorithm. • ours w/ uniform : strategies by AMP replace with  uniform layer assignment.
Revised: 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.

Original: • ours w/ parameter .

Original: strategies by AMP replace with uniform parameter assignment.

Original: Results are presented in Table 4.
Revised: Results are presented in Table 5.

Original: Our algorithm finds  1 .
Revised: AMP finds a 1 .

Original: 5 × faster strategies compared to Megatronbaseline because it does not enforce strategies to use large data parallelism degree.
Revised: 77 × faster strategy than Megatron ones. We analyze the speedup in two dimensions. First, if we only consider the micro-batch size and the parallelism degrees aspect, the strategy found by AMP is 1 . 51 × faster than the one by Megatron. The reason is similar to the one in the heterogeneous cluster setting.

Original: By enabling ourpipeline layer assignment algorithm, the strategy gains an additional 1 . 17 ×speedup, which gives anend-to-end 1 . 76 × speedup.
Revised: Second, if we additionally consider the pipeline layer assignment aspect, the strategy gains an additional 1 . 17 × speedup.

Original: To study the relationship between heterogeneity in  model and the system performance, we fix theparallelism degrees and vary only the model configuration.
Revised: 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.

Original: Results are presented in 4.
Revised: Results are presented in 3.

Original: Since we onlyuse transformer layers with two different hidden size, a simple definition of heterogeneity is:
Revised: Since we only use transformer layers with two different hidden sizes, we simply define heterogeneity as:

Original: [Equation]
Revised: [Equation]

Revised: Cost model accuracy To correctly select the top strategies, AMP must rely on an accurate cost model.

Original: [Figure]
Revised: 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 to a lower actual f ( · ) .

Original: [Figure]
Revised: More importantly, strategies with low actual f ( · ) are ranked among the top. Concretely, the best strategy is ranked 3 rd , 2 nd , and 1 st 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 launchreal 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).

Original: [Figure]
Revised: [Figure]

Original:  methods.
Revised: two Megatron layer assignment methods.

Original: Numbers obtained when using the best parallelism strategy: mp=1, pp=4, dp=4.
Revised: Numbers obtained when a top parallelism strategy: mp=1, pp=4, dp=4.

Original: Optimization cost.
Revised: Optimization cost.

Original: The dynamic programming algorithm takes the major time. In experiments with heterogeneous models with 24 transformer layers, our methodtakes 140.7 seconds for profiling and 235.6 seconds for65 iterations of the optimization procedure.
Revised: In the heterogeneous model experiment, AMP takes 140.7 seconds for profiling and 235. seconds for 65 iterations of the optimization procedure.

Original: In particular,230.8 seconds are spent in the dynamic programming algorithm.
Revised: In particular, 230.8 seconds are spent in the dynamic programming algorithm.

Original: We also launch several real trials for top predictedcandidates to select the final best one.
Revised: We also launch several real trials for top predicted candidates to select the final best one.

Original: In total, the costof our optimization is around 5 minutes with several realtrials (around 1 − 2 minutes per trial).
Revised: In total, the cost of our optimization is around 5-10 minutes with several real trials (around 1 − 2 minutes per trial).

Original: Comparison with Piper.

Original: Piper planner [23] also explores the 3D parallelism space but with a different problem setup.

Original: They assume that the cluster has a uniformbandwidth.

Original: Our method attains its most speedup when thecluster has non-uniform bandwidth, which is out of thedomain of Piper input.

Original: Additionally, Piper reports resultson simulation, while we report numbers on real trials based on DeepSpeed, which can not supportmost of Piper’s solutions.

Original: Limitations.
Revised: Limitations.

Original: AMP’s main limitation is its unawareness of memory footprint.
Revised: AMP does not model memory footprint.

Revised: Consequently, it can suggest strategies that do not fit in the device’s memory, and launch unnecessary real trials.

Original: In our experiments,a small number of top strategies are out of memory ( i.e., 7 th and 9 th of  top 10 strategies in the TransGAN experiment).
Revised: In our experiments, a small number of top strategies are out of memory ( i.e., 7 th and 9 th of the top 10 strategies in the TransGAN experiment).

Revised: Additionally, the solution for the pipeline assignment problem is not efficient enough.

Original: We defer memory footprint modeling to future work.
Revised: 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 O ( L 4 ) . In the future, we aim to either improve the dynamic programming solution itself, or combine the solution with techniques such as operator clustering to keep L low [2, 29].

Original:  Conclusion
Revised:  Conclusion

Original: We present AMP, an automatic algorithm that can find model parallel training strategies with highsystem throughput.
Revised: We present AMP, an automatic algorithm that can find model parallel training strategies with high system throughput.

Original: AMP is equipped with an expert-design cost model that considers cluster andmodel configurations.
Revised: AMP is equipped with an expert-design cost model that considers cluster and model configurations.

Original: We show that when heterogeneity exists in cluster or model setup, currentheuristics are sub-optimal.
Revised: We show that when heterogeneity exists in cluster or model setup, current heuristics are sub-optimal.

Original: Our automatic algorithm can find strategies with 1 . 37 × and 1 .
Revised: Our automatic algorithm can find strategies with 1 . 51 × and 1 .

Original: 76 ×speedup within similar budgets.
Revised: 76 × speedup within similar budgets.

Original: Checklist
Revised: Acknowledgements

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: We would like to thank Sangkeun Choe and Aurick Qiao for their comments in the early stage of the paper.

Original: contributions and scope? [Yes] (b) Did you describe the limitations of your work? [Yes]
Revised: We thank Tailin Zhang and Tianyi Yu for their help in deriving the pipeline objective and the dynamic programming solution.

Original: (c) Did you discuss any potential negative societal impacts of your work? [No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: We thank all reviewers for their invaluable comments and feedback. This research was supported by NSF IIS1563887, NSF CCF1629559, NSF IIS1617583, NGA

Original: them?
Revised: HM

Original: [Yes] If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: We consider overlapping in the P2P communication as: the sender sends the message to the receiver, while continuing its computation.

Original: [N/A](b) Did you include complete proofs of all theoretical results?
Revised: The receiver needs to wait for the message before continuing its computation.

Original: [N/A]
Revised: This overlapping is illustrated in Figure 4.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: We present the rest details of Algorithm 1. We briefly describe how each sub-procedure is implemented here:

Original: mental results (either in the supplemental material or as a URL)?
Revised: placement() generates the device placement using the heuristic in Megatron.

Original: [Yes] Codes and how to reproduce results are included in the supplementary materials. (b) Did you specify all  the training details (e.g., data splits, hyperparameters,
Revised: enumerate_degrees() takes in the cluster C information, and outputs all possible parallelism degrees, each with format ( pp, dp, tmp ) with constraints that pp × dp × tmp = |D| . pipe_ast() takes in the model W information, the number of stages d.pp .

Original: how they
Revised: It generates the per layer cost, and per edge cost using formulas in section 3.4, and uses the dynamic programming algorithm in section 3.5 to solve for an optimal layer assignment.

Original: were chosen)? [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experi-
Revised: 4. estimate() takes in the optimal layer assignment, and generates the final estimate time using Equation 4. 6.3 Randomized Optimization

Original: ments multiple times)?
Revised: In the layer assignment problem, we leverage the structure that only continuous layers can be assigned to the same stage, which enables a solution with optimality and polynomial runtime.

Original: [No] We do not have randomness in our main experiments. (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: However, other aspects exhibit less structure that allows us to leverage a deterministic algorithm. Moreover, these dimensions compose a large space that prohibits us from simply enumerating all possibilities.

Original: [Yes] We include cluster setup in § 4.3,and total time used in § 4.
Revised: Specifically, there are exponentially many possible device placements [17].

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: At a high level, we would like to optimize a cost function over a discrete domain with a large support. One effective family of algorithms to solve this problem is known as Simulated Annealing (SA) [3], where it explores states (strategies in our semantics) in a neighbor-to-neighbor fashion and gradually improves the quality. A neighbor of a state is produced by conservatively changing the current state. Specifically, we use a SA algorithm that varies only dp or tmp at a time (i.e. strategies with the same dp or tmp are neighbors), guided by the cost model in section 3.4 to iteratively find a good strategy.

Revised: To optimize the device placement aspect given tmp and dp , we consider a problem setup similar to the Domino Tiling problem: we place devices as a 2d mesh and tile it using dominos of size tmp × dp either horizontally or vertically.

Revised: Each domino represents ranks with the same pipeline stage..

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: An example domino tiling scheme is shown in Figure 5. The described optimization procedure is presented at algorithm 2. 6.4 Proof of the dynamic programming solution At a cut with position i’, equation 6 can be rewritten as

Original: [Yes] We cite the underlying
Revised: [Equation] The associate cost stored at dp [ i ′ ] [ j − 1][ max { t 2 ,i ′ , m } ] is: [Equation]

Original: system Deepspeed in § 4.
Revised: Claim: [Equation] Prove claim by enumerate all possibility: <

Original: (b) Did
Revised: < m +

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re [Table] using/curating? [No] We focus on system performance under different setups, which
Revised: 13 = t

Original: [Table]
Revised: [Equation]

Original: does not need to use actual data.

Original: [Table]
Revised: [Equation]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable
Revised: m

Original: [Table]
Revised: [Equation]

Original: information or offensive content?
Revised: t

Original: [No] Since we focus on the system performance,
Revised: =

Original: we
Revised: - m

Original: [Table]
Revised: 13 =

Original: do not release model weights. [Table] If you used crowdsourcing or conducted research with human subjects...
Revised: <

Original: [Table]
Revised: < m <

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: <

Original: [Table]
Revised: , Equation 13 =

Original: applicable? [N/A]
Revised: Using the claim: [Equation]

Original: [Table]
Revised: In this section, we provide detailed experiment setup and results to help understand the performance of each training strategy, and how different methods find different top strategies. In particular, we provide the Top 5 strategies proposed by Megatron and AMP. 6.5.1 Homogeneous setting

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: Model architecture We use a 24 layers GPT-2 model, where layer 3-26 are transformer layers, and the rest are lambda functions or embedding layers. We use hyper-parameters: hidden size 1024, sequence length 1024, vocabulary size 52256, and batch size 32. mbs pipeline layer assignment tmp pp run time [0, 30]  1. [Table]

Original: [Table]
Revised: [Table]

Original: Board (IRB) approvals, if applicable?
Revised: Model architecture We use the same model configuration as in the homogeneous setup.

Original: [N/A] [Table]
Revised: [Table] [Table]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: Model architecture We use a TransGAN Generator with 24 transformer layers, where layer 4-15, 42-53 are transformer layers, and the rest are lambda functions or embedding layers. We use hyperparameters: hidden size 1024, bottom width 9, batch size 64, 12 transformer layers for stage 1, and 12 transformer layers for stage 6.

Original: [Table]
Revised: [Table]

Original: spent on participant compensation?
Revised: [Table] [Table] [Table]

Original: [N/A]
Revised: [Table] [Table] [Table] [Table]

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]