Delayed Gradient Averaging: Tolerate the Communication Latency for Federated Learning

1 Introduction . Federated Learning [ 18 , 27 ] has gained growing attention as it enables multi-clients distributed training without exposing the data from private users . During the training , only the model updates are exchanged between clients and servers , thus private training data never leave local devices , enhancing privacy . Many successful applications such as next-word prediction [ 10 ] , voice recognition [ 38 ] , and health care applications [ 50 ] have been derived under the framework . A significant difference between federated scenarios and typical in-center distributed settings [ 6 , 8 ] is the networking condition . Unlike high-end in-cluster network infrastructures where high bandwidth ( 100⇠Gbps ) and low latency ( 1ms ) network is available , edge devices are usually connected through wireless and long-distance connection , thus the bandwidth and latency are strictly limited . This dwarfs the performance of federated systems and slows the development of related applications . While the bandwidth constraints have been efficiently addressed by gradient compression [ 25 ] , lowrank updates [ 18 ] and quantization techniques [ 43 ] , the issue related to network latency is rarely studied in the recent literature [ 12 , 15 ] . However , high network latency is inevitable in federated settings because of ( 1 ) wireless connections and ( 2 ) long-distance transmission ( Figure . 1. ( b ) ) . On the one hand , the high-density urban office and home environments create a lot of contention , as dozens of devices compete for the same radio frequency . On the other hand , the multi-geographic located data entails a minimum latency cost : even with the speed of light , it requires hundreds of milliseconds to send a packet across the world . In either case , the high latency is a hard barrier introduced by physical limits thus inevitable . If not handled specifically , such communication lag , in the magnitude of hundreds of milliseconds , or even seconds , will significantly degrade the scalability of the learning algorithm as shown in the Figure . 1 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 16x4 = 64 cards , batch size 32-2 latency ( ms ) Scalability speed 1us 1.00 504 501 5 1ms 0.99 501 503 10 10ms 0.94 472 472 20 100ms 0.70 352 447 50 500ms 0.29 145 399 100 1s 0.17 87 352 200 234 500 145 1000 87 2000 39 5000 Sp ee du p 0.00 0.25 0.50 0.75 1.00 1us 1ms 10ms 100ms 500ms 98x worse4x worse Across the World > 100ms Wireless Connection ~ 20ms Within a Rack < 1us Within a Data Center ~ 1ms ( a ) Conventional Distributed Settings ( b ) Federated Distributed Settings ( < 1us ) Within a Rack ( ~1ms ) Within a Data Center ( a ) Conventional Distributed Settings ( b ) Federated Distributed Settings Sp ee du p 0.00 0.25 0.50 0.75 1.00 1us 1ms 10ms 100ms 500ms 98x worse 4x worse ( c ) System Performance Comparison Within a Rack ( < 1us ) Within a Data Center ( ~1ms ) World-wide ( > 100ms ) Wireless ( ~ 20ms ) ( a ) Conventional Distributed Settings ( b ) Federated Distributed Settings ( ~20ms ) Across the World ( > 100ms ) Wireless ConnectionWirel ss Connection ( ~20ms ) 1 In this paper , we propose Delayed Gradient Aggregation ( DGA ) to address the latency bottleneck . The key idea is to delay the gradient averaging to a future iteration thus the communication can be pipelined with computation . By accepting stale average gradients for model updates , DGA allows the communic tion to execute par llel with the computation , thus scalable even under extreme latency . We prove that our DGA shares the same convergence as FedAvg and provide extensive experiments on image and language tasks in both i.i.d . and non-i.i.d settings . We demonstrate that ( i ) DGA speeds up FedAvg by a factor 2.6⇥ to 4.1⇥ over various datasets ; ( ii ) no accuracy drop under extreme communication lag ( e.g. , > 1 second ) . We further set up a Raspberry Pi cluster to simulate the real-world scenario and demonstrate that DGA is robust against network stragglers . Our contributions can be summarized as follows : • We propose DGA , a novel distributed optimization method to tolerate the communication latency for federated learning . To our best knowledge , our algorithm is the first work that can achieve scalable federated training under high latency ( > 1s ) . • We theoretically justify the convergence of DGA . We show that under reasonable delay interval , it shares the same convergence rate as FedAvg [ 51 ] . We also discuss its extension with momentum update to better suit modern federated optimizations . • We empirically evaluate the accuracy on diverse datasets and benchmark the speed on different latency setups . Under an extremely high latency , DGA can show impressive improvement over previous algorithms while preserving similar performance on both i.i.d and non-i.i.d partitions . • We build up a Raspberry Pi cluster of 16 devices to evaluate our algorithm in a real-world setting . Even with unpredictable packet loss and latency fluctuation , DGA still shows consistent speedup . 2 Delayed Gradient Averaging . 2.1 Problem Setting , Preliminaries and Related Works . In this work , we formalize the problem as the minimization of sum of stochastic functions , min w f ( w ) = 1 N NX i=1 fi ( w ) with fi ( w ) = E⇣i [ Fi ( w , ⇣i ) ] , where N denotes the number of clients and fi represents the loss function on i-th client . In the setting of empirical risk minimization , fi could be further expressed as finite sums and the random variable ⇣i corresponds to a mini-batch sample . Computation and Communication cost Each client represents a computational resource that can either perform local gradient updates or exchange information with others . We model the computation and communication cost by : • We encounter a computational cost ⌧g for each gradient/mini-batch evaluation . • The communication suffers a cost ⌧c per transmission . We further define the normalized com- munication cost : D = ⇠ ⌧c ⌧g ⇡ . 2 In other words , one communication round takes the same amount of time as D gradient updates . This normalized cost D serves as the main parameter in our discussion on the computation/communication trade-off . Note that the communication cost usually depends on the bandwidth , amount of bits transmitted and the latency , we start revisiting FedAvg [ 27 ] to illustrate the effect on latency . FedAvg and its variants One of the key components in Federated Learning is the step of parameter averaging . This step aggregates the model parameter across clients , without revealing personal training data , enhancing privacy . However , the averaging step is synchronous by design [ 3 ] , in the sense that no local updates are allowed when the averaging takes place . Specifically , FedAvg [ 27 ] alternate the local updates and averaging step : FedAvg : averaging ( communication ) ! K⇥local updates ! averaging ( communication ) · · · where the parameter K indicates the number of mini-batches ( gradients ) locally updated between two averaging steps . The alternating structure of the update naturally decomposes the algorithm into rounds , with K local updates and one communication per round . The total running time of FedAvg after T rounds is given by TFedAvg = T ( K⌧g + ⌧c ) = T ✓ K + ⌧c ⌧g ◆ ⌧g ⇡ T ( K +D ) ⌧g . ( 1 ) As we can see , the communication cost D inflates the total run time by a non-negligible factor , especially when D is large . For instance , if D K , half of the running time is indeed wasted , where clients sit and wait on the synchronization of averages , making the algorithm inefficient . The same drawback persists in all recent extensions of FedAvg [ 17 , 21 , 22 , 28 , 51 ] , its momentum variants [ 5 , 26 , 47 ] , variance reduced variant SCAFFOLD [ 16 ] and adaptive variant [ 34 ] , etc . Compression/Quantization techniques in FL To improve communication efficiency , extensive efforts have been devoted to reducing the bits on gradient exchanges in large-scale distributed training . Techniques such as gradient quantization QSGD [ 1 ] , Onebit [ 39 ] , Tengrad [ 49 ] and gradient compression DGC [ 25 ] , Sparsification [ 48 ] , DenseCommu [ 43 ] , DoubleSqueeze [ 44 ] can safely reduce the amount of information transmitted by a factor of 1000 , while maintaining the high performance of the model . These techniques have been successfully applied to accelerate FL [ 9 , 24 , 35 , 36 ] , enabling training on mobile devices with limited bandwidth connection . Therefore , bandwidth is no longer a critical bottleneck on edge learning . On the other hand , the latency issue is less discussed in the current FL literature [ 15 ] as it is a hard physical barrier . To focus on such inherent and non-improvable subject , we assume that the bandwidth is sufficient , * and the communication cost ⌧c is dominated by network ’ s latency . In this case , the running time of FedAvg and its variants in Eq.1 can still be significantly inflated by the latency , even compression/quantization techniques are applied . The ineffectiveness of the existing approaches to deal with the high latency setting motivates us to ask the following question : is it possible to design an algorithm that could scale up under high latency settings while keeping the accuracy ? We provide a positive answer by introducing the delayed gradient averaging . 2.2 A step-by-step walk through Delayed Gradient Averaging . The main idea of our algorithm is to allow local updates during the communication of the averaging . In FedAvg , clients send its parameter to each other at the end of each round , wait until the averaging ( communication ) ended , later starts the next round . In our algorithm , the averaging barrier is delayed to a later iteration so that clients can immediately start the next round and an gradient correction term is designed to compensate the staleness : 1 . Clients send updates to each other at the end of t-th round ; 2 . Clients continue performing local updates using the latest local parameter ; 3 . When other ’ s t-th round info arrived , the client has already performed D extra local updates . 4 . Delayed averaging step : replace the local gradients at the t-th round by the received averaging ( we will detail this step right away ) . * otherwise apply gradient compression techniques . 3 Algorithm 1 Delayed Gradient Averaging ( DGA ) 1 : Initialize each worker with wi1,1 = w1 for i 2 [ 1 , N ] , the number of local update K , the delayed parameter D 1 . Define s = ( D 1 ) //K as the integer quotient . 2 : for rounds t = 1 , · · · , T do 3 : for client i in parallel do 4 : Set wit,1 = wit 1 , K+1 as the last iterate at round ( t 1 ) 5 : for k = 1 , · · · , K do 6 : Sample the stochastic gradient git , k at the previous iterate w i t , k and update w i t , k+1 = ( w i t , k ⌘git , k if k 6⌘ D ( mod K ) or t 1 s < 1 ; w i t , k ⌘ ( git , k mit 1 s +mt 1 s ) if k ⌘ D ( mod K ) and t 1 s 1. where mit 1 s is the accumulated gradient ( see line 8 ) at the earlier round t 1 s , mt 1 s is the average of mit 1 s among all clients , i.e . mt 1 s = 1 N P i m i t 1 s. 7 : end for 8 : Send the t-th round accumulated gradient mit = PK k=1 g i t , k to all other clients . † 9 : end for 10 : end for 11 : Return wT = 1N PN i=1 w i T , K+1 . The benefit is plain to see : we no longer freeze the local computation power during the communication . To make the discussion more explicit , we denote the parameter on the i-th client at the k-th iteration within the t-th round by wit , k and the corresponding stochastic gradient as g i t , k . In the very first round t = 1 , solely local updates are performed , hence the last iterate in the first round can be expressed as w i 1 , K+1 = w i 1 , K ⌘gi1 , K = · · · = w1 ⌘ KX k=1 g i 1 , k | { z } : =mi1 . As the first round ’ s computation has completed , we send the accumulated gradient mi1 : = PK k=1 g i 1 , k to all the clients , i.e . perform averaging . Right after the gradients are sent , we immediately proceed the second round ’ s local updates , leaving the first round averages in transmission . By the time the average is received , we already performed D extra local updates in the second round , starting from the last iterate of first round wi1 , K+1 : w i 2 , D ⌘gi2 , D| { z } last iterate before averaging = · · · = wi1 , K+1 ⌘ DX k=1 g i 2 , k | { z } D updates = w1 ⌘ KX k=1 g i 1 , k | { z } 1st round ⌘ DX k=1 g i 2 , k | { z } 2nd round . ( 2 ) At this point , the average of the first round arrives m1 = 1 N NX i=1 m i 1 = 1 N NX i=1 KX k=1 g i 1 , k = KX k=1 1 N NX i=1 g i 1 , k ! : = KX k=1 g1 , k , ( accumulated averages ) where g1 , k represents the average gradient at k-th iteration in the 1st round . With the averaging m1 in hand , we substitute all the first round local gradients in Eq.2 by their averages , leading to : w i 2 , D+1 = w1 ⌘ KX k=1 g1 , k | { z } common on all clients ⌘ DX k=1 g i 2 , k | { z } 2nd round unchanged ( delayed gradients averaging ) †The accumulated gradient is indeed implemented as mt mt + git , k under the for loop on k , which do not require extra memory to store gt , k . 4 Computation Communication Synchronization Barrier Increased latency slows the training ! Latency is amortized by multiple steps , but still slow the training ! No longer blocks the training ! ( 1 ) Vanilla Synchronous SGD for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) send ( grad [ i ] , iter=i ) avg_grad = recv ( iter=i ) W = W - lr * avg_grad for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) send ( grad [ i ] , iter=i ) stale_avg_grad = recv ( iter=i-d ) W = W - lr * ( grad [ i ] - grad [ i-d ] + stale_avg_grad ) for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) if ( I+d ) % k == 0 : send ( grad [ i ] , iter=i ) If I % k == 0 stale_avg_grad = recv ( iter=i-d ) W = W - lr * ( grad [ i ] - grad [ i-d ] + stale_avg_grad ) K 统⼀⼤写 ( 2 ) Federated Averaging ( K=2 ) ( 4 ) DGA ( D=2 , K=2 ) Latency is amortized by multiple steps , but still slow the training ! High latency longer blocks the training ! Delay the averaging High latency longer blocks the training ! ( a ) Federated Averaging ( K=1 ) ( c ) Delayed Averaging ( D=2 , K=1 ) ( b ) Federated Averaging ( K=2 ) ( d ) Delayed Averaging ( D=2 , K=2 ) ( b ) Delayed Gradient Averaging : D indicates how many steps the averaging receiving timing are delayed . Computation Communication Synchronization Barrier Increased latency slows the training ! Latency is amortized by multiple steps , but still slow the training ! No longer blocks the training ! ( 1 ) Vanilla Synchronous SGD for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) send ( grad [ i ] , iter=i ) avg_grad = recv ( iter=i ) W = W - lr * avg_grad for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) send ( grad [ i ] , iter=i ) stale_avg_grad = recv ( iter=i-d ) W = W - lr * ( grad [ i ] - grad [ i-d ] + stale_avg_grad ) for i in iters : L = loss ( net , data , label ) grad [ i ] = L.backward ( ) if ( I+d ) % k == 0 : send ( grad [ i ] , iter=i ) If I % k == 0 stale_avg_grad = recv ( iter=i-d ) W = W - lr * ( grad [ i ] - grad [ i-d ] + stale_avg_grad ) K 统⼀⼤写 ( 2 ) Federated Averaging ( K=2 ) ( 4 ) DGA ( D=2 , K=2 ) Latency is amortized by multiple steps , but still slow the training ! High latency longer blocks the training ! Delay the averaging High latency longer blocks the training ! ( a ) Federated Averaging ( K=1 ) ( c ) Delayed Averaging ( D=2 , K=1 ) ( b ) Federated Averaging ( K=2 ) ( d ) Delayed Averaging ( D=2 , K=2 ) Rewrite the update rule using the accumulated gradients mi1 and m1 gives the compact form : w i 2 , D+1 = w i 2 , D ⌘ ( gi2 , D mi1 +m1| { z } gradient correction ) , ( compact DGA update ) serving as the main building block in our Algorithm 1 . In the most general setting , the delayed gradients could take several rounds before arriving . This happens when the latency is high , i.e . D > K. To take care of such cases , we express the delay parameter in the quotient/remainder form D = sK + r , with s 0 and r 2 [ 1 , K ] . The integer quotient s indicates the number of rounds skipped due to the latency . In this scenario , the gradients sent at t 1 s round is received in the middle of the t-th round , which justifies our update rule in Algorithm 1 . The example we discussed earlier in ( compact DGA update ) is a special case when s = 0 , i.e . D K , where the gradients sent in the last round are received in the current round . We call such operation as Delayed Gradients Averaging ( DGA ) , as every round ’ s averaging has been delayed to later rounds.‡ With our designed gradient correction term in ( compact DGA update , blue part in Figure . 2b ) , each update ensures that all clients replace the previous round ’ s local gradients by their averages , thus shares the same ( t D ) -round gradients . In another word , different workers only differ on the most recent D gradients . Last but not the least , in the ideal case when there is no delay , i.e . D = 0 , our algorithm DGA recovers the original Federated Averaging [ 27 ] . When D 1 , our algorithm DGA allows the clients to pipeline local updates with communication , and the divergence of Federated Average and DGA is bounded by a constant number . To further clarify , we now conduct a sketch of the convergence analysis by bounding the staleness between different clients . ‡DGA requires to store the recent D copies of gradients to perform the update . This can bring challenges to edge platforms when the model size is very large . 5

Summary: The work proposes to overlap the local gradient computation step on each client in a federated ML setup with the latency cost associated with communicating the averaged gradients from all clients. This is done by allowing the clients to proceed with computation of local gradients on K mini-batches for a new round (t+1) using the latest local parameter until the gradient average from the previous round (t) is received with a delay (aka, *Delayed Gradient Averaging*). In doing so, the client is said to have performed D extra local updates which are combined with the delayed gradient average. The work is said to be the first of its kind to consider high latency (>1s) situations to perform scalable federated learning without sacrificing performance.