Communication-Efficient Distributed Blockwise Momentum SGD with Error-Feedback

Introduction

Deep neural networks have been highly successful in recent years . To achieve state-of-the-art performance, they often have to leverage the computing power of multiple machines during training . Popular approaches include distributed synchronous SGD and its momentum variant SGDM, in which the computational load for evaluating a mini-batch gradient is distributed among the workers. Each worker performs local computation, and these local informations are then merged by the server for final update on the model parameters. However, its scalability is limited by the possibly overwhelming cost due to communication of the gradient and model parameter . Let $d$ be the gradient/parameter dimensionality, and $M$ be the number of workers. $64Md$ bits need to be transferred between the workers and server in each iteration.

To mitigate this communication bottleneck, the two common approaches are gradient sparsification and gradient quantization. Gradient sparsification only sends the most significant, information-preserving gradient entries. A heuristic algorithm is first introduced in , in which only the large entries are transmitted. On training a neural machine translation model with 4 GPUs, this greatly reduces the communication overhead and achieves 22% speedup . Deep gradient compression is another heuristic method that combines gradient sparsification with other techniques such as momentum correction, local gradient clipping, and momentum factor masking, achieving significant reduction on communication cost. Recently, a stochastic sparsification method was proposed in that balances sparsity and variance by solving a constrained linear programming. MEM-SGD combines top-$k$ sparsification with error correction. By keeping track of the accumulated errors, these can be added back to the gradient estimator before each transmission. MEM-SGD converges at the same rate as SGD on convex problems, whilst reducing the communication overhead by a factor equal to the problem dimensionality.

On the other hand, gradient quantization mitigates the communication bottleneck by lowering the gradient’s floating-point precision with a smaller bit width. 1-bit SGD achieves state-of-the-art results on acoustic modeling while dramatically reducing the communication cost . TernGrad quantizes the gradients to ternary levels $\{-1,0,1\}$. QSGD employs stochastic randomized rounding to ensure unbiasedness of the estimator. Error-compensated quantized SGD (ECQ-SGD) was proposed in , wherein a similar stochastic quantization function used in QSGD is employed, and an error bound is obtained for quadratic loss functions. Different from the error-feedback mechanism proposed in MEM-SGD, ECQ-SGD requires two more hyper-parameters and its quantization errors are decayed exponentially. Thus, error feedback is limited to a small number of iterations. Also, ECQ-SGD uses all-to-all broadcast (which may involve large network traffic and idle time), while we consider parameter-server architecture. Recently, Bernstein et al. proposed signSGD with majority vote , which only transmits the 1-bit gradient sign between workers and server. A variant using momentum, called signum with majority vote, is also introduced though without convergence analysis . Using the majority vote, signSGD achieves a notion of Byzantine fault tolerance . Moreover, it converges at the same rate as distributed SGD, though it has to rely on the unrealistic assumptions of having a large mini-batch and unimodal symmetric gradient noise. Indeed, signSGD can diverge in some simple cases when these assumptions are violated . With only a single worker, this divergence issue can be fixed by using the error correction technique in MEM-SGD, leading to SGD with error-feedback (EF-SGD) .

While only a single worker is considered in EF-SGD, we study in this paper the more interesting distributed setting. An extension of MEM-SGD and EF-SGD with parallel computing was proposed in for all-to-all broadcast. Another related architecture is allreduce. Compression at the server can be implemented between the reduce and broadcast steps in tree allreduce, or between the reduce-scatter and allgather steps in ring allreduce. However, allreduce requires repeated gradient aggregations, and the compressed gradients need to be first decompressed before they are summed. Hence, heavy overheads may be incurred.

Our contributions are: (i) We provide a bound on dist-EF-SGD with general stepsize schedule for a class of compressors (including the commonly used sign-operator and top-$k$ sparsification). In particular, without relying on the unrealistic assumptions in , we show that dist-EF-SGD with constant/decreasing/increasing stepsize converges at an $\mathcal{O}(1/\sqrt{MT})$ rate, which matches that of distributed synchronous SGD; (ii) We study gradient compression with Nesterov’s momentum in a parameter server. For dist-EF-SGDM with constant stepsize, we obtain an $\mathcal{O}(1/\sqrt{MT})$ rate. To the best of our knowledge, these are the first convergence results on two-way gradient compression with Nesterov’s momentum; (iii) We propose a general blockwise compressor and show its theoretical properties. Experimental results show that the proposed algorithms are efficient without losing prediction accuracy. After our paper has appeared, we note a similar idea was independently proposed in . Different from ours, they do not consider changing stepsize, blockwise compressor and Nesterov’s momentum.

Related Work: SGD with Error-Feedback

Recently, Karimireddy et al. introduced SGD with error-feedback (EF-SGD), which combines gradient compression with error correction (Algorithm 1). A single machine is considered, which keeps the gradient difference that is not used for parameter update in the current iteration. In the next iteration $t$, the accumulated residual $e_{t}$ is added to the current gradient. The corrected gradient $p_{t}$ is then fed into an $\delta$-approximate compressor.

Examples of $\delta$-approximate compressors include the scaled sign operator $\mathcal{C}(v)=\|v\|_{1}/d\cdot\text{sign}(v)$ and top-$k$ operator (which only preserves the $k$ coordinates with the largest absolute values) . One can also have randomized compressors that only satisfy Definition 1 in expectation. Obviously, it is desirable to have a large $\delta$ while achieving low communication cost.

Distributed Blockwise Momentum SGD with Error-Feedback

In the following, we investigate the convergence of dist-EF-SGD. We make the following assumptions, which are common in the stochastic approximation literature.

The full gradient $\nabla F$ is uniformly bounded: $\|\nabla F(x_{t})\|_{2}^{2}\leq\omega^{2}$.

Suppose that Assumptions 1-3 hold. Assume that $0<\eta_{t}<3/(2L)$ for all $t$. For the $\{x_{t}\}$ sequence generated from Algorithm 2, we have

where $o\in\{0,\dots,T-1\}$ is an index such that $P(o=k)=\frac{\eta_{k}\left(3-2L\eta_{k}\right)}{\sum_{t=0}^{T-1}\eta_{t}\left(3-2L\eta_{t}\right)}$, $\forall k=0,\dots,T-1$.

Let stepsize $\eta=\min(\frac{1}{2L},\frac{\gamma}{\sqrt{T}/\sqrt{M}+(1-\delta)^{1/3}\left(1/\delta^{2}+16/\delta^{4}\right)^{1/3}T^{1/3}})$ for some $\gamma>0$. Then,

In comparison, under the same assumptions, distributed synchronous SGD achieves

Thus, the convergence rate of dist-EF-SGD matches that of distributed synchronous SGD (with full-precision gradients) after $T\geq O(1/\delta^{2})$ iterations, even though gradient compression is used. Moreover, more workers (larger $M$) leads to faster convergence. Note that the bound above does not reduce to that of EF-SGD when $M=1$, as we have two-way compression. When $M=1$, our bound also differs from Remark 4 in in that our last term is $O((1-\delta)^{1/3}/(\delta^{4/3}T^{2/3}))$, while theirs is $O((1-\delta)/(\delta^{2}T))$ (which is for single machine with one-way compression). Ours is worse by a factor of $O(T^{1/3}\delta^{2/3}/(1-\delta)^{2/3})$, which is the price to pay for two-way compression and a linear speedup of using $M$ workers. Moreover, unlike signSGD with majority vote , we achieve a convergence rate of $\mathcal{O}(1/\sqrt{MT})$ without assuming a large mini-batch size ($=T$) and unimodal symmetric gradient noise.

Theorem 1 only requires $0<\eta_{t}<3/(2L)$ for all $t$. This thus allows the use of any decreasing, increasing, or hybrid stepsize schedule. In particular, we have the following Corollary.

Let $\eta_{t}=\frac{\gamma}{((t+1)T)^{1/4}/(\sqrt{M})+(1-\delta)^{1/3}\left(1/\delta^{2}+16/\delta^{4}\right)^{1/3}T^{1/3}}$ (decreasing stepsize) with $T\geq 16L^{4}\gamma^{4}M^{2}$ or $\eta_{t}=\frac{\gamma\sqrt{t+1}}{T/\sqrt{M}+(1-\delta)^{1/3}\left(1/\delta^{2}+16/\delta^{4}\right)^{1/3}T^{5/6}}$ (increasing stepsize) with $T\geq 4L^{2}\gamma^{2}M$. Then, dist-EF-SGD converges to a stationary point at a rate of $\mathcal{O}(1/\sqrt{MT})$.

To the best of our knowledge, this is the first such result for distributed compressed SGD with decreasing/increasing stepsize on nonconvex problems. These two stepsize schedules can also be used together. For example, one can use an increasing stepsize at the beginning of training as warm-up, and then a decreasing stepsize afterwards.

2 Blockwise Compressor

Compared to using only the sign operator as in signSGD, the factor $\|v\|_{1}/d$ can preserve the gradient’s magnitude. However, as shown in , its $\delta$ in Definition 1 is $\|v\|_{1}^{2}/(d\|v\|_{2}^{2})$, and can be particularly small when $v$ is sparse. When $\delta$ is closer to $1$, the bound in Corollary 1 becomes smaller and thus convergence is faster. In this section, we achieve this by proposing a blockwise extension of (1).

Specifically, we partition the compressor input $v$ into $B$ blocks, where each block $b$ has $d_{b}$ elements indexed by $\mathcal{G}_{b}$. Block $b$ is then compressed with scaling factor $\|v_{\mathcal{G}_{b}}\|_{1}/d_{b}$ (where $v_{\mathcal{G}_{b}}$ is the subvector of $v$ with elements in block $b$), leading to: $\mathcal{C}_{B}(v)=[\|v_{\mathcal{G}_{1}}\|_{1}/d_{1}\cdot\text{sign}(v_{\mathcal{G}_{1}}),\dots,\|v_{\mathcal{G}_{B}}\|_{1}/d_{B}\cdot\text{sign}(v_{\mathcal{G}_{B}})]$. A similar compression scheme, with each layer being a block, is considered in the experiments of . However, they provide no theoretical justifications. The following Proposition first shows that $\mathcal{C}_{B}(\cdot)$ is also an approximate compressor.

Let $[B]=\{1,2,\dots,B\}$. $\mathcal{C}_{B}$ is a $\phi(v)$-approximate compressor, where $\phi(v)=\min_{b\in[B]}\frac{\|v_{\mathcal{G}_{b}}\|_{1}^{2}}{d_{b}\|v_{\mathcal{G}_{b}}\|_{2}^{2}}\geq\min_{b\in[B]}\frac{1}{d_{b}}$.

The resultant algorithm will be called dist-EF-blockSGD (Algorithm 3) in the sequel. As can be seen, this is a special case of Algorithm 2. By replacing $\delta$ with $\phi(v)$ in Proposition 1, the convergence results of dist-EF-SGD in Section 3.1 can be directly applied.

The per-iteration communication costs of the various distributed algorithms are shown in Table 1. Compared to signSGD with majority vote , dist-EF-blockSGD requires an extra $64MB$ bits for transmitting the blockwise scaling factors (each factor $\|v_{\mathcal{G}_{b}}\|_{1}/d_{b}$ is stored in float32 format and transmitted twice in each iteration). By treating each vector/matrix/tensor parameter as a block, $B$ is typically in the order of hundreds. For most problems of interest, $64MB/(2Md)<10^{-3}$. The reduction in communication cost compared to full-precision distributed SGD is thus nearly 32x.

3 Nesterov’s Momentum

Momentum has been widely used in deep networks . Standard distributed SGD with Nesterov’s momentum and full-precision gradients uses the update: $m_{t,i}=\mu m_{t-1,i}+g_{t,i},\forall i\in[M]$ and $x_{t+1}=x_{t}-\eta_{t}\frac{1}{M}\sum_{i=1}^{M}(\mu m_{t,i}+g_{t,i})$, where $m_{t,i}$ is a local momentum vector maintained by each worker $i$ at time $t$ (with $m_{0,i}=0$), and $\mu\in[0,1)$ is the momentum parameter. In this section, we extend the proposed dist-EF-SGD with momentum. Instead of sending the compressed $g_{t,i}+\frac{\eta_{t-1}}{\eta_{t}}e_{t,i}$ to the server, the compressed $\mu m_{t,i}+g_{t,i}+\frac{\eta_{t-1}}{\eta_{t}}e_{t,i}$ is sent. The server merges all the workers’s results and sends it back to each worker. The resultant procedure with blockwise compressor is called dist-EF-blockSGDM (Algorithm 4), and has the same communication cost as dist-EF-blockSGD. The corresponding non-block variant is analogous.

Suppose that Assumptions 1-3 hold. Let $\eta_{t}=\eta$ for some $\eta>0$. For any $\eta\leq\frac{(1-\mu)^{2}}{2L}$, and the $\{x_{t}\}$ sequence generated from Algorithm 4, we have

Compared to Theorem 1, using a larger momentum parameter $\mu$ makes the first term (which depends on the initial condition) smaller but a worse variance term (second term) and error term due to gradient compression (last term). Similar to Theorem 1, a larger $\eta$ makes the third term larger. The following Corollary shows that the proposed dist-EF-blockSGDM achieves a convergence rate of $\mathcal{O}(((1-\mu)[F(x_{0})-F_{*}]+\sigma^{2}/(1-\mu))/\sqrt{MT})$.

Experiments

In this experiment, we demonstrate that the proposed dist-EF-blockSGDM and dist-EF-blockSGD ($\mu=0$ in Algorithm 4), though using fewer bits for gradient transmission, still has good convergence. For faster experimentation, we use a single node with multiple GPUs (an AWS P3.16 instance with 8 Nvidia V100 GPUs, each GPU being a worker) instead of a distributed setting.

Experiment is performed on the CIFAR-100 dataset, with 50K training images and 10K test images. We use a 20-layer ResNet . Each parameter tensor/matrix/vector is treated as a block in dist-EF-blockSGD(M). They are compared with (i) distributed synchronous SGD (with full-precision gradient); (ii) distributed synchronous SGD (full-precision gradient) with momentum (SGDM); (iii) signSGD with majority vote ; and (iv) signum with majority vote . All the algorithms are implemented in MXNet. We vary the mini-batch size per worker in $\{8,16,32\}$. Results are averaged over 5 repetitions. More details of the experiments are shown in Appendix A.1.

Figure 2 shows convergence of the testing accuracy w.r.t. the number of epochs. As can be seen, dist-EF-blockSGD converges as fast as SGD and has slightly better accuracy, while signSGD performs poorly. In particular, dist-EF-blockSGD is robust to the mini-batch size, while the performance of signSGD degrades with smaller mini-batch size (which agrees with the results in ). Momentum makes SGD and dist-EF-blockSGD faster with mini-batch size of $16$ or $32$ per worker, particularly before epoch $100$. At epoch 100, the learning rate is reduced, and the difference is less obvious. This is because a larger mini-batch means smaller variance $\sigma^{2}$, so the initial optimality gap $F(x_{0})-F_{*}$ in (2) is more dominant. Use of momentum ($\mu>0$) is then beneficial. On the other hand, momentum significantly improves signSGD. However, signum is still much worse than dist-EF-blockSGDM.

2 Distributed Training on ImageNet

In this section, we perform distributed optimization on ImageNet using a 50-layer ResNet. Each worker is an AWS P3.2 instance with 1 GPU, and the parameter server is housed in one node. We use the publicly available codehttps://github.com/PermiJW/signSGD-with-Majority-Vote in , and the default communication library Gloo communication library in PyTorch. As in , we use its allreduce implementation for SGDM, which is faster.

As momentum accelerates the training for large mini-batch size in Section 4.1, we only compare the momentum variants here. The proposed dist-EF-blockSGDM is compared with (i) distributed synchronous SGD with momentum (SGDM); and (ii) signum with majority vote . The number of workers $M$ is varied in $\{7,15\}$. With an odd number of workers, a majority vote will not produce zero, and so signum does not lose accuracy by using 1-bit compression. More details of the setup are in Appendix A.2.

Figure 3 shows the testing accuracy w.r.t. the number of epochs and wall clock time. As in Section 4.1, SGDM and dist-EF-blockSGDM have comparable accuracies, while signum is inferior. When 7 workers are used, dist-EF-blockSGDM has higher accuracy than SGDM (76.77% vs 76.27%). dist-EF-blockSGDM reaches SGDM’s highest accuracy in around 13 hours, while SGDM takes 24 hours (Figure 3(b)), leading to a $46\%$ speedup. With 15 machines, the improvement is smaller (Figure 3(e)). This is because the burden on the parameter server is heavier. We expect comparable speedup with the 7-worker setting can be obtained by using more parameter servers. In both cases, signum converges fast but the test accuracies are about $4\%$ worse.

Figures 3(c) and 3(f) show a breakdown of wall clock time into computation and communication time.Following , communication time includes the extra computation time for error feedback and compression. All methods have comparable computation costs, but signum and dist-EF-blockSGDM have lower communication costs than SGDM. The communication costs for signum and dist-EF-blockSGDM are comparable for 7 workers, but for 15 workers signum is lower. We speculate that it is because the sign vectors and scaling factors are sent separately to the server in our implementation, which causes more latency on the server with more workers. This may be alleviated if the two operations are fused.

Conclusion

References

Appendix A Experimental Setup

where $\lambda$ is the weight decay parameter. In the experiment, the sign is mapped to $\{-1,1\}$ and takes 1 bit. Note that the gradient sign has zero probability of being zero.

Each algorithm is run for 200 epochs. We only tune the initial stepsize, using a validation set with 5K images that is carved out from the training set. For dist-EF-blockSGD (resp. dist-EF-blockSGDM), we use the stepsize tuned for SGD (resp. SGDM). The stepsize with the best validation set performance is used to run the algorithm on the full training set. The stepsize is divided by $10$ at the $100$-th and $150$-th epochs. The weight decay parameter is fixed to $0.0005$, and the momentum parameter $\mu$ is $0.9$. When mini-batch size is $16$ per worker, for both SGD and SGDM, the stepsize is tuned from $\{0.05,0.1,0.5,1\}$, and for signSGD and signum, the stepsize is chosen from $\{0.0005,0.001,0.005,0.01\}$. When we obtain the best stepsize $\eta_{0}$ tuned with mini-batch size $B=16$ per worker, for $B=8$, the best stepsize is selected from $\{\eta_{0}/2,\eta_{0}\}$; whereas for $B=32$, it is selected from $\{\eta_{0},2\eta_{0}\}$. The best stepsizes obtained are shown in Table 2

A.2 Setup: Distributed Training on ImageNet

We use the default hyperparameters for SGDM and signum in the code base, which have been tuned for the ImageNet experiment in . Specifically, the momentum parameter $\mu$ is $0.9$, and weight decay parameter is $0.0001$. A mini-batch size of 128 per worker is employed.

For SGDM, we use $\eta=0.1M$ (used for SGDM on the ImageNet experiment in the code base). For signum, $\eta=0.0001$ (used for signum on the ImageNet experiment in the code base) on 7 workers and $\eta=0.0002$ on 15 workers. For dist-EF-blockSGDM, we also use $\mu=0.9$ and a weight decay of $0.0001$. Its stepsize $\eta$ is $0.1$ for 7 workers,We observe that $\eta=0.1M$ is too large for dist-EF-blockSGDM, while SGDM with $\eta=0.1$ performs worse than SGDM with $\eta=0.1M$. and $0.2$ for 15 workers.

Appendix B Proof of Lemmas 1 and 3

The Lemmas 1 and 3 hold by substituting $z_{t,i}=g_{t,i}$ and $z_{t,i}=\mu m_{t,i}+g_{t,i}$, respectively. ∎

Appendix C Proof of Theorem 1

By the smoothness of the function $F$, we have

where the second inequality follows from Young’s inequality with $\rho>0$. The last inequality follows from the smoothness of the function $F$. Let $\rho=1/2$. Taking total expectation and using Lemma 6 with $\mu=0$, we get

Assume that $\eta_{t}<3/(2L)$ for all $t$. Rearranging the terms, taking summation, and dividing by $\sum_{k=0}^{T-1}\frac{\eta_{k}}{4}\left(3-2L\eta_{k}\right)$ gives

Let $o\in\{0,\dots,T-1\}$ be an index such that

Appendix D Proof of Corollary 1

Let $\eta=\min\left(\frac{1}{2L},\frac{\gamma}{\frac{\sqrt{T}}{\sqrt{M}}+\frac{(1-\delta)^{1/3}}{\delta^{2/3}}\left(1+\frac{16}{\delta^{2}}\right)^{1/3}T^{1/3}}\right)$ for some $\gamma>0$, then $3-2L\eta\geq 2$. Substituting this into (3), we get

The bound on full-precision distributed SGD follows similar proof. For completeness, we present proof here. By the smoothness of the function $F$, we have

Let $\eta_{t}=\eta$. Taking total expectation, rearranging terms, and averaging over $T$, we obtain

Substituting $\eta=\min\left(\frac{1}{2L},\frac{\gamma\sqrt{M}}{\sqrt{T}}\right)$, we get

Appendix E Proof of Corollary 2

Let $\eta_{t}=\frac{\gamma}{\frac{((t+1)T)^{1/4}}{\sqrt{M}}+\frac{(1-\delta)^{1/3}}{\delta^{2/3}}\left(1+\frac{16}{\delta^{2}}\right)^{1/3}T^{1/3}}$. The following implies that $\eta_{t}\leq 1/(2L)$ for all $0\leq t\leq T-1$.

Using the fact that $\sum_{t=1}^{T}t^{\alpha-1}\leq\int_{0}^{T}x^{\alpha-1}dx=\frac{T^{\alpha}}{\alpha}$, for any $0<\alpha<1$, we have

Substituting the above results into Theorem 1, we obtain

Similarly, let $\eta_{t}=\frac{\gamma\sqrt{t+1}}{\frac{T}{\sqrt{M}}+\frac{(1-\delta)^{1/3}}{\delta^{2/3}}\left(1+\frac{16}{\delta^{2}}\right)^{1/3}T^{5/6}}$. We obtain

Using the fact that $\sum_{t=1}^{T}t^{\alpha}\leq\int_{1}^{T+1}x^{\alpha}dx\leq\frac{(T+1)^{\alpha+1}}{\alpha+1}$ for any $\alpha>0$, we also have

Assuming that $T\geq 4L^{2}\gamma^{2}M$, we have $\eta_{t}\leq 1/(2L)$ for all $0\leq t\leq T-1$. Substituting the above results into Theorem 1, we obtain

Appendix F Proof of Proposition 1

Appendix G Proof of Theorem 2

where in the first inequality we use Jensen’s inequality. In the second-to-last equality, we apply Assumptions 2 and 3. The last inequality follows from the sum of a geometric series. ∎

Now, we can consider two terms separately. For the second term, we have

where the first inequality follows from the definition of the compressor $\mathcal{C}$. The second inequality follows from Young’s inequality with any $\beta>0$, and the third inequality follows from Lemma 5. The third equality follows from the definition of $p_{t,i}$ and the assumption $\eta_{t}=\eta$. The last inequality follows from the sum of a geometric series. Let $\beta=\frac{\delta}{2(1-\delta)}$, then $1+1/\beta=(2-\delta)/\delta\leq 2/\delta$. We get

Then, combining (4), (6) and (8), we obtain

In the sequel, we assume $\eta_{t}=\eta$ for some $\eta>0$. Let us introduce the following virtual iterate:

By the smoothness of the function $F$, we get

where in the last inequality we use Lemma 6. Let $A_{t-1}=\sum_{k=0}^{t-1}\mu^{t-1-k}=\frac{1-\mu^{t}}{1-\mu}$. Then, we bound the last term:

where the first inequality follows from Jensen’s inequality. Then, combining (9), (10), (11), and (12), we obtain

Taking total expectation and telescoping this inequality from to $T-1$, we obtain

Let $\eta\leq\frac{(2-\rho)(1-\mu)^{2}}{2L}$ and $\rho$ is selected such that $\rho\geq(2-\rho)\mu^{3}$, we get

Hence, combining (13) and (14), and dividing by $T$,

Appendix H Proof of Corollary 3

Let $\eta=\frac{\gamma}{\frac{\sqrt{T}}{\sqrt{M}}+\frac{(1-\delta)^{1/3}}{\delta^{2/3}}\left(1+\frac{16}{\delta^{2}}\right)^{1/3}T^{1/3}}$ for some $\gamma>0$. As $T\geq\frac{4\gamma^{2}L^{2}M}{(1-\mu)^{4}}$, we have $\eta\leq\frac{(1-\mu)^{2}}{2L}$ and