# Training DNN Models over Heterogeneous Clusters with Optimal Performance

Chengyi Nie  Stony Brook UniversityStony BrookNYUSA ,  Jessica Maghakian
Stony Brook UniversityStony BrookNYUSA and  Zhenhua Liu  Stony Brook
UniversityStony BrookNYUSA

(2023)

###### Abstract.

Adjusting batch sizes and adaptively tuning other hyperparameters can
significantly speed up deep neural network (DNN) training. Despite the
ubiquity of heterogeneous clusters, existing adaptive DNN training techniques
solely consider homogeneous environments. Optimizing distributed DNN training
over heterogeneous clusters is technically challenging, and directly adapting
existing techniques results in low utilization and poor performance. To solve
this problem, we introduce Cannikin – a novel data-parallel distributed
training system. Cannikin achieves efficient and near optimal performance by
accurately modeling the optimal system performance and predicting adaptive
batch size training metrics for DNNs in heterogeneous clusters. We implemented
Cannikin in PyTorch and conducted experiments over 16 GPUs in Chameleon.
Empirical results show that Cannikin reduces DNN training in heterogeneous
clusters by up to $52\%$ compared to the state-of-art adaptive training system
and up to $85\%$ compared to native PyTorch DistributedDataParallel.

machine learning systems, heterogeneous system, load balancing, distributed
training

††copyright: acmcopyright††journalyear: 2023††doi:
XXXXXXX.XXXXXXX††conference: XXXXXX; XXXXX; XXXXX††ccs: Computing
methodologies Distributed computing††ccs: Computing methodologies Machine
learning††ccs: Computing methodologies Heterogeneous computing

##  1\. Introduction

With the explosive increase of deep learning (DL) applications in fields such
as image classification (Deng et al., 2009; Krizhevsky, 2009), natural
language processing (Rajpurkar et al., 2016; Wu et al., 2016), and recommender
systems (Harper and Konstan, 2015; Zhang et al., 2019), the demand for deep
neural network (DNN) training resources is doubling every six months (Sevilla
et al., 2022). In order to achieve efficient DNN training, practitioners rely
on accelerators (Jouppi et al., 2017; Zhang et al., 2018; Nurvitadhi et al.,
2018), hyper-parameter tuning (McCandlish et al., 2018), and distributed
training (Li et al., 2020; Abadi et al., 2016). Meanwhile, the short hardware
update cycle results in newly released accelerators that significantly
outperform previous models within a short time (Schaller, 1997). Table 1 shows
the evolution of NVIDIA data center GPUs released in recent years. Each new
flagship model is over two times faster than the preceding flagship data
center GPU. When companies and research institutes upgrade their machine
learning systems, newly released GPUs are installed before older models
retire, so homogeneous environments cannot always be guaranteed when running
distributed training jobs. To enhance the utilization of computing resources
and speed up DNN model training in heterogeneous environments, specialized
methods are required for DNN training in heterogeneous environments.

Table 1. Evolution of NVIDIA data center GPUs Model | Year | Archit. | CUDA | Memory | FP16
---|---|---|---|---|---
Cores | (GB) | (TFLOPS)
Tesla P100 | 2016 | Pascal | 3584 | 16 | 21.2
Tesla V100 | 2017 | Volta | 5120 | 16/32 | 31.4
A100 | 2020 | Ampere | 6912 | 40/80 | 77.97
H100 | 2022 | Hopper | 16896 | 80 | 204.9

Previous work on specialized distributed training for heterogeneous
environments focuses on two major schemes: data parallelism (Li et al., 2020;
Sergeev and Balso, 2018) and model parallelism (Shoeybi et al., 2019; Park et
al., 2020). For data-parallelism heterogeneous distributed training systems,
HetSeq (Ding et al., 2021) manually tunes the local mini batch size for each
node to balance workloads, while LB-BSP (Chen et al., 2020) and DLB (Ye et
al., 2022) improve performance by iteratively tuning the workloads assigned to
each worker based on the computing time of each node. On the other hand,
BlueConnect (Cho et al., 2019) boosts data-parallelism distributed training by
optimizing communication. Existing data-parallelism systems do not jointly
consider the computing and communication models for heterogeneous clusters,
resulting in suboptimal performance. Model-parallelism systems (Park et al.,
2020; Harlap et al., 2018) pipeline the DNN model in heterogeneous clusters,
which is near optimal for resource utilization with fine-tuning. But model
parallelism requires a specific configuration of each node in a cluster, thus
limiting scalability. Furthermore, existing model-parallelism and data-
parallelism methods cannot manage the sudden changes of resources that occur
in clusters with dynamic resource allocation (Peng et al., 2018; Qiao et al.,
2021).

Another efficient DNN training method, adaptive batch size training, tunes
hyper-parameters such as batch size and learning rate during training,
significantly speeding up convergence. Previous work (Qiao et al., 2021;
Devarakonda et al., 2017; McCandlish et al., 2018; Mai et al., 2020) focuses
on adaptive batch size training in homogeneous environments. These approaches
continuously monitor system metrics and optimize the batch size according to
their adaption policies. However, the adaptive batch size metrics, algorithms,
and adaption policies are all designed for homogeneous environments. Directly
adopting existing methods in heterogeneous clusters will cause a large margin
of error for adaptive batch size training metrics measurement and prediction.
To the best of our knowledge, no specialized adaptive training system for
heterogeneous environments has been developed so far.

There are three main challenges in designing an automatic near-optimal
training system. First, each node has a different performance model in
heterogeneous clusters, which causes complexity in predicting the optimal
performance and the corresponding cluster configuration. The system overhead
will significantly affect the training performance for larger clusters.
Second, considering data parallelism distributed training in heterogeneous
clusters, given a new total batch size to the cluster in the adaptive batch
size training, the optimal configuration of each node will change. Given the
total batch size, previous work (Chen et al., 2020; Ye et al., 2022)
iteratively tunes each node’s configuration to approach the optimal
performance, which is inefficient in adaptive batch size training. Third, in
heterogeneous data parallelism training systems, different local batch sizes
are assigned to different GPUs. This introduces challenges in accurately
modeling gradient noise (McCandlish et al., 2018) across the heterogeneous
clusters.

Based on these insights and challenges, we study the performance model of
data-parallel distributed training for heterogeneous GPU clusters and propose
the optimal performance OptPerf. For the prediction of OptPerf and the
corresponding configuration, we design Cannikin, an efficient near-optimal
data-parallel distributed training system. Taking both computing and
communication models into consideration, Cannikin has accurate modeling and
prediction of the performance models for DL in heterogeneous clusters. With
the cluster performance model learned online, Cannikin can predict OptPerf and
configuration with low overhead when the cluster is given a new total batch
size. Cannikin optimizes the measurement of system parameters using inverse
variance weighting of different observations from each node in the cluster. We
also prove that in heterogeneous clusters we can correctly model adaptive
batch size training metrics, just like systems (Qiao et al., 2021; McCandlish
et al., 2018) designed for homogeneous clusters. While developed for single
DNN training, Cannikin can be readily integrated with adaptive batch size
training engines (Qiao et al., 2021; Mai et al., 2020) and dynamic resource
allocation schedulers (Qiao et al., 2021; Peng et al., 2018) for multiple
jobs. The main contributions of this paper are:

  * •

We are the first to consider training DNN models using adaptive batch sizes in
heterogeneous clusters.

  * •

We deduce and predict the optimal performance, denoted as OptPerf, along with
its corresponding configuration and the optimal total batch sizes during DL
training in heterogeneous clusters.

  * •

We design and implement Cannikin that can be easily integrated with the state-
of-art adaptive batch size training systems to train the DNN models in
heterogeneous clusters optimally.

  * •

We evaluate the performance of Cannikin in two real heterogeneous clusters
using multiple popular DNN workloads and optimizers. Results highlight that
Cannikin reduces DNN training by up to $52\%$ and $85\%$ in heterogeneous
clusters compared to AdaptDL (Qiao et al., 2021) and PyTorch
DistributedDataParallel(DDP) (Li et al., 2020).

##  2\. Background

###  2.1. Data-parallel Distributed Training

DNN workloads are highly computing-intensive (You et al., 2018). Distributed
deep learning accelerates DNN model training with multiple GPUs, by either
using model parallelism and data parallelism. This paper focuses on data-
parallelism distributed learning, which trains the same model on multiple GPUs
and each GPU uses different data samples.

During data-parallelism distributed training, node $i$ first trains its local
mini batch by forward and backward passes for the local gradient estimation.
Node $i$’s local gradient $g_{i}$ is:

(1)
$g_{i}=\frac{1}{b_{i}}\sum_{j=0}^{b_{i}-1}\nabla_{\theta}L_{x_{j}}(\theta),$

where $\theta$ is the vector of DNN weight parameters, $b_{i}$ is node $i$’s
local mini-batch size, $x_{j}$ is a sample in local mini batch, and $L$ is the
loss function.

Upon the completion of the local gradient estimation, node $i$ will aggregate
its own gradient with all the local gradients calculated by other nodes in the
cluster to get the gradient of the total batch (full batch of the DNN model):

(2) $g=\frac{1}{N}\sum_{i=0}^{N-1}g_{i},$

where N is the count of total nodes (GPUs) that train the model. Finally, each
node will use the averaged gradient to update the weight parameters $\bf{w}$
for the next batch. In a DNN training system, the gradient averaging can be
handled using backends such as NCCL (Nvidia, [n. d.]), MPI (mpi, 1993), and
Gloo (incubator, [n. d.]).

###  2.2. Adaptive Batch Size Training

A deep learning model usually consists of millions to trillions of parameters
(He et al., 2015; Vaswani et al., 2017; Barham et al., 2022), requiring a long
time for training. The selection of hyperparameters such as batch size in
different convergence phases significantly impacts training efficiency. Recent
work on adaptive batch size training (Qiao et al., 2021; Devarakonda et al.,
2017; McCandlish et al., 2018; Mai et al., 2020) greatly speeds up deep
learning model training by dynamically tuning hyperparameters such as batch
size and learning rate according to the gradient noise (McCandlish et al.,
2018), data throughput, and other customized metrics.

To determine the most statistically efficient batch size for a training
iteration, McCandlish et al. (2018) propose the _gradient noise scale_ , an
estimation of the signal-to-noise ratio of the stochastic gradient. This
metric can be used to predict the most statistically efficient batch size
during training. When the gradient noise is low, a small batch size can
achieve a great contribution to the convergence. When the gradient noise is
large, the error of the gradient estimated by a small batch would be
significant. In this situation, a larger batch size could reduce the training
time with little reduction of statistical efficiency.

However the optimal convergence progress is not guaranteed by using the most
statistically efficient batch size, because the most statistically efficient
batch size often comes with relatively low data throughput. Pollux (Qiao et
al., 2021) introduces _goodput_ , the product of the system throughput and
statistical efficiency modeled by the gradient noise scale. Goodput optimizes
training by balancing the trade-off between data throughput and statistical
efficiency.

##  3\. The _OptPerf_ of GPU Clusters

To improve the performance of a cluster, first we need to know the optimal
performance the cluster can achieve. In this section, we define and deduce the
optimal performance OptPerf of a heterogeneous GPU cluster using the metrics
collected from each GPU in a general case.

###  3.1. The Definition of _OptPerf_

For data-parallelism distributed training with a given total batch size $B$,
OptPerf is the optimal batch processing time a heterogeneous GPU cluster can
achieve by ideally tuning each node’s local mini-batch size. Consider
heterogeneous GPU Cluster $A$ with a set of nodes $\mathcal{N}$,
$|\mathcal{N}|=n$. For the standalone training of local mini batch size
$b_{i}$ at node $i$, let the batch processing time without gradient
synchronization among nodes be denoted by $t_{i}^{b_{i}}$. Due to
hetergoeneity, we assume $t_{i}^{b}\neq t_{j}^{b}$ for any $b$ and $i\neq j$.
The local mini batch sizes satisfy $\sum_{i\in\mathcal{N}}b_{i}=B$. Define
local mini batch size ratio ${\bf r}=[r_{0},r_{1},\dots,r_{n-1}]$, where
$r_{i}=b_{i}/B$. In this paper, we only consider synchronized data-parallelism
distributed training, meaning all nodes synchronize their updates after each
batch’s training. For this training model, fast nodes in the cluster always
wait for the stragglers to finish local gradient estimation, hence the batch
processing time
$T=max\\{t_{0}^{b_{0}},t_{1}^{b_{1}},\dots,t_{n-1}^{b_{n-1}}\\}$. We can infer
that there exists an optimal local mini batch size ratio
$\mathbf{r}_{\text{opt}}$ that minimizes the batch processing time $T$. We
define the minimized batch processing time to be OptPerf. To determine OptPerf
and $\mathbf{r}_{\text{opt}}$ for heterogeneous cluster $A$ with total batch
size $B$, we model the performance of GPUs in the cluster as a function of
$T$, $B$, and ${\bf r}$.

###  3.2. Performance Modeling

When considering the performance model of a heterogeneous cluster, rather than
simultaneously modeling the performance of heterogeneous nodes, we can instead
model the performance of each GPU separately and then combine all the GPU
performance models to determine the cluster’s performance. In data parallel
distributed training, the batch processing time is composed of the computing
time for local gradient estimation and the communication time for gradient
synchronization across nodes.

####  3.2.1. Computing Time

The batch computing time can be separated into data loading, forward
propagation, backward propagation and parameter updating. For any node $i$,
the computing time $t_{compute}^{i}$ is a linear function of local batch size
$b_{i}$ (Luo et al., 2022; Qiao et al., 2021). Furthermore, the parameter
updating time remains constant regardless of variations in the local batch
size. In contrast, data loading time, forward propagation time, and
backpropagation time exhibit a linear relationship with batch sizes $b_{i}$.
So for all nodes in Cluster $A$, the computing time can be expressed as:

(3) $\displaystyle t_{compute}^{i}=a_{i}+{P}_{i},\ \forall i\in\mathcal{N},$
$\displaystyle{a}_{i}=q_{i}b_{i}+s_{i},\ \forall i\in\mathcal{N},$
$\displaystyle{P}_{i}=k_{i}b_{i}+m_{i},\ \forall i\in\mathcal{N},$

where $a_{i}$ is the total time of parameter updating, data loading, and
forward propagation, with the corresponding coefficients $q_{i}$ and $s_{i}$.
${P}_{i}$ is the backpropagation time, while $k_{i}$ and $m_{i}$ are
coefficients related to GPU types and DL jobs. Note that within a
heterogeneous GPU cluster, different GPU models exhibit varying pairs of
$q_{i}$ and $s_{i}$, as well as $k_{i}$ and $m_{i}$, even when performing the
same DL job. If cluster $A$ has $n$ different types of GPUs, there are $n$
different pairs of linear functions corresponding to each type of GPUs.

####  3.2.2. Gradient Synchronization Time

We focus on the ring All-reduce mechanism (Patarasuk and Yuan, 2009a) adopted
by Pytorch DistributedDataParallel (Li et al., 2020). Ring All-reduce is a
synchronized communication method that starts the gradient synchronization
when all nodes are ready to synchronize and ends the synchronization when all
nodes finish the gradient synchronization. The gradient synchronization time
$T_{comm}$ is dependent on model size (size of gradient parameters) and
network status. In the scenario that the network and allocated resources are
stable in a heterogeneous cluster, even though each node’s network performance
varies, the gradient synchronization time $T_{comm}$ is a learnable constant
when we train the same job with different batch sizes.

####  3.2.3. Computing and Communication Overlap

Modern distributed training frameworks (Li et al., 2020; Sergeev and Balso,
2018) support the overlap between gradient computing and synchronization by
separating the locally-computed gradients into buckets (Li et al., 2020).
Rather than synchronizing after all nodes finish computing the full gradient
at the end of backpropagation, each gradient bucket starts synchronization
when all nodes finish computing the gradient of the same bucket. In batch
processing, only the last bucket cannot overlap its synchronization with its
gradient computation. So Cluster $A$’s per batch gradient synchronization time
$T_{comm}$ is the sum of $T_{u}$, the last gradient bucket synchronization
time, and $T_{o}$, the gradient synchronization time for all the other
gradient buckets: $T_{comm}=T_{o}+T_{u}$.

For all nodes in Cluster $A$, the gradient size (model size) is determined
when training starts. Different types of GPUs have the same gradient computing
procedure (gpu, 2011). Even though the local batch size of each node is
different, the gradient bucket will be ready for synchronization at a fixed
proportion of each node’s backpropagation time $P_{i}$. We can infer that the
starting point of the gradient synchronization of node $i$ is:

(4) $\textit{syncStart}_{i}=a_{i}+\gamma P_{i},$

where the overlap ratio $\gamma$ is the ratio of the first bucket computing
time to the total backpropagation time. The first bucket computing time is the
period from the start point of the backward pass to the first bucket ready for
synchronization point. The first bucket computing cannot be overlapped with
the gradient synchronization.

Although varying the local mini batch size $b_{i}$ changes the backpropagation
time $P_{i}$, $T_{comm}$ is a fixed, learnable constant. So the computing and
communication overlap pattern differs when the local mini batch size varies.
The first gradient bucket’s ready-for-synchronization point is determined by
the computing time along with $\gamma$, which is a constant that can be
accurately measured through all the nodes in the cluster. The following
buckets’ synchronization can be blocked by the previous buckets’
synchronization. To eliminate the effect of measurement error and system
overhead, we only measure the first bucket’s ready-for-synchronization point
syncStart and assume all buckets’ computing time and communication time are
evenly distributed in the rest of gradient computing time and communication
time. There are two possible overlap patterns of computing and communication.

Figure 1. A node running in the computing-bottleneck situation. PU, DL, FP,
and BP are the parameter update, data loading, forward propagation, and
backpropagation.

When computing is the bottleneck. Figure 1 shows a computing-bottleneck node.
The dashed lines show the buckets’ ready-for-synchronization points. When
$(1-\gamma)P_{i}\geq T_{o}$, node $i$’s bottleneck is the gradient
computation. In this case, the gradient synchronization of each bucket
finishes before the next gradient bucket is ready for synchronization, so
$T_{o}$ fully overlaps with the gradient computing. The total processing time
of one batch for node $i$ in cluster $A$ is:

(5) $T=t_{compute}^{i}+T_{u}.$

Figure 2. A communication-bottleneck node.

When communication is the bottleneck. Figure 2 shows the communication-
bottleneck pattern. If $(1-\gamma)P_{i}<T_{o}$, node $i$’s bottleneck is the
gradient communication. Although the unfinished bucket synchronization won’t
block gradient computing (due to the All-reduce mechanism), the
synchronization of previous unfinished buckets will block the synchronization
of the following bucket. In this situation, the total processing time of one
batch for node $i$ is:

(6) $T=\textit{syncStart}_{i}+T_{comm}.$

###  3.3. Expression of OptPerf

To predict OptPerf for heterogeneous clusters, with learned performance models
of all GPUs, first we deduce cluster $A$’s batch processing time $T$. In
cluster $A$,

(7)
$T=\max\left\\{\max\limits_{i\in\mathcal{N}}\\{t_{compute}^{i}+T_{u}\\},\max\limits_{i\in\mathcal{N}}\\{\textit{syncStart}_{i}+T_{comm}\\}\right\\},$

where $\mathcal{N}$ is the set of nodes in Cluster $A$. Since each node’s
bottleneck is unknown, minimizing $T$ is a mixed integer linear programming
problem. Rather than solve an NP-hard problem, we instead provide the criteria
for different situations to achieve OptPerf.

When optimizing the performance of Cluster $A$, the batch processing time
$T_{i}$ of stragglers during training can be reduced by adjusting the local
mini batch sizes for all nodes. Intuitively, the cluster’s fast nodes should
be assigned larger local mini batches, while the stragglers should have
smaller batches. With the computing and communication overlap patterns of each
node in Section 3.2.3, we first look into two special scenarios.

All nodes are computing-bottleneck. Since $(1-\gamma)P_{i}\geq T_{o}$,
$\forall i\in\mathcal{N}$, all the nodes’ performance models are Equation (5).
OptPerf is achieved when all nodes in cluster $A$ have the same computing time
$t_{compute}$. The proof is in Appendix A.1.

All nodes are communication-bottleneck. If $(1-\gamma)P_{i}<T_{o}$, $\forall
i\in\mathcal{N}$, all nodes’ performance models are Equation (6). OptPerf will
be achieved when all nodes start the first gradient bucket synchronization at
the same time. The proof is in Appendix A.2.

Figure 3. An example of the general case, where GPU 0 and GPU 1 are
communication-bottleneck, GPU 2 and GPU 3 are computing-bottleneck. FT is the
abbreviation of fixed time.

The general case. In the general case, some nodes’ bottlenecks are computing
while others’ are communication (see Figure 3).

The computing-bottleneck nodes’ performance models follow Equation (5), and
the communication-bottleneck nodes’ performance models follow Equation (6).
OptPerf is achieved when all computing-bottleneck nodes have the same
computing time $t_{compute}$ and all communication-bottleneck nodes start the
first bucket synchronization at the same time. Moreover, the computing and
communication bottleneck nodes simultaneously get ready for the last bucket
synchronization. The proof is located in Appendix A.3.

With OptPerf, we can determine the optimal performance of a heterogeneous
cluster with different batch sizes using the parameters measured and learned
during training.

##  4\. System Design of Cannikin

In this section, we give details on the workflow of Cannikin and describe how
Cannikin configures the cluster before the start of each epoch, optimizes the
measurement of learnable parameters, guarantees the gradient quality, and how
Cannikin integrates with existing adaptive batch size training systems.

###  4.1. Workflow of Cannikin

Figure 4 shows the overview of the workflow of Cannikin. In a heterogeneous
environment, after a user submits a training Job $J$ to the dynamic resource
job scheduler, the job scheduler allocates a number of (possibly
heterogeneous) GPUs to form Cluster $A$ to initialize Job $J$. Before the
start of each epoch, the adaptive batch size engine enumerates the total batch
size candidates from the batch size range (Qiao et al., 2021). The optimizer
uses performance models learned by the analyzer to predict OptPerf with its
corresponding total batch size and $\mathbf{r}_{\text{opt}}$ for the next
training epoch, then loads each node’s local mini batch based on
$\mathbf{r}_{\text{opt}}$ and starts the next training epoch. During the
training epoch, each node continually collects performance metrics and learns
the performance models locally. After each training epoch, the analyzer
gathers the updated performance models of all nodes.

Figure 4. The overall workflow of Cannikin.

###  4.2. OptPerf Optimizer

Cannikin continuously learns the computing and communication models for all
nodes during the training process. Despite having performance models for all
nodes, the overlap state of each node in the cluster remains unknown as it is
contingent on the total batch size. For example, a larger total batch size can
lead to a higher number of nodes experiencing computational bottlenecks. To
address this challenge, we have developed a novel search algorithm aimed at
unveiling the overlap state for all nodes across the cluster.

Determine the overlap state. Given an enumerated total batch size, Cannikin
uses Algorithm 1 to determine the overlap state, then predict OptPerf with
$\mathbf{r}_{\text{opt}}$ for each node.

If all nodes are computing-bottleneck or communication-bottleneck, we can use
Check 1 and Check 2 to verify.

Input: Total batch size $B=\sum_{i=0}^{n-1}b_{i}$.

Given: $\gamma,T_{o},T_{u},\\{a_{i},k_{i},m_{i}\\},i=0,1,\dots,n-1.$

$\triangleright$ Overlap ratio, two parts of communication time, computing
coefficients.

/*Check 1: All nodes are computing-bottleneck*/

Solve $t_{compute}=t_{compute}^{0}=t_{compute}^{1}=\dots=t_{compute}^{n-1}$.

if  _$\forall(1-\gamma)P_{i}\geq T_{o}$_ then

/*If all nodes are computing-bottleneck. */

OptPerf $=t_{compute}+T_{u};$

return OptPerf, $b_{i},i=0,1,\dots,n-1.$

break

/*Check 2: If nodes are communication-bottleneck.*/

Solve syncStart $=$ syncStart${}_{0}=\dots=$ syncStart${}_{n-1}.$

if  _$\forall(1-\gamma)P_{i} <T_{o}$_ then

/*All nodes are communication-bottleneck.*/

OptPerf $=$ syncStart $+T_{comm};$

return OptPerf$,b_{i},i=0,1,\dots,n-1.$

break

/*The nodes are mixed-bottleneck.*/

while  _$(\exists$ syncStart${}_{i}>$ syncStart${}^{\prime})\lor(\exists\
t_{compute}>t_{compute}^{\prime})$_ do

/*Search the overlap state.*/

beg $=$ outlier${}_{min};$

end $=$ outlier${}_{max};$

C$=($beg$+$end$)/2;$

/*Node $0$ to Node $C-1$ are computing-bottleneck, Node $C$ to Node $n-1$ are
comm-bottleneck.*/

Solve $T_{comb}=t_{compute}^{\prime}=$ syncStart${}^{\prime}+T_{o}$.

OptPerf$=T_{comb}+T_{u}.$

return OptPerf$,b_{i},i=0,1,\dots,n-1.$

Algorithm 1 Overlap state and OptPerf configuration

However, when nodes are mixed-bottleneck, the overhead of the enumeration
method to determine each node’s overlap pattern is relatively high. Algorithm
1 addresses this issue in the following steps: If node $i$ is a computing
(communication) bottleneck node in check 1 and check 2, then node $i$ is also
a computing (communication) bottleneck node in the mixed-bottleneck situation.
For all other outliers that have different overlap patterns in check 1 and
check 2, we use a binary-search-like algorithm to determine the computing and
communication bottleneck nodes. We rank all the intermediate nodes in
increasing order based on the fixed processing time and then set a
hypothetical bottleneck boundary node that separates the computing-bottleneck
and communication-bottleneck nodes. For any overlap state, the mixed-
bottleneck OptPerf solver from Section 3.3 will indicate $(\forall\
$syncStart$\leq$ syncStart${}^{\prime})$ $\land(\forall\ t_{compute}\leq
t_{compute}^{\prime})$ if the overlap pattern is correct. Thus we can
iteratively set the middle element to be the boundary node until we find the
correct overlap pattern.

Since the time complexity of checks 1 and 2 is at most
$O\left((n+1)^{3}\right)$ while solving linear equations (Bojańczyk, 1984) and
the time complexity of the mixed-bottleneck search algorithm is at most
$O(\log n)$, Algorithm 1 is $O\left((n+1)^{3}\log n\right)$. In Section 4.5,
we improve the time complexity of Algorithm 1 to $O\left((n+1)^{3}\right)$.

Approaching OptPerf when no available performance models. As the computation
time scales linearly with the local batch size of each GPU, deriving the
computing time models $a_{i}$, $P_{i}$ to $b_{i}$ necessitates the execution
of at least two distinct local batch sizes per GPU. Consequently, during the
initial two epochs, there will be no available performance model to predict
OptPerf. In this scenario, we employ the inverse proportion of the sample
computation time for each node to determine their respective local batch sizes
for the next epoch. Assume the per sample computing time of node $i$ at the
previous epoch is $t_{sample}^{i}=\frac{t_{compute}^{i}}{b_{current}^{i}}$,
where $b_{current}^{i}$ is the local batch size of node $i$ in the previous
epoch. The local batch size of node $i$ for the next epoch can be expressed
as:

(8)
$b_{next}^{i}=\frac{\sum_{i\in\mathcal{N}}t_{sample}^{i}}{t_{sample}^{i}}(\sum_{i\in\mathcal{N}}\frac{\sum_{i\in\mathcal{N}}t_{sample}^{i}}{t_{sample}^{i}})^{-1}B,$

where $B$ is the total batch size for the upcoming epoch. With this method,
each node can experiment with various local mini-batch sizes necessary for
performance model learning, and Cannikin iteratively approaches OptPerf when
no available performance models. Note that the primary purpose of this method
is to adjust each node’s local batch size for performance model learning,
rather than relying on the less efficient iterative method to find OptPerf.
Once the performance models are established, they are employed to predict
OptPerf before each epoch.

###  4.3. Optimized Gradient Aggregation

In homogeneous environments, the local gradient of each node can be aggregated
by the cluster via averaging. This procedure guarantees each training sample
has an identical weight in the global gradient after synchronization. However
local gradient averaging cannot be utilized for adaptive local batch training
in heterogeneous clusters due to the variety of local batch sizes. Simply
averaging each node’s local gradient results in over-representation of
training samples from smaller local batches in the global gradient. To address
this problem, LB-BSP (Chen et al., 2020) introduced proportional-weighted
gradient aggregation. By weighting each local gradient proportionally to the
local batch size, samples assigned to different nodes have identical weights
in the global gradient. Cannikin computes the global gradient $g$ using:

(9) $g=\sum_{i\in\mathcal{N}}r_{i}g_{i},$

where $g_{i}$ is the local gradient in Equation (1) and $r_{i}$ is the local
mini batch ratio of node $i$. For _i.i.d._ data, $g$ is equivalent to the
averaged gradients in homogeneous environments.

###  4.4. Gradient Noise Scale in Heterogeneous Environment

Adaptive batch size training uses the gradient noise scale (GNS) (McCandlish
et al., 2018) to model the convergence efficiency (statistical efficiency).
The GNS, $\mathcal{B}_{noise}$, measures how large the gradient is compared to
its variance: $\mathcal{B}_{noise}=\text{tr}(\Sigma)/|G|^{2}$, where $\Sigma$
is the covariance matrix and $G$ is the noiseless true gradient. Since
$\text{tr}(\Sigma)$ and $|G|^{2}$ are not available in practice, standard
methods instead rely on good estimators of these two quantities. Previous work
has only considered how to estimate $\text{tr}(\Sigma)$ and $|G|^{2}$ (and
thus $\mathcal{B}_{noise}$) in homogeneous clusters.

To compute the GNS, we first construct local estimates $\mathcal{G}_{i}$ and
$\mathcal{S}_{i}$ of $|G|^{2}$ and $\text{tr}(\Sigma)$ for each node $i$:

(10) $\displaystyle\mathcal{G}_{i}$
$\displaystyle=\frac{1}{B-b_{i}}(B|g|^{2}-b_{i}|g_{i}|^{2}),\quad\mathcal{S}_{i}=\frac{b_{i}B}{B-b_{i}}(|g_{i}|^{2}-|g|^{2})$

where the estimates incorporate local and global gradient information. For any
batch of size $b$, the expected gradient norm $\mathbb{E}[|g_{est}|^{2}]$
satisfies the equality
$\mathbb{E}[|g_{est}|^{2}]=|G|^{2}+\frac{1}{b}\text{tr}(\Sigma)$ (McCandlish
et al., 2018). Using this equation, we can prove that $\mathcal{G}_{i}$ and
$\mathcal{S}_{i}$ are unbiased estimators of $|G|^{2}$ and
$\text{tr}(\Sigma)$, respectively. Aggregating the local estimates
$\mathcal{G}_{i}$ and $\mathcal{S}_{i}$ across all nodes can provide high
quality, unbiased estimates of $|G|^{2}$ and $\text{tr}(\Sigma)$ that have
improved, lower variance. The variance of these estimators is crucial since
the standard ratio estimator used for the GNS is inherently biased (McCandlish
et al., 2018).

In homogeneous clusters, it is optimal to separately aggregate
$\mathcal{G}_{i}$ and $\mathcal{S}_{i}$ via averaging. However in Lemma B.1,
we prove that the variance of both $\mathcal{G}_{i}$ and $\mathcal{S}_{i}$
depend on the local mini batch size. Furthermore, the local estimates of
$\text{tr}(\Sigma)$ and $|G|^{2}$ for different nodes are correlated via
dependence on $|g|^{2}$. As a result, aggregating $\mathcal{G}_{i}$ and
$\mathcal{S}_{i}$ is more challenging for heterogeneous clusters. The
following theorem states the optimal weighted combination of the local
estimators, with the proof located in Appendix B.

###### Theorem 4.1.

$\mathcal{G}=\sum_{i\in\mathcal{N}}w^{\mathcal{G}}_{i}\mathcal{G}_{i}$ and
$\mathcal{S}=\sum_{i\in\mathcal{N}}w^{\mathcal{S}}_{i}\mathcal{S}_{i}$ are
minimum variance, unbiased linear estimators of $|G|^{2}$ and
$\text{tr}(\Sigma)$ when:

(11)
$\displaystyle\mathbf{w}^{\mathcal{G}}=\frac{\mathbf{1}^{T}A_{\mathcal{G}}^{-1}}{\mathbf{1}^{T}A_{\mathcal{G}}^{-1}\mathbf{1}},\quad\quad\mathbf{w}^{\mathcal{S}}=\frac{\mathbf{1}^{T}A_{\mathcal{S}}^{-1}}{\mathbf{1}^{T}A_{\mathcal{S}}^{-1}\mathbf{1}},$

where $\mathbf{1}$ is an $n$-dimensional column vector of ones and both
$A_{\mathcal{G}}$ and $A_{\mathcal{S}}$ are $n\times n$ matrices with
respective entries $a_{\mathcal{G}}(i,j)$ and $a_{\mathcal{S}}(i,j)$:

$a_{\mathcal{G}}(i,i)=\dfrac{B+2b_{i}}{B^{2}-Bb_{i}},\quad
a_{\mathcal{G}}(i,j)=\dfrac{B^{2}-b_{i}^{2}-b_{j}^{2}}{B(B-b_{i})(B-b_{j})}\text{
for }i\neq j$ $a_{\mathcal{S}}(i,i)=\dfrac{Bb_{i}}{B-b_{i}},\quad
a_{\mathcal{S}}(i,j)=\dfrac{b_{i}b_{j}(B-b_{i}-b_{j})}{(B-b_{i})(B-b_{j})}\text{
for }i\neq j$

Cannikin delivers a novel method to estimate the GNS $\mathcal{B}_{noise}$ in
heterogeneous clusters. First, each node estimates the sum of the variances of
the individual gradient components and the global norm of the gradient using
(10). We optimally aggregate the local estimates $\mathcal{G}_{i}$ and
$\mathcal{S}_{i}$ using (B), and then take the ratio of the resulting terms to
get the global gradient noise scale
$\mathcal{B}_{noise}=\mathcal{S}/\mathcal{G}$. Despite the added challenge of
heterogeneity, Figure 5 shows Cannikin’s convergence is comparable to the
homogeneous baseline AdaptDL with the same training epochs, which means the
larger batch size chosen by Cannikin won’t harm the convergence efficiency.

Figure 5. (a) shows the total batch size selected by Cannikin and AdaptDL in a
heterogeneous cluster for the training of CIFAR10. In most epochs, Cannikin
trains with a larger batch size compared with AdaptDL due to the throughput
improvement. (b) shows with the same training epochs, Cannikin achieves the
same accuracy compared with AdaptDL. (c) shows the accuracy in the convergence
process of Cannikin and AdaptDL.

###  4.5. Implementation

Cannikin is implemented as a PyTorch library based on AdaptDL (Qiao et al.,
2021) that can be imported into DNN training scripts. Cannikin introduces the
HeteroDataLoader class, which unevenly loads local mini batches to each node
based on the OptPerf prediction. The implementation of Cannikin addresses the
following concerns to improve efficiency.

Parameter learning. During each epoch, Cannikin collects the backpropagation
time ($P_{i}$) and the total time of data loading, optimization steps, forward
propagation($a_{i}$) for each local batch size. With the collected data from
two epochs using different local mini-batch sizes, each node can construct the
computing time model to the local mini batch size by solving linear equations.
In the subsequent epochs, employing more different local batch sizes, allows
for the refinement of the computation time model, making it increasingly
accurate.

When it comes to learning the communication time ($T_{comm}$) and $\gamma$,
it’s important to note that $T_{comm}$ and $\gamma$ remain constant across
different local and total batch sizes. Cannikin collects the overlap ratio
($\gamma$) as well as communication times ($T_{o}$ and $T_{u}$) for each node
in the cluster. We proceed to optimize the learning of $\gamma$ and $T_{comm}$
as follows.

Optimized parameter measurement in the cluster. Cannikin collects observation-
based parameters such as overlap ratio $\gamma$ and communication time $T$
from all nodes in the cluster. However since different nodes can have
different levels of noise in their measurement of $\gamma$, simply averaging
across all nodes causes significant error. Figure 6 shows the measured
$\gamma$ for ResNet18 training on Cifar10 on different GPUs. The stochastic
nature of distributed DNNs results in randomness when measuring $\gamma$ (You
et al., 2022), which can lead to error when learning $\gamma$.

Figure 6. The measured overlap ratio $\mathbf{\gamma}$ of different local mini
batch sizes validated on different types of GPUs (Duplyakin et al., 2019).

We adopt the inverse-variance weighting (Marín-Martínez and Sánchez-Meca,
2010) to adjust the weighting of each node’s measurement of $\gamma$:

(12)
$\gamma=\sum_{i\in\mathcal{N}}\frac{\gamma_{i}}{\hat{\sigma}_{\gamma_{i}}^{2}}\left(\sum_{i\in\mathcal{N}}\frac{1}{\hat{\sigma}_{\gamma_{i}}^{2}}\right)^{-1},$

where $\gamma_{i}$ is the overlap ratio estimation of node $i$ and
$\hat{\sigma}_{\gamma_{i}}^{2}$ is the sample variance of $\gamma_{i}$. This
estimation of $\gamma$ is optimal when observation errors are uncorrelated
across nodes.

The communication time $T$ is also a fixed value across all nodes because of
the ring All-reduce mechanism (Patarasuk and Yuan, 2009a). However, each node
reports different $T_{i}$ during training in heterogeneous clusters because of
the wait-for-synchronization time of some nodes. In Cannikin’s implementation,
we use $T=\min_{i\in\mathcal{N}}\\{T_{i}\\}$ to eliminate the error in
communication time measurement because this value typically corresponds to the
slowest node that does not need to wait for any other node for
synchronization.

Total batch size selection. Although the search algorithm is efficient in
finding the overlap pattern, the overhead can be significant if we determine
the overlap pattern for each total batch size candidate determined by the
adaptive batch size engine (Qiao et al., 2021) before every epoch. Cannikin
instead calculates $\textit{OptPerf}_{init}$ for all batch size candidates
after the initial epoch. In the upcoming epochs, since OptPerf is unrelated to
the training progress, Cannikin uses $\textit{OptPerf}_{init}$ and the updated
gradient noise scale to choose the total batch size. Then Cannikin determines
OptPerf along with $\mathbf{r}_{\text{opt}}$ according to the updated
performance metrics. If the overlap pattern has changed from the initial
pattern, Cannikin will start over to determine the pattern for each candidate
again to choose the total batch size. Otherwise Cannikin will update
$\textit{OptPerf}_{init}$ for the corresponding total batch size candidate.
With this strategy, in most epochs Cannikin only needs to determine OptPerf
for one total batch size.

Overlap state searching. In the initialization epoch, Cannikin goes through
all the total batch size candidates and calculates OptPerf for each candidate.
When the total batch size increases, more cluster nodes will be computing-
bottleneck nodes. Hence in the total batch size enumeration from small to
large in sequence, the search starting point of an enumerated candidate is the
overlap pattern of the previous one. In the following epochs, the search
starting point of an enumerated candidate is its overlap state in
$\textit{OptPerf}_{init}$.

Integer batch sizes. The batch size of each node is required to be an integer.
To avoid the overhead of integer programming, Cannikin instead computes the
optimal $b_{i}$ without enforcing the integer constraint and then rounds any
decimal local mini batch sizes. In our evaluation, the error caused by the
approximation is insignificant.

##  5\. Evaluation

We evaluate the effectiveness of Cannikin in optimal distributed DNN training
using convergence time, batch processing time, prediction accuracy, and the
system overhead. Key results are as follows:

  * •

Cannikin reduced the overall convergence time by up to $85\%$ and $52\%$ in
heterogeneous clusters compared with PyTorch and the adaptive batch size
training system AdaptDL.

  * •

Compared to the state-of-art data parallel distributed deep learning
strategies for heterogeneous clusters, Cannikin reduces the batch processing
time by up to $18\%$ compared to the baseline.

  * •

Cannikin predicted OptPerf for heterogeneous clusters within $7\%$ error with
low overhead less than $4\%$.

###  5.1. Experimental Setup

Testbed. We conduct our experiments in two different clusters: cluster $A$ and
cluster $B$. Cluster $A$ is a heterogeneous 3-node cluster with different
types of NVIDIA GPUs specified in Table 2. Cluster $B$ is a heterogeneous
10-server cluster consisting of $16$ GPUs, as detailed in the table 3. Note
that in Cluster $B$, each GPU is a node for data-parallelism distributed DL
training.

Table 2. Hardware specification of cluster _A_ in evaluation Node | Node | GPU | GPU | Main | CPU
---|---|---|---|---|---
type | count | model | Count | Memory
a5000 | 1 | RTX A5000 | 1 | 32GB | i9-10980XE
a4000 | 1 | RTX A4000 | 1 | 32GB | Xeon W-2255
p4000 | 1 | Quadro P4000 | 1 | 32GB | Xeon W-2102
Table 3. Hardware specification of cluster _B_ in evaluation Node | Node | GPU | GPU | Main | CPU
---|---|---|---|---|---
type | count | model | Count | Memory
a100 | 1 | A100 | 4 | 512GB | Xeon Plati. 8380*2
v100 | 1 | V100 | 4 | 128GB | Xeon Gold 6230*2
rtx | 8 | RTX6000 | 1 | 192GB | Xeon Gold 6126*2
Table 4. The training models, with the datasets ImageNet, CIFAR-10, LibriSpeech, SQuAD, and MovieLens. The model size (parameters in the model), optimizer, learning rate scaler, original batch size $\mathbf{B_{0}}$, and the target metrics used for Cannikin’s evaluation. Note that we use fine-tuning BERT for evaluation. Task | Dataset | Model | Size | Optimizer | LR scaler | $\mathbf{B_{0}}$ | Target
---|---|---|---|---|---|---|---
Image Classification | ImageNet (Deng et al., 2009) | ResNet-50 (He et al., 2015) | 25.6M | SGD | Adascale | 100 | 75% Top1 acc.
Image Classification | CIFAR-10 (Krizhevsky, 2009) | ResNet-18 (He et al., 2015) | 11M | SGD | Adascale | 64 | 94% Top1 acc.
Speech Recognition | LibriSpeech (Panayotov et al., 2015) | DeepSpeech2 (Amodei et al., 2015) | 52M | SGD | Adascale | 12 | WER = 40.0%
Question Answering | SQuAD (Rajpurkar et al., 2016) | BERT (Devlin et al., 2018) | 110M | AdamW | Square-Root | 9 | F1 = 88%
Recommendation | MovieLens (Harper and Konstan, 2015) | NeuMF (He et al., 2017) | 5.2M | Adam | Square-Root | 64 | Hit rate = 69%

Workloads. Evaluated workloads are listed in Table 4. The range of batch sizes
is determined by each GPU’s memory; the initial batch size is relatively small
(Qiao et al., 2021) and configured by users. For the optimizer, learning rate
scaler, and target metrics choices, we adopt the canonical setting for each
training task, with the philosophy of testing different models and optimizers
on applications of various sizes.

Baselines. We evaluate Cannikin by comparing it with the state-of-art adaptive
batch size training system, data-parallelism heterogeneous distributed DL
training system, and PyTorch DDP:

  * •

AdaptDL (Qiao et al., 2021): The state-of-art adapted distributed DNN training
system for homogeneous clusters.

  * •

LB-BSP (Chen et al., 2020): LB-BSP is a data-parallelism distributed training
system for heterogeneous GPU clusters, which recurrently tune each node’s
local mini batch size for efficient model training. We set step size
$\Delta=5$ in our experiments, which is identical to the original paper.

  * •

PyTorch DistributedDataParallel (Li et al., 2020): Pytorch DDP is one of the
most efficient distributed training libraries for homogeneous clusters.

###  5.2. Performance with Heterogeneous GPUs

####  5.2.1. Overall convergence performance

For the overall convergence performance evaluation, we compared Cannikin with
the baselines in cluster $B$. Figure 7 shows the convergence processes of
example tasks Cifar10 and Imagenet training in cluster $B$. Due to the
weighted gradient aggregation, each gradient descent step of Cannikin is
equivalent to the homogeneous gradient descent given the same total batch
size, This equivalence guarantees that the convergence is not compromised.
With this precondition, Cannikin achieves a convergence speed up from
throughput improvement and the improved prediction of optimal total batch size
in heterogeneous clusters, thus increasing the adaptive training system’s
goodput. Our results show Cannikin reduces the convergence time by $52\%$ and
$29\%$ for CIFAR-10 and ImageNet compared with AdaptDL.

Figure 7. Convergence process of ResNet-18 on CIFAR-10 (left) and ResNet-50 on
ImageNet (right). Figure 8. Normalized convergence time of all training tasks.

The normalized overall convergence time for each evaluated workload is
depicted in Figure 8. In the case of PyTorch DDP, which trains deep learning
models with fixed total batch size and distributes local batch sizes evenly
across heterogeneous clusters, the speedup achieved by Cannikin is primarily
attributed to its optimized prediction of total/local batch sizes during
training. Unlike PyTorch DDP, which uses fixed batch sizes, AdaptDL evenly
distributes local batch sizes across the cluster and predicts the optimal
total batch size in homogeneous environments. In the context of AdaptDL, the
speedup observed with Cannikin results from the optimized selection of local
batch sizes to maximize the utilization of heterogeneous GPUs and the improved
prediction of total batch sizes in heterogeneous environments. LB-BSP
iteratively tunes local batch sizes for each GPU within heterogeneous clusters
which improves the utilization of the heterogeneous GPUs. However, LB-BSP only
supports the DL training with a fixed total batch size and LB-BSP hasn’t
consider the communication and communication overlap. The speedup with
Cannikin compared to LB-BSP primarily arises from the faster determination of
the optimal local batch sizes considering communication and computing overlap
and optimized total batch size selection during training. The results indicate
Cannikin significantly enhances in overall convergence time to achieve the
target accuracy, with improvements of up to $85\%$, $52\%$, and $82\%$
compared to PyTorch DDP, AdaptDL, and LB-BSP respectively.

####  5.2.2. Batch processing time for heterogeneous clusters

We evaluate Cannikin using two methodologies for batch processing time. The
first is the fixed total batch size training, i.e. classical DNN training. The
second is the adaptive batch size situation when the total batch size of the
cluster varies in each training epoch. Since AdaptDL’s batch processing time
in heterogeneous clusters is equivalent to Pytorch DDP, we don’t consider
AdaptDL in this section.

With fixed batch size. We fix the total batch size of the cluster and each
node’s optimal local mini batch size ratio $\mathbf{r}_{\text{opt}}$ . Figure
9 shows an example of Cannikin and LB-BSP training ResNet-50 with ImageNet in
cluster $A$. Given the total batch size of 128, Cannikin and LB-BSP initialize
training by evenly assigning local batch size for each node. Cannikin approach
OptPerf as early as the third epoch, because Cannikin requires two epochs to
learn the performance models discussed in Section 4.2. However, LB-BSP
requires more than ten epochs to reach its best performance.

Figure 9. Cluster $\mathbf{A}$’s batch size processing time when training
ImageNet from evenly assigned local mini batch size initialization given fixed
total batch size 128.  Figure 10. Each evaluation task’s normalized batch
processing time to the total batch size. The normalized performances of LB-BSP
in fixed batch size situations and adapted batch size situations will be
approached when the batch size is large enough. Because when the batch size is
large, the optimal local mini batch size ratio to the total batch size of each
node will approach a constant. This is also why when the batch size is large,
Cannikin’s normalized performance speedup will approach a constant.

Assume Cannikin and each compared method have reached their best batch
processing time for the given batch size. Figure 10 shows OptPerf of cluster
$B$ compared with the baselines. We can observe that OptPerf is at most $18\%$
faster compared with LB-BSP, and up to 53% faster than the Pytorch DDP. Note
that when the batch size is large enough, all nodes become computing
bottleneck. The OptPerf will be achieved when all nodes have the same
computing time $T_{compute}$. The performance of LB-BSP will approach OptPerf
because asymptotically the two have the same target that all nodes have the
same $T_{compute}$.

With adaptive batch size. Since only data-parallelism distributed systems are
sensitive to batch size changes, we evaluate Cannikin with LB-BSP for the
adaptive batch size situation. Assuming Cannikin and LB-BSP have already
achieved their best performance for the previous batch size, now given a new
batch size that is 10% of the total batch size range larger than the previous
one, the batch processing time of LS-BSP will become sub-optimal because the
$\mathbf{r}_{\text{opt}}$ has changed. In the meantime, Cannikin can still
accurately predict the OptPerf for the newly assigned batch size, just like
the fixed batch size situation. Figure 10 shows the batch processing time of
LS-BSP in cluster $B$ for adaptive batch size training.

###  5.3. OptPerf Prediction

In cluster $A$, we evaluate the prediction of OptPerf with and without inverse
variance weighting in measurement compared to the manually tuned OptPerf.
Results show that without inverse variance weighting, the maximum error of
OptPerf prediction can reach up to 21%. With the inverse variance weighting
method introduced in Section 4.5, Cannikin’s prediction of OptPerf in small
and medium models like NeuMF, ResNet-18, and ResNet-50 have a maximum 3%
error. For larger models like BERT and DeepSpeech2, larger model sizes lead to
more gradient buckets to synchronize, which increases the probability of
contingency in gradient synchronization, so the maximum error in the
prediction of OptPerf is 7% in the batch size range. However, Cannikin trains
with varying batch sizes, so the maximum 7% error of the OptPerf prediction is
only used in a fraction of the entire training process.

###  5.4. Overhead and Scalability of Cannikin

Table 5 shows the overhead of Cannikin for each task we deployed in the large-
scale test cluster $B$. The overhead encompasses the time required to evaluate
each candidate total batch size alongside its corresponding OptPerf, as well
as the configuration time for each node’s local batch size and local training
data index. For all the medium and large applications, the configuring time
for OptPerf of Cannikin before each epoch is much less than $1\%$ of the total
epoch training time across all candidate batch sizes in the range specified by
(Qiao et al., 2021), which is insignificant for the entire training process.
For small applications like CIFAR-10 and MovieLens, the overhead of Cannikin
will reach up to $9\%$ and $12\%$ when the system runs with batch sizes around
the upper limit of the batch size range. However, during the training
progress, the system will use the batch sizes near the upper limit only when
the model almost converges. The period using the batch sizes near the upper
limit for training is a minority part of training time. Considering the entire
training progress, the overheads of CIFAR-10 and MovieLens are $2.7\%$ and
$3.9\%$.

Table 5. The maximum overhead for an epoch and the overall overhead for the complete training process of Cannikin. Dataset | Model | Max | Overall
---|---|---|---
Overhead | Overhead
ImageNet | ResNet-50 | $\ll 1\%$ | $\ll 1\%$
LibriSpeech | DeepSpeech2 | $\ll 1\%$ | $\ll 1\%$
SQuAD | BERT | $\ll 1\%$ | $\ll 1\%$
CIFAR-10 | ResNet-18 | $9\%$ | $2.7\%$
MovieLens | NeuMF | $12\%$ | $3.9\%$

##  6\. Discussion

Impact of varying heterogeneity. The performance improvement compared with the
baseline depends on the heterogeneity of the cluster. Generally speaking, a
cluster with more heterogeneity will get more benefits from Cannikin. In
homogeneous clusters, the performance of Cannikin is identical to AdaptDL. In
Cluster $B$, the fastest GPU A100 is about $3.42$ times faster compared with
RTX6000 which is the slowest GPU. The degree of heterogeneity we evaluated in
this paper generally exists in today’s computing platform. As shown in Table
1, after two years, the H100 GPU is faster than A100 by more than $4$ times.
We know that the A100 GPUs are not outdated and are still prevalent within the
majority of deep learning clusters.

Potentials with Sharing-caused heterogeneity. The heterogeneity can arise not
only from hardware differences but also from resource sharing. Recent studies
(Wu et al., 2023; Xiao et al., 2020) have introduced GPU-sharing mechanisms
that enable the sharing of a single GPU’s resources among multiple instances.
In this context, even when the same GPU type is present within a cluster, the
resources available at each node can still exhibit heterogeneity during
distributed training.

We create Cluster $C$, a 16-node homogeneous cluster in Chameleon Cloud
(Keahey et al., 2020). Each node is equipped with one NVIDIA RTX6000 GPU. We
use the container’s constraint to construct the heterogeneous environment. We
adopt docker containers (Merkel, 2014) for cluster $C$ to configure the
heterogeneous environment. In each node, we start two docker containers. The
first container runs Cannikin distributed training workloads, and the second
docker container runs a local dummy GPU workload to share the same GPU’s
computing power and memory with Cannikin. To tune each node’s computing power,
we manually adjust the local dummy GPU workload’s batch size to change the
computing power and memory of Cannikin workloads.

For the fastest node in cluster $C$, we allocate the entire RTX6000 GPU to the
Cannikin container, while for the slowest node in cluster $C$, we restrict the
performance of the node’s Cannikin container to approximately one-fourth of an
RTX6000 GPU. For the other intermediate nodes, we assign the batch size of
each dummy workload to make them evenly distributed between the batch sizes of
the fastest and slowest nodes. For example, in Cluster $C$ the batch size of
the dummy workload in the slowest node is $150$, and in the fastest node is
$0$. Then for the intermediate nodes, the dummy batch sizes are
$10,20,\dots,140$. The results indicate that Cannikin’s performance in Cluster
$C$ aligns with that of Cluster $A$ and $B$. This brief experiment
demonstrates the potential of Cannikin in addressing heterogeneity induced by
resource sharing.

Memory limitation. In the adaptive batch size training, the system will be
initialized with a small total batch size (Qiao et al., 2021) because when
training starts, the relatively small batch size could guarantee statistical
efficiency with less demand for the hardware resources. In Cannikin, this is a
benefit to our system. In the experiments, we observe that powerful GPUs often
outfit large GPU memory. On the other hand, weak GPUs’ memory is usually
relatively small. Since Cannikin evenly initializes each node’s local mini
batch size, if the initial total batch size is large, the nodes with weak GPU
often reach their memory limitation, which causes training failure. Then in
the following epoch, the total batch size increases. Cannikin will assign
larger local mini batch sizes to the fast nodes and smaller local mini batch
sizes to the slow nodes, which could avoid the GPU memory limit problem. We
have also established local batch size constraints for each node to prevent
out-of-memory failures.

Adapt to schedulers for heterogeneous clusters. Existing dynamic resource
allocating schedulers (Qiao et al., 2021; Peng et al., 2018; Zhang et al.,
2017) only support the scheduling of homogeneous clusters. Sia (Jayaram
Subramanya et al., 2023) is a scheduler with heterogeneity awareness, however,
at each job level, the resource allocated for each job is still homogeneous.
With Cannikin, we fully utilize the computing resources in heterogeneous
clusters. When we design schedulers in future work, the scheduler should be
able to allocate a heterogeneous cluster for each job, which can significantly
increase resource utilization.

With the performance metrics of Cannikin, the scheduler optimizes multi-job
performance and reallocates resources for each job between epochs. In the
multiple jobs scenario, if the scheduler removes nodes for Job $J$, Cannikin
can easily use the learned computing models of remaining nodes along with the
communication time model proposed by Qiao et al. (2021) to predict and
configure OptPerf for Job $J$. When adding new nodes to Cluster $A$, Cannikin
will re-initialize the cluster for job $J$ with two epochs. Then in the
following epochs, Job $J$ will achieve OptPerf.

##  7\. Other Related Work

Performance modeling of DNN training. The importance of performance modeling
in deep neural network training has been shown in recent research work. Paleo
(Qi et al., 2016) proposed the DNN training model by studying the NN topology.
Clockwork (Gujarati et al., 2020) model the GPU runtime with tracing. The All-
reduce communication between nodes is studied and modeled (Patarasuk and Yuan,
2009b). For the cloud-based DNN training (Luo et al., 2022), the accurate
performance modeling and prediction significantly increase the training
efficiency and reduce the clients’ costs. From the scheduler perspective, an
accurate DNN performance modeling (Qiao et al., 2021; Peng et al., 2018; Le et
al., 2020; Narayanan et al., 2020) increases the resource utilization and
improves the fairness of multiple jobs execution. However, before Cannikin,
none of the previous work synthetically considered the computing and
communication model with the overlap in heterogeneous environments.

Systems for adapted batch size training. The philosophy of adaptive batch size
training is choosing the batch size with the most convergence efficiency
during training. Previous work (Lee et al., 2019; Devarakonda et al., 2017)
gradually increase the batch size and learning rate during training. Anytime
Minibatch (AMB) (Ferdinand et al., 2018) automatically tunes batch size by
setting a fixed computing time for each batch, and each node computes as many
samples as possible. However, AMB has only been evaluated on small
applications and cannot support DNN workloads so far. The state-of-art
adaptive batch size training system Pollux (Qiao et al., 2021) can
automatically configure the batch size by optimizing goodput during training.
However, Pollux is designed for homogeneous environments. Cannikin is the
first adapted batch size training system for heterogeneous environments.

Accelerating ML on heterogeneous environments. Most previous work about ML
acceleration for heterogeneous clusters is on the scheduler level for multiple
jobs. Hare (Chen et al., 2022) optimized inter- and intra-job parallelism to
increase the utilization of heterogeneous clusters. EasyScale (Li et al.,
2022) dynamically assigned workers to scale distributed training for
heterogeneous GPUs. Gandiva (Chaudhary et al., 2020) designed a strategy for
heterogeneous GPU allocation that guarantees fairness and improves
utilization. GPUlet (Choi et al., 2022) designed an efficient scheduler by
adopting spatial partitioning of GPU resources. However, looking into the
single job level for the schedulers, the training strategy is still
homogeneous. For job-level optimization on heterogeneous clusters, SnuHPL (Kim
et al., 2022) improved the training in a heterogeneous HPC system by
optimizing the data distribution for a given cluster configuration. BytePS
(Jiang et al., 2020) accelerated DNN training by leveraging CPU and bandwidth
resources, however, BytePS focuses on the heterogeneity between CPU/GPU.
Cannikin is a job-level-optimized system designed for heterogeneous GPU
clusters. It automatically explores and determines the optimal local batch
sizes assigned to each GPU and total batch sizes.

##  8\. Conclusion

By modeling the GPU computing and communication jointly, we derive the optimal
performance OptPerf for heterogeneous clusters in the general case. Then we
design Cannikin, which automatically predicts and configures OptPerf for the
heterogeneous cluster with high efficiency. Cannikin is the first adaptive
distributed training system for heterogeneous clusters with near-optimal
performance and high scalability by overcoming challenges such as optimal
scenario determination and metrics measurement caused by heterogeneity.
Cannikin outperforms state-of-art training systems for diverse workloads in
real heterogeneous clusters. We believe Cannikin will inspire future DL
training systems and DNN job schedulers in heterogeneous environments.

## References

  * (1)
  * mpi (1993) 1993\.  MPI: A message passing interface. In _Supercomputing ’93:Proceedings of the 1993 ACM/IEEE Conference on Supercomputing_. 878–883.  https://doi.org/10.1145/169627.169855
  * gpu (2011) 2011\.  Chapter 4 - The CUDA Execution Model.  In _CUDA Application Design and Development_ , Rob Farber (Ed.). Morgan Kaufmann, Boston, 85–108.  https://doi.org/10.1016/B978-0-12-388426-8.00004-5
  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek G. Murray, Benoit Steiner, Paul Tucker, Vijay Vasudevan, Pete Warden, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. 2016.  TensorFlow: A system for large-scale machine learning.  https://doi.org/10.48550/ARXIV.1605.08695
  * Amodei et al. (2015) Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Ng, Sherjil Ozair, Ryan Prenger, Jonathan Raiman, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Yi Wang, Zhiqian Wang, Chong Wang, Bo Xiao, Dani Yogatama, Jun Zhan, and Zhenyao Zhu. 2015.  Deep Speech 2: End-to-End Speech Recognition in English and Mandarin.  https://doi.org/10.48550/ARXIV.1512.02595
  * Barham et al. (2022) Paul Barham, Aakanksha Chowdhery, Jeff Dean, Sanjay Ghemawat, Steven Hand, Daniel Hurt, Michael Isard, Hyeontaek Lim, Ruoming Pang, Sudip Roy, Brennan Saeta, Parker Schuh, Ryan Sepassi, Laurent Shafey, Chandu Thekkath, and Yonghui Wu. 2022.  Pathways: Asynchronous Distributed Dataflow for ML. In _Proceedings of Machine Learning and Systems_ , D. Marculescu, Y. Chi, and C. Wu (Eds.), Vol. 4. 430–449.  https://proceedings.mlsys.org/paper/2022/file/98dce83da57b0395e163467c9dae521b-Paper.pdf
  * Bojańczyk (1984) A. Bojańczyk. 1984.  Complexity of Solving Linear Systems in Different Models of Computation.  _SIAM J. Numer. Anal._ 21, 3 (1984), 591–603.  http://www.jstor.org/stable/2157070
  * Chaudhary et al. (2020) Shubham Chaudhary, Ramachandran Ramjee, Muthian Sivathanu, Nipun Kwatra, and Srinidhi Viswanatha. 2020.  Balancing Efficiency and Fairness in Heterogeneous GPU Clusters for Deep Learning. In _Proceedings of the Fifteenth European Conference on Computer Systems_ (Heraklion, Greece) _(EuroSys ’20)_. Association for Computing Machinery, New York, NY, USA, Article 1, 16 pages.  https://doi.org/10.1145/3342195.3387555
  * Chen et al. (2020) Chen Chen, Qizhen Weng, Wei Wang, Baochun Li, and Bo Li. 2020.  Semi-Dynamic Load Balancing: Efficient Distributed Learning in Non-Dedicated Environments. In _Proceedings of the 11th ACM Symposium on Cloud Computing_ (Virtual Event, USA) _(SoCC ’20)_. Association for Computing Machinery, New York, NY, USA, 431–446.  https://doi.org/10.1145/3419111.3421299
  * Chen et al. (2022) Fahao Chen, Peng Li, Celimuge Wu, and Song Guo. 2022.  Hare: Exploiting Inter-Job and Intra-Job Parallelism of Distributed Machine Learning on Heterogeneous GPUs. In _Proceedings of the 31st International Symposium on High-Performance Parallel and Distributed Computing_ (Minneapolis, MN, USA) _(HPDC ’22)_. Association for Computing Machinery, New York, NY, USA, 253–264.  https://doi.org/10.1145/3502181.3531462
  * Cho et al. (2019) Minsik Cho, Ulrich Finkler, David Kung, and Hillery Hunter. 2019.  BlueConnect: Decomposing All-Reduce for Deep Learning on Heterogeneous Network Hierarchy. In _Proceedings of Machine Learning and Systems_ , A. Talwalkar, V. Smith, and M. Zaharia (Eds.), Vol. 1. 241–251.  https://proceedings.mlsys.org/paper/2019/file/9b8619251a19057cff70779273e95aa6-Paper.pdf
  * Choi et al. (2022) Seungbeom Choi, Sunho Lee, Yeonjae Kim, Jongse Park, Youngjin Kwon, and Jaehyuk Huh. 2022.  Serving Heterogeneous Machine Learning Models on Multi-GPU Servers with Spatio-Temporal Sharing. In _2022 USENIX Annual Technical Conference (USENIX ATC 22)_. USENIX Association, Carlsbad, CA, 199–216.  https://www.usenix.org/conference/atc22/presentation/choi-seungbeom
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. 2009.  ImageNet: A large-scale hierarchical image database. In _2009 IEEE Conference on Computer Vision and Pattern Recognition_. 248–255.  https://doi.org/10.1109/CVPR.2009.5206848
  * Devarakonda et al. (2017) Aditya Devarakonda, Maxim Naumov, and Michael Garland. 2017.  AdaBatch: Adaptive Batch Sizes for Training Deep Neural Networks.  https://doi.org/10.48550/ARXIV.1712.02029
  * Devlin et al. (2018) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2018.  BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding.  https://doi.org/10.48550/ARXIV.1810.04805
  * Ding et al. (2021) Yifan Ding, Nicholas Botzer, and Tim Weninger. 2021.  HetSeq: Distributed GPU Training on Heterogeneous Infrastructure.  _Proceedings of the AAAI Conference on Artificial Intelligence_ 35, 17 (May 2021), 15432–15438.  https://doi.org/10.1609/aaai.v35i17.17813
  * Duplyakin et al. (2019) Dmitry Duplyakin, Robert Ricci, Aleksander Maricq, Gary Wong, Jonathon Duerig, Eric Eide, Leigh Stoller, Mike Hibler, David Johnson, Kirk Webb, Aditya Akella, Kuangching Wang, Glenn Ricart, Larry Landweber, Chip Elliott, Michael Zink, Emmanuel Cecchet, Snigdhaswin Kar, and Prabodh Mishra. 2019.  The Design and Operation of CloudLab. In _Proceedings of the USENIX Annual Technical Conference (ATC)_. 1–14.  https://www.flux.utah.edu/paper/duplyakin-atc19
  * Ferdinand et al. (2018) Nuwan Ferdinand, Haider Al-Lawati, Stark Draper, and Matthew Nokleby. 2018.  ANYTIME MINIBATCH: EXPLOITING STRAGGLERS IN ONLINE DISTRIBUTED OPTIMIZATION. In _International Conference on Learning Representations_.
  * Gujarati et al. (2020) Arpan Gujarati, Reza Karimi, Safya Alzayat, Wei Hao, Antoine Kaufmann, Ymir Vigfusson, and Jonathan Mace. 2020.  Serving DNNs like Clockwork: Performance Predictability from the Bottom Up. In _Proceedings of the 14th USENIX Conference on Operating Systems Design and Implementation_ _(OSDI’20)_. USENIX Association, USA, Article 25, 20 pages.
  * Harlap et al. (2018) Aaron Harlap, Deepak Narayanan, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur, Greg Ganger, and Phil Gibbons. 2018.  PipeDream: Fast and Efficient Pipeline Parallel DNN Training.  https://doi.org/10.48550/ARXIV.1806.03377
  * Harper and Konstan (2015) F. Maxwell Harper and Joseph A. Konstan. 2015.  The MovieLens Datasets: History and Context.  _ACM Trans. Interact. Intell. Syst._ 5, 4, Article 19 (dec 2015), 19 pages.  https://doi.org/10.1145/2827872
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2015.  Deep Residual Learning for Image Recognition.  https://doi.org/10.48550/ARXIV.1512.03385
  * He et al. (2017) Xiangnan He, Lizi Liao, Hanwang Zhang, Liqiang Nie, Xia Hu, and Tat-Seng Chua. 2017.  Neural Collaborative Filtering.  https://doi.org/10.48550/ARXIV.1708.05031
  * incubator ([n. d.]) Facebook incubator. [n. d.].  Facebookincubator/gloo: Collective Communications Library with various primitives for multi-machine training.  https://github.com/facebookincubator/gloo.
  * Jayaram Subramanya et al. (2023) Suhas Jayaram Subramanya, Daiyaan Arfeen, Shouxu Lin, Aurick Qiao, Zhihao Jia, and Gregory R. Ganger. 2023.  Sia: Heterogeneity-Aware, Goodput-Optimized ML-Cluster Scheduling. In _Proceedings of the 29th Symposium on Operating Systems Principles_ (Koblenz, Germany) _(SOSP ’23)_. Association for Computing Machinery, New York, NY, USA, 642–657.  https://doi.org/10.1145/3600006.3613175
  * Jiang et al. (2020) Yimin Jiang, Yibo Zhu, Chang Lan, Bairen Yi, Yong Cui, and Chuanxiong Guo. 2020.  A unified architecture for accelerating distributed $\\{$DNN$\\}$ training in heterogeneous $\\{$GPU/CPU$\\}$ clusters. In _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_. 463–479.
  * Jouppi et al. (2017) Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Daniel Killebrew, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le, Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. 2017.  In-Datacenter Performance Analysis of a Tensor Processing Unit.  _SIGARCH Comput. Archit. News_ 45, 2 (jun 2017), 1–12.  https://doi.org/10.1145/3140659.3080246
  * Keahey et al. (2020) Kate Keahey, Jason Anderson, Zhuo Zhen, Pierre Riteau, Paul Ruth, Dan Stanzione, Mert Cevik, Jacob Colleran, Haryadi S. Gunawi, Cody Hammock, Joe Mambretti, Alexander Barnes, François Halbach, Alex Rocha, and Joe Stubbs. 2020.  Lessons Learned from the Chameleon Testbed.  In _Proceedings of the 2020 USENIX Annual Technical Conference (USENIX ATC ’20)_. USENIX Association.
  * Kim et al. (2022) Jinpyo Kim, Hyungdal Kwon, Jintaek Kang, Jihwan Park, Seungwook Lee, and Jaejin Lee. 2022.  SnuHPL: High Performance LINPACK for Heterogeneous GPUs. In _Proceedings of the 36th ACM International Conference on Supercomputing_ (Virtual Event) _(ICS ’22)_. Association for Computing Machinery, New York, NY, USA, Article 18, 12 pages.  https://doi.org/10.1145/3524059.3532370
  * Krizhevsky (2009) Alex Krizhevsky. 2009.  Learning Multiple Layers of Features from Tiny Images.
  * Le et al. (2020) Tan N. Le, Xiao Sun, Mosharaf Chowdhury, and Zhenhua Liu. 2020.  AlloX: Compute Allocation in Hybrid Clusters. In _Proceedings of the Fifteenth European Conference on Computer Systems_ (Heraklion, Greece) _(EuroSys ’20)_. Association for Computing Machinery, New York, NY, USA, Article 31, 16 pages.  https://doi.org/10.1145/3342195.3387547
  * Lee et al. (2019) Sunwoo Lee, Qiao Kang, Sandeep Madireddy, Prasanna Balaprakash, Ankit Agrawal, Alok Choudhary, Richard Archibald, and Wei-keng Liao. 2019.  Improving Scalability of Parallel CNN Training by Adjusting Mini-Batch Size at Run-Time. In _2019 IEEE International Conference on Big Data (Big Data)_. 830–839.  https://doi.org/10.1109/BigData47090.2019.9006550
  * Li et al. (2022) Mingzhen Li, Wencong Xiao, Biao Sun, Hanyu Zhao, Hailong Yang, Shiru Ren, Zhongzhi Luan, Xianyan Jia, Yi Liu, Yong Li, Depei Qian, and Wei Lin. 2022.  EasyScale: Accuracy-consistent Elastic Training for Deep Learning.  https://doi.org/10.48550/ARXIV.2208.14228
  * Li et al. (2020) Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, et al. 2020\.  Pytorch distributed: Experiences on accelerating data parallel training.  _arXiv preprint arXiv:2006.15704_ (2020).
  * Luo et al. (2022) Liang Luo, Peter West, Pratyush Patel, Arvind Krishnamurthy, and Luis Ceze. 2022.  SRIFTY: Swift and Thrifty Distributed Neural Network Training on the Cloud. In _Proceedings of Machine Learning and Systems_ , D. Marculescu, Y. Chi, and C. Wu (Eds.), Vol. 4. 833–847.  https://proceedings.mlsys.org/paper/2022/file/f457c545a9ded88f18ecee47145a72c0-Paper.pdf
  * Mai et al. (2020) Luo Mai, Guo Li, Marcel Wagenländer, Konstantinos Fertakis, Andrei-Octavian Brabete, and Peter Pietzuch. 2020.  $\\{$KungFu$\\}$: Making Training in Distributed Machine Learning Adaptive. In _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_. 937–954.
  * Marín-Martínez and Sánchez-Meca (2010) Fulgencio Marín-Martínez and Julio Sánchez-Meca. 2010.  Weighting by Inverse Variance or by Sample Size in Random-Effects Meta-Analysis.  _Educational and Psychological Measurement_ 70, 1 (2010), 56–73.  https://doi.org/10.1177/0013164409344534 arXiv:https://doi.org/10.1177/0013164409344534
  * McCandlish et al. (2018) Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAI Dota Team. 2018.  An empirical model of large-batch training.  _arXiv preprint arXiv:1812.06162_ (2018).
  * Merkel (2014) Dirk Merkel. 2014.  Docker: lightweight linux containers for consistent development and deployment.  _Linux journal_ 2014, 239 (2014), 2.
  * Narayanan et al. (2020) Deepak Narayanan, Keshav Santhanam, Fiodar Kazhamiaka, Amar Phanishayee, and Matei Zaharia. 2020.  Heterogeneity-Aware Cluster Scheduling Policies for Deep Learning Workloads. In _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_. USENIX Association, 481–498.  https://www.usenix.org/conference/osdi20/presentation/narayanan-deepak
  * Nurvitadhi et al. (2018) Eriko Nurvitadhi, Jeffrey Cook, Asit Mishra, Debbie Marr, Kevin Nealis, Philip Colangelo, Andrew Ling, Davor Capalija, Utku Aydonat, Aravind Dasu, and Sergey Shumarayev. 2018.  In-Package Domain-Specific ASICs for Intel® Stratix® 10 FPGAs: A Case Study of Accelerating Deep Learning Using TensorTile ASIC. In _2018 28th International Conference on Field Programmable Logic and Applications (FPL)_. 106–1064.  https://doi.org/10.1109/FPL.2018.00027
  * Nvidia ([n. d.]) Nvidia. [n. d.].  Nvidia/NCCL: Optimized Primitives for collective multi-gpu communication.  https://github.com/NVIDIA/nccl.
  * Oehlert (1992) Gary W Oehlert. 1992.  A note on the delta method.  _The American Statistician_ 46, 1 (1992), 27–29.
  * Panayotov et al. (2015) Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. 2015.  Librispeech: An ASR corpus based on public domain audio books. In _2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_. 5206–5210.  https://doi.org/10.1109/ICASSP.2015.7178964
  * Park et al. (2020) Jay H. Park, Gyeongchan Yun, Chang M. Yi, Nguyen T. Nguyen, Seungmin Lee, Jaesik Choi, Sam H. Noh, and Young ri Choi. 2020.  HetPipe: Enabling Large DNN Training on (Whimpy) Heterogeneous GPU Clusters through Integration of Pipelined Model Parallelism and Data Parallelism. In _2020 USENIX Annual Technical Conference (USENIX ATC 20)_. USENIX Association, 307–321.  https://www.usenix.org/conference/atc20/presentation/park
  * Patarasuk and Yuan (2009a) Pitch Patarasuk and Xin Yuan. 2009a.  Bandwidth Optimal All-Reduce Algorithms for Clusters of Workstations.  _J. Parallel Distrib. Comput._ 69, 2 (feb 2009), 117–124.  https://doi.org/10.1016/j.jpdc.2008.09.002
  * Patarasuk and Yuan (2009b) Pitch Patarasuk and Xin Yuan. 2009b.  Bandwidth Optimal All-Reduce Algorithms for Clusters of Workstations.  _J. Parallel Distrib. Comput._ 69, 2 (feb 2009), 117–124.  https://doi.org/10.1016/j.jpdc.2008.09.002
  * Peng et al. (2018) Yanghua Peng, Yixin Bao, Yangrui Chen, Chuan Wu, and Chuanxiong Guo. 2018.  Optimus: An Efficient Dynamic Resource Scheduler for Deep Learning Clusters. In _Proceedings of the Thirteenth EuroSys Conference_ (Porto, Portugal) _(EuroSys ’18)_. Association for Computing Machinery, New York, NY, USA, Article 3, 14 pages.  https://doi.org/10.1145/3190508.3190517
  * Qi et al. (2016) Qi, Evan R. Sparks, and Ameet S. Talwalkar. 2016.  Paleo: A Performance Model for Deep Neural Networks. In _International Conference on Learning Representations_.
  * Qiao et al. (2021) Aurick Qiao, Sang Keun Choe, Suhas Jayaram Subramanya, Willie Neiswanger, Qirong Ho, Hao Zhang, Gregory R Ganger, and Eric P Xing. 2021.  Pollux: Co-adaptive cluster scheduling for goodput-optimized deep learning. In _15th $\\{$USENIX$\\}$ Symposium on Operating Systems Design and Implementation ($\\{$OSDI$\\}$ 21)_.
  * Rajpurkar et al. (2016) Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. 2016.  SQuAD: 100,000+ Questions for Machine Comprehension of Text. In _Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing_. Association for Computational Linguistics, Austin, Texas, 2383–2392.  https://doi.org/10.18653/v1/D16-1264
  * Schaller (1997) R.R. Schaller. 1997.  Moore’s law: past, present and future.  _IEEE Spectrum_ 34, 6 (1997), 52–59.  https://doi.org/10.1109/6.591665
  * Sergeev and Balso (2018) Alexander Sergeev and Mike Del Balso. 2018.  Horovod: fast and easy distributed deep learning in TensorFlow.  _CoRR_ abs/1802.05799 (2018).  arXiv:1802.05799 http://arxiv.org/abs/1802.05799
  * Sevilla et al. (2022) Jaime Sevilla, Lennart Heim, Anson Ho, Tamay Besiroglu, Marius Hobbhahn, and Pablo Villalobos. 2022.  Compute Trends Across Three Eras of Machine Learning. In _2022 International Joint Conference on Neural Networks (IJCNN)_. 1–8.  https://doi.org/10.1109/IJCNN55064.2022.9891914
  * Shoeybi et al. (2019) Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019.  Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.  https://doi.org/10.48550/ARXIV.1909.08053
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. 2017.  Attention is All you Need. In _Advances in Neural Information Processing Systems_ , I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (Eds.), Vol. 30. Curran Associates, Inc.  https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf
  * Wu et al. (2023) Bingyang Wu, Zili Zhang, Zhihao Bai, Xuanzhe Liu, and Xin Jin. 2023.  Transparent GPU Sharing in Container Clouds for Deep Learning Workloads. In _20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)_. USENIX Association, Boston, MA, 69–85.  https://www.usenix.org/conference/nsdi23/presentation/wu
  * Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Łukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. 2016\.  Google’s Neural Machine Translation System: Bridging the Gap between Human and Machine Translation.  https://doi.org/10.48550/ARXIV.1609.08144
  * Xiao et al. (2020) Wencong Xiao, Shiru Ren, Yong Li, Yang Zhang, Pengyang Hou, Zhi Li, Yihui Feng, Wei Lin, and Yangqing Jia. 2020.  AntMan: Dynamic Scaling on GPU Clusters for Deep Learning. In _14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)_. USENIX Association, 533–548.  https://www.usenix.org/conference/osdi20/presentation/xiao
  * Ye et al. (2022) Qing Ye, Yuhao Zhou, Mingjia Shi, Yanan Sun, and Jiancheng Lv. 2022.  DLB: A Dynamic Load Balance Strategy for Distributed Training of Deep Neural Networks.  _IEEE Transactions on Emerging Topics in Computational Intelligence_ (2022), 1–11.  https://doi.org/10.1109/TETCI.2022.3220224
  * You et al. (2022) Jie You, Jae-Won Chung, and Mosharaf Chowdhury. 2022.  Zeus: Understanding and Optimizing GPU Energy Consumption of DNN Training.  https://doi.org/10.48550/ARXIV.2208.06102
  * You et al. (2018) Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. 2018.  ImageNet Training in Minutes. In _Proceedings of the 47th International Conference on Parallel Processing_ (Eugene, OR, USA) _(ICPP 2018)_. Association for Computing Machinery, New York, NY, USA, Article 1, 10 pages.  https://doi.org/10.1145/3225058.3225069
  * Zhang et al. (2017) Haoyu Zhang, Logan Stafman, Andrew Or, and Michael J. Freedman. 2017.  SLAQ: Quality-Driven Scheduling for Distributed Machine Learning. In _Proceedings of the 2017 Symposium on Cloud Computing_ (Santa Clara, California) _(SoCC ’17)_. Association for Computing Machinery, New York, NY, USA, 390–404.  https://doi.org/10.1145/3127479.3127490
  * Zhang et al. (2019) Shuai Zhang, Lina Yao, Aixin Sun, and Yi Tay. 2019.  Deep Learning Based Recommender System: A Survey and New Perspectives.  _ACM Comput. Surv._ 52, 1, Article 5 (feb 2019), 38 pages.  https://doi.org/10.1145/3285029
  * Zhang et al. (2018) Xiaofan Zhang, Junsong Wang, Chao Zhu, Yonghua Lin, Jinjun Xiong, Wen-mei Hwu, and Deming Chen. 2018.  DNNBuilder: an Automated Tool for Building High-Performance DNN Hardware Accelerators for FPGAs. In _2018 IEEE/ACM International Conference on Computer-Aided Design (ICCAD)_. 1–8.  https://doi.org/10.1145/3240765.3240801

##  Appendix A Proof of Optimality Conditions

###  A.1. Compute-Bottleneck Scenario

###### Proof.

When computing is the bottleneck for all nodes, the total processing time of
one batch for the cluster is
$\max_{i\in\mathcal{N}}\\{t_{compute}^{i}+T_{u}\\}$. Since $T_{u}$ is the same
across all nodes, we can just consider
$\max_{i\in\mathcal{N}}\\{a_{i}+k_{i}b_{i}+m_{i}\\}$. Minimizing this quantity
is equivalent to solving the following optimization problem:

$\displaystyle\min\quad$ $\displaystyle\mu$ s.t. $\displaystyle
a_{i}+k_{i}b_{i}+m_{i}-\mu\leq 0,\quad\forall i\in\mathcal{N}$ $\displaystyle
B-\sum_{i\in\mathcal{N}}b_{i}=0.$

This optimization problem has corresponding Lagrangian:

$L(\mu,\mathbf{b},\mathbf{\lambda},\nu)=\mu+\sum_{i\in\mathcal{N}}\lambda_{i}(a_{i}+k_{i}b_{i}+m_{i}-\mu)+\nu\left(B-\sum_{i\in\mathcal{N}}b_{i}\right)$

Using the complimentary slackness conditions, we can solve and get
$\lambda_{i}=(1/k_{i})(\sum_{i\in\mathcal{N}}1/k_{i})^{-1}$. The Karush-Kuhn-
Tucker (KKT) conditions state that the optimal solution $\mu^{*}$ to this
problem must satisfy $\lambda_{i}(a_{i}+k_{i}b_{i}+m_{i}-\mu^{*})=0$. Since
$\lambda_{i}$ is strictly positive, $a_{i}+k_{i}b_{i}+m_{i}=\mu^{*},\forall
i\in\mathcal{N}$ so $t_{compute}^{i}=t_{compute}^{j},\forall
i,j\in\mathcal{N}$. ∎

###  A.2. Communication-Bottleneck Scenario

###### Proof.

When communication is the bottleneck for all nodes, the total processing time
of one batch is $\max_{i\in\mathcal{N}}\\{syncStart_{i}+T_{comm}\\}$. Since
$T_{comm}$ is the same across all nodes, we can just consider
$\max_{i\in\mathcal{N}}\\{a_{i}+\gamma(k_{i}b_{i}+m_{i})\\}$. Using the same
technique as in the previous proof, we now get that for the optimal solution
$\mu^{*}$, $a_{i}+\gamma(k_{i}b_{i}+m_{i})=\mu^{*},\forall i\in\mathcal{N}$.
Thus OptPerf must have $syncStart_{i}=syncStart_{j},\forall
i,j\in\mathcal{N}$. ∎

###  A.3. General Optimal Scenario

###### Proof.

Let $\mathcal{N}_{1}$ be the set of computation-bottleneck nodes in
$\mathcal{N}$ and let $\mathcal{N}_{2}$ be the set of communication-bottleneck
nodes in $\mathcal{N}$. Minimizing the total cluster processing time of one
batch is equivalent to the following optimization problem:

$\displaystyle\min\quad$ $\displaystyle\mu$ s.t. $\displaystyle
a_{i}+k_{i}b_{i}+m_{i}-\mu\leq 0,\quad\forall i\in\mathcal{N}_{1}$
$\displaystyle a_{i}+\gamma(k_{i}b_{i}+m_{i})+T_{o}-\mu\leq 0,\quad\forall
i\in\mathcal{N}_{2}$ $\displaystyle B-\sum_{i\in\mathcal{N}}b_{i}=0.$

If we construct the Lagrangian for this problem, we see that by solving the
complimentary slackness equations, the coefficients $\lambda_{i}$ are strictly
positive for all $i$. Thus the optimal solution $\mu^{*}$ satisfies
$a_{i}+k_{i}b_{i}+m_{i}=\mu^{*}$ for $i\in\mathcal{N}_{1}$ and satisfies
$a_{i}+\gamma(k_{i}b_{i}+m_{i})+T_{o}=\mu^{*}$ for $i\in\mathcal{N}_{2}$,
giving the desired result. ∎

##  Appendix B The GNS in Heterogeneous Clusters

Theorem 4.1.
_$\mathcal{G}=\sum_{i\in\mathcal{N}}w^{\mathcal{G}}_{i}\mathcal{G}_{i}$ and
$\mathcal{S}=\sum_{i\in\mathcal{N}}w^{\mathcal{S}}_{i}\mathcal{S}_{i}$ are
minimum variance, unbiased linear estimators of $|G|^{2}$ and
$\text{tr}(\Sigma)$ when:_

$\displaystyle\mathbf{w}^{\mathcal{G}}=\frac{\mathbf{1}^{T}A_{\mathcal{G}}^{-1}}{\mathbf{1}^{T}A_{\mathcal{G}}^{-1}\mathbf{1}},\quad\quad\mathbf{w}^{\mathcal{S}}=\frac{\mathbf{1}^{T}A_{\mathcal{S}}^{-1}}{\mathbf{1}^{T}A_{\mathcal{S}}^{-1}\mathbf{1}},$

_where $\mathbf{1}$ is an $n$-dimensional column vector of ones and both
$A_{\mathcal{G}}$ and $A_{\mathcal{S}}$ are $n\times n$ matrices with
respective entries $a_{\mathcal{G}}(i,j)$ and $a_{\mathcal{S}}(i,j)$:_

$a_{\mathcal{G}}(i,i)=\dfrac{B+2b_{i}}{B^{2}-Bb_{i}},\quad
a_{\mathcal{G}}(i,j)=\dfrac{B^{2}-b_{i}^{2}-b_{j}^{2}}{B(B-b_{i})(B-b_{j})}\text{
for }i\neq j$ $a_{\mathcal{S}}(i,i)=\dfrac{Bb_{i}}{B-b_{i}},\quad
a_{\mathcal{S}}(i,j)=\dfrac{b_{i}b_{j}(B-b_{i}-b_{j})}{(B-b_{i})(B-b_{j})}\text{
for }i\neq j$

###### Proof.

Since $\mathcal{G}_{i}$ is an unbiased estimator of $|G|^{2}$, $\mathcal{G}$
is an unbiased estimator when $\sum_{i\in\mathcal{N}}w_{i}^{\mathcal{G}}=1$.
Furthermore $\mathcal{G}$ is the minimum variance, unbiased linear estimator
when $\mathbf{w}$ minimizes the quadratic form of $\Sigma(\mathcal{G}_{i})$,
where $\Sigma(\mathcal{G}_{i})$ is the correlation matrix of the estimators
$\mathcal{G}_{i}$. Using Lagrange multipliers, we get:

$\mathbf{w}^{\mathcal{G}}=\frac{\mathbf{1}^{T}\Sigma(\mathcal{G}_{i})^{-1}}{\mathbf{1}^{T}\Sigma(\mathcal{G}_{i})^{-1}\mathbf{1}}$

Similarly, we get that $\mathcal{S}$ is the minimum variance, unbiased linear
estimator of $\text{tr}(\Sigma)$ when:

$\mathbf{w}^{\mathcal{S}}=\frac{\mathbf{1}^{T}\Sigma(\mathcal{S}_{i})^{-1}}{\mathbf{1}^{T}\Sigma(\mathcal{S}_{i})^{-1}\mathbf{1}}$

where $\Sigma(\mathcal{S}_{i})$ is the covariance matrix of the estimators
$\mathcal{S}_{i}$.

To compute $\mathbf{w}^{\mathcal{G}}$, we require $\Sigma(\mathcal{G}_{i})$.
By definition, the matrix’s diagonal elements are
$\text{Var}(\mathcal{G}_{i})$ and the off-diagonal elements are
$\text{Cov}(\mathcal{G}_{i},\mathcal{G}_{j})$ for $i\neq j$. Lemma B.1 gives
us $\text{Var}(\mathcal{G}_{i})$ and Lemma B.2 gives us
$\text{Cov}(\mathcal{G}_{i},\mathcal{G}_{j})$. Since all terms of
$\Sigma(\mathcal{G}_{i})$ have a common factor of $4|G|^{2}\text{tr}(\Sigma)$,
this factor will cancel for $\mathbf{w}^{\mathcal{G}}$, so we can equivalently
solve for $\mathbf{w}^{\mathcal{G}}$ using the matrix $A_{\mathcal{G}}$
instead of $\Sigma(\mathcal{G}_{i})$, where:

$a_{\mathcal{G}}(i,i)=\dfrac{B+2b_{i}}{B^{2}-Bb_{i}},\quad
a_{\mathcal{G}}(i,j)=\dfrac{B^{2}-b_{i}^{2}-b_{j}^{2}}{B(B-b_{i})(B-b_{j})}\text{
for }i\neq j$

We can use a similar argument for $\mathcal{S}$ with Lemmas B.4 and B.3, where
rather than using $\Sigma(\mathcal{S}_{i})$ we can use the matrix
$A_{\mathcal{G}}$ with entries:

$a_{\mathcal{S}}(i,i)=\dfrac{Bb_{i}}{B-b_{i}},\quad
a_{\mathcal{S}}(i,j)=\dfrac{b_{i}b_{j}(B-b_{i}-b_{j})}{(B-b_{i})(B-b_{j})}\text{
for }i\neq j$

∎

###### Lemma B.1.

The estimators $\mathcal{G}_{i}$ and $\mathcal{S}_{i}$ have variances:

$\displaystyle\text{Var}(\mathcal{G}_{i})$
$\displaystyle=4|G|^{2}\text{tr}(\Sigma)\left(\frac{B+2b_{i}}{B^{2}-Bb_{i}}\right)$
$\displaystyle\text{Var}(\mathcal{S}_{i})$
$\displaystyle=4|G|^{2}\text{tr}(\Sigma)\left(\frac{Bb_{i}}{B-b_{i}}\right)$

where $\text{tr}(\Sigma)$ is the sum of the variances of the individual
gradient components and $|G|^{2}$ is the global norm of the gradient.

###### Proof.

First we compute the variance of $\mathcal{G}_{i}$:

$\displaystyle\text{Var}(\mathcal{G}_{i})=\text{Var}(\frac{B}{B-b_{i}}|g|^{2}-\frac{b_{i}}{B-b_{i}}|g_{i}|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\left(\frac{B}{B-b_{i}}\right)^{2}\text{Var}(|g|^{2})+\left(\frac{b_{i}}{B-b_{i}}\right)^{2}\text{Var}(|g_{i}|^{2})-$
$\displaystyle-2\left(\frac{B}{B-b_{i}}\right)\left(\frac{b_{i}}{B-b_{i}}\right)\text{Cov}(|g|^{2},|g_{i}|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\left(\frac{B}{B-b_{i}}\right)^{2}\cdot\frac{4|G|^{2}\text{tr}(\Sigma)}{B}+\left(\frac{b_{i}}{B-b_{i}}\right)^{2}\cdot\frac{4|G|^{2}\text{tr}(\Sigma)}{b_{i}}-$
$\displaystyle-\frac{2Bb_{i}}{(B-b_{i})^{2}}\text{Cov}(|g|^{2},|g_{i}|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(3)}}{{=}}\frac{B+b_{i}}{(B-b_{i})^{2}}\left(4|G|^{2}\text{tr}(\Sigma)\right)-\frac{2Bb_{i}}{(B-b_{i})^{2}}\cdot\frac{4b_{i}|G|^{2}\text{tr}(\Sigma)}{B^{2}}=$
$\displaystyle=4|G|^{2}\text{tr}(\Sigma)\left(\frac{B+b_{i}}{(B-b_{i})^{2}}-\frac{2b_{i}^{2}}{B(B-b_{i})^{2}}\right)=$
$\displaystyle=4|G|^{2}\text{tr}(\Sigma)\left(\frac{B+2b_{i}}{B^{2}-Bb_{i}}\right)$

where $(1)$ follows from the variance of sums of random variables, $(2)$
follows from Lemma B.4 and $(3)$ follows from Lemma B.5. We can similarly
compute the variance of $\mathcal{S}_{i}$:

$\displaystyle\text{Var}(\mathcal{S}_{i})=\text{Var}(\frac{b_{i}B}{B-b_{i}}|g_{i}|^{2}-\frac{b_{i}B}{B-b_{i}}|g|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\left(\frac{b_{i}B}{B-b_{i}}\right)^{2}\left(\text{Var}(|g_{i}|^{2})+\text{Var}(|g|^{2})-2\text{Cov}(|g|^{2},|g_{i}|^{2})\right)=$
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\left(\frac{b_{i}B}{B-b_{i}}\right)^{2}\left(\frac{4|G|^{2}\text{tr}(\Sigma)}{b_{i}}+\frac{4|G|^{2}\text{tr}(\Sigma)}{B}-2\text{Cov}(|g|^{2},|g_{i}|^{2})\right)=$
$\displaystyle\stackrel{{\scriptstyle(3)}}{{=}}4|G|^{2}\text{tr}(\Sigma)\left(\frac{b_{i}B}{B-b_{i}}\right)^{2}\left(\frac{1}{b_{i}}+\frac{1}{B}-\frac{2b_{i}}{B}\right)=$
$\displaystyle=4|G|^{2}\text{tr}(\Sigma)\left(\frac{Bb_{i}}{B-b_{i}}\right)$

where again $(1)$ follows from the variance of sums of random variables, $(2)$
follows from Lemma B.4 and $(3)$ follows from Lemma B.5. ∎

###### Lemma B.2.

The estimators $\mathcal{G}_{i}$ and $\mathcal{G}_{j}$ have covariance:

$\text{Cov}(\mathcal{G}_{i},\mathcal{G}_{j})=4|G|^{2}\text{tr}(\Sigma)\frac{B^{2}-b_{i}^{2}-b_{j}^{2}}{B(B-b_{i})(B-b_{i})}$

###### Proof.

Using the definition of $\mathcal{G}_{i}$ and $\mathcal{G}_{j}$:

$\displaystyle\text{Cov}(\mathcal{G}_{i},\mathcal{G}_{j})=\text{Cov}\left(\frac{B|g|^{2}-b_{i}|g_{i}|^{2}}{B-b_{i}},\frac{B|g|^{2}-b_{j}|g_{j}|^{2}}{B-b_{j}}\right)=$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\frac{B^{2}\text{Var}(|g|^{2})}{(B-b_{i})(B-b_{j})}-\frac{Bb_{i}\text{Cov}(|g|^{2},|g_{i}|^{2})}{(B-b_{i})(B-b_{j})}-\frac{Bb_{j}\text{Cov}(|g|^{2},|g_{j}|^{2})}{(B-b_{i})(B-b_{j})}=$
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\frac{4|G|^{2}\text{tr}(\Sigma)}{(B-b_{i})(B-b_{j})}\left(\frac{1}{B}-\frac{b_{i}^{2}}{B}-\frac{b_{j}^{2}}{B}\right)$

where (1) follows from the covariance of linear combinations of random
variables and the independence of $g_{i}$ and $g_{j}$, and (2) follows from
Lemma B.4 and Lemma B.5. ∎

###### Lemma B.3.

The estimators $\mathcal{S}_{i}$ and $\mathcal{S}_{j}$ have covariance:

$\text{Cov}(\mathcal{S}_{i},\mathcal{S}_{j})=4|G|^{2}\text{tr}(\Sigma)\frac{b_{i}b_{j}(B-b_{i}-b_{j})}{(B-b_{i})(B-b_{j})}$

###### Proof.

Using the definition of $\mathcal{S}_{i}$ and $\mathcal{S}_{j}$:

$\displaystyle\text{Cov}(\mathcal{S}_{i},\mathcal{S}_{j})=\text{Cov}\left(\frac{Bb_{i}}{B-b_{i}}(|g_{i}|^{2}-|g|^{2}),\frac{Bb_{j}}{B-b_{j}}(|g_{j}|^{2}-|g|^{2})\right)=$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\frac{B^{2}b_{i}b_{j}}{(B-b_{i})(B-b_{j})}\left(\text{Var}(|g|^{2})-\text{Cov}(|g|^{2},|g_{i}|^{2})-\text{Cov}(|g|^{2},|g_{i}|^{2})\right)=$
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\frac{4|G|^{2}\text{tr}(\Sigma)B^{2}b_{i}b_{j}}{(B-b_{i})(B-b_{j})}\left(\frac{1}{B}-\frac{b_{i}}{B^{2}}-\frac{b_{j}}{B^{2}}\right)$

where (1) follows from the covariance of linear combinations of random
variables and the independence of $g_{i}$ and $g_{j}$, and (2) follows from
Lemma B.4 and Lemma B.5. ∎

###### Lemma B.4.

For any estimated gradient $g_{est}$ with corresponding batch size $b$, the
variance of the gradient norm satisfies:

$\text{Var}(|g_{est}|^{2})\approx\frac{4|G|^{2}\text{tr}(\Sigma)}{b}$

###### Proof.

Error propagation using Taylor’s rule (Oehlert, 1992) (also known as the delta
method) gives us the approximation:

$\text{Var}(|g_{est}|^{2})\approx
4\mathbb{E}[|g_{est}|]^{2}\text{Var}(|g_{est}|)=4|G|^{2}\cdot\frac{1}{b}\text{tr}(\Sigma)$

∎

###### Lemma B.5.

For local gradient $g_{i}$ at node $i$ with batch size $b_{i}$ and global
gradient $g$ with batch size $B$ and computed using Eq (9), the covariance of
the two gradient norms is:

$\text{Cov}(|g|^{2},|g_{i}|^{2})=\frac{4b_{i}|G|^{2}\text{tr}(\Sigma)}{B^{2}}$

###### Proof.

The global gradient norm $|g|^{2}$ can be written in terms of $g_{i}$ and
non-$g_{i}$ components:

$\displaystyle|g|^{2}$ $\displaystyle=\frac{1}{B^{2}}\sum_{j\in
b_{i}}(\nabla_{\theta}L_{x_{j}}(\theta))^{2}+\frac{1}{B^{2}}\sum_{j\notin
b_{i}}(\nabla_{\theta}L_{x_{j}}(\theta))^{2}=$
$\displaystyle=\frac{b_{i}^{2}}{B^{2}}|g_{i}|^{2}+\frac{1}{B^{2}}\sum_{j\notin
b_{i}}(\nabla_{\theta}L_{x_{j}}(\theta))^{2}$

Rewriting the covariance using this expression:

$\displaystyle\text{cov}(|g|^{2},|g_{i}|^{2})=\text{cov}(\frac{b_{i}^{2}}{B^{2}}|g_{i}|^{2}+\frac{1}{B^{2}}\sum_{j\notin
b_{i}}(\nabla_{\theta}L_{x_{j}}(\theta))^{2},|g_{i}|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(1)}}{{=}}\frac{b_{i}^{2}}{B^{2}}\text{cov}(|g_{i}|^{2},|g_{i}|^{2})+\frac{1}{B}\text{cov}(\sum_{j\notin
b_{i}}(\nabla_{\theta}L_{x_{j}}(\theta))^{2},|g_{i}|^{2})=$
$\displaystyle\stackrel{{\scriptstyle(2)}}{{=}}\frac{b_{i}^{2}}{B^{2}}\text{Var}(|g_{i}|^{2})\stackrel{{\scriptstyle(3)}}{{=}}\frac{b_{i}^{2}}{B^{2}}\cdot\frac{4|G|^{2}\text{tr}(\Sigma)}{b_{i}}=\frac{4b_{i}|G|^{2}\text{tr}(\Sigma)}{B^{2}}$

where $(1)$ follows from the covariance of linear combinations of random
variables, $(2)$ follows from the variance-covariance relationship and the
independence of $g_{i}$ and $g_{j}$ for $i\neq j$ and $(3)$ follows from Lemma
B.4. ∎