IntSGD: Adaptive Floatless Compression of Stochastic Gradients

1 INTRODUCTION . Many recent breakthroughs in machine learning were made possible due to the introduction of large , sophisticated and high capacity supervised models whose training requires days or even weeks of computation ( Hinton et al. , 2015 ; He et al. , 2016 ; Huang et al. , 2017 ; Devlin et al. , 2018 ) . However , it would not be possible to train them without corresponding advances in parallel and distributed algorithms capable of taking advantage of modern hardware . Very large models are typically trained on vast collections of training data stored in a distributed fashion across a number of compute nodes that need to communicate throughout the training process . In this scenario , reliance on efficient communication protocols is of utmost importance . Communication in distributed systems . The training process of large models relies on fast synchronization of gradients computed in a parallel fashion . Formally , to train a model , we want to solve the problem of parallel/distributed minimization of the average of n functions : min x∈Rd [ f ( x ) def = 1n n∑ i=1 fi ( x ) ] , fi ( x ) def = Eξ [ fi ( x ; ξ ) ] , ( 1 ) where we will compute the gradients of stochastic realizations fi ( x ; ξ ) . The two dominating protocols for gradient synchronization are all-reduce and all-gather aggregation , which may use either Parameter Server or all-to-all communication under the hood . The core difference between them lies in that all-gather communicates all vectors , whereas all-reduce only outputs their average . As shown in previous works , current distributed deep learning algorithms predominantly use all-reduce as it scales much better than all-gather ( Vogels et al. , 2019 ; Agarwal et al. , 2021 ) . A popular way to reduce the communication cost of both all-reduce and all-gather primitives is to use lossy compression of gradients ( Ramesh et al. , 2021 ) . To study the benefit of lossy compression , large swaths of recent literature on distributed training attribute the cost of sending a single vector from a worker to the server to the number of bits needed to represent it . Based on this abstraction , various elaborate vector compression techniques ( see Table 1 in ( Beznosikov et al. , 2020 ; Xu et al. , 2020 ; Safaryan et al. , 2020 ) ) and algorithms have been designed for higher and higher compression ratios . However , in real systems , the efficiency of sending a vector is not fully characterized by the number of bits alone , because : • First , many compressors with high compression ratio ( e.g. , natural compression ( Horváth et al. , 2019 ) , quantization ( Alistarh et al. , 2017 ) , top-k sparsification , sign ( Bernstein et al. , 2018 ) ) are not compatible with the efficient all-reduce primitive and require all-gather implementation . • Secondly , some compressors rely on expensive operations such as low-rank decomposition ( Wang et al. , 2018 ; Vogels et al. , 2019 ) or bit-level operations ( Horváth et al. , 2019 ) , whose computation overhead may outweigh the benefits of reduced communication load . • Thirdly , algorithms with biased compressors such as Top-k SGD , SignSGD , PowerSGD ( Vogels et al. , 2019 ) , require the error-feedback ( EF-SGD ) mechanism ( Stich et al. , 2018 ; Karimireddy et al. , 2019 ) to ensure the convergence . Alas , error feedback needs extra sequences that may not fit the low memory budget of GPUs . Moreover , to the best of our knowledge , no convergence guarantee has been established for EF-SGD on the non-smooth objectives with multiple workers . SwitchML . Another approach to combating long communication times is to improve the hardware itself . The recently proposed SwitchML is an alternative to the NVIDIA Collective Communications Library ( NCCL ) on real-world hardware ( Sapio et al. , 2021 ) . The first key component of SwitchML is the in-network aggregation ( INA ) by a programmable switch . INA reduces the communication cost and latency because the execution can be paralleled and pipelined . To be specific , it splits the vector to aggregate into chunks and processes them individually by the switch pipeline . The advantages of INA over parameter server and all-reduce in terms of latency and communication cost have been theoretically and empirically justified by Sapio et al . ( 2021 ) . The second key component of SwitchML is stochastic gradient descent with integer rounding and aggregation . Instead of reducing the data volume to exchange , the goal of integer rounding in SwitchML is to fit the limited computation capability of the modern programmable switch , which only supports integer additions or logic operations . To increase the rounding precision , the gradient gki on device imultiplies a positive scaling factor αk known to every worker and then rounded to an integer number Int ( αk ◦ gki ) . As there is no additional scaling or decompression before aggregating the communicated vectors , their sums can be computed on the fly . Then , each worker can divide the aggregated gradient by nαk to update the model . However , Sapio et al . ( 2021 ) remark that the choice of the scaling factor αk requires special care . In their presentation1 , one of the authors notes : “ A bad choice of scaling factor can reduce the performance. ” To this end , they propose a heuristic-based profiling step that is executed before the gradient aggregation and keeps the rounded integers small to fit in 32 bits . We refer to their algorithm including the profiling step as Heuristic IntSGD . Unfortunately , no convergence guarantee for that algorithm has been established . This is where our theory comes to the rescue . By rigorously and exhaustively analyzing integer rounding based on scaling , we find adaptive rules for the scaling factor αk that do not require the profiling employed by Sapio et al . ( 2021 ) . As we will show in the remainder of the paper , our algorithm is perfectly suited for both in-network aggregation ( INA ) of SwitchML and for other efficient primitives such as all-reduce . 1.1 CONTRIBUTIONS . We summarize the key differences of our algorithm and prior work in Table 1 , and we also list our main contributions below . • Adaptive IntSGD . We develop a family of computationally cheap adaptive scaling factors for provably convergent IntSGD . It is a better alternative to the Heuristic IntSGD in ( Sapio et al. , 2021 ) that requires expensive operations and does not ensure convergence . • Rates . We obtain the first analysis of the integer rounding and aggregation for distributed machine learning . For all of the proposed variants , we prove convergence rates of IntSGD that match those of full-precision SGD up to constant factors . Our results are tight and apply to both convex and nonconvex problems . Our analysis does not require any extra assumption compared to those typically invoked for SGD . In contrast to other compression-based methods , IntSGD has the same rate as that of full-precision SGD even on non-smooth problems . • IntDIANA . We observe empirically that IntSGD struggles when the devices have heterogeneous ( non-identical ) data—an issue it shares with vanilla SGD—and propose an alternative method , IntDI- 1https : //youtu.be/gBPHFyBWVoM ? t=606 ANA , that can provably alleviate this issue . We also show that our tools are useful for extending the methodology beyond SGD methods , for example , to variance reduced methods ( Johnson & Zhang , 2013 ; Allen-Zhu & Hazan , 2016 ; Kovalev et al. , 2020 ; Gower et al. , 2020 ) with integer rounding . Please refer to Appendix A.2 for theoretical results and Appendix C.5 for the empirical verification . 2 ADAPTIVE INTEGER ROUNDING AND INTSGD . By randomized integer rounding we mean the mapping Int : R→ Z defined by Int ( t ) def= { [ t ] + 1 , with probability pt def = t− [ t ] , [ t ] , with probability 1− pt , where [ t ] denotes the floor of t ∈ R , i.e. , [ t ] = k ∈ Z , where k is such that k ≤ t < k+ 1 . Note that E [ Int ( t ) ] = ( t− [ t ] ) ( [ t ] + 1 ) + ( [ t ] + 1− t ) [ t ] = t. We extend this mapping to vectors x ∈ Rd by applying in element-wise : Int ( x ) i def = Int ( xi ) . 2.1 ADAPTIVE INTEGER ROUNDING . Given a scaling vector α ∈ Rd with nonzero entries , we further define the adaptive integer rounding operator Q : Rd → Rd by Q ( x ) def = 1α ◦ Int ( α ◦ x ) , ( 2 ) where a ◦ b def= ( a1b1 , . . . , adbd ) ∈ Rd denotes the Hadamard product of two vectors a = ( a1 , . . . , ad ) ∈ Rd and b = ( b1 , . . . , bd ) ∈ Rd . As we show below , the adaptive integer rounding operator ( 2 ) has several properties which will be useful in our analysis . In particular , the operator is unbiased , and its variance can be controlled by choice of a possibly random scaling vector α ∈ Rd++ . Lemma 1 . For any x ∈ Rd and α ∈ Rd++ , we have 1 α ◦ E [ Int ( α ◦ x ) ] = x , ( 3 ) E [ ∥∥ 1 α ◦ Int ( α ◦ x ) − x ∥∥2 ] ≤ d∑ j=1 1 4α2j , ( 4 ) Algorithm 1 IntSGD . Default setting for the tested problems : β = 0.9 , ε = 10−8 . 1 : Params : Stepsizes ηk , scaling vectors αk ∈ Rd 2 : Init : x0 ∈ Rd , x1 = x0 − η0 1n ∑n i=1 g 0 i 3 : for k = 1 , 2 , . . . do 4 : for each device i = 1 , 2 , . . . , n do 5 : Compute stochastic gradient gki ( E [ gki | xk ] ∈ ∂fi ( xk ) ) 6 : Maintain the moving average : rk = βrk−1 + ( 1− β ) ‖xk − xk−1‖2 { rk is a scalar } 7 : Compute the adaptive scaling factor : αk = √ d√ 2nrk/η2k+ε 2 8 : Scale and round the local gradient Q ( gki ) = Int ( αk ◦ gki ) 9 : end for 10 : Aggregate Q ( gki ) by either all-reduce or in-network aggregation ( INA ) 11 : for each device i = 1 , 2 , . . . , n do 12 : Compute the ( sub ) gradient estimator : g̃k = 1nαk ∑n i=1Q ( g k i ) 13 : Update the model parameter xk+1 = xk − ηkg̃k 14 : end for 15 : end for The expectations above are taken with respect to the randomness inherent in the rounding operator . 2.2 NEW ALGORITHM : INTSGD . We are ready to present our algorithm , IntSGD . At iteration k , each device i computes a stochastic ( sub ) gradient vector gki , i.e. , a vector satisfying E [ gki | xk ] ∈ ∂fi ( xk ) . ( 5 ) Prior to communication , each worker i rescales its stochastic ( sub ) gradients gki using the same vector αk ∈ Rd++ , and applies the randomized rounding operator Int . The resulting vectors Int ( αk ◦ gki ) are aggregated to obtain ∑n i=1 Int ( αk ◦ gki ) , which is also an integer . Each device subsequently performs division by n and inverse scaling to decode the message , obtaining the vector g̃k def = 1nαk ◦ n∑ i=1 Int ( αk ◦ gki ) = 1n n∑ i=1 1 αk ◦ Int ( αk ◦ gki ) ( 2 ) = 1n n∑ i=1 Q ( gki ) . Here αk is a random adaptive scaling factor calculated based on the historical information . We left the design of αk to Section 4 . By combining ( 5 ) and ( 3 ) , we observe that gk is a stochastic ( sub ) gradient of f at xk . Finally , all devices perform in parallel an SGD-type step of the form xk+1 = xk−ηkg̃k and the process is repeated . Our IntSGD method is formally stated as Algorithm 1 with the suggested rule of α . Relation to QSGD ( Alistarh et al. , 2017 ) . QSGD bears some similarity to the IntSGD : Both of them scale gki by a factor before the quantization ( the scaling factor in QSGD is 1/‖gki ‖ for normalization ) . However , some key difference makes the communication efficiency of IntSGD much better than that of QSGD . It is worth noting that the normalization factors 1/‖gki ‖ in QSGD are different for various workers . Then , the quantized values of various workers need to be gathered and decompressed before aggregation . On the contrary , the scaling factor αk in our IntSGD is the same for all workers such that the sum of integers can be computed on the fly . Thus , IntSGD supports the efficient allreduce primitive while QSGD does not . As seen in the experimental results in Section 5 , this makes a big difference in empirical performance . Moreover , the proof technique for IntSGD is also intrinsically different from that of QSGD . Please see the next section for the details .

Summary: The paper introduces a randomized compress-to-integer operator with a shared scaling factor for use in data communication in distributed SGD.  The resulting algorithm is provably convergent, matches the behavior of SGD up to constant factors, and works well with the all-reduce primitive.  The authors claim three main contributions: the IntSGD algorithm itself, the analysis of convergence rates, and the IntDIANA variant for heterogeneous data distributions (though this is only discussed in the appendices).