Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices

1 Introduction . Many recent influential discoveries in deep learning were enabled by the trend of scaling model and dataset size . Over the last decade , computer vision has grown from training models with 60 million parameters [ 1 ] on 1.3 million images [ 2 ] to 15 times more parameters [ 3 ] and 200 times more training data [ 4 ] . In natural language processing , the state-of-the-art language models [ 5 ] with 175 billion parameters are trained on over 570GB of texts , and even this does not saturate the model quality [ 6 ] . Training these large models can take years even with a top-of-the-line GPU server [ 7 ] . As a result , researchers and practitioners often have to run distributed training with multiple machines [ 8 ] . The dominant approach to distributed deep learning is data-parallel training [ 9 ] , where each worker processes a fraction of the training batch and then exchanges its gradients with peers . If done naïvely , the gradient exchange step can overload the network as the number of workers increases . To combat this issue , modern distributed training algorithms take advantage of communication-efficient protocols , such as all-reduce [ 10 ] . These protocols allow workers to collectively compute the global average gradient with a constant communication overhead , regardless of the total number of peers . ∗Equal contribution . Correspondence to mryabinin0 @ gmail.com . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . However , this efficiency makes the protocols more fragile : if any single participant fails or takes too long to process its batch , all other nodes are stalled . Therefore , scaling all-reduce protocols beyond a couple of servers requires specialized infrastructure with dedicated ultra-high bandwidth networking [ 8 ] . This kind of infrastructure is notoriously expensive compared to regular GPU servers or preemptible cloud VMs ( see Appendix A for details ) . Hence , it is tempting to consider distributed training on cheap unreliable instances as a cost-efficient alternative . A similar scenario arises in federated learning [ 11 ] , where a single model is trained on heterogeneous devices due to privacy concerns . In both scenarios , workers use a shared network , where both latency and bandwidth can vary drastically due to interference from other users [ 12 ] . Furthermore , compute nodes are also subject to failure ( or preemption ) caused by factors beyond the protocol ’ s control . Running large-scale distributed training in these circumstances requires fault- and latency-tolerant algorithms [ 14 , 15 ] . Most of these algorithms replace all-reduce averaging with gossip : each participant periodically downloads the latest parameters from their neighbors in a sparsely connected communication graph and averages the results . The updates gradually propagate through the graph over multiple rounds of averaging . However , the communication required to perform gossip grows linearly with the number of neighbors . Hence , when scaling to hundreds of peers , decentralized SGD has to keep the communication graph sparse , slowing down the convergence . In this work , we propose an alternative approach . Instead of relying on a predefined communication graph , participants dynamically organize themselves into groups using a fully decentralized matchmaking algorithm called Moshpit All-Reduce . This strategy allows us to use communicationefficient all-reduce protocols that significantly reduce the network load compared to gossip-based averaging , while still being able to operate in unreliable hardware and network conditions . Our contributions can be summarized as follows : • We propose Moshpit All-Reduce — a novel decentralized averaging protocol for large-scale training with unreliable communication-constrained devices . According to our analysis , this method has exponential convergence rate independent of network topology and size . • Armed with this averaging protocol , we develop Moshpit SGD for distributed optimization . We derive convergence rates for this algorithm and establish its equivalence to Centralized ( Local ) SGD in terms of iteration complexity under realistic assumptions . • Our experiments demonstrate that Moshpit All-Reduce is significantly more efficient under network latency in realistic conditions . In particular , we train ResNet-50 on ImageNet to 75 % accuracy 1.3 times faster than existing decentralized training algorithms and pretrain ALBERT-large 1.5 times faster on preemptible cloud VMs.2 2 Related Work . 2.1 Data parallel training . The most popular way to accelerate neural network training with multiple devices is data-parallel training [ 9 , 16 , 17 ] . On each optimization step , this strategy splits the training batch among participants . Each participant then runs forward and backward passes to obtain gradients of the objective function on their part of the training batch . After that , we can aggregate the gradients from workers and perform an optimization step . There are two main strategies for this aggregation . Historically , the first solution to gradient aggregation was to use Parameter Server ( PS ) [ 18 ] : a separate process or a dedicated server that keeps track of model parameters and optimizer statistics . After each round , the PS accumulates the gradients from each worker and updates the model parameters using SGD or any other optimizer , such as Adam [ 19 ] . Finally , the server distributes the updated model parameters to workers . This strategy is robust and easy to implement , but it requires the server to regularly download full model gradients from every single worker . As a result , the parameter server can quickly become a bottleneck for large-scale training [ 20 ] . Since the original PS , researchers have proposed several modifications that reduce the communication load : accumulating multiple batches [ 22 ] , compression [ 23 , 24 ] , server sharding [ 25 , 26 ] . A more detailed overview is given in Appendix B . 2Implementation and code of experiments are at github.com/yandex-research/moshpit-sgd . In turn , many practical distributed training systems have instead switched to averaging with AllReduce [ 16 , 27 , 28 , 17 ] . This name refers to a collection of protocols originally developed for HPC applications . Workers can follow these protocols to collectively compute the average3 gradient more efficiently than with a central server . 2.2 Communication-efficient All-Reduce . There are several all-reduce protocols optimized for different network topologies . The simplest one is known as Butterfly All-Reduce [ 10 ] . Each of N participants splits its local vector into N chunks . Then , i-th worker aggregates i-th chunk of data from all peers and sends back the averaged chunk . As long as the vector size s is greater than N , this protocol uses O ( s× N−1N ) total bandwidth on each worker . However , it requires all-to-all communication , which is not always practical for the HPC infrastructure due to network contention [ 10 ] . As a result , real-world systems typically use Ring or Tree All-Reduce , where each worker only communicates with a small subset of its peers . These protocols enable highly efficient and scalable averaging with O ( 1 ) or O ( logN ) total communication per worker , but they also share a common drawback : they can not tolerate node failures or network instability . If any single participant fails to execute its part or takes long to respond , this paralyzes all other workers . 2.3 Distributed training in unstable conditions . Some distributed training applications must deal with unstable network bandwidth and/or unreliable workers . This issue is most prevalent in federated learning [ 11 , 29 , 30 ] . When dealing with privacysensitive data distributed across multiple actors , such as hospital servers [ 31 , 32 ] or mobile phones [ 33 , 34 ] , one must train the model using whichever hardware and network available to those actors . Another important motivational factor is cost : HPC-grade infrastructure can be prohibitively expensive , pushing researchers and practitioners towards commodity servers or preemptible cloud VMs that are significantly cheaper ( see Appendix A ) . Another solution is to use volunteer computing [ 35 , 36 ] with abundant , but even less reliable , compute resources . Training under these conditions requires specialized strategies . At a small scale , one can deploy one or a few reliable parameter servers to aggregate the updates from workers . This strategy can tolerate individual node failures [ 37 ] , but scales poorly due to the reasons discussed in Section 2.1 . 2.4 Decentralized training . If there are too many participants for PS , it can be advantageous to use decentralized SGD via gossip-based averaging [ 38 , 39 , 14 ] . In this scenario , participants form a sparse graph : each worker periodically downloads parameters from its neighbors and mixes them with local parameters . In essence , gossip-based averaging removes the communication bottlenecks of PS at the cost of using different local parameters on each peer . That said , gossip-based optimization algorithms can match , and sometimes even outperform , their centralized counterparts in terms of training speed [ 40 , 41 , 42 , 14 , 43 ] . However , the convergence properties of gossip averaging and gossipbased optimization methods significantly depend on the communication graph through the spectral properties of the mixing matrix [ 44 , 42 ] or the Laplacian matrix of the network [ 45 , 46 ] . 3All-Reduce works with any commutative associative operation , such as min , max , or product . Consequently , as the number of peers increases , gossip-based averaging has to either increase the number of neighbors ( hence more communication ) or accept slower convergence speed . Because of this , gossip is less communication-efficient than all-reduce algorithms reviewed in Section 2.2 . However , gossip-based algorithms are more robust to changes , which makes them applicable to time-varying networks [ 47 , 48 , 49 , 50 ] and federated learning [ 51 , 52 , 53 ] . 3 Moshpit SGD . Large-scale training with unreliable participants requires a protocol that is both communicationefficient and fault-tolerant . Unfortunately , existing methods have only provide one of these properties . To better address our conditions , we propose Moshpit All-Reduce — a fully decentralized averaging protocol that combines the efficiency of all-reduce and the fault tolerance of gossip-based averaging . The rest of this section is organized as follows : • Section 3.1 describes the protocol and proves its correctness and communication efficiency ; • Section 3.2 provides the analysis of the protocol and proves exponential convergence rate for averaging and the rate matching the one of centralized Local-SGD for optimization ; • Section 3.3 contains implementation details for training with heterogeneous compute nodes . 3.1 Moshpit All-Reduce . The core idea of Moshpit All-Reduce is that workers perform averaging in small independent groups . That way , a single failed participant would only affect his current group . In turn , the composition of each group should be chosen dynamically to converge in the least number of steps . Ideally , if there are 9 peers with local parameters θ , we can average them in 2 rounds , as demonstrated in Figure 2 : First round Second round A θA θA θA θB θB θB θC θC θC Group C Group B Group Average θ in groups θ3θ1 θ2 θ6θ4 θ5 θ9θ8θ7 Algorithm 1 Moshpit All-Reduce ( for i-th peer ) Input : parameters { θj } Nj=1 , number of peers N , d , M , number of iterations T , peer index i θ0i : = θi C0i : = get_initial_index ( i ) for t ∈ 1 . . . T do DHT [ Ct−1i , t ] .add ( addressi ) Matchmaking ( ) // wait for peers to assemble peerst : = DHT.get ( [ C t−1 i , t ] ) θti , c t i : = AllReduce ( θ t−1 i , peerst ) Cti : = ( C t−1 i [ 1 : ] , c t i ) // same as eq . ( 1 ) To achieve this in a decentralized system , we use Distributed Hash Tables ( DHT ) — a decentralized key-value storage ; Appendix B contains its more detailed description . On each averaging round : • Each worker computes its group key Ci ; • Workers add their network addresses to the DHT key corresponding to Ci ; • Each worker can now fetch a full list of peers that have the same Ci and run All-Reduce with those peers . Unfortunately , the averaging structure from Figure 2 is impossible to maintain when participants are constantly joining , leaving , and failing . However , we can achieve equivalent results without global structure using a simple rule : if two peers were in the same group in round t , they must choose different groups in round t+1 . A natural way to enforce this rule is to take advantage of the chunk indices from Butterfly All-Reduce ( see Figure 1 ) . Recall that each worker accumulates a unique chunk of parameters defined by an index ci . By setting Ci : = ci , we can guarantee that any workers that were in the same group at a round t will have different group indices in round t+1 . This averaging scheme can be generalized to more than two dimensions in order to fit a larger number of peers or reduce the group size . For a d-dimensional hypercube , nodes should find groups of peers that they have not communicated with during d−1 previous rounds . To that end , we define Ci as tuples containing chunk indices from d−1 previous rounds ( t denotes the communication round ) : Cti : = ( c t−d+1 i , c t−d+2 i , . . . , c t i ) . ( 1 ) The above intuition can be formalized with Algorithm 1 . Here , N peers form a virtual d-dimensional grid with M peers per row and average their parameters θi over T rounds . DHT [ · ] is a shortcut for using the DHT to add or retrieve values for a given key . The Matchmaking step corresponds to the decentralized matchmaking procedure that organizes active workers with the same index into groups , described in detail in Appendix E. In turn , AllReduce denotes running all-reduce to compute the average θ in a given group . The get_initial_index function takes the peer index i and returns d−1 integers in range [ 0 , M ) such that the size of initial groups does not exceed M . This way , the groups formed on subsequent rounds will also have at most M participants . One possible strategy is : get_initial_index ( i ) = ( ⌊i/Md−1⌋ modM ) j∈ { 1 , ... , d } ( 2 ) If N=Md and there are no node/network failures , Algorithm 1 is equivalent to Torus All-Reduce [ 54 ] , achieving the exact average after d rounds of communication ( see Appendix C.1 ) . However , our typical use case is far from this perfect scenario ; for example , some groups can have less than M members . Furthermore , a peer might fail during all-reduce , causing its groupmates to skip a round of averaging . Still , Moshpit All-Reduce is applicable even in these conditions : Theorem 3.1 ( Correctness ) . If all workers have a non-zero probability of successfully running a communication round and the order of peerst is random , then all local vectors θ t i converge to the global average with probability 1 : ∀i , ∣∣∣∣∣∣θti − 1N ∑ i θ0i ∣∣∣∣∣∣2 2 −−−→ t→∞ 0 . ( 3 ) Proof ( sketch , complete in Appendix C.2 ) . Running all-reduce with a subset of peers preserves the invariant 1N ∑ i θ t i = 1 N ∑ i θ t−1 i and reduces the deviation of θ t i from the overall average . Complexity . The matchmaking protocol is implemented over Kademlia DHT [ 55 ] , meaning that each read and write operation needs at most O ( logN ) requests and O ( M ) bandwidth to load peerst . After the matchmaking is over , each group runs a single all-reduce round to compute the average . In principle , Moshpit Averaging can use any general-purpose all-reduce protocol . We opted for a butterfly-like version ( Figure 1 ) , as it is simpler than Ring All-Reduce while still being communicationefficient . The communication complexity of this algorithm is O ( max ( s , M ) × M−1M ) , where s is the size of vector θ . Thus , the total time complexity of Algorithm 1 becomes : O ( T × [ log2 N +M +max ( s , M ) × M − 1 M ] ) . ( 4 ) This compares favorably to gossip , where network load grows linearly with the number of neighbors .

Summary: This paper introduces Moshpit All-Reduce, a fully decentralized iterative averaging protocol for large-scale training under unstable communication conditions and/or with unreliable workers. The proposed protocol allows an exponential convergence to the global average. It is shown that Moshpit SGD, built on top of Moshpit All-Reduce, is equivalent to the centralized SGD in terms of iteration complexity.