DANA: Scalable Out-of-the-box Distributed ASGD Without Retuning

1 INTRODUCTION . Modern deep neural networks are comprised of millions of parameters , which require massive amounts of data and time to learn . Steady growth of these networks over the years has made it impractical to train them from scratch on a single GPU . Distributing the computations over several GPUs can drastically reduce this training time . Unfortunately , stochastic gradient descent ( SGD ) , typically used to train these networks , is an inherently sequential algorithm . As a result , training deep neural networks on multiple workers ( computational devices ) is difficult , especially when trying to maintain high efficiency , scalability and final accuracy . Data Parallelism is a common practice for distributing computation : data is split across multiple workers and each worker computes over its own data . Synchronous SGD ( SSGD ) is the most straightforward method to distribute the training process of neural networks : each worker computes the gradients over its own separate mini-batches , which are then aggregated to update a single model . The result is identical to multiplying the batch size B by the number of workers N , so the effective batch size is B ·N . This severely limits scalability and reduces the model accuracy if not handled carefully ( Smith et al. , 2018 ; Devarakonda et al. , 2017 ; Goyal et al. , 2017 ) . Furthermore , synchronization limits SSGD progress to the slowest worker : all workers must finish their current mini-batch and update the parameter server before any can proceed to the next mini-batch . Asynchronous SGD ( ASGD ) addresses these drawbacks by removing synchronization between the workers ( Dean et al. , 2012 ) . Unfortunately , it suffers from gradient staleness : gradients sent by workers are often based on parameters that are older than the master ’ s ( parameter server ) current parameters . Hence , distributed ASGD suffers from slow convergence and reduced final accuracy , and may not converge at all if the number of workers is high ( Zhang et al. , 2016b ) . Several works attempt to address these issues ( Zheng et al. , 2017 ; Zhang et al. , 2015 ; 2016b ; Dean et al. , 2012 ) , but none has managed to overcome these problems when scaling to a large number of workers . More crucially , many ASGD algorithms require re-tuning of hyperparameters when scaling to different numbers of workers , and several even introduce new hyperparameters that must also be tuned ( Zheng et al. , 2017 ; Zhang et al. , 2015 ; 2016b ) . In practice , the vast number of potential hyperparameter configurations means that tuning is often done in parallel , with each worker independently evaluating a single configuration using standard SGD . Once the optimal hyperparameters are selected , training is completed on larger clusters of workers . Any additional tuning for ASGD can thus be computationally expensive and time-consuming . Though many algorithms have been proposed to reduce the cost of tuning ( Bergstra & Bengio , 2012 ; Li et al. , 2017 ; Klein et al. , 2017 ; Hazan et al. , 2018 ; Snoek et al. , 2015 ) , hyperparameter search remains a significant obstacle , and many practitioners can not afford to re-tune hyperparameters for distributed training . Our contribution : We propose Distributed Accelerated Nesterov ASGD ( DANA ) , a new distributed ASGD algorithm that works out of the box : it achieves state-of-the-art accuracy on existing architectures without any additional hyperparameter tuning or changes to the training schedule , while scaling as well or better than existing ASGD approaches , and without any additional overhead . Our DANA implementation achieves state-of-the-art accuracy on ImageNet when training ResNet-50 with 16 and even 32 workers , as well as on CIFAR-10 and CIFAR-100 . 2 BACKGROUND . The goal of SGD is to minimize an optimization problem J ( θ ) where J is the objective function ( i.e. , loss ) and the vector θ ∈ Rk is the model ’ s parameters . Let ∇J be the gradient of J with respect to its argument θ . Then the update rule of SGD for the given problem with learning rate η is : gt = ∇θJ ( θt ) θt+1 = θt − ηgt ( 1 ) Momentum Momentum ( Polyak , 1964 ) has been demonstrated to accelerate SGD convergence and reduce oscillation ( Sutskever et al. , 2013 ) . Momentum can be compared to a heavy ball rolling downhill that accumulates speed on its way towards the minima . Mathematically , the momentum update rule is obtained by adding a fraction γ of the previous update vector vt−1 to the current update vector vt : gt = ∇J ( θt ) vt+1 = γvt + gt θt+1 = θt − ηvt+1 ( 2 ) When successive gradients have similar direction , momentum results in larger update steps ( higher speed ) , yielding up to quadratic speedup in convergence rate for stochastic and standard gradient descent ( Loizou & Richtárik , 2017b ; a ) . Algorithm 1 ASGD : worker Receive parameters θt from master Compute gradients gt ← ∇J ( θt ) Send gt to master Algorithm 2 ASGD : master Receive gradients gt from worker i ( at iteration t+ τ ) θt+τ+1 ← θt+τ − ηgt Send parameters θt+τ+1 to worker i Nesterov Continuing the analogy of a heavy ball rolling downhill , higher speed might make the heavy ball overshoot the bottom of the valley ( the local or global minima ) if it does not slow down in time . Nesterov ( 1983 ) proposed Nesterov Accelerated Gradient ( NAG ) , which gives the ball a “ sense ” of where it is going , allowing it to slow down in advance . Formally , NAG approximates θ̂t , the future value of θt , using the previous update vector vt : θ̂t = θt − ηγvt , and computes the gradients on the parameters ’ approximated future value θ̂ instead of their current value θ . This allows NAG to slow down near the minima before overshooting the goal and climbing back up the hill . We call this look-ahead since it allows us a peek at θ ’ s future position . The NAG update rule is identical to Equation 2 , except that the gradient gt is computed on the approximated future parameters θ̂t instead of θt : gt = ∇J ( θ̂t ) . It is then applied to the original parameters θt via vt as in Equation 2 . Equation 3 shows that the difference between the updated parameters θt+1 and the approximated future position θ̂t is only affected by the newly computed gradients gt , and not by vt . Hence , NAG can accurately estimate future gradients even when the update vector vt is large . θt+1 − θ̂t = θt − ηvt+1 − θt + ηγvt = ηγvt − η ( γvt + gt ) = −ηgt ( 3 ) 3 GRADIENT STALENESS AND MOMENTUM . In ASGD training , each worker i pulls up-to-date parameters θt from the master and computes gradients of a single batch ( Algorithm 1 ) . Once computation has finished , worker i sends the computed gradient gt back to the master . The master ( Algorithm 2 ) then applies the gradient gt to its current set of parameters θt+τ , where τ is the lag : the number of updates the master has received from other workers while worker i was computing gt . In other words , gradient gt is stale : it was computed from parameters θt but applied to θt+τ . This gradient staleness is major obstacle to scaling ASGD : the lag τ increases as the number of workers N grows , decreasing gradient accuracy , and ultimately reducing the accuracy of the trained model . From Lag to Gap We denote by ∆θ = θt − θt+τ the difference between the master and worker parameters , and define the gap as the sum of layer-wise RMSE : G ( ∆θ ) = ∑ ψ∈layers RMSE ( ψ ) , where for each model layer ψ with m parameters , RMSE ( ψ ) = ‖ψ‖/ √ m. Ideally , there should be no difference between θt and θt+τ : when ∆θ = 0 , gradients are computed on the same parameters they will be applied to . This is the case for sequential and synchronous methods such as SGD and SSGD . In asynchronous methods , however , more workers result in an increased lag τ and thus a larger gap , as demonstrated by Figure 1 ( a ) . A larger gap means less accurate gradients , since they are computed on parameters that differ significantly from those they will be applied to . Conversely , a smaller gap means that gradients are likely to be more accurate . The Effect of Momentum While momentum and Nesterov methods improve SGD convergence and accuracy of trained models , they make scaling to more workers more difficult . As Figure 1 ( b ) shows , adding NAG to ASGD exacerbates gradient staleness , even though the lag τ is unchanged . Put differently , NAG and momentum increase the gap G ( ∆θ ) . Let xi be the variable x for worker i ( for the master , i = 0 ) and xit be the value of that variable at the worker ’ s t iteration . For ASGD without momentum or NAG , ∆θ is the sum of gradients1 , ∆ASGDθ = η ∑N i=1 g i t , whereas in the case of ASGD with NAG , ∆θ is the sum of update vectors : ∆NAG-ASGDθ = ∑N i=1 v i t. Recall that 1To simplify analysis , we assume that all workers have equal computation power and are applying their gradients to the master in a round-robin order . These assumptions can be removed without loss of generality by keeping track of worker updates and weighting them accordingly . by design , momentum and NAG increase the magnitude of updates to θt : ‖vt‖ ≥ ‖gt‖ . Moreover , if the distribution of training data in each worker is the same ( e.g. , the common case of assigning data to workers uniformly at random ) , then the directions of updates vit are approximately similar , since the loss surfaces are similar . Applying the identity ‖a + b‖2 = ‖a‖2 + ‖b‖2 + 2〈a , b〉 and the triangle inequality , it follows that in general the gap with NAG is larger than the gap without it : G ( ∆NAG-ASGDθ ) ≥ G ( ∆ASGDθ ) . Figure 1 ( b ) shows that the gap for ASGD with NAG is substantially larger than for ASGD without it . Conversely , DANA-Zero , detailed in the next section , maintains a low gap throughout training even though it also uses NAG . 4 DANA : DISTRIBUTED ACCELERATED NESTEROV ASGD . DANA is a distributed optimizer that converges without hyperparameter tuning even when training with momentum on large clusters . It reduces the gap G ( ∆θ ) by computing the worker ’ s gradients on parameters that more closely resemble the master ’ s future position θt+τ . We extend NAG to the common distributed setting with N workers and one master , obtaining similar look-ahead to the traditional method with a single worker . This means that for the same lag τ , DANA suffers from a reduced gap and therefore suffers less from gradient staleness . 4.1 THE DANA-ZERO UPDATE RULE . In DANA-Zero , the master maintains a separate update vector vi for each worker , which is updated with the worker ’ s gradients gi using the same update rule as in classic SGD with momentum ( Equation 2 ) . Since the master updates each vi only with the gradients from worker i , we can apply look-ahead using the most recent update vectors of the other workers . We know that vit−1 will move the master ’ s parameters θ0 on iteration t of worker i by ηγvit−1 . Thus , computing θ 0 t − ηγvit−1 gives us an approximation of the next position of the master ’ s parameters after worker i has sent its gradients . Instead of sending the master ’ s current parameters θ0 to the worker , DANA-Zero sends the estimated future position of the master ’ s parameters after N updates , one for each worker : θ̂DANA , θ 0 − ηγ N∑ i=1 viprev ( i ) ( 4 ) where prev ( i ) denotes the last iteration where worker i sent gradients to the master . Algorithm 3 shows the DANA-Zero master algorithm ; the worker code is the same as in ASGD ( Algorithm 1 ) . Given the update rule , we calculate the gap of DANA-Zero , G ( ∆DANAθ ) , similarly to Equation 3 : ∆DANAθ = η N∑ i=1 vit − ηγ N∑ i=1 vit−1 = η N∑ i=1 ( vit+1 − γvit ) = η N∑ i=1 git ( 5 ) Algorithm 3 DANA-Zero master . DANA-Zero uses the standard ASGD worker ( Algorithm 1 ) . Receive gradients gi from worker i Update worker ’ s momentum vi ← γvi + gi Update master ’ s weights θ0 ← θ0 − ηvi Send estimate θ̂ = θ0 − ηγ ∑N j=1 v j to worker i Algorithm 4 DANA worker i. DANA uses the standard ASGD master ( Algorithm 2 ) . Receive parameters Θi from master Compute gradients gi ← ∇J ( Θi ) Update momentum vi ← γvi + gi Send update step γvi + gi to master Equation 5 shows that DANA-Zero has the same gap as ASGD without momentum . Figure 1 ( b ) demonstrates this empirically : ASGD with momentum has a larger gap than ASGD throughout the training process , whereas DANA-Zero ’ s gap is similar to ASGD despite also using momentum . Additionally , when running with one worker ( N = 1 ) , DANA-Zero reduces to a single standard NAG optimizer : with one worker , θ1t = θ 0 t − ηγ , so merging the master and the worker algorithms yields the Nesterov update rule ( see Appendix A for more details ) .

Summary: The paper addresses an important problem in distributed training of deep learning models, i.e., the gradient staleness vs the parallel performance. Keeping the gradient up-to-date in distributed training is important in order to achieve a low test error and high accuracy, but that comes at a cost: the overhead of more communication and synchronization. Asynchronous methods to update the gradient have been proposed, but they usually suffer from staleness, i.e., the communication latency between the master and the slaves impacts the accuracy and training time since the accumulated gradient already is "old" in relation to the model parameters when distributed to the slaves.