Asynchronous Decentralized SGD with Quantized and Local Updates

1 Introduction . Decentralized optimization has recently emerged as a promising approach for scaling the distributed training of machine learning models , in particular via stochastic gradient descent ( SGD ) [ Lian et al. , 2017 , Tang et al. , 2018 , Koloskova et al. , 2019a ] . Its key advantage is that it removes the need for a central coordinator node in distributed training , and therefore it can allow for high scaling . The general decentralized optimization setting is the following : we are given n nodes , each with a subset of data from some distribution , which can communicate over some underlying graph topology . In each global round , each node samples some local data , performs a local gradient step , and it is paired with a neighbor , which may be chosen randomly . The nodes exchange model information pairwise , and then update their models , often via direct model averaging . Variants of this setting have been analyzed since pioneering work by Tsitsiklis [ 1984 ] , for various estimation and optimization algorithms [ Xiao and Boyd , 2004 , Nedic and Ozdaglar , 2009 , Johansson et al. , 2009 , Shamir and Srebro , 2014 ] and have seen renewed interest given its applicability to training deep neural networks ( DNNs ) at scale , e.g . [ Lian et al. , 2017 , 2018 , Assran et al. , 2018 ] . Recently , there has been significant focus on reducing the synchronization overheads for decentralized training , usually employing three approaches : 1 ) implementing faster non-blocking communication between communication partners at a round [ Lian et al. , 2018 , Assran et al. , 2018 ] , which may cause them to see stale versions of their models , 2 ) allowing nodes to take local steps in between 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . their communication rounds [ Wang and Joshi , 2018 , Koloskova et al. , 2020 ] , and 3 ) applying quantization to the communication [ Lu and De Sa , 2020 , Tang et al. , 2018 , Koloskova et al. , 2019a , b ] . The above impressive line of work contributes a rich set of algorithmic and analytic ideas ; however , one common limitation is that the algorithms are usually set in the synchronous gossip model , which requires all nodes to perform their communication in lock-step rounds , and share a common notion of time , thus reducing their practicality . To mitigate this fact , some references , e.g . [ Lian et al. , 2018 , Assran et al. , 2018 , Lu and De Sa , 2020 ] partially relax this requirement , although they do so at the cost of additional assumptions , or reduced guarantees , as we discuss in related work . Another relative limitation is that the analyses are usually customized to the bespoke communication-reduced methods being applied , and therefore are hard to generalize to other methods . Our Contribution . In this paper , we consider decentralized SGD-based optimization in the simpler , but harder to analyze , asynchronous gossip model [ Xiao and Boyd , 2004 ] , in which communication occurs in discrete , randomly chosen pairings among nodes , and does not require a common notion of time . We prove that a new variant of SGD we call SwarmSGD converges in this setting , even though it supports all three communication-reduction approaches mentioned above in conjunction . Our analysis generalizes to heterogeneous data distributions and communication topologies . At a high level , SwarmSGD works as follows . Each node i maintains a local model estimate Xi based on which gradients are generated , and a shared buffer where quantized models are stored for communication with other nodes . In each step , node i first computes a sequence of H local gradient steps , which it does not yet apply . Next , the node chooses communication partner j , uniformly at random among its neighbors . Then , node i reads from its own communication buffer and from the communication buffer of j , obtaining quantized models Qi and Qj . A subtlety here is that Qi is not necessarily the quantized version of the model Xi , since other nodes can write concurrently to i ’ s buffer . The node i then averages Qi with Qj , and updates the neighbor ’ s remote buffer to the quantized average . Finally , it applies its local gradient steps to the resulting average , adopts this as its next model Xi , and a writes quantized version of it in its own shared buffer . This procedure can be implemented in a deadlock-free , non-blocking manner , by using either shared-memory or the remote direct-memory access ( RDMA ) calls supported by MPI [ Woodall et al. , 2006 ] . Importantly , the communication partner j does not need to block its computation during communication , and may be contacted by more than one interaction partner during a single local step , although we do assume that individual reads and writes are performed atomically . A key component of this procedure is the quantization scheme : directly using an unbiased quantizer , e.g . [ Alistarh et al. , 2017 ] would destroy convergence guarantees , as the quantization error would be proportional to the model norm , which may not be bounded . Instead , we use a customized variant of the quantization scheme of Davies et al . [ 2021 ] , whose error depends on the distance between the point being quantized ( the model ) , and an arbitrary reference point , provided as a parameter . We prove that each node can reliably use its own model as a reference point to quantize and de-quantize messages placed in its buffer by other nodes . In turn , this requires care in the analysis . Specifically , the key observation behind our analysis is exactly in showing that the nodes ’ local models stay well-enough concentrated around their mean throughout optimization to allow for correct decoding of quantized models , which in turn implies joint convergence by the nodes towards a point of vanishing gradient . This concentration follows via a non-trivial super-martingale argument . If nodes take a constant number of local SGD steps between communication steps , then SwarmSGD has Θ ( √ n ) speedup to convergence for non2-convex objectives . This matches results from previous work which considered decentralized dynamics but with global synchronization [ Lian et al. , 2017 ] . Experimental Validation . We apply SwarmSGD to train deep neural networks on image classification and machine translation ( NMT ) tasks , deployed on the Piz Daint supercomputer [ Piz , 2019 ] . Experiments confirm the intuition that the average synchronization cost of SwarmSGD per iteration is low : it stays at less than 10 % of the batch computation time , and remains constant as we increase the number of nodes . For example , using SwarmSGD , we are able to train a TransformerXL [ Vaswani et al. , 2017 ] model on WMT17 ( En-Ge ) 1.5× faster than a highly-optimized largebatch SGD baseline , and to slightly higher accuracy , without additional hyper-parameter tuning . At the same time , due to the reduced communication frequency , Swarm also improves upon the speed of the previous practical decentralized methods , e.g . [ Lian et al. , 2017 , 2018 , Assran et al. , 2018 ] . Importantly , we also note that , in less overparametrized settings such as training residual CNNs [ He et al. , 2016 ] on ImageNet [ Russakovsky et al. , 2015 ] , nodes do need to perform more iterations over the dataset relative to the baseline in order to recover full accuracy . This is predicted by the analysis , and confirms similar findings in previous work [ Assran et al. , 2018 ] . Overall , our method does appear well-suited to training large modern models at node counts where global synchronization among all nodes is prohibitively expensive . Related Work . Decentralized optimization has a long history [ Tsitsiklis , 1984 ] , and is related to the study of gossip algorithms , e.g . [ Kempe et al. , 2003 , Xiao and Boyd , 2004 , Boyd et al. , 2006 ] . Gossip is usually studied in one of two models [ Boyd et al. , 2006 ] : synchronous , structured in global rounds , where each node interacts with a randomly chosen neighbor , forming a matching , and asynchronous , where each node wakes up at random times , e.g . given by a Poisson clock , and picks a random neighbor to interact with . Several classic optimization algorithms have been analyzed in the asynchronous gossip model [ Nedic and Ozdaglar , 2009 , Johansson et al. , 2009 , Shamir and Srebro , 2014 ] . In this paper , we focus on analyzing decentralized SGD in this model . As mentioned , the growing line of work on decentralized optimization for machine learning has mostly focused on variants of the synchronous gossip model . Specifically , Lian et al . [ 2017 ] considered this setting in the context of DNN training , while and Tang et al . [ 2018 ] and Koloskova et al . [ 2019b ] also analyzed decentralized optimization with quantization in the synchronous model . Wang and Joshi [ 2018 ] and Koloskova et al . [ 2020 ] provided analysis frameworks for synchronous decentralized SGD with local updates , and possibly changing topologies . Lian et al . [ 2018 ] and Assran et al . [ 2018 ] focused specifically on reducing synchronization costs in this setting , and proposed algorithms with partially non-blocking communication , in which nodes may read a stale version of the interaction partner ’ s information , modelling e.g . a communication buffer . However , the maximum staleness must be bounded by a global variable τ , which must be enforced throughout the execution . As observed by Assran et al . [ 2018 ] , enforcing this bound can cause blocking , and therefore the authors of these works propose to implement a relaxed roundbased model , in which nodes interact once per round in perfect matchings . Their algorithms provide O ( 1/ √ Tn ) convergence rates , under analytical assumptions . Upon careful examination , we find that their analysis approach can be extended to the asynchronous gossip model we consider , by defining the “ contact matrices ” to correspond to pairwise interactions . However , this introduces two significant limitations . First , the analysis will not support local gradient updates to models nor quantized communication . If we remove these practical relaxations , our technique yields better bounds , as our potential analysis is specifically-tailored to this dynamic interaction model . Second , as we detail in the Appendix , some of their technical conditions imply existence of global synchronization . For Assran et al . [ 2018 ] , as we detail in the Appendix , their analysis would not guarantee any non-trivial speedup due to parallelization in the asynchronous gossip model . We describe these issues in detail and present a systematic comparison in Appendix A. Lu and De Sa [ 2020 ] provided a novel approach to analyze decentralized SGD with quantization and limited asynchrony : specifically , their algorithm requires blocking communication , i.e . nodes have to synchronize explicitly during interactions , but may see old versions of eachothers ’ models . More precisely , during each interaction , both parties are responsible for updating their local models , meaning that once node is woken up ( we call it initiator node ) and chooses interaction partner it has to block until the partner is woken up as well . In our case , initiator can update both its local model and the local model of its partner and proceed to the next step without blocking . Koloskova et al . [ 2019a ] use a similar update rule in the synchronous model . Zhang and You [ 2021 ] recently proposed a decentralized algorithm which is fully-asynchronous as long as node activation rates and message delays are bounded . As noted earlier , bounding activation rates does imply blocking ; however , tolerating ( bounded ) message delays does improve over our approach of updating models using atomic writes . The setting further differs in that they assume that nodes compute full ( nonstochastic ) gradients , as well as that the loss function satisfies the PL condition . In sum , we are the first to explicitly consider the asynchronous gossip model , and the impact of local updates , asynchrony , and quantization used in conjunction together with decentralized SGD . Our technique is new , relies on a fine-grained analysis of individual interactions , and can yield improved bounds even in the case where H = 1 . Further , our algorithm is the first to allow for both communication-compression and non-blocking communication . From the implementation perspective , the performance of our algorithm matches or improves that of previous methods , notably D-PSGD [ Lian et al. , 2017 ] , AD-PSGD [ Lian et al. , 2018 ] and SGP [ Assran et al. , 2018 ] .

Summary: This paper proposes a decentralized optimization method, SwarmSGD, for reducing the communication overhead in distributed deep learning. The proposed method combines asynchronous pair-wise gossip, quantized communication, and local updates into a single method and a unified analysis. Convergence guarantees are provided for smooth non-convex function, using the quantizer from Davies et al., 2021, and a bounded second moment assumption. Numerical results are provided on image classification (CIFAR, ImageNet) and machine translation (WMT En-Ge) tasks.