Local SGD Meets Asynchrony

1 INTRODUCTION . In this paper , we consider the classic problem of minimizing an empirical risk , defined simply as min x∈Rd ∑ i∈ [ I ] fi ( x ) , ( 1 ) where d is the dimension , x ∈ Rd denotes the set of model parameters , [ I ] is the training set , and fi ( x ) : Rd → R is the loss on the training sample i ∈ [ I ] . Stochastic gradient descent ( SGD ) ( Robbins & Monro , 1951 ) is an extremely popular iterative approach to solving this problem : xk+1 = xk − αk∇fBk ( xk ) , ( 2 ) where∇fBk ( xk ) = 1|Bk| ∑ i∈Bk ∇fi ( xk ) is the sum of gradients computed over samples , typically selected uniformly and randomly as a minibatch Bk ⊆ [ I ] , and αk is the learning rate at iteration k . 1.1 BACKGROUND ON DECENTRALIZED DATA-PARALLEL SGD . For better or worse , SGD and its variants currently represent the computational backbone for many large-scale optimization tasks , most notably the training of deep neural networks ( DNNs ) . Arguably the most popular SGD variant is minibatch SGD ( MB-SGD ) ( Bottou ( 2012 ) ) . In a distributed setting with decentralized workers q ∈ [ Q ] , it follows the iteration xk+1 = xk − αk 1 Q Q∑ q=1 ∇fBqk , ( 3 ) where Bqk ⊆ [ I ] is a local minibatch selected by worker q ∈ [ Q ] at iteration k. This strategy is straightforward to scale in a data-parallel way , as each worker can process a subset of the samples in parallel , and the model is then updated by the average of the workers ’ gradient computations . For convenience , we assume the same batch size per worker . This approach has achieved tremendous popularity recently , and there has been significant interest in running training with increasingly large batch sizes aggregated over a large number of GPUs , e.g . Goyal et al . ( 2017 ) . An alternative approach is parallel or local SGD ( L-SGD ) ( Zinkevich et al . ( 2010 ) ; Zhang et al . ( 2016c ) ; Lin et al . ( 2020 ) ) : xqj , t+1 = x q j , t − αj , t∇fBqj , t , 0 ≤ t < Kj ; x q j+1,0 = 1 Q ∑ q xqj , Kj , ( 4 ) where xqj , t denotes the local model at worker q ∈ [ Q ] after j synchronization rounds followed by t local gradient updates and Bqj , t is the local minibatch sampled at the same iteration . Kj denotes the number of local gradient update steps before the jth synchronization . Essentially , workers run SGD without any communication for several local steps , after which they globally average the resulting local models . This method is intuitively easy to scale , since it reduces the frequency of the communication . Recently , a variant called post local SGD ( PL-SGD ) ( Lin et al . ( 2020 ) ) , was introduced to address the issue of loss in generalization performance of L-SGD , wherein the averaging frequency during the initial phase of training is high and is reduced later when optimization stabilizes . Method Bloc Train Loss Train Acc . Test Loss Test Acc Time ( Sec ) Quality/ Perf . MB-SGD 128 0.016 99.75 0.234 92.95 1754 Baseline MB-SGD 1024 0.023 99.51 0.293 91.38 1201 OK PL-SGD 128 0.018 99.69 0.245 92.98 1603 Good PL-SGD 1024 0.154 94.69 0.381 87.81 1159 Poor Bloc , where Bloc = 128 , Q is the number of workers , 2 here , and α0 = 0.1 . In PL-SGD , we average the model after each gradient update for first 150 epochs and thereafter averaging frequencyK is set to 16 as in Lin et al . ( 2020 ) ; other HPs are identical to theirs . The listed results are average of 3 runs with different seeds . with PL-SGD . Clearly , these methods can not tolerate a larger Bloc , though the GPUs can support them . This shortcoming of the existing methods in harnessing the growing data-parallelism is also identified via empirical studies ( Golmant et al . ( 2018 ) ; Shallue et al . ( 2019 ) ) existing in literature . To our knowledge no effective remedy ( yet ) exists to address this challenge . Notice that , here our core target is maximally harnessing the local data-parallelism and therefore the larger local batch size , as against the existing trend in the literature wherein large number of GPUs are deployed to have a large aggregated global batch size with a relatively small Bloc . For example , refer to the performance of MB-SGD and PL-SGD as listed in Table 1 of Lin et al . ( 2020 ) . Notice that with 16 GPUs , each with Bloc = 128 , thus totaling the minibatch size as 2048 , identical to the one with 2 GPUs each with Bloc = 1024 as above , with exactly the same LR scaling and warmup strategy , both MB-SGD and PL-SGD do not face generalization degradation . However , unfortunately , such an implementation setting would incur excessive wastage of available data-parallel compute resources on each of the GPUs . Indeed , the existing specific techniques such as LARS ( You et al . ( 2017 ) ) to address the issue of poor generalization for global large batch training are insufficient for the larger local minibatch size ; we empirically describe it in Section 3 ( Table 11 ) . 1.2 LOCALLY-ASYNCHRONOUS PARALLEL SGD . Now , consider an implementation scheme as the following : 1 . In a decentralized setting of L-SGD , i.e . wherein each worker q ∈ [ Q ] has a local model xq undergoing local SGD updates as described earlier , multiple local concurrent processes u ∈ Uq share the model xq . Processes u ∈ Uq perform asynchronous concurrent gradient updates locally . 2 . The workers average their models whenever any one of them would have had at least Kj local shared updates , where Kj is as that in Equation 4 . The averaging is performed asynchronously and in a non-blocking way by the ( averaging- ) processes aq on behalf of each worker q ∈ [ Q ] . Essentially , the decentralized workers run shared-memory-based asynchronous SGD locally and periodically synchronize in a totally non-blocking fashion . More formally , consider Algorithm 1 . The model xq on a GPU q ∈ [ Q ] is shared by the processes p ∈ P q = { { aq } ∪Uq } locally . The processes p ∈ P q also maintain a shared counter Sq , initialized to 0 . The operation read-and-inc implements an atomic ( with lock ) read and increment of Sq , whereas , read provides an atomic read . Sq essentially enables ordering the shared gradient updates . In turn , this order streamlines the synchronization among workers , thereby determines the averaging rounds j . The ( updater ) processes u ∈ Uq asynchronously and lock-freely update xq with gradients computed over a non-blocking , potentially inconsistent , snapshot va , q of xq , essentially going Hogwild ! ( Recht et al . ( 2011 ) ) , see Algorithm 1a . 1 Initialize s = 0 ; 2 while s ≤ T do 3 vu , q [ i ] : = xq [ i ] , ∀ 1 ≤ i ≤ d ; 4 s : = read-and-inc ( S ) ; 5 Compute∇fBqs ( v u , q ) ; 6 xq [ i ] −= αs∇fBqs ( v u , q ) [ i ] , ∀ 1 ≤ i ≤ d ; ( a ) Local asynchronous gradient update by process u ∈ Uq . 1 Initialize scur = spre = |Uq| , j = 0 ; 2 while scur ≤ T do 3 scur : = read ( S ) ; Compute j corresponding to scur ; 4 if scur − spre ≥ Kj then 5 va , qj [ i ] : = x q [ i ] , ∀ 1 ≤ i ≤ d ; 6 Synchronize across ar , r ∈ [ Q ] \ { q } to compute vj : = 1 Q ∑ q∈ [ Q ] v a , q j ; 7 Compute ∆vqj = vj − v a , q j ; spre : = scur ; 8 xq [ i ] += ∆vqj [ i ] , ∀ 1 ≤ i ≤ d ; j = j + 1 ; ( b ) Asynchronous non-blocking in-place averaging . Algorithm 1 : Locally-asynchronous Parallel SGD ( LAP-SGD ) The process aq , which performs averaging for the worker q ∈ [ Q ] , concurrently keeps on atomically reading Sq , see Algorithm 1b . As soon as it notices an increment Kj in Sq , i.e . xq got concurrently updated with Kj number of gradients , it takes a non-blocking snapshot v a , q j of x q and synchronizes with ar of peers r ∈ [ Q ] /q to compute the average vj of the snapshots . Thereafter , aq adds the difference of the average with the snapshot va , qj to the model x q without blocking the concurrent asynchronous local gradient updates . We call this method locally-asynchronous parallel SGD ( LAP-SGD ) . This method closely resembles Hogwild++ ( Zhang et al . ( 2016a ) ) , which targets the heterogeneous NUMA based multi-core machines , though there are key differences which we describe in Section 4 . Results of the same training task as before by LAP-SGD is given in Table 2 . The distinction of this implementation is that it harnesses the compute power of the GPUs not by increasing the size of Bloc but by concurrently computing many minibatch gradients . Evidently , LAP-SGD provides speed-up without losing the quality of optimization in comparison to the baseline . Recently , Kungurtsev et al . ( 2019 ) presented a shared-memory based method wherein they showed that partitioned gradient updates for some iterations during the course of training over a shared model can reasonably save on total computation cost by means of restricted backpropagation without necessarily losing on optimization quality . Their method is limited to a centralized sharedmemory setting . Moreover , aiming to establish convergence under non-smoothness assumption , they ensure write consistency under a model-wide lock . Having designed our asynchronous parallel SGD , it inspires us to adapt the partitioned gradient update strategy to our lock-free decentralized setting . More specifically , building on LAP-SGD , we consider locally partitioned gradient computation along with asynchronous lock-free updates . Essentially , we partition the model xq to { xqi ( u ) } for u ∈ Uq , i ( u ) ∩ i ( w ) = ∅ , ∀u , w ∈ Uq ( i.e. , non-overlapping block components of the vector x ) . With that , a partitioned gradient computation will amount to computing ∇i ( u ) fBqs ( vq , u ) , the minibatch gradient with respect to the partition xqi ( u ) at line 5 in Figure 1a . Accordingly , the update step at line 6 in Algorithm 1a transforms to xq [ i ] −= αs∇fBqs ( vq , u ) [ i ] , ∀ i ∈ i ( u ) . It is to be noted that we do not use write lock for iterations at any stage . Having devised a partitioned update scheme , we propose locally-partitionedasynchronous parallel SGD ( LPP-SGD ) as described below . 1 . Processes u ∈ Uq maintain a process-local variable last iter which can take two values PARTITIONED and FULL . Each u ∈ Uq initializes last iter as FULL . 2 . While s ≤ Tst , each process u ∈ Uq performs LAP-SGD updates as lines 3 to 6 of Algorithm 1a . 3 . If Tst < s ≤ T , each process u ∈ Uq performs ( a ) a partitioned gradient computation and update : xq , u [ i ] −= αs∇fBqs ( vu , q ) [ i ] , ∀i ∈ i ( u ) if last iter = FULL , and sets last iter = PARTITIONED . ( b ) an LAP-SGD update if last iter = PARTITIONED , and sets last iter = FULL . Essentially , after some initial stabilizing epochs each process u ∈ Uq alternates between a full and a partitioned lock-free asynchronous gradient updates to the model xq . Our experiments showed that Tst = T 10 was almost always sufficient to obtain a competitive optimization result . The results of a sample implementation of LPP-SGD are available in Table 3 . It is clear that LPP-SGD handsomely speeds up the computation and provides equally competitive optimization results .

Summary: In this paper, the authors argue that the mini-batch method and local SGD method suffers generalization performance degradation for large local mini-batch size. An asynchronous method is proposed to improve the generalization performance. A sublinear convergence rate is provided for the non-convex objective. As there are some missing definitions and little explanation of the proposed method, the reviewer finds the paper hard to read.