A Communication-Efficient Distributed Gradient Clipping Algorithm for Training Deep Neural Networks
Anonymous Author(s)
AffiliationAddress email
Abstract
In distributed training of deep neural networks, people usually run Stochastic
Gradient Descent (SGD) or its variants on each machine and communicate withother machines periodically.
However, SGD might converge slowly in trainingsome deep neural networks (e.g., RNN, LSTM) because of the exploding gradientissue.
Gradient clipping is usually employed to address this issue in the singlemachine setting, but exploring this technique in the distributed setting is still inits infancy: it remains mysterious whether the gradient clipping scheme can takeadvantage of multiple machines to enjoy parallel speedup.
The main technicaldifficulty lies in dealing with nonconvex loss function, non-Lipschitz continuousgradient, and skipping communication rounds simultaneously.
In this paper, weexplore a relaxed-smoothness assumption of the loss landscape which LSTM wasshown to satisfy in previous works, and design a communication-efficient gradientclipping algorithm.
This algorithm can be run on multiple machines, where eachmachine employs a gradient clipping scheme and communicate with other machinesafter multiple steps of gradient-based updates.
Our algorithm is proved to have
O (cid:0)
Nϵ 4 (cid:1) iteration complexity and O ( 1 ϵ 3 ) communication complexity for findingan ϵ -stationary point in the homogeneous data setting, where N is the number ofmachines.
This indicates that our algorithm enjoys linear speedup and reducedcommunication rounds.
Our proof relies on novel analysis techniques of estimatingtruncated random variables, which we believe are of independent interest.
Ourexperiments on several benchmark datasets and various scenarios demonstrate thatour algorithm indeed exhibits fast convergence speed in practice and thus validatesour theory.
Introduction
Deep learning has achieved tremendous successes in many domains including computer vision [23,15], natural language processing [5] and game [39].
To obtain good empirical performance, peopleusually need to train large models on a huge amount of data and it is usually very computationallyexpensive.
To speed up the training process, distributed training becomes indispensable [4].
Forexample, Goyal et al.
[12] trained a ResNet-50 on ImageNet dataset by distributed SGD withminibatch size 8192 on 256 GPUs in only one hour, which not only matches the small minibatchaccuracy but also enjoys parallel speedup, and hence improves the running time.
Recently, local
SGD [40, 47], as a variant of distributed SGD, achieved tremendous attention in federated learningcommunity [29].
The local SGD algorithm runs multiple steps of SGD on each machine beforecommunicating with other clients.
Despite the empirical success of distributed SGD and its variants (e.g., local SGD) in deep learning,they may not exhibit good performance when training some neural networks (e.g., Recurrent Neural
Networks, LSTMs), due to the exploding gradient problem [33, 34].
To address this issue, Pascanuet al.
[34] proposed to use the gradient clipping strategy, and it has become a standard technique
[Table] [Table]
when training language models [8, 35, 31].
There are some recent works trying to theoreticallyexplain gradient clipping from the perspective of non-convex optimization [51, 50].
These works arebuilt upon an important observation made in [51]: for certain neural networks such as LSTM
[Table]
thegradient does not vary uniformly over the loss landscape (i
[Table]
e., the gradient is not Lipschitz continuouswith a uniform constant)
[Table]
and the gradient Lipschitz constant can scale linearly with respect to thegradient norm.
This is referred to as the relaxed smoothness condition (i.e., ( L 0 , L 1 ) -smoothnessdefined in Definition 2), which generalizes but strictly relaxes the usual smoothness condition (i.e.,
L -smoothness defined in Definition 1).
Under the relaxed smoothness condition, Zhang et al.
[51, 50]proved that gradient clipping enjoys polynomial-time iteration complexity for finding the first-orderstationary point in the single machine setting, and it can be arbitrarily faster than fix-step gradientdescent.
In practice, both distributed learning and gradient clipping are important techniques toaccelerate neural network training.
However, the theoretical analysis of gradient clipping is onlyrestricted to the single machine setting [51, 50].
Hence it naturally motivates us to consider thefollowing question:
Is it possible that the gradient clipping scheme can take advantage of multiple machines toenjoy parallel speedup in training deep neural networks, with only infrequent communication?
In this paper, we give an affirmative answer to the above question.
Built upon the relaxed smoothnesscondition as in [51, 50], we design a communication-efficient distributed gradient clipping algorithm.
The key characteristics of our algorithm are: (i) unlike naive parallel gradient clipping algorithm 1which requires averaging model weights and gradients from all machines for every iteration, ouralgorithm only aggregates weights with other machines after a certain number of local updateson each machine; (ii) our algorithm clips the gradient according to the norm of the local gradienton each machine, instead of the norm of the averaged gradients across machines as in the naiveparallel algorithm.
These key features make our algorithm amenable to the distributed setting withinfrequent communication, and it is nontrivial to establish desired theoretical guarantees (e.g., linearspeedup, reduced communication complexity).
The main difficulty in the analysis lies in dealing withthe nonconvex objective function, non-Lipschitz continuous gradient, and skipping communicationrounds simultaneously.
Our main contribution is summarized as the following:
• We design a novel communication-efficient distributed stochastic local gradient clippingalgorithm,
namely CELGC, for solving a nonconvex optimization problem under the relaxedsmoothness condition.
The algorithm only needs to clip the gradient according to the localgradient’s magnitude and globally averages the weights on all machines periodically.
To thebest of our knowledge, this is the first work proposing communication-efficient distributedstochastic gradient clipping algorithms under the relaxed smoothness condition.• Under the relaxed smoothness condition, we prove iteration and communication complexityresults of our algorithm for finding an ϵ -stationary point, when each machine has access to the same data distribution.
First, comparing with [50], we prove that our algorithm enjoyslinear speedup, which means that the iteration complexity of our algorithm is reduced by afactor of N (the number of machines).
Second, comparing with the naive parallel version ofthe algorithm of [50], we prove that our algorithm enjoys better communication complexity.
The detailed comparison over existing algorithms under the same relaxed smoothnesscondition is described in Table 1.
To achieve this result, we introduce a novel techniqueof estimating truncated random variables, which is of independent interest and could beapplied in related problems involving truncation operations such as gradient clipping.• We empirically verify our theoretical results by conducting experiments on different neuralnetwork architectures on benchmark datasets.
The experimental results demonstrate that ourproposed algorithm indeed exhibits speedup in practice.
Related Work
Gradient Clipping/Normalization Algorithms
In deep learning literature, gradient clipping(normalization) technique was initially proposed by [34] to address the issue of exploding gradientproblem in [33], and it has become a standard technique when training language models [8, 35, 31].Itis shown that gradient clipping is robust and can mitigate label noise [30].
Recently gradientnormalization techniques [45, 46] were applied to train deep neural networks on the very large batchsetting.
For example, You et al.
[45] designed the LARS algorithm to train a ResNet50 on ImageNetwith batch size 32k, which utilized different learning rates according to the norm of the weightsand the norm of the gradient.
In optimization literature, gradient clipping (normalization) was usedin early days in the field of convex optimization [7, 1, 38].
Nesterov [32] and Hazan et al.
[14]considered normalized gradient descent for quasi-convex functions in deterministic and stochasticcases respectively.
Gorbunov et al.
[10] designed an accelerated gradient clipping method to solveconvex optimization problems with heavy-tailed noise in stochastic gradients.
Mai and Johansson [26]established the stability and convergence of stochastic gradient clipping algorithms for convex andweakly convex functions.
In nonconvex optimization, Levy [24] showed that normalized gradientdescent can escape from saddle points.
Cutosky and Mehta [3] showed that adding a momentumprovably improves the normalized SGD in nonconvex optimization.
Zhang et al.
[51] and Zhanget al.
[50] analyzed the gradient clipping for nonconvex optimization under the relaxed smoothnesscondition rather than the traditional L -smoothness condition in nonconvex optimization [9].
However,all of them only consider the algorithm in the single machine setting or the naive parallel setting, andnone of them can apply to the distributed setting where only limited communication is allowed.
Communication-Efficient Algorithms in Distributed and Federated Learning In large-scalemachine learning, people usually train their model using first-order methods on multiple machinesand these machines communicate and aggregate their model parameters periodically.
When thefunction is convex, there is a scheme named one-shot averaging [55, 28, 53, 37], in which everymachine runs a stochastic approximation algorithm and averages the model weights across machinesonly at the very last iteration.
The one-shot averaging scheme is communication-efficient and enjoysstatistical convergence with one pass of the data [53, 37, 16, 22], but the training error may notconverge in practice.
Mcmahan et al.
[29] considered the Federated Learning setting where thedata is decentralized and might be non-i.i.d.
across devices and communication is expensive anddesigned the very first algorithm for federated learning (a.k.a., FedAvg).
Stich [40] considered aconcrete case of FedAvg, namely local SGD, which runs SGD independently in parallel on differentworks and averages the model parameters only once in a while.
Stich [40] also showed that local
SGD enjoys linear speedup for strongly-convex objective functions.
There are also some worksanalyzing local SGD and its variants on convex [6, 20, 19, 43, 44, 11, 49] and nonconvex smoothfunctions [54, 47, 48, 17, 41, 25, 2, 13, 19].
Recently, Woodworth et al.
[43, 44] analyzed advantagesand drawbacks of local SGD compared with minibatch SGD for convex objectives.
Woodworthet al.
[42] proved hardness results for distributed stochastic convex optimization.
Reddi et al.
[36]introduced a general framework of federated optimization and designed several federated versionsof adaptive optimizers.
Zhang et al.
[52] considered employing gradient clipping to optimize Lsmooth functions and achieve differential privacy.
Koloskova et al.
[21] developed a unified theory ofdecentralized SGD with changing topology and local updates for smooth functions.
Due to a vastamount of literature on federated learning and limited space, we refer readers to [18] and referencestherein.
However, all of these works either assume the objective function is convex or L -smooth.
Tothe best of our knowledge, our algorithm is the first communication-efficient algorithm that does notrely on these assumptions but still enjoys linear speedup.
Preliminaries, Notations and Problem Setup
Preliminaries and Notations Denote ∥ · ∥ by the Euclidean norm.
We denote f : R d → R asthe overall loss function, and F : R d → R as the loss function on i -th machine, where i ∈ [ N ] :={ 1 , . . .
, N } .
Denote ∇ h ( x ) as the gradient of h evaluated at the point x , and denote ∇ h ( x ; ξ ) as thestochastic gradient of h calculated based on sample ξ .
Definition 1 ( L -smoothness) .
A function h is L -smooth if ∥∇ h ( x ) − ∇ h ( y ) ∥
≤ L ∥ x − y ∥ for allx , y ∈ R d .
Definition 2 ( ( L 0 , L 1 ) -smoothness) .
A second order differentiable function h is ( L 0 , L 1 ) -smooth if∥∇ 2 h ( x ) ∥ ≤ L 0 + L 1 ∥∇ h ( x ) ∥ for any x ∈ R d .
Definition 3 ( ϵ -stationary point) .
x ∈ R
d is an ϵ -stationary point of the function h if ∥∇ h ( x ) ∥ ≤ ϵ .
Remark: ( L 0 , L 1 ) -smoothness is strictly weaker than L -smoothness.
First, we know that L -smoothfunctions is ( L 0 , L 1 ) -smooth with L 0 =
L and L 1 = 0 .
However the reverse is not true.
For example,consider the function h
( x ) = x 4 , we know that the gradient is not Lipschitz continuous and henceh is not L -smooth, but | h ′′ ( x ) | = 12 x 2 ≤ 12 + 3 × 4 | x | 3 = 12 + 3 | h
( x ) | , so h ( x ) = x 4 is(12 , 3) -smooth.
Zhang et al.
[51] empirically showed that the ( L 0 , L 1 ) -smoothness holds for the
AWD-LSTM [31].
In nonconvex optimization literature [9, 50], the goal is to find an ϵ -stationarypoint since it is NP-hard to find a global optimal solution for a general nonconvex function.
Problem Setup
In this paper, we consider the following optimization problem using N machines:
[Equation]
where D stands for the data distribution which each machine has access to, and f is the populationloss function.
We make the following assumptions throughout the paper.
Assumption 1.
(i) f
( x ) is ( L 0 , L 1 ) -smooth, i.e., ∥∇ 2 f ( x ) ∥
≤ L 0 + L 1 ∥∇ f ( x ) ∥ , for ∀ x ∈
R d .
(ii) There exists ∆ > 0 such that f ( x 0 ) − f ∗ ≤ ∆ , where f ∗ is the global optimal value of f .
(iii)
For all x ∈ R d , E ξ ∼D [ ∇ F ( x ; ξ )] = ∇ f ( x ) , and ∥∇ F ( x ; ξ ) − ∇ f ( x ) ∥ ≤ σ almost surely.
(iv) The distribution of ∇ F ( x ; ξ ) is symmetric around its mean ∇ f ( x ) , and the density ismonotonically decreasing over the ℓ 2 distance between the mean and the value of random
variable.
Remark:
The Assumption 1 (i) means that the loss function satisfies the relaxed-smoothnesscondition, and it holds when training a language model with LSTMs.
We consider the homogeneousdistributed learning setting, where each machine has access to same data distribution as in [44].
Assumption 1 (ii) and (iii) are standard assumptions in nonconvex optimization [9, 51].
Note thatit is usually assumed that the stochastic gradient is unbiased and has bounded variance [9], but inthe relaxed smoothness setting, we follow [51] to assume we have unbiased stochastic gradient withalmost surely bounded deviation σ .
Assumption 1 (iv) assumes the noise is unimodal and symmetricaround its mean, which we empirically verify in Appendix E. Examples satisfying (iii) and (iv)include truncated Gaussian distribution, truncated student’s t-distribution, etc.
Algorithm and Theoretical Analysis
4.1 Main Difficulty and the Algorithm Design
We briefly present the main difficulty in extending the single machine setting [50] to the distributedsetting.
In [50], they split the contribution of decreasing objective value by considering two cases:clipping large gradients and keeping small gradients.
If communication is allowed at every iteration,then we can aggregate gradients on each machine and determine whether we should clip or keepthe averaged gradient or not.
However, in our setting, communicating with other machines at everyiteration is not allowed.
This would lead to the following difficulties: (i) the averaged gradient maynot be available to the algorithm if communication is limited, so it is hard to determine whetherclipping operation should be performed or not; (ii) the model weight on every machine may not be
1: for t = 0 , . . .
, T do 2:
Each node i samples its stochastic gradient ∇ F ( x it ; ξ it ) , where ξ it ∼ D .
Each node i updates it local solution in parallel :
[Equation]
4: if t is a multiple of I then 5:
Each worker resets the local solution as the averaged solution across nodes:
[Equation]
6: end if 7: end for the same when communication does not happen at the current iteration; (iii) the loss function is not
L -smooth, so the usual local SGD analysis for L -smooth functions cannot be applied in this case.
To address this issue, we design a new algorithm, namely Communication-Efficient Local Gradient
Clipping (CELGC), which is presented in Algorithm 1.
The algorithm calculates a stochastic gradientand then performs multiple local gradient clipping steps on each machine in parallel, and aggregatesmodel parameters on all machines after every I steps of local updates.
We aim to establish iterationand communication complexity for Algorithm 1 for finding an ϵ -stationary point when I > 1 .
4.2 A Lemma for Truncated Random Variables
As indicated in the previous subsection, Algorithm 1 clips gradient on each local machine.
Thisfeature greatly complicates the analysis: it is difficult to get an unbiased estimate of the stochasticgradient when its magnitude is not so large such that clipping does not happen.
The reason is due tothe dependency between random variables (i.e., stochastic gradient and the indicator of clipping).
To get around of this difficulty, we introduce the following Lemma for estimating truncated randomvariables.
Lemma 1.
Denote by g ∈ R d a random vector.
Suppose the distribution of g is symmetric around itsmean, and the density is monotonically decreasing over the ℓ 2 distance between the mean and thevalue of random variable, then there exists Λ = diag ( c 1 , . .
c d ) and 0 <
c i ≤ 1 with i = 1 , . . .
, dsuch that
[Equation]
where α > 0 is a constant, I ( · ) is the indicator function.
The proof of Lemma 1 is included in Appendix B.
This Lemma provides an unbiased estimate fora truncated random variable.
In the subsequent of this paper, we regard g in the Lemma as thestochastic gradient and α as the clipping threshold.
In addition, we define c min = min( c 1 , . . .
, c d ) ,and c max = max( c 1 , . . .
, c d ) .
We have 0 < c min ≤ c max ≤ 1 .
4.3 Main Results
Theorem 1.
Suppose Assumption 1 holds and σ ≥ 1 .
Take ϵ ≤ min( AL 0 BL 1 , 0 . 1) be a smallenough constant and N ≤ min( 1 ϵ , 14 AL 0 5 BL 1 ϵ ) .
In Algorithm 1, choose I ≤ (cid:113) 1 c min σNϵ , γ ≤c min
Nϵ 28 σ min { ϵAL 0 , 1 BL 1 } and the fixed ratio γη = 5 σ , where A ≥ 1 and B ≥ 1 are constantswhich will be specified in the proof, and run Algorithm 1 for T
= O (cid:16)
∆ L 0 σ 2 Nϵ 4 (cid:17) iterations.
Define¯ x t = 1 N (cid:80) Ni
=1 x it .
Then we have 1 T (cid:80)
Tt =
1 E ∥∇ f
(¯ x t ) ∥ ≤ 9 ϵ.
Remark: We have some implications of Theorem 1.
When the number of machines is not large(i.e.
, N ≤ O (1 /ϵ ) ) and the number of skipped communications is not large (i.e., I ≤ O ( σ/ϵN ) ), then with proper setting of the learning rate, we have following observations.
First, our algorithmenjoys linear speedup, since the number of iterations we need to find an ϵ -stationary point is dividedby the number of machines N when comparing the single machine algorithm in [50].
Second, ouralgorithm is communication-efficient, since the communication complexity (a.k.a., number of rounds)is T/I
= O (cid:0) ∆ L 0 σϵ − 3 (cid:1) , which provably improves the naive parallel gradient clipping algorithmof [50] with O (∆ L 0 σ 2 / ( Nϵ 4 )) communication complexity when N ≤ O (1 /ϵ ) .
Another important fact is that both iteration complexity and communication complexity only dependon L 0 while being independent of L 1 and the gradient upper bound M .
This indicates that ouralgorithm does not suffer from slow convergence even if these quantities are large.
This is in linewith [50] as well.
In addition, local gradient clipping is a good mechanism to alleviate the bad effectsbrought by a rapidly changing loss landscape (e.g., some language models such as LSTM).
4.4 Sketch of the Proof of Theorem 1
In this section, we present the sketch of our proof of Theorem 1.
The detailed proof can be found in
Appendix C. The key idea in our proof is to establish the descent property of the sequence { f (¯ x t ) }
Tt =0in the distributed setting under the relaxed smoothness condition, where ¯ x t = 1 N (cid:80) ti =1 x it is theaveraged weight across all machines at t -th iteration.
The main challenge is that the descent propertyof ( L 0 , L 1 ) -smooth function in the distributed setting does not naturally hold, which is in sharpcontrast to the usual local SGD proof for L -smooth functions.
To address this challenge, we need tocarefully study whether the algorithm is able to decrease the objective function in different situations.
Our main technical innovations in the proof are listed as the following.
First, we monitor the algorithm’s progress in decreasing the objective value according to some novelmeasures.
The measures we use are the magnitude of the gradient evaluated at the averaged weightand the magnitude of local gradients evaluated at the individual weights on every machine.
To thisend, we introduce Lemma 3, whose goal is to carefully inspect how much progress the algorithmmakes, according to the magnitude of local gradients calculated on each machine.
The reason is thatthe local gradient’s magnitude is an indicator of whether the clipping operation happens or not.
Foreach fixed iteration t , we define J ( t ) =
{ i ∈ [ N ] : ∥∇ F ( x it , ξ it )
∥ ≥ γ/η } and ¯ J ( t ) =
[ N ] \ J ( t ) .
Briefly speaking, J ( t ) contains all machines that perform clipping operation at iteration t and ¯ J ( t )is the set of machines that do not perform clip operation at iteration t .
In Lemma 3, we performthe one-step analysis and consider all machines with different clipping behaviors at the iteration t .
The proof of Lemma 3 crucially relies on the lemma for estimating truncated random variables (i.e.,
Lemma 1) to calculate the expectation of non-clipped gradients.
Second, Zhang et al.
[50] inspect their algorithm’s progress by considering the magnitude of thegradient at different iterations, so they treat every iteration differently.
However, this approach doesnot work in the distributed setting with infrequent communication since one cannot get access to theaveraged gradient across machines at every iteration.
Instead, we treat every iteration of the algorithmas the same but consider the progress made by each machine.
Third, by properly choosing hyperparameters ( η, γ, I ) and using an amortized analysis, we prove thatour algorithm can decrease the objective value by a sufficient amount, and the sufficient decrease ismainly due to the case where the gradient is not too large (i.e., clipping operations do not happen).
This important insight allows us to better characterize the training dynamics.
Now we present how to proceed with the proof in detail.
Lemma 2 characterizes the ℓ 2 error between averaged weight and individual weights at t -th iteration.
Lemma 2.
Under Assumption 1, for any i and t , Algorithm 1 ensures ∥ ¯ x t − x it ∥ ≤ 2 γI holds almostsurely.
Lemma 5 and Lemma 6 (in Appendix A) are properties of ( L 0 , L 1 )
-smooth functions we need touse.
To make sure they work, we need 2 γI ≤ c/L 1 for some c > 0 .
This is proved in the proof of
Theorem 1 (in Appendix C).
Let J ( t ) be the index set of i such that ∥∇ F ( x it , ξ it )
∥ ≥ γη at fixed iteration t , i.e., J ( t ) =
{ i ∈[ N ] | ∥∇ F ( x it ; ξ it )
∥ ≥ γ/η } .
Lemma 3 characterizes how much progress we can get in oneiteration of Algorithm 1, which is decomposed into contributions from every machine (note that
J ( t ) ∪ ¯ J ( t ) = { 1 , . . .
, N } for every t ).
Lemma 3.
Let J ( t ) be the set defined as above.
If 2 γI ≤ c/L 1 for some c > 0 , AL 0 η ≤ 1 / 2 , and 262 γ/η = 5 σ , then we have
E [ f (¯ x t +1 ) − f (¯ x t )]
[Equation]
[Table]
[Equation]
[Table]
where A = 1 + e c − e c − 1 c and B = e c − 1 c .
Lemma 4 quantifies an upper bound of the averaged ℓ 2 error between the local gradient evaluated atthe local weight and the gradient evaluated at the averaged weight.
Lemma 4.
Suppose Assumption 1 holds.
When 2 γI ≤ c/L 1 for some c > 0 , the following inequalityholds for every i almost surely with A = 1 + e c − e c − 1 c and B = e c − 1 c :
[Equation]
Putting all together Suppose our algorithm runs T iterations.
Taking summation on bothsides of Lemma 3 over all t = 0 , . .
, T − 1 , we are able to get an upper bound of 270 (cid:80) T −t =
E [ f (¯ x t +1 ) − f (¯ x t )]
=
E [ f (¯ x T ) − f
(¯ x 0 )] .
Note that E [ f (¯ x T ) − f
(¯ x 0 )] ≥ − ∆ due to
Assumption 1, so we are able to get an upper bound of gradient norm.
For details, please refer to theproof of Theorem 1 in Appendix C.
Experiments
We conduct extensive experiments to validate the merits of our algorithm on various tasks (e.g.,image classification, language modeling).
In the main text below, we consider the homogeneousdata setting where each machine has the same data distribution and every machine participate thecommunication at each round.
In the Appendix D, we also test our algorithm in the general federatedlearning setting (e.g., partial participation of machines), ablation study on small/large batch-sizes,and other experiments, etc.
We conduct each experiment in two nodes with 4 Nvidia-V100 GPUs on each node.
In our experiments, one “machine” corresponds to one GPU, and we use the word “GPU” and “machine” in thissection interchangeably.
We compared our algorithm with the baseline across three deep learningbenchmarks: CIFAR-10 image classification with ResNet, Penn Treebank language modeling with
LSTM, and Wikitext-2 language modeling with LSTM.
All algorithms and the training frameworkare implemented in Pytorch 1.4.
Due to limited computational resources, for our algorithms, wechoose same hyperparameters like the clipping thresholds according to the best-tuned baselines unlessotherwise specified.
We compare our algorithm of different I with the baseline, which is the naive parallel version of thealgorithm in [50].
We want to re-emphasize that the difference is that the baseline algorithm needs toaverage the model weights and local gradients at every iteration while ours only requires averagingthe model weights after every I iterations.
We find that our algorithm with a range of values of I canmatch the results of the baseline in terms of epochs on different models and data.
This immediatelysuggests that our algorithm will gain substantial speedup in terms of the wall clock time, which isalso supported by our experiments.
5.1 Effects of Skipping Communication
We focus on one feature of our algorithm: skipping communication rounds.
Theorem 1 says thatour algorithm enjoys reduced communication complexity since every node only communicates withother nodes periodically with node synchronization interval length I .
To study how communicationskipping affects the convergence of Algorithm 1, we run it with I ∈ { 2 , 4 , 8 , 16 , 32 } .
CIFAR-10 classification with ResNet-56.
We train the standard 56-layer ResNet [15] architecture on
CIFAR-10.
We use SGD with clipping as the baseline algorithm with a stagewise decaying learningrate schedule, following the widely adopted fashion on training the ResNet architecture.
Specifically,
[Figure] [Figure]
we use the initial learning rate η = 0 .
3 , the clipping threshold γ = 1 . 0 , and decrease the learningrate by a factor of 10 at epoch 80 and 120 .
The local batch size at each GPU is 64 .
These parametersettings follow that of [47].
The results are illustrated in Figure 1.
Figure 1a shows the convergence of training loss and testaccuracy v.s.
the number of epochs that are jointly accessed by all GPUs.
This means that, if thex-axis value is 8, then each GPU runs 1 epoch of training data.
The same convention applied to allother figures for multiple GPU training in this paper.
Figure 1b verifies our algorithm’s advantageof skipping communication by plotting the convergence of training loss and test accuracy v.s. thewall clock time.
Overall, we can clearly see that our algorithm matches the baseline epoch-wise butgreatly speeds up wall-clock-wise.
Language modeling with LSTM on Penn Treebank.
We adopt the 3-layer AWD-LSTM [31] to dolanguage modeling on Penn Treebank (PTB) dataset [27](word level).
We use SGD with clippingas the baseline algorithm with the clipping threshold γ = 7 . 5 .
The local batch size at each GPUis 3.
These parameter settings follow that of [50].
We fine-tuned the initial learning rate η for allalgorithms (including the baseline) by choosing the one giving the smallest final training loss in therange { 0 . 1 , 0 . 5 , 1 , 5 , 10 , 20 , 30 , 40 , 50 , 100 } .
We report the results in Figure 2.
It can be seen that we can match the baseline in both trainingloss and validation perplexity epoch-wise while gaining substantial speedup (4x faster for I = 4 )wall-clock-wise.
Language modeling with LSTM on Wikitext-2.
We adopt the 3-layer AWD-LSTM [31] to dolanguage modeling on Wikitext-2 dataset [27](word level).
We use SGD with clipping as the baselinealgorithm with the clipping threshold γ = 7 . 5 .
The local batch size at each GPU is 10.
Theseparameter settings follow that of [31].
We fine-tuned the initial learning rate η for all algorithms(including the baseline) by choosing the one giving the smallest final training loss in the range{ 0 . 1 , 0 . 5 , 1 , 5 , 10 , 20 , 30 , 40 , 50 , 100 } .
We report the results in Figure 3.
We can match the baseline in both training loss and validationperplexity epoch-wise, but we again obtain large speedup (4x faster for I = 4 ) wall-clock-wise.
This,together with the above two experiments, clearly show our algorithm’s effectiveness in speeding upthe training in distributed settings.
Another observation is that Algorithm 1 can allow relatively large
I without hurting the convergence behavior.
[Figure]
[Figure]
in which clipping is triggered v.s. epochs showing clipping is very frequent.
5.2 Verifying Parallel Speedup
Figure 4 show the training loss and test accuracy v.s. the number of iterations.
In the distributedsetting, one iteration means running one step of Algorithm 1 on all machines; while in the singlemachine setting, one iteration means running one step of SGD with clipping.
In our experiment, weuse minibatch size 64 on every GPU in the distributed setting to run Algorithm 1, while we also use64 minibatch size on the single GPU to run SGD with clipping.
In the left two panels of Figure 4, wecan clearly find that even with I > 1 , our algorithm still enjoys parallel speedup, since our algorithmrequires less number of iterations to converge to the same targets (e.g., training loss, test accuracy).
This observation is consistent with our iteration complexity results in Theorem 1.
5.3 Clipping Operation Happens Frequently
Figure 5 reports the proportion of iterations in each epoch that clipping is triggered.
We observethat for our algorithm, clipping happens more frequently than the baseline, especially for NLP tasks.
We conjecture that this is because we only used local gradients in each GPU to do the clippingwithout averaging them across all machines as the baseline did.
This leads to more stochasticityof the norm of the gradient in our algorithm than the baseline, and thus causes more clippings tohappen.
This observation highlights the importance of studying clipping algorithms in the distributedsetting.
Another interesting observation is that clipping happens much more frequently when traininglanguage models than image classification models.
Hence this algorithm is presumably more effectivein training deep models in NLP tasks.
Conclusion
In this paper, we design a communication-efficient distributed stochastic local gradient clippingalgorithm to train deep neural networks.
By exploring the relaxed smoothness condition which wasshown to be satisfied for certain neural networks, we theoretically prove both the linear speedupproperty and the improved communication complexity when the data distribution across machinesis homogeneous.
Our empirical studies show that our algorithm indeed enjoys parallel speedupand greatly improves the runtime performance in various federated learning scenarios (e.g., partialclient participation).
One limitation of our work is that our convergence analysis is only applicablefor homogeneous data, and it would be interesting to analyze the settings of heterogeneous datatheoretically in the future.