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

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 distributed system expertise to carefully design model-parallel execution strategies that suit 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, a framework that automatically derives such strategies.
Revised: In this paper, we develop AMP, a framework that automatically derives such strategies.

Original: 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.
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 cluster setups 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 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.
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 finds strategies with 1 . 54 × and 1 .
Revised: On heterogeneous clusters or models with heterogeneous architectures, AMP finds strategies with 1 . 54 × and 1 .

Original: 77 × higher throughput than state-of-the-art model-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 [6, 4], multimodal learning [1], and many other ML applications 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 language models [4, 14] 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: 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.
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 [28, 5, 15], model parallelism considers a much larger 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 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]).
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 memory requirements, these model parallelism dimensions have different applicable scenarios, which may change 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 communication volumes, hence is favored on devices connected with high bandwidth, but may become less effective with 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, then tuning the configurations in each parallelism dimension to align with the heterogeneity of the model and 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 or two 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 example, the 3D parallelism strategy adopted in Megatron-LM [23] and
Revised: For example, the 3D parallelism strategy adopted in Megatron-LM [23] and

Original: 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].
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 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 ).
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 the core 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 for more 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 is difficult 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, 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.
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-based methods [11] 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 capture the heterogeneity in the model or cluster (e.g. [24] 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 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.
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 makes the 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 heterogeneous settings.
Revised: We empirically show that existing systems output sub-optimal strategies in heterogeneous settings.

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

Original: 77 × 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 for composing 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 each worker holds an entire model replica [15, 22, 9].
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 gradients over 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 models with 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 [23], is a popular model-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 two consecutive layers are partitioned row-wise ( i.e., input dimension) first, then column-wise ( i.e., output dimension) [23].
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 first layer, but requires heavy cross-device communications afterward.
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 with data parallelism to increase the training throughput [18].
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-batch is 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 pipelined across 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 device idle ( i.e., pipeline bubbles) [10, 18].
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 [10],
Revised: Despite synchronous pipelining schedules such as GPipe [10],

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

Original: TeraPipe develops token-level pipelining schedules in a single training sequence for auto-regressive models [16].
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 automatic computation graph partitioning over a few typical cluster device setups [11, 18, 7], 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 desired strategy over the search space.
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.

Original: 3D parallelism.
Revised: 3D parallelism.

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

Original: 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: DeepSpeed carefully designs a 3D parallelism strategy to train a massive-scale language model with 17 billion parameters [21].
Revised: DeepSpeed carefully designs a 3D parallelism strategy to train a massive-scale language model with 17 billion parameters [21].

Original: [19] manually tunes the degrees of DP, MP, PP to scale to 3072 GPUs – both of 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 [24], which leverages a cost model 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, 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.
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, 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).
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 switches emerge rapidly, thus it is common nowadays that multiple types of devices are deployed together in 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 clusters introduces 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 effects easily; heterogeneous memory capacity among GPUs introduces new constraints on model and data partitioning.
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 of heterogeneity.
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 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.
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 3.1 Problem Formulation
Revised:  Methods 3.1 Problem Formulation

Original: Formally, the inputs of AMP are (i) the model, (ii) the cluster, and (iii) the global batch size ( 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 a set 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 edge v W ,i =
Revised: { i, · · · , | C W |} represents a layer, and each edge v W ,i =

Original: ( c W ,i ; c W ,i +1 )
Revised: ( 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 an edge 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 set of 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 bandwidth between device
Revised: B ∈ R |D|×|D| is a symmetric matrix where each matrix element b ij represents bandwidth between device

Original: d i and d j .
Revised: d i and d j .

Original: The goal of AMP is to output a high throughput parallel training strategy s given 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: We formulate 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 the given
Revised: where f ( · ) denotes per iteration running time as a function of the strategy s , conditioned on the given

Original: ML model, cluster, and user-specified global batch size.
Revised: 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 system performance (§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 difficulty in solving this optimization with 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 real trials, but this is prohibitively expensive because the entire optimization procedure may require evaluating a large number of strategies.
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 discrete nature 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, evaluating all strategies in the search space S 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 optimize the objective is to exploit structure in the strategy space S , and prune areas with no promising candidates.
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 prunes out 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: 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 .
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: 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
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: 1.
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 degrees of each affect the system performance.
Revised: Since PP, DP, and TMP have different characteristics, different degrees of each affect the system performance.

Original: We 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: This tuple generates pp × dp × tmp virtual ranks that carry out corresponding tasks.
Revised: This tuple generates pp × dp × tmp virtual ranks that carry out corresponding tasks.

Original: For instance, ranks in the same data parallelism group communicate with each other to finish the gradient update.
Revised: For instance, ranks in the same data parallelism group communicate with each other to finish the gradient update.

Original: These ranks are bijectively mapped to available devices D for execution.
Revised: These ranks are bijectively mapped to available devices D for execution.

Original: [Figure]
Revised: [Figure]

Original: Assuming
Revised: Assuming

Original: uniform bandwidth,
Revised: uniform bandwidth,

Original: uniform
Revised: uniform

Original: layer assignment cuts
Revised: layer assignment cuts

Original: Device placement.
Revised: Device placement.

Original: 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.
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 high bandwidth, the per iteration time will be shorter.
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: Formally, We denote a device placement as a bijective function from virtual 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 batch size a device proceeds at a time in the pipeline parallelism execution (§2).
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, devices usually have different utilization [18, 26].
Revised: When mbs varies, devices usually have different utilization [18, 26].

Original: For instance, a larger micro-batch size usually leads to better GPU utilization because the GPU can better parallel the computation.
Revised: For instance, a larger micro-batch size usually leads to better GPU utilization because the GPU can better parallel the computation.

Original: However, a larger microbatch size also leads to a larger pipeline bubble, which makes selecting mbs hard.
Revised: 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: Different assignments of layers to each device 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 next device 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 device is assigned a much higher load than other devices, it blocks all other devices and thus incurs a longer iteration 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 maps layers 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 tuple:
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: Since launching real trials to obtain the actual f ( · ) is expensive, we develop a 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 ( · ) .

Original: In this section, we introduce the proposed cost model.
Revised: In this section, we introduce the proposed cost model.

Original: Communication primitive.
Revised: Communication primitive.

Original: We consider the ring topology for all-reduce collective communication in both DP and TMP [22].
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 communication group of n workers with uniform bandwidth B , the communication time of a single ring all-reduce operation can be modeled as [13]:
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 group with non-uniform bandwidth, we approximate the bandwidth with the lowest bandwidth
Revised: In a group with non-uniform bandwidth, we approximate the bandwidth with the lowest bandwidth

Original: B ′ .
Revised: B ′ .

Original: For point-to-point (P2P) communication operations between consecutive pipeline stages,
Revised: For point-to-point (P2P) communication operations between consecutive pipeline stages,

Original: we
Revised: we

Original: [Table]
Revised: [Table]

Original: assume the communication operation can use the full bandwidth:
Revised: assume the communication operation can use the full bandwidth:

Original: e = vb [11].
Revised: e = vb [11].

Original: Here v is the activation volume and b is the bandwidth between these two GPUs.
Revised: 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-down view.
Revised: We now present our cost model designed from a top-down view.

Original: First, DP is at the top 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 their schedules.
Revised: It synchronizes gradients when all pipelines finish their schedules.

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

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

Original: [Equation]
Revised: [Equation]

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 to synchronize gradients with an all-reduce operation.
Revised: T dpsyncd denotes the time for a device d to synchronize gradients with an all-reduce operation.

Original: 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.
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 pipeline executes a training schedule that includes activation/gradients exchange at a layer level.
Revised: At a lower granularity, each pipeline executes a training schedule that includes activation/gradients exchange at a layer level.

Original: It does not need to consider how each layer is executed, which is carried out by MP workers.
Revised: It does not need to consider how each layer is executed, which is carried out by MP workers.

Original: 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 .
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: 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 )=
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: We model the execution time for a single pipeline as 2 :
Revised: We model the execution time for a single pipeline as 2 :

Original: [Equation]
Revised: [Equation]

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

Original: If we assume gas = 1 , Equation 5 reduces to T
Revised: If we assume gas = 1 , Equation 5 reduces to T

Original: pp =
Revised: pp =

Original: (cid:80)
Revised: (cid:80)

Original: pp − 1 i =

Revised: pp − 1 i =

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

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

Original: 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.
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: 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.
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 heuristic is often adopted as the default layer assignment solution in current systems[23].
Revised: This heuristic is often adopted as the default layer assignment solution in current systems[23].

Original: At the bottom granularity, MP workers consider how a single layer is executed.
Revised: At the bottom granularity, MP workers consider how a single layer is executed.

Original: 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.
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: One way to estimate t layerj is using the number of floating point operations (FLOPs) for computation time, and Equation 3 for communication time.
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, we find this inaccurate in our setting.
Revised: However, we find this inaccurate in our setting.

Original: Concretely, when mp grows, the arithmetic intensity of each matrix multiplication becomes lower.
Revised: Concretely, when mp grows, the arithmetic intensity of each matrix multiplication becomes lower.

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

Original: Thus, we simply obtain t layerj by profiling several tensor model parallelism settings and store results in a lookup table.
Revised: Thus, we simply obtain t layerj by profiling several tensor model parallelism settings and store results in a lookup table.

Original: This profiling is feasible because it only needs to be done once in the entire optimization.
Revised: This profiling is feasible because it only needs to be done once in the entire optimization.

Original: [Table]
Revised: [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 to the same stage.
Revised: A key observation in the layer assignment problem is that only consecutive layers can be assigned to the same stage.

Original: In particular, 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: This pattern fits neatly in a dynamic programming solution.
Revised: This pattern fits neatly in a dynamic programming solution.

Original: To handle the max term in objective 5, we use an auxiliary variable tolerance m .
Revised: To handle the max term in objective 5, we use an auxiliary variable tolerance m .

Original: m takes value from all possible execution time for a single stage.
Revised: m takes value from all possible execution time for a single stage.

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

Original: However, it can not be that of the first and the third layer because they are not consecutive.
Revised: However, it can not be that of the first and the third layer because they are not consecutive.

Original: Using m , we rewrite an objective for the dynamic programming algorithm:
Revised: Using m , we rewrite an objective for the dynamic programming algorithm:

Original: [Equation]
Revised: [Equation]

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

Original: 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 ] .
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, then dp [ L ][ k ][0] stores the solution.
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 the sum 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 positions of the last cut 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 value stored 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 ′ , the communication cost between the ( j − 1) th and j th device at layer i ′
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 ′

Original: as e i ′ ,j .
Revised: as e i ′ ,j .

Original: The cost associated with 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 entry dp [ 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: Running 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 is
Revised: Thus, our problem size is

Original: O ( 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 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) .
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: In this section, we introduce the overall optimization procedure to find the best strategies defined at § 3.3.
Revised: In this section, we introduce the overall optimization procedure to find the best strategies defined at § 3.3.

Original: First, we enumerate all possible parallelism degrees and micro-batch sizes.
Revised: First, we enumerate all possible parallelism degrees and micro-batch sizes.

Original: For each parallelism degree and micro-batch size, we deterministically solve for the optimal pipeline layer assignment using the algorithm 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 [19] to solve the device placement problem: prioritize placing model parallelism workers in the same node.
Revised: We adopt the heuristics at [19] to solve the device placement problem: prioritize placing model parallelism workers in the same node.

Original: If there is more available intra-node bandwidth, then place data parallelism workers and pipeline parallelism workers.
Revised: If there is more available intra-node bandwidth, then place data parallelism workers and pipeline parallelism workers.

Original: The procedure is written in Algorithm 1.
Revised: The procedure is written in Algorithm 1.

Original: More details can be found in § 6 .
Revised: More details can be found in § 6 .

Original: in § 6.
Revised: in § 6.

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

Original: 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 ) //
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 ) //

Original: Device placement heuristics a = pipe_ast( W , d.pp , p ) //
Revised: Device placement heuristics a = pipe_ast( W , d.pp , p ) //

Original: Optimal pipeline assignment s = ( a , mbs , d , p ) //
Revised: Optimal pipeline assignment s = ( a , mbs , d , p ) //

Original: current strategy cost = estimate( s ) //
Revised: current strategy cost = estimate( s ) //

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

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

Original: run(record, budget) 11 return best_s
Revised: run(record, budget) 11 return best_s

Original: [Table]
Revised: [Table]

Original: Complexity of the overall optimization.
Revised: Complexity of the overall optimization.

Original: The number of factors of an integer N can be loosely upper bounded by O ( N 1 / 2 ) .
Revised: The number of factors of an integer N can be loosely upper bounded by O ( N 1 / 2 ) .

Original: Thus, the number of possible degrees can be upper bounded by
Revised: Thus, the number of possible degrees can be upper bounded by

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

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

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

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

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

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

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

Original: We find that the actual optimization time is near linear with respect to the cluster size.
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 AMP on different types of clusters and models.
Revised: In this section, we evaluate AMP on different types of clusters and models.

Original: Besides a homogeneous setup, we study two heterogeneous setups, one with a heterogeneous cluster and one with a heterogeneous model.
Revised: Besides a homogeneous setup, we study two heterogeneous setups, one with a heterogeneous cluster and one with a heterogeneous model.

Original: We use the state-of-the-art Megatron heuristics as our baseline [19].
Revised: We use the state-of-the-art Megatron heuristics as our baseline [19].

Original: 4.1 Baseline
Revised: 4.1 Baseline

Original: Megatron [23] is a runtime engine 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 suggests several 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 heuristics in the following procedure.
Revised: We summarize these heuristics in the following procedure.

Original: Mgetraon explores all possible micro-batch sizes as in AMP.
Revised: Mgetraon explores all possible micro-batch sizes as in AMP.

Original: For each micro-batch size, it forces tmp to be smaller than the smallest number of devices in a node.
Revised: For each micro-batch size, it forces tmp to be smaller than the smallest number of devices in a node.

Original: Then it selects parallelism degrees with the smallest tmp × pp .
Revised: Then it selects parallelism degrees with the smallest tmp × pp .

Original: It solves the device placement problem by prioritizing tensor model parallelism workers 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 data parallelism workers next.
Revised: If there are additional available intra-node connections, it places data parallelism workers next.

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

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

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

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

Original: Thus, We hypothesize under such as homogeneous setup, it is (near) optimal.
Revised: Thus, We hypothesize under such as homogeneous setup, it is (near) optimal.

Original: To validate this hypothesis, we conduct experiments on a homogeneous setup ( § 4.3).
Revised: To validate this hypothesis, we conduct experiments on a homogeneous setup ( § 4.3).

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

Original: Concretely, we ask two questions.
Revised: Concretely, we ask two questions.

Original: What is optimal when the cluster is equipped with
Revised: What is optimal when the cluster is equipped with

Original: 8 Estimated Per Iter.
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 . )

Original: spearman correlation = 0.
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 . )

Original: Spearman Correlation = 0.
Revised: Spearman Correlation = 0.

Original: (b) heterogeneous cluster 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 Estimated Per Iter.
Revised: (b) heterogeneous cluster 1.0 1.5 2.0 2.5 3.0 3.5 4.0 4.5 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 . )

Original: Spearman Correlation = 0.
Revised: Spearman Correlation = 0.

Original: [Figure]
Revised: [Figure]

Original: heterogeneous model different 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 heterogeneity in models?
Revised: What is the effect of heterogeneity in models?

Original: To answer these, we design two sets of experiments under a heterogeneous cluster and a heterogeneous in
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: As described in Section 3.6, we adopt the heuristic in Megatron to solve the device placement algorithm in our main 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: During the development time, we also explore an alternative device placement optimization method.
Revised: During the development time, we also explore an alternative device placement optimization method.

Original: Concretely, we adopt a Domino-tiling style device placement representation.
Revised: Concretely, we adopt a Domino-tiling style device placement representation.

Original: Since this representation introduces much more possible strategies, enumerating them is costly.
Revised: Since this representation introduces much more possible strategies, enumerating them is costly.

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

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

Original: Details can be found in § 6.3.
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 )
Revised: We conduct experiments using GPT-2 ( L = 24 , H = 1024 )

Original: [20] on 4 AWS EC g4dn.12xlarge
Revised: [20] on 4 AWS EC g4dn.12xlarge

Original: nodes with a global batch size 32.
Revised: 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: In all experiments, we place instances in the same AWS placement group to make sure they have the full 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 different bandwidths, we ensemble a cluster using 3 AWS EC2 p3.8xlarge instances (V100 GPUs) andg4dn.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 with approximately 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: We evaluate the performance 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 [12], which stacks transformer layers in the way that later layers have 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 uselayers with dimensions 1024, 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 the running time per iteration.
Revised: We evaluate strategies by the running time per iteration.

Original: Specifically, we run each strategy byiterations and take the average iteration time for the latest 20 iterations [27].
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 (built on top of the Megatron engine) [21] 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.

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

Original: Heuristics, (2) AMP : our method described in § 3.6, and (3) SA : optimize the device placement aspect using the Simulated Annealing.
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 setup in § 4.3, we hypothesize that Megatron heuristics 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.

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

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

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

Original: our method.
Revised: our method.

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

Original: We find that Megatron is indeed near-optimal and
Revised: We find that Megatron is indeed near-optimal and

Original: AMP can find such strategies (in fact slightly better).
Revised: AMP can find such strategies (in fact slightly better).

Original: We analyze the selection procedure below.
Revised: We analyze the selection procedure below.

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

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

Original: 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.
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: In fact, AMP finds that dp
Revised: In fact, AMP finds that dp

Original: = 4 ,
Revised: = 4 ,

Original: pp =with the smallest micro-batch size 1 is a better choice.
Revised: pp =with the smallest micro-batch size 1 is a better choice.

Original: We also find that exploring more device placement strategies with SA finds worse results.
Revised: We also find that exploring more device placement strategies with SA finds worse results.

Original: This is because considering more device placement has limited influence in the homogeneous setup.
Revised: This is because considering more device placement has limited influence in the homogeneous setup.

Original: For instance, whether the tensor model parallelism communication happens within a node or cross node shall 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 much more strategies that make the optimization hard.
Revised: However, it introduces much more strategies that make the optimization hard.

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

Original: In the heterogeneous cluster setup, 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 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).
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: We are also interested in how much heterogeneity can influence the performance of strategies.
Revised: We are also interested in how much heterogeneity can influence the performance of strategies.

Original: Concretely, we examine how much
Revised: Concretely, we examine how much

Original: Megatron and AMP improve compared to random strategies in the space.
Revised: 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
Revised: In this experiment, we have four methods to compare: Megatron, AMP, SA, and

Original: • RS : Randomly selecting strategies without heuristics and a cost model.
Revised: • 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 Table 4.
Revised: Results are presented in Table 4.

Original: We find that the mean of random strategies is 1 .
Revised: We find that the mean of random strategies is 1 .

Original: 16 × worse than that of Megatron, and 1 .
Revised: 16 × worse than that of Megatron, and 1 .

Original: 76 × worse than that of ours.
Revised: 76 × worse than that of ours.

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

Original: Specifically, the heuristic 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 microbatch size 1, the best strategy shall be dp = 2 , pp = 8 , which results in a 1.28 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 the model can be fit in a single GPU, Megatron selects dp=16, which results in 2.72 second/iteration performance.
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 because some GPUs have a low bandwidth (10Gbps) to others.
Revised: This strategy is much slower because some GPUs have a low bandwidth (10Gbps) to others.

Original: On the contrary, by modeling each GPU’s bandwidth in
Revised: On the contrary, by modeling each GPU’s bandwidth in

Original: Equation 3, our cost model penalizes this strategy heavily.
Revised: Equation 3, our cost model penalizes this strategy heavily.

Original: We let the cost model to estimate f ( · ) for all strategies and find that AMP ranks this strategy in the 21 st place.
Revised: 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 the model.
Revised: Heterogeneity in the model.

Original: 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.
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: In such a setup, we expect our pipeline assignment can outperform these two.
Revised: In such a setup, we expect our pipeline assignment can outperform these two.

Original: 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.
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: Results are presented in Table 5.
Revised: Results are presented in Table 5.

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

Original: 77 × faster strategy than Megatron ones.
Revised: 77 × faster strategy than Megatron ones.

Original: We analyze the speedup in two dimensions.
Revised: We analyze the speedup in two dimensions.

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

Original: The reason is similar to the one in the heterogeneous cluster setting.
Revised: The reason is similar to the one in the heterogeneous cluster setting.

Original: Second, if we additionally consider the pipeline layer assignment aspect, the strategy gains an additional 1 . 17 × 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 the model and the system performance, we fix the parallelism degrees and vary only the pipeline layer assignment.
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 3.
Revised: Results are presented in 3.

Original: Since we only use transformer layers with two different hidden sizes, we simply define heterogeneity as:
Revised: Since we only use transformer layers with two different hidden sizes, we simply define heterogeneity as:

Original: [Equation]
Revised: [Equation]

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

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

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

Original: The results are presented in Figure 2.
Revised: The results are presented in Figure 2.

Original: We find that generally a lower estimate f ( · ) corresponds to a lower actual f ( · ) .
Revised: We find that generally a lower estimate f ( · ) corresponds to a lower actual f ( · ) .

Original: More importantly, strategies with low actual f ( · ) are ranked among the top.
Revised: More importantly, strategies with low actual f ( · ) are ranked among the top.

Original: Concretely, the best strategy is ranked 3 rd , 2 nd , and 1 st in the three settings.
Revised: Concretely, the best strategy is ranked 3 rd , 2 nd , and 1 st in the three settings.

Original: As a result, the best strategy can be selected because AMP runs several real trials at the end of the optimization procedure (Algorithm 1).
Revised: As a result, the best strategy can be selected because AMP runs several real trials at the end of the optimization procedure (Algorithm 1).

Original: 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).
Revised: 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: two Megatron layer assignment methods.
Revised: two Megatron layer assignment methods.

Original: Numbers obtained when a top 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: In the heterogeneous model experiment, AMP takes 140.7 seconds for profiling and 235.
Revised: In the heterogeneous model experiment, AMP takes 140.7 seconds for profiling and 235.

Original: seconds for 65 iterations of the optimization procedure.
Revised: 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 predicted candidates 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 cost of our optimization is around 5-10 minutes with several real trials (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: Limitations.
Revised: Limitations.

Original: AMP does not model memory footprint.
Revised: AMP does not model memory footprint.

Original: Consequently, it can suggest strategies that do not fit in the device’s memory, and launch unnecessary real trials.
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 the 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).

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

Original: In our experiments, it takes the major time of the whole optimization procedure.
Revised: In our experiments, it takes the major time of the whole optimization procedure.

Original: When the number of layers L goes up, the running time is increasing in the O
Revised: When the number of layers L goes up, the running time is increasing in the O

Original: ( L 4 ) .
Revised: ( L 4 ) .

Original: 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].
Revised: 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 high system 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 and model 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, current heuristics 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 . 51 × 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: Acknowledgements
Revised: Acknowledgements

Original: We would like to thank Sangkeun Choe and Aurick Qiao for their comments in the early stage of the paper.
Revised: We would like to thank Sangkeun Choe and Aurick Qiao for their comments in the early stage of the paper.

Original: We thank Tailin Zhang and Tianyi Yu for their help in deriving the pipeline objective and the dynamic programming solution.
Revised: We thank Tailin Zhang and Tianyi Yu for their help in deriving the pipeline objective and the dynamic programming solution.

Original: We thank all reviewers for their invaluable comments and feedback.
Revised: We thank all reviewers for their invaluable comments and feedback.

Original: This research was supported by NSF IIS1563887, NSF CCF1629559, NSF IIS1617583, NGA
Revised: This research was supported by NSF IIS1563887, NSF CCF1629559, NSF IIS1617583, NGA

Original: HM
Revised: HM

Original: Supplementary Materials 6.1 Computation and communication overlap

Original: We consider overlapping in the P2P communication as: the sender sends the message to the receiver, while continuing its computation.

Original: The receiver needs to wait for the message before continuing its computation.

Original: This overlapping is illustrated in Figure 4.

Original: 6.2 Optimization procedure details

Original: We present the rest details of Algorithm 1.

Original: We briefly describe how each sub-procedure is implemented here:

Original: placement() generates the device placement using the heuristic in Megatron.

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

Original: pipe_ast() takes in the model W information, the number of stages d.pp .

Original: 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: 4. estimate() takes in the optimal layer assignment, and generates the final estimate time using Equation 4.

Original: 6.3 Randomized Optimization

Original: 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: However, other aspects exhibit less structure that allows us to leverage a deterministic algorithm.

Original: Moreover, these dimensions compose a large space that prohibits us from simply enumerating all possibilities.

Original: Specifically, there are exponentially many possible device placements [17].

Original: At a high level, we would like to optimize a cost function over a discrete domain with a large support.

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

Original: A neighbor of a state is produced by conservatively changing the current state.

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

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

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

Original: An example domino tiling scheme is shown in Figure 5.

Original: The described optimization procedure is presented at algorithm 2.

Original: 6.4 Proof of the dynamic programming solution

Original: At a cut with position i’, equation 6 can be rewritten as

Original: [Equation]

Original: The associate cost stored at dp [ i ′ ]

Original: [ j − 1][ max { t 2 ,i ′ , m } ] is:

Original: [Equation]

Original: Claim:

Original: [Equation]

Original: Prove claim by enumerate all possibility:

Original: <

Original: < m

Original: +

Original: 13 = t

Original: [Equation]

Original: [Equation]

Original: m

Original: [Equation]

Original: t

Original: =

Original: - m

Original: 13 =

Original: <

Original: < m <

Original: <

Original: , Equation 13 =

Original: Using the claim:

Original: [Equation]

Original: 6.5 Experiment Details

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

Original: In particular, we provide the Top 5 strategies proposed by Megatron and AMP.

Original: 6.5.1 Homogeneous setting

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

Original: We use hyper-parameters: hidden size 1024, sequence length 1024, vocabulary size 52256, and batch size 32.

Original: mbs pipeline layer assignment tmp pp run time [0, 30]  1.

Original: [Table]

Original: [Table]

Original: 6.5.2 Heterogeneous Cluster

Original: Model architecture We use the same model configuration as in the homogeneous setup.

Original: [Table]

Original: [Table]

Original: 6.5.3 Heterogeneous Model

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

Original: 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]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]