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

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AffiliationAddress email
Revised: 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 distributedsystem expertise to carefully design model-parallel execution strategies that suitwith 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, aframework to automatically derive 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 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.

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

Original: On heterogeneous clusters or models with heterogeneous architectures,
Revised: On heterogeneous clusters or models with heterogeneous architectures,

Original: AMP demonstrates up to 1 .
Revised: AMP demonstrates up to 1 .

Original: 37 × and 1 .
Revised: 37 × and 1 .

Original: 76 × higher throughput than state-of-the-artmodel-parallel systems, respectively.
Revised: 76 × higher throughput than state-of-the-artmodel-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 [4, 2], multimodal learning [1], and many other MLapplications 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 languagemodels [2, 12] 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: 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.

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 [23, 3, 13], model parallelism considers a muchlarger 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 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) [21]).

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 memoryrequirements, these model parallelism dimensions have different applicable scenarios, which maychange 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 communicationvolumes, hence is favored on devices connected with high bandwidth, but may become less effectivewith 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, thentuning the configurations in each parallelism dimension to align with the heterogeneity of the modeland 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 ortwo 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 examples, the 3D parallelism strategy adopted in Megatron-LM [21] 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 [19] is manually specialized for transformer-based language models, assuming a fixedcluster setup and the model having the same layer repeated [21, 19, 17].

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

Original: This motivates thecore questions in this paper:
Revised: This motivates thecore 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 formore 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 isdifficult 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, 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.

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-basedmethods [9] 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 capturethe heterogeneity in the model or cluster (e.g. [22] 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 largenumber of strategies within a limited budget, we develop a new cost model which serves as a cheaperproxy 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 makesthe following contributions:

Original: • We identify factors that influence model-parallel performance in heterogeneous settings, and
Revised: • We identify factors that influence model-parallel performance in heterogeneous settings, and

Original: define strategies based on these factors.
Revised: define strategies based on these factors.

Original: • We present a new cost model that handles heterogeneity in both the model and the cluster to
Revised: • We present a new cost model that handles heterogeneity in both the model and the cluster to

Original: evaluate the quality of model parallel strategies with a minimal amount of real trials.
Revised: 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
Revised: • We develop a dynamic programming approach to handle the load imbalance issue from

Original: heterogeneous models in the pipeline layer assignment problem.
Revised: heterogeneous models in the pipeline layer assignment problem.

Original: • We empirically show that existing systems
Revised: • We empirically show that existing systems

Original: output sub-optimal strategies in heterogeneoussettings.
Revised: output sub-optimal strategies in heterogeneoussettings.

Original: In contrast, AMP achieves 1 .
Revised: In contrast, AMP achieves 1 .

Original: 37 × speedup with heterogeneous clusters, and 1 .
Revised: 37 × speedup with heterogeneous clusters, and 1 .

Original: 76 ×speedup on heterogeneous model architectures compared to state-of-the-art heuristics.
Revised: 76 ×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 forcomposing 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 eachworker holds an entire model replica [13, 20, 7].

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 gradientsover 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 modelswith 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 [21], is a popularmodel-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 twoconsecutive layers are partitioned row-wise ( i.e., input dimension) first, then column-wise ( i.e., outputdimension) [21].

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 firstlayer, but requires heavy cross-device communications afterwards.

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

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, different layers are placed across GPUs, the trainingmini-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 thenpipelined 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 suffersfrom device idle ( i.e., pipeline bubbles) [8, 16].

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

Original:  PipeDream [17] proposes an 1F1B asynchronous schedule to reduce the pipeline bubble.
Revised: GPipe [8], PipeDream [16] 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-regressivemodels [14].

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 automaticcomputation graph partitioning over a few typical cluster device setups [9, 16, 5], 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 desiredstrategy over the search space.3D parallelism.

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

Original: DP, MP, and PP have to be used together (which is also known as 3D parallelism ).
Revised: DP, MP, and 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: DeepSpeedcarefully designs a 3D parallelism strategy to train a massive scale language model with 17 billionparameters [19].

Original: [18] manually tunes the degrees of DP, MP, PP to scale to 3072 GPUs – bothof them are not automatic .
Revised: [17] manually tunes the degrees of DP, MP, PP to scale to 3072 GPUs – bothof 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 [22], which leverages a simple costmodel 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, Piperworks only with transformer-based models, and assumes that same type of devices in the cluster arehomogeneously 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, 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).

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 switchesemerge rapidly, thus it is common nowadays that multiple types of devices are deployed togetherin a training cluster.

Original: Conducting distributed training over such types of heterogeneous clustersintroduces extra challenges.
Revised: Conducting distributed training over such types of heterogeneous clustersintroduces 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 effectseasily; heterogeneous memory capacity among GPUs introduces new constraints on model and datapartitioning.

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

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 sizesand computing characteristics, e.g, e.g., typical convolution networks have their layer width growingwith depth [6], which would require distinct partitioning and placement strategies for each layerduring model parallelism.

Original:  Methods
Revised:  Methods

Original: 3.1 Problem Formulation
Revised: 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 input of AMP contains: (i) a given model, (ii) a given cluster, and (iii) the globalbatch size ( gbs ).

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

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

Original: { i, · · · , | CW |} represents a layer, and each edgev W ,i = ( c W ,i ; c W ,i +1 )
Revised: { i, · · · , | C |} 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 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 116 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 117 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 betweendevice 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 thehardware 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 formulatethis 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 thegiven 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 systemperformance (§3.3).

Original: However, there are two major problems in solving this optimization problemwith these identified factors.
Revised: However, there are two major problems in solving this optimization problemwith 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 given that we possibly need to evaluate 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 ofthe optimization problem disallows us to leverage gradient-based optimization methods.

Original: However, enumerating all strategies is slow even with the cost model.
Revised: However,enumerating all strategies 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 theobjective 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). TMP DP PP

Original: 3.3 Representation of a Strategy
Revised: GPU:2 GPU: GPU:0 GPU: GPU: GPU:

Revised: GPU:

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: GPU: [Figure] Representation of a Strategy Intuitively, a strategy s defines how to train model { C W , V W } on a given cluster {D , B } with thebatch 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: A strategy s that composes DP, TMP, and PP strategies can be represented by afew 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

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:

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 of them affect the system performance

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

Original: The tuple( pp, dp, tmp ) defines a three-dimensional grid, which satisfies the constraint pp × dp × tmp =|D| ; pp, dp, tmp ∈ { 1 , . .
Revised: 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).
Revised: ( i.e., this tuple generates pp × dp × tmp virtual ranks that will be bijectively mapped to available devices).

Original: Two ranks are 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). [Figure]
Revised: Two ranks are in the same DP(TMP/PP) communication groups if they have different values in the DP(TMP/PP) dimension and the same values in others ( e.g.,

Original: Assuming a uniform bandwidth, a uniform-layer assignment cuts balanced
Revised: rank (0 , 0 , 0) and rank (0 , 0 , 1) forms a TMP group)

Original: Device placement.
Revised: Device placement.

Original: Since physical devices may have different connection speeds, i.e., bandwidth to each other, placing physical devices in different virtual ranks leads to dif- ferent performance.
Revised: Since physical devices may have different connection speeds, i.e., bandwidthto each other, 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, We denote a device placement as a bijective function fromvirtual ranks to devices p : { 1 , . .
Revised: Forinstance, if ranks in a communication-heavy group are all allocated with physical devices with highbandwidth, the per iteration time will be shorter Formally, We denote a device placement as a bijectivefunction 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 holds at a time.

Original: When mbs varies, devicesgenerally have different utilization.
Revised: When mbsvaries, devices generally have different utilization.

Original: This effect is empiricallyverified in [18], where the observation is that different microbatch sizes lead to noticeably different in system throughput.4.
Revised: This effect is empirically verified in [17], wherethe observation is that different micro-batch sizes lead to noticeably different in system throughput.4.

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

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

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 deviceneeds 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 isassigned a much higher load than other devices, it blocks all other devices and thus incurs a longeriteration 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 mapslayers 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 space:

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

Original: Communication primitive.
Revised: Communication primitive.

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

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 time cost of a single ring all-reduce operation can bemodeled as[12]:
Revised: In a cluster of n workers with uniformbandwidth B , the communication time cost of a single ring all-reduce operation can be modeled as

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

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: For point-to-point communicationoperations between consecutive pipeline stages, we assume the communication operation can alwaysenjoy the full bandwidth: e = vb where v is the activation volume and b is the bandwidth betweenthese 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-downview.

Original: First, DP stands at the highest granularity, where there are dp pipelines running in parallel.
Revised: First, DP stands at the highest 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.

Original: At the DP level, the system runtime can bedecomposed into two terms:
Revised: At the DP level, the system runtime can bedecomposed 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 tocommunicate data parallelism gradient updates.

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: 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.

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 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).

Original: Assume that a single pipeline consists of k stages.
Revised: Assume that a single pipeline consists of k stages.

Original: The communication time between the i th and( i + 1) th is denoted as e i .
Revised: The communication time between the i th and( i + 1) th is denoted as e 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: 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 )=

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

Original: Assuming that the P2P communication and computation can be overlapped1 , the runtime for a singlepipeline can be modeled as:[Equation]
Revised: Assuming that the P2P communication and computation can be overlapped , the runtime for a singlepipeline can be modeled as:

Original: [Equation]
Revised: [Equation]

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

Original: When gas
Revised: When 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.
Revised: = 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: In this case, T pp given at equation 5 reduces to (cid:80) pp − 1 i =

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

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

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: 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.

Original: This heuristics is oftenadopted as the default layer assignment in current systems[22].
Revised: This heuristics is oftenadopted as the default layer assignment in current systems[21].

Original: At the lowest granularity, MP carries out results of individual layers: t layerj =
Revised: At the lowest granularity, MP carries out results of individual layers: t layerj =

Original: c j tmp + T mpsyncj .
Revised: c j tmp + T mpsyncj .

Original: T mpsyncj is the all-reduce communication time, which can also be calculated by equation 3.
Revised: T mpsyncj is the all-reduce communication timewhich can also be calculated by equation 3.

Original: Oneway to estimate the execution cost for a single layer c
Revised: Oneway to estimate the execution cost for a single layer c

Original: j is by counting the number of floating-pointoperations (FLOPs).
Revised: j is by counting the number of floating-pointoperations (FLOPs)

Original: However, when matrix size is small,its arithmetic intensity is low, and GPUs canhardly parallel efficiently.
Revised: However, when matrix size is small, arithmetic intensity is low, and GPUs canhardly parallel efficiently.

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

Original: To handle general deep learning models, we simply obtain t layerj by profiling withseveral tensor model parallelism settings.
Revised: To handle general deep learning models, we simply obtain c j by profiling withseveral tensor model parallelism settings.

Original: This profiling only needs to  run once, which is negligiblecompared to model training time.
Revised: This profiling only needs to be run once, which is negligiblecompared to model training time.

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 tothe same stage.

Original: Intuitively, computation up to layer i can be reused for that of layer i + 1 .
Revised: Intuitively, 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: 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:

Original: [Equation]
Revised: [Equation]

Revised: m }

Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: Note that when m = 0 , objective 6 reduces to the desired objective 5.
Revised: Note that when m = 0 , objective 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 record theoptimal assignments and the corresponding cost up to layer i with j stages and variable m in tabledp[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 and k be the number of stages, then dp[L][k][0] stores thesolution.

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 thesum 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 position for lastcut 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 valuestored 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 ′ , thecommunication 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 associatedwith 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 entrydp [ i ][ j ][ m ] is updated as: ( a i ′ opt ∪ ( i − i ′ opt ) , g

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

Original: Run time.
Revised: Run time.

Original: There are [-(cid:80) Li =
 i-] = O
Revised: There are [-(cid:80) Li =
 i-]1 + 2 + · · · + L = 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 subproblem 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 run time is O ( kL 4 ) , while a bruteforce 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: To find best strategies defined at § 3.3, we enumerate all possible parallelism degrees and microbatch sizes.

Original: We deterministically solve for the optimal pipeline layer assignment using the dynamicprogramming approach developed at § 3.5.
Revised: We deterministically solve for the optimal pipeline layer assignment using the dynamicprogramming approach developed 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 [17] to solve the deviceplacement problem: prioritize placing model parallelism workers in the same node, over dataparallelism workers, and pipeline parallelism workers.

Original: We also explore an alternative optimizationfor device placement at § 4.2.
Revised: We also explore an alternative optimizationfor device placement at § 4.2.

Original: The concrete algorithms and complexity analysis can be found inappendix.

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 our algorithm and several of its variants on different types of clusters andmodels.

Original: Specifically, we compare against the state-of-the-art Megatron heuristics [18].
Revised: Specifically, we compare against the state-of-the-art Megatron heuristics [17].

Original: 4.1 Baseline
Revised: 4.1 Baseline

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

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

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

Original: Megatron heuristics loops over all possible micro-batch sizes.
Revised: Megatron heuristics loops over all possible micro-batch sizes.

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 selectsparallelism degrees with the smallest tmp × pp , whose tmp is smaller than the number of devices inthe smallest node.

Original: It solves the device placement problem by prioritizing tensor model parallelismworkers in the same node. If there are additional available intra-node connections, it places data2 proof can be found at the supplementary materials.1.5 2.0 2.5 3.0 3.5 4.0 4.5 Estimated Per Iter. Time (sec.)
Revised: It solves the device placement problem by prioritizing tensor model parallelism1 proof can be found at the supplementary materials.

Original: P e r I t e r .

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

Original: (a) homogeneous 1. 3.0 Estimated Per Iter. Time (sec.) P e r I t e r .
Revised: workers in the same node.

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

Original: (b) heterogeneous cluster 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 5.0 Estimated

Original: [Figure]

Original: Time (sec.) 1. P e r I t e r . T i m e ( s e c . ) (c) heterogeneous model parallelism workers next.
Revised: If there are additional available intra-node connections, it places dataparallelism workers next.

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

Original: Megatron heuristics assume a homogeneous cluster and model setup.
Revised: Megatron heuristics assume a homogeneous cluster and model setup.

Original: The model is composedof simple stacks of identical layers.
Revised: The model is composedof simple stacks of identical layers.

Original: The cluster is equipped with the same type of devices withsimilar bandwidths to each other.
Revised: 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.
Revised: We hypothesize that under this setup, it is (near) optimal.

Original: Tovalidate this hypothesis, we conduct experiments on a homogeneous setup (section 4.3).
Revised: Tovalidate this hypothesis, we conduct experiments on a homogeneous setup (section 4.3).

Original: Moreover,we are interested in how Megatron heuristics perform when the homogeneous assumption breaks.
Revised: Moreover,we are interested in how Megatron heuristics perform when the homogeneous assumption breaks.

Original: Specifically, we want to answer two questions.
Revised: Specifically, we want to answer two questions.

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

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

Original: Thus, we design two sets of experiments under heterogeneous cluster and heterogeneousmodel in section 4.3.
Revised: Thus, we design two sets of experiments under heterogeneous cluster and heterogeneousmodel in section 4.3.

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: Since the number of possible device placements is large and often difficult to optimize efficiently[15], we adopt Megatron heuristics in our method.

Original: To study the optimality of this heuristic, we use a
Revised: To study the optimality of this heuristic, we use a

Original: Simulated Annealing(SA) algorithm to explore candidates with different placements.
Revised: Simulated Annealing(SA) algorithm to explore candidates with different placements.

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

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

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 ) [18] on 4 AWS EC2g4dn.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.

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: We ensemble a cluster using 3 AWS EC2 p3.8xlarge instances and 1g4dn.12xlarge 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 approximately170Gbps NVLink intra-node connection and 10 Gbps inter-node bandwidth.

Original: evaluateperformance on GPT-2{+We evaluate performanceon GPT-
+} ( L = 24 , H = 1024 ) with a global batch size 32 .
Revised: {+We evaluate performanceon GPT-
+} ( 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 [10], which stackstransformer layers with gradually smaller hidden sizes.

Original: Specifically, we use 12 layers with dimensions1024, and 12 layers with dimensions 16.
Revised: Specifically, we use 12 layers with dimensions1024, 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 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 by 40 iterations andtake the average iteration time for the latest 20 iterations.

Original: The underlying system is Deepspeed (builton top of Megatron system) [20] with fp16 optimization enabled.
Revised: The underlying system is Deepspeed (builton top of Megatron system) [19] with fp16 optimization enabled.

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

Original: Homogeneous setup.
Revised: Homogeneous setup.

Original: Under homogeneous setup as specified in 4.3, we hypothesize that Megatronheuristics is (near) optimal.
Revised: Under homogeneous setup as specified in 4.3, we hypothesize that Megatronheuristics 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.

Original: [Table]

Original: Additionally, we conjecture that our device placement algorithm can not propose better strategiesthan the default one from Megatron.
Revised: Additionally, we conjecture that our device placement algorithm can not propose better strategiesthan the default one from Megatron.

Original: In this experiment, we compare three methods:
Revised: In this experiment, we compare three methods:

Original: • Megatron : Megatron Heuristics.
Revised: • Megatron : Megatron Heuristics.

Original: • AMP : our cost model that enumerates parallelism degrees and micro-batch size, and uses dynamic programming to solve the pipeline layer assignment.
Revised: • AMP : our cost model that enumerates parallelism degrees and micro-batch size, and usesdynamic programming to solve the pipeline layer assignment.

Original: • SA : additionally optimize device placement using the simulated Annealing algorithm.
Revised: • SA : additionally optimize device placement using the simulated Annealing algorithm.

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

Original: To understand thedifference in strategies between our proposal and the Megatron proposal, we evaluate the top 10strategies from our method.
Revised: To understand thedifference in strategies between our proposal and the Megatron proposal, we evaluate the top 10strategies from our method.

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

Original: AMP finds slightly better strategiesthan Megatron Heuristics.
Revised: AMP finds slightly better strategies2 Codes are available in the supplementary materials. [Table] than Megatron Heuristics.

Original: Since the model can be fit into a single GPU, Megatron heuristics selectsan extreme configuration - dp = 16 .
Revised: Since the model can be fit into a single GPU, Megatron heuristics selectsan extreme configuration - dp = 16 .

Original: Our method automates the strategies selection process andfinds that dp = 4 is a better choice.
Revised: Our method automates the strategies selection process andfinds that dp = 4 is a better choice.

Original: By enabling our SA algorithm to explore device placement,
Revised: By enabling our SA algorithm to explore device placement,

Original: SA finds slightly worse results.
Revised: SA finds slightly 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: The reason is that in such a cluster setup where devices have thesame connection speeds, device placements have limited influence on the system performance.

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

Original: However, it introduces more strategies that make optimizationhard.
Revised: However, it introduces more strategies that make optimizationhard.

Original: This effect is illustrated in Figure 3.
Revised: This effect is illustrated in Figure 2.

Original: [Figure]
Revised: [Figure]

Original: top

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

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

Original: In the setup at section 4.3,devices have different connectivity to each other, and different execution speeds.
Revised: In the setup at section 4.3, devices have different connectivity toeach 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 canlead to bad performance, but our method should be able to automatically address the heterogeneityin (especially) bandwidth since we explicitly model it (3.4).

Original: Withheterogeneity in the cluster, we are also interested in examining how much Megatron heuristics and our method improvecompared to average strategies in the space.
Revised: With heterogeneity in the cluster, weare also interested in examining how much Megatron heuristics and our method improve comparedto average 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 2.

Original: We found that average strategies have a high cost, butrandomly evaluating them can find better strategies than Megatron.
Revised: We found that average strategies havea high cost, but randomly evaluating them can find better strategies than Megatron.

Original: However, this comes with a muchhigher cost (50 vs 10 real trials), and cannot reliably produce good results.
Revised: However, thiscomes with a much higher cost (50 vs 10 real trials), and cannot reliably produce good results.

Original: By using our cost model,we can find strategies up to 1 .
Revised: Byusing our cost model, we can find strategies up to 1 .

Original: 37 × speedup compared to Megatron.
Revised: 37 × speedup compared to Megatron.

Original: Specifically, the heuristics of Megatron that minimize mp × pp is sub-optimal.
Revised: Specifically,the heuristics 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 micro-batch size1, the best strategy shall be dp=4 and pp=4, which results in a 1.44 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.72second/iteration performance.

Original: This strategy is much slower because it 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 it requires all GPUs to performa data-parallel synchronization with large message size (the whole model parameter), while some GPUs have low bandwidth (10 Gbps).

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: On the contrary, our cost model is aware of such a slowconnection speed in the data-parallel communication group and punishes such strategies heavily.

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

Original: When model is no longer homogeneous, layers have different executiontime.
Revised: When model is no longer homogeneous, layers have different executiontime.

Original: Simply balancing the number of layers or parameters can not balance the workload.
Revised: Simply balancing the number of layers or parameters can not balance the workload.

Original: We expect Megatron-LM [Table] our pipeline assignment is better than those in Megatron.
Revised: We expectour pipeline assignment is better than those in Megatron.

Original: In this experiment, We compare fourmethods:
Revised: In this experiment, We compare fourmethods:

Original: • Megatron : Megatron heuristics with uniform layer assignment.
Revised: • Megatron : Megatron heuristics with uniform layer assignment.

Revised: [Table]

Original: • AMP : our algorithm with dynamic programming layer assignment algorithm. • ours w/ uniform : strategies by AMP replace with uniform layer assignment.
Revised: • AMP : our algorithm with dynamic programming layer assignment algorithm.• ours w/ uniform : strategies by AMP replace with uniform layer assignment.

Original: • ours w/ parameter .
Revised: • ours w/ parameter .

Original: strategies by AMP replace with uniform parameter assignment.
Revised: strategies by AMP replace with uniform parameter assignment

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

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

Original: 5 × faster strategies compared to Megatronbaseline because it does not enforce strategies to use large data parallelism degree.
Revised: 5 × faster strategies compared to Megatronbaseline because it does not enforce strategies to use large data parallelism degree.

Original: By enabling ourpipeline layer assignment algorithm, the strategy gains an additional 1 . 17 × speedup, which gives anend-to-end 1 .
Revised: By enabling ourpipeline layer assignment algorithm, the strategy gains an additional 1 . 17 × speedup, which gives anend-to-end 1 .

Original: 76 × speedup.
Revised: 76 × 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 model and the system performance, we fix theparallelism degrees and vary only the model configuration.

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 onlyuse transformer layers with two different hidden size, a simple definition of heterogeneity is:heterogeneity

Original: [Equation]
Revised: [Equation]

Original:  [Figure] 
Revised: 16 layers time [Equation] [Figure] [Equation]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: methods.
Revised: methods.

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

Original: Optimization cost.
Revised: Optimization cost.

Original: The dynamic programming algorithm takes the major time.
Revised: The dynamic programming algorithm takes the major time in our searching algorithm.

Original: 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: Inexperiments with heterogeneous models with 24 transformer layers, our method takes 140.7 seconds for profiling and 235.6 seconds for 65 iterations of the optimizationprocedure.

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

Original: We also launch several real trials for top predictedcandidates to select the final best one. In total, the costof our optimization is around 5 minutes with several realtrials (around 1 − 2 minutes per trial).
Revised: The additional cost(around 5 minutes in total) introduced by our method isnegligible compared to usually days to weeks of trainingtime in real-world DNNs applications.

Original: Comparison with Piper.
Revised: Comparison with Piper.

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

Original: They assume that the cluster has a uniformbandwidth.
Revised: 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.
Revised: Our method attains its most speedup whenthe cluster has non-uniform bandwidth, which is out ofthe domain 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.
Revised: Additionally, Piper reportsresults on simulation, while we report numbers on realtrials based on DeepSpeed, which can not support most of Piper’s solutions.

Original: Limitations.
Revised: Limitations.

Original: AMP’s main limitationis its unawareness of  memory footprint.
Revised: The main limitation of our method is that we do not model memory footprint.

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: Duringthe experiments, we discover that a small number of top strategies are out of memory (i.e. 7 th and9 th of top 10 strategies in the TransGAN experiment).

Original: We defer memory footprint modeling to  future work.
Revised: We plan to incorporate a memory model inour future work.

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 highsystem 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 andmodel 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, currentheuristics are sub-optimal.

Original: Our automatic algorithm can find strategies with 1 .
Revised: Our automatic algorithm can find strategies with 1 .

Original: 37 × and 1 .
Revised: 37 × and 1 .

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

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] (c) Did you discuss any potential negative societal impacts of your work?
Revised: [No](c) Did you discuss any potential negative societal impacts of your work?

Original: [No]
Revised: [No]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [N/A](b)
Revised: [N/A]

Original:  Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [N/A]
Revised: [N/A]3.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] Codes and how to reproduce results are included in the supplementary materials.
Revised: [Yes] Codes and howto reproduce results are included in the supplementary materials.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes]
Revised: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-
Revised: (c) Did you report error bars (e g., with respect to the random seed after running experi-

Original: ments multiple times)?
Revised: ments multiple times)?

Original: [No] We do not have randomness in our main experiments.
Revised: [No] We do not have randomness in our main experiments.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type
Revised: (d) Did you include the total amount of compute and the type of resources used (e

Original: of GPUs, internal cluster, or cloud provider)?
Revised: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] We include cluster setup in § 4.3,and total time used in § 4.
Revised: [Yes] We include cluster setup in § 4.3,and total time used in § 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
Revised: If you are using existing assets (e g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We cite the underlying
Revised: [Yes] We cite the underlying

Original: system Deepspeed in § 4.
Revised: system Deepspeed in § 4.4(b) Did you mention the license of the assets?

Original: (b) Did 
Revised: [No] (c) Did you include any new assets either in the supplemental material or as a URL? [No]No new assets are introduced.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re[Table]
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [No] We focus on system performance under different setups, which[Table]
Revised: [No] We focus on system performance under different setups, which

Original: does not need to use actual data.
Revised: does not need to use actual data.

Original: [Table]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable [Table] information or offensive content?
Revised: (e) Did you discuss whether the data you are using /curating contains personally identifiableinformation or offensive content?

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

Original: do not release model weights.

Original: [Table]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: [Table] (a) Did you include the full text of instructionsgiven to participants and screenshots, if [Table] applicable?
Revised: instructions

Original: [N/A]

Original: [Table]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: [Table]

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: [Table]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: [Table]
Revised: [Table]

Original: spent on participant compensation?

Original: [N/A]

Original: [Table]