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 presents a new collective algorithm and associated implementation of SGD training that is optimized for distributed systems where some ranks will not participate in each training step.  The observation is that parameter servers are not scalable, and that most collective based (all-reduce) distributed algorithms are either not fault tolerant, or have significant performance impacts when there is load imbalance or system noise.  The proposed Moshpit all-reduce and SGD strive to overcome these challenges by breaking down the collective communication domains into many more "sub-domains" that can tolerate delays or absences of its participants.  However, it does so in a structured manner that is more communication efficient than a standard gossip, random, algorithm.  They show that despite stale gradients for each participant, the algorithm still converges and does so faster than SOTA algorithms.