Evaluation and Optimization of Gradient Compression for Distributed Deep Learning

S-SGD. The S-SGD is one of the most popular algorithms used to accelerate DNN training with data parallelism [1, 2]. Each iteration of S-SGD consists of two main phases: gradient computation and gradient aggregation. During the gradient computation phase, the local gradient at each worker is computed via feed-forward and back-propagation (FF&BP) process, followed by the gradient aggregation. As shown in Fig. 1(a), local gradients are synchronously communicated (summed up) among all workers, so that each worker has the aggregated global gradient for the model update. Due to the large size of deep models, millions to billions of parameters are communicated among workers, leading to significant performance bottlenecks [3, 4, 5].

System Optimizations. To alleviate the communication overheads, on one hand, many system level optimization techniques [11] have been proposed to improve the training efficiency of S-SGD. In this paper, we introduce three commonly used system advances : ring all-reduce primitive [6], wait-free back-propagation [7], and tensor fusion [8, 25].

Another line of work has proposed gradient compression methods to mitigate gradient communication costs [29]. Gradient compression methods aim to reduce the communication traffic by compressing the gradients, as demonstrated in Fig 1(c), in the following three categories.

Methods: we compare the performance of different gradient compression methods with S-SGD which is well optimized atop off-the-shelf implementation of PyTorch-DDP [25]. Following the prior work [36], we choose three most scalable gradient compression methods in each category: Sign-SGD with majority vote [17], Top-$k$ SGD with multiple sampling [21] ²²2The exact Top-$k$ selection is very computationally inefficient in GPUs, instead, multiple sampling uses binary search to find a close top-k threshold. , and Power-SGD [24]. For Power-SGD, we use the all-reduce collective, while for Sign-SGD and Top-$k$ SGD, we use the all-gather collective [34]. The gradients are packed together to be compressed and communicated for better performance [36].

Models: following the work [36], we choose four popular DNN models: ResNet-50 [13], ResNet-152 [13], BERT-Base [39], and BERT-Large [39], with per-GPU batch size of 64, 32, 32, and 8, respectively. The input image size is set as $3\times 224\times 224$ for ResNets, and the input sequence length is set as $64$ for BERTs. The model statistics and compression ratios are given in Table I. We are optimistic to use $1000\times$ compression ratio in Top-$k$ SGD. For Power-SGD, we use a rank $r=4$ for ResNets, and a higher rank $r=32$ for BERTs to achieve competitive results as suggested in [24].

Testbed: we conduct our experiments on a 32-GPU cluster of 8 nodes. Each node has 4 Nvidia RTX2080Ti GPUs (11GB RAM) connected by PCIe3.0x16. The interconnect between nodes is 10Gb/s Ethernet (10GbE). The common software includes PyTorch-1.12.1, CUDA-11.3, cuDNN-8.3.2, and NCCL-2.10.3.

Metric: we use the metric of average iteration time in running 120 iterations (exclude the first 20 iterations). Each gradient compression method mainly consists of gradient computation (i.e., FF&BP), gradient compression and decompression, and gradient communication costs. For the communication cost, we only measure its non-overlapped overhead.

We report iteration time of S-SGD, Sign-SGD, Top-$k$ SGD, and Power-SGD in Fig. 2. It shows that Sign-SGD and Top-$k$ SGD usually perform worse than S-SGD, while they are able to compress the gradients by $32\times$ and $1000\times$ times, respectively. For example, Sign-SGD and Top-$k$ SGD take $1.70\times$ and $1.66\times$ higher iteration time than S-SGD for training a ResNet-50 model. For training the largest BERT-Large model, the Top-$k$ SGD runs faster than S-SGD, while Sign-SGD runs out of memory due to its increased memory requirement. Among three gradient compression methods, Power-SGD has achieved the best performance over all four models. However, Power-SGD outperforms S-SGD only on large models (BERT-Base and BERT-Large), and it performs worse or closely than S-SGD on small models (ResNet-50 and ResNet-152). Overall, it is disappointing to find that three gradient compression methods, with $16\times$ to $1000\times$ high compression ratios, fail to achieve speedup over S-SGD in many cases on 10GbE bandwidth, which is very common in public cloud clusters [40, 21].

To understand the performance bottleneck, we dive into time breakdowns of gradient compression methods. Specifically, we measure the FF&BP computation time, compression and decompress time, and non-overlapped communication time. The time breakdowns of ResNet-50 and BERT-Base are shown in Fig. 3. The results of ResNet-152 and BERT-Large are similar to ResNet-50 and BERT-Base, respectively.

First, the communication overhead of S-SGD on BERT-Base is much larger than that on ResNet-50, because BERT-Base has more parameters, as well as a higher communication-to-computation ratio than ResNet-50. Thus, the communication overhead cannot be well hidden on BERT-Base, leaving more speedup space for gradient compression methods. Second, Sign-SGD has cheaper compression cost than Top-$k$ SGD and Power-SGD, however, its communication cost is even higher than uncompressed S-SGD. We contribute it to the fact that Sign-SGD has at most $32\times$ compression ratio, but it requires all-gather for communication, which is less efficient than all-reduce used in S-SGD. On the other hand, even though Top-$k$ SGD uses all-gather for communication, its communication overhead is very low due to its up to $1000\times$ compression ratio. However, Top-$k$ SGD suffers from the computation bottleneck of Top-$k$ operations, for example, its takes $4\times$ compression time to achieve $7\times$ communication speedup than Sign-SGD when training BERT-Base. It is partly because our multiple sampling top-k selection implemented atop PyTorch is less efficient than its highly-optimized CUDA version [21], which however is not publicly available. Third, Power-SGD achieves competitive communication performance by using all-reduce, while having a mild compression cost. For large model BERT-Base, Power-SGD is $1.4\times$ faster than S-SGD, showing it is one of the most promising gradient compression methods to outperform S-SGD.

In summary, existing gradient compression methods require non-negligible compression and/or communication overheads, resulting in poor scalability. It is worth noting that the S-SGD baseline has been well optimized in the system-level, while gradient compression methods are not. We will discuss the limitations of current gradient compression methods to integrate these techniques.

For S-SGD, the ring all-reduce primitive, wait-free back-propagation, and tensor fusion optimization techniques have been widely applied to accelerate distributed training, such as in Horovod [8] and PyTorch-DDP [25]. However, these techniques are not readily applied into aforementioned gradient compression methods.

We first summarize two good properties of S-SGD to enable these system optimization techniques as follows:

• •

  additive communication: the gradient communication operation in each DNN layer is to sum up the local gradient tensor from each workers. The gradient addition enables efficient aggregation with ring all-reduce primitive to distribute the reduced gradient result to all workers.

• •

  non-blocking communication: the gradient communication operation in each DNN layer is independent of backward operations, without blocking the subsequent back-propagation. It enables wait-free back-propagation and tensor fusion techniques very friendly to overlap communication with computation tasks.

However, these two properties are not both satisfied in current gradient compression methods. For Sign-SGD and Top-$k$ SGD, each DNN layer requires gradient calculation and gradient quantitation/sparsification tasks, followed by communicating the compressed gradient, so the communication will not block the subsequent back-propagation [41]. However, two quantized gradients cannot be simply added, e.g., the result of adding two $+1$ values is out of the range in Sign-SGD. Two sparsified gradient tensors are not additive as well, because their selected gradient elements (i.e., Top-$k$ elements) have different coordinates. As a compromise, Top-$k$ SGD and Sign-SGD typically utilize the all-gather primitive [34] to collect quantized and/or sparsified gradient results. Specifically, the communication complexity comparison is given in Table II, showing that the communication complexity of Sign-SGD and Top-$k$ SGD (with all-gather) is linear to the number of workers, which is system inefficient compared to S-SGD (with all-reduce) [21, 36]. Sign-SGD lacks compatibility with all-reduce, as studied in Fig. 3, so its actual communication overhead (with $32\times$ compression ratio) is however $24\%$ higher than S-SGD when training BERT-Base.

On the other hand, the compressed gradients of Power-SGD are still dense matrices, which allows Power-SGD to use efficient ring all-reduce primitive to aggregate these smaller matrices (i.e., $N_{c}\ll N$ as shown in Table II). However, in Power-SGD, each DNN layer needs to calculate the gradient $M$ and low-rank $P$, followed by aggregating $P$, and then it needs to calculate $Q$ based on the result of aggregated $P$, followed by aggregating $Q$, as shown in Fig. 4(a). In other words, the communication of aggregating $P$ will block the subsequent operations of computing and aggregating $Q$, and causes trouble for overlapping communication tasks (i.e., aggregating $P$ and $Q$) with back-propagation. As demonstrated in Fig. 4(b), Power-SGD with WFBP will overlap the whole gradient compression and communication process with back-propagation. By doing so, gradient compression and communication tasks (e.g., $P_{2}$, $AP_{2}$, $Q_{2}$) are performed in parallel with gradient computation tasks (e.g., $M_{1}$). However, as gradient compression and gradient computation are both compute intensive tasks (see Table II, the compress complexity of Power-SGD), they will compete for compute resources on the GPU, leading to performance slowdown [36], i.e., affecting computing tasks of $M_{1}$ and $P_{2}$ in Fig. 4(b). For example, Power-SGD with WFBP causes an overall of $13\%$ slowdown than Power-SGD without WFBP, when training ResNet-50 on one GPU (with only computation tasks).

Motivated by these limitations, we propose a novel gradient compression algorithm that can satisfy two good properties of S-SGD, providing opportunities for system optimizations.

Power-SGD. The Power-SGD algorithm [24] uses a single step of power iteration to decompose the large gradient matrix into two low-rank matrices as $M\approx PQ^{T}$. As shown in Algorithm 1, it requires one right multiplication, one left multiplication, and an orthogonalization operation to compute low-rank matrices $P$ and $Q$. Besides, it calls two all-reduce operations to aggregate $P$ and $Q$, respectively.

In each power iteration, the previous result at the last step (i.e., $Q_{t-1}$) will be re-used to update the low-rank matrices at the current step (i.e., $P_{t}$ and $Q_{t}$). This query reuse trick helps Power-SGD approximate the stochastic gradient matrix at each step more accurately [24]. In the very beginning, $Q$ is initialized from an i.i.d. standard normal distribution.

However, Power-SGD involves interleaving computing and communication tasks (i.e., compute $P\rightarrow$ aggregate $P\rightarrow$ compute $Q\rightarrow$ aggregate $Q$) in each DNN layer. That is, the communication tasks will block the subsequent back-propagation operations. And if we simply pipeline the gradient compression and aggregation tasks with gradient computing tasks, gradient compression and gradient computation will compete for compute resources on the GPU, leading to the performance interference and slowdown [36].

Alternate compression. To avoid this, we propose an alternate compressed Power-SGD (ACP-SGD) algorithm (Algorithm 1), which compresses the gradient into $P$ and $Q$ alternately. It only requires computing $P$ (or $Q$), followed by aggregating $P$ (or $Q$) in each iteration. By doing so, the communication of ACP-SGD is not blocking anymore, leading to the benefits of overlapping computing and communication tasks like S-SGD. We defer the details of system optimization to the next subsection.

In each power iteration, ACP-SGD reuses the previous low-rank results ($P_{t-1}$ or $Q_{t-1}$) to approximate the stochastic gradient. Consider two consecutive iterations with gradients $M_{t-1}$ and $M_{t}$, alternate compression is equal to performing a complete step of power iteration, except that it computes $P$ with $M_{t-1}$ but computes $Q$ with $M_{t}$. Under a small update stepsize, one can expect that $M_{t}$ is close to $M_{t-1}$ [24]. Thus, query reuse is helpful to ensure the approximation quality.

In addition, another benefit of ACP-SGD is that it can halve the gradient compression and communication costs compared to Power-SGD. That is, ACP-SGD performs one orthogonalization and one matrix multiplication with the computation complexity of $O(\frac{n+m}{2}r^{2}+nmr)$, and one all-reduce primitive with the communication volume of $O(\frac{n+m}{2}r)$.

Error-feedback. The error-feedback mechanism computes the difference between one worker’s gradient and the compressed gradient (i.e., error), and adds it back to the next gradient (i.e., feedback) [14, 16, 42, 24]. It has been shown to be crucial to improve the convergence performance of biased compression methods (i.e., compressed gradient is not equal to the original one in expectation), such as Sign-SGD, Top-$k$ SGD, and Power-SGD. Our ACP-SGD compression (derived from Power-SGD) is biased, therefore, it requires error-feedback to achieve good performance as shown in Algorithm 2.

Take computing and aggregating $P$ as an example, we apply the error-feedback into ACP-SGD as follows: 1) incorporate the previous local error $E_{t-1}$ into the local gradient $M_{t}$ before compression, i.e., $P_{t}\leftarrow(M_{t}+E_{t-1})Q_{t}$, and 2) update the local error via $E_{t}\leftarrow M_{t}+E_{t-1}-P_{t}Q_{t}^{T}$ before aggregation. ACP-SGD with error-feedback performs the layer-wise computation (including compute $M_{t}$, orthogonalize $Q_{t-1}$, compute $P_{t}$, and update $E_{t}$), followed by communication of aggregating $P_{t}$, whose communication operation is still non-blocking. In the very beginning, $E_{0}$ is initialized as zeros, and low-rank matrices $P_{0}$ and $Q_{0}$ are initialized randomly from standard normal distribution.

Wait-free Back-propagation (WFBP). Apart from using the ring all-reduce primitive for aggregating compressed tensors, the good property of non-blocking gradient communication allows ACP-SGD to apply WFBP to overlap communication tasks with computing tasks. As shown in Fig. 4(c), one can aggregate the compressed tensor of layer 2 (i.e., $AP_{2}$) immediately after the tasks of calculating and compressing the gradient (i.e., $M_{2}$ and $P_{2}$) have been finished. By doing WFBP, the communication task of $AP_{2}$ can be overlapped with the computing tasks of layer 1 (i.e., $M_{1}$). In the next iteration, WFBP can be equally applied to overlap the computing and communication tasks for $Q$. Meanwhile, ACP-SGD with WFBP only overlap all-reduce operations that are communication intensive and do not compete for the compute resources.

Tensor Fusion (TF). If the compressed tensors $P$ and $Q$ are layer-wisely aggregated via ring all-reduce operations, the communication overheads are easily dominated by start-up costs [9]. This is because the compressed tensors in ACP-SGD are much smaller than uncompressed tensors in S-SGD. For example, in Fig. 5, we report the cumulative distribution functions (CDFs) of the number of parameters in two models. After low-rank decomposing the gradients of ResNet-50 (or BERT-Base), we observe that there is a $30\%$ increase in the proportion of small tensors having less than $10^{4}$ (or $10^{5}$) parameters.

Thus, all-reducing these small and compressed tensors separately can be very communication inefficient. To reduce start-up costs, one shall apply the TF technique to merge several small tensors to be communicated together via one all-reduce operation. For example, on our 10GbE platform, all-reducing uncompressed gradients of ResNet-50 takes $243$ms, while all-reducing them together takes $169$ms (with $1.4\times$ speedup). However, in ACP-SGD, all-reducing compressed tensors separately takes $55.9$ms, and all-reducing them together only takes $2.3$ms (with $24.3\times$ speedup). Although TF is very helpful to reduce start-up costs in ACP-SGD, fusing all tensors for one communication requires waiting the completion of gradient computation of BP, which will lose the opportunity of WFBP to hide communication overheads.

To enable both WFBP and TF, it is suggested to merge a subset of compressed tensors of nearby layers. For example, given a three-layer DNN, during the BP from layer $3$ to $1$, one can merge and communicate the compressed gradients of layer $3$ and $2$ immediately when they are available, so that the fused communication task can be overlapped with the subsequent gradient computation task of layer $1$.

Buffer Size. It is non-trivial to determine which tensors to be merged in order to minimize the communication overhead for different model and hardware configurations. In this work, we choose to allocate the buffer with a pre-defined buffer size, select the available tensors to fit in the buffer, and execute the all-reduce operation on the buffer. This has shown to be very effective in S-SGD [8, 25]. For example, the default buffer size in PyTorch-DDP is 25MB [25], and it will batch the gradient tensors of a ResNet-50 model (with $97.5$MB parameters) into $4$ buffers. The selected tensors are copied to one buffer, and the buffer is aggregated as a whole when it is ready.

The buffer size is critical to control the trade-off between WFBP and TF. When the buffer size is too small, each tensor is communicated immediately to maximize the overlap (optimal WFBP), but it loses any opportunity of fusing tensors (no TF). When the buffer size is too large, all tensors will be merged together to be communicated at the end of BP (optimal TF), losing any change for overlapping (no WFBP).

However, the compressed tensors ($P$ and $Q$) to be communicated in ACP-SGD are smaller than the gradient tensors ($M$), as shown in Fig. 5. This means the default buffer size (i.e., 25MB) used in S-SGD is not suitable for ACP-SGD anymore. For example, a ResNet-50 involves only $0.63$MB and $1.04$MB parameters in $P$ and $Q$, respectively, for ACP-SGD with a rank of $4$. To reduce the trouble of tuning the hyper-parameter of buffer size for different models and ranks, we choose to configure the compressed buffer size by scaling the default buffer size with the compression rate. For instance, the compression rates of ACP-SGD in ResNet-50 are $0.64\%$ and $1.07\%$ for $P$ and $Q$, giving compressed buffer sizes of $0.16$MB and $0.27$MB, respectively. Thus, it will batch $P$ tensors into 4 buffers, and batch $Q$ tensors into another 4 buffers for aggregation. Note that fusion results for $P$ and $Q$ are different to each other. In this work, we use the default buffer size of 25MB, and the compressed buffer sizes are derived based on the compression rates. We notice that buffer size can be automatically tuned using e.g. Bayesian optimization technique [43], but we do not adopt it in this work as the default buffer size can generally provide nearly optimal performance as studied later in Fig. 10.

We implement our ACP-SGD prototype atop PyTorch. It wraps the SGD optimizer to cope with the underling gradient compression and communication operations. We use native APIs, i.e., torch.linalg.qr and torch.distributed.all_reduce, to perform reduced QR decomposition for efficient orthogonalization (which is required for compression), and ring all-reduce for compressed gradient aggregation. The vector-shaped parameters (e.g., biases) requires no compression, while other parameters are reshaped into matrices for compression.

To support WFBP, we register a hook function for each learnable parameter tensor during the back-propagation, and the registered hook function will be called when each gradient is ready. In the hook function, we implement the ACP-SGD with EF to compress the gradient $M$ into $P$ and $Q$ alternately. To support TF, the compressed tensors are copied into one of the buffers, and the ready buffer will be aggregated asynchronously via the all-reduce operation. At the end of back-propagation, we synchronize the all-reduce operations on the $P$ (or $Q$) buffers.

For convergence experiments, we compare the performance of ACP-SGD with S-SGD and Power-SGD. We choose VGG-16 [44] and ResNet-18 [13] models on the Cifar-10 [45] dataset. We run each algorithm for $300$ epochs on $4$ GPUs. We use per-GPU batch size of $128$ and learning rate of $0.1$, with a gradual warmup in the first $5$ epochs, and multiple learning rate decays (by a factor of $10$) at $150$ and $220$ epochs [1]. Each algorithm uses a momentum of $0.9$, and Power-SGD and ACP-SGD use a rank of $4$.

For time efficiency related experiments, we compare the iteration time of ACP-SGD with S-SGD and two versions of Power-SGD. The Power-SGD is the original implementation of [24], which packs gradients after back-propagation for compression and communication. The Power-SGD with WFBP and TF optimizations (i.e., Power-SGD*) is implemented on PyTorch’s communication hook [46], which overlaps gradient compression and communication tasks with back-propagation. For a fair comparison, two Power-SGD baselines use reduced QR decomposition for orthogonalization like ACP-SGD. Note that S-SGD, Power-SGD*, and ACP-SGD are all optimized in the system-level with WFBP and TF, using a buffer size of $25$MB. We do not compare with other gradient compression methods, as we have shown that Power-SGD performed better than them in all cases (see §III). For other model and testbed settings, we follow the same configurations in §III-A.

We first verify the convergence performance of our ACP-SGD on training two different models VGG-16 and ResNet-18 on the Cifar-10 dataset. The results are given in Fig. 6, showing that ACP-SGD can achieve very close accuracy results (i.e., $94.1\%$ for VGG-16, and $94.6\%$ for ResNet-18) compared with S-SGD and Power-SGD. We see that gradient compression methods Power-SGD and ACP-SGD converge slightly slower than S-SGD in the early training stage, which however does not affect their final achieved model accuracy.

We contribute the good convergence performance of ACP-SGD to the usage of EF and query reuse mechanisms. To validate it, we perform an ablation study by disabling one of them, and report their results in Fig. 7. It demonstrates that ACP-SGD without EF and reuse perform poorly, verifying that both of them are key components to ACP-SGD.

We compare the performance of our ACP-SGD with S-SGD [25], Power-SGD [24], and Power-SGD* (i.e., Power-SGD optimized with WFBP and TF) [46]. We report the iteration time (mean$\pm$std) in Table III, where the best results are in bold.

From Table III, it is observed that ACP-SGD consistently outperforms other methods in all tested cases. Specifically, ACP-SGD achieves an average of $4.06\times$, $1.34\times$, and $1.51\times$ speedups over S-SGD, Power-SGD, and Power-SGD* methods, respectively. And ACP-SGD can provide a significant speedup over S-SGD, e.g., training BERT-Large up to $9.42\times$ faster than S-SGD. However, two gradient compression baselines (Power-SGD and Power-SGD*) do not guarantee improved performance over S-SGD, for example, they run about $11\%$ slower than S-SGD in training ResNet-50.

For Power-SGD* with system optimizations, it performs better than Power-SGD in small models (ResNet-50 and ResNet-152), but performs worse than Power-SGD in large models (BERT-Base and BERT-Large). This is because larger models typically require more compute-intensive gradient computation and compression tasks, and hence cause more severe performance interference for Power-SGD*. In contrast, ACP-SGD with system optimizations will not lead to resource competition, which in turns can achieve $1.60\times$ speedup than Power-SGD when training the BERT-Large model. The effects of system optimizations for Power-SGD and ACP-SGD are studied later in Fig. 9.

To further investigate the performance of ACP-SGD, we give the time breakdowns in Fig. 8 for two representative models (ResNet-50 and BERT-Base). It shows that ACP-SGD has very low gradient compression and communication overheads, and it achieves almost a linear scalability. Meanwhile, S-SGD can hide communication overhead very well for the small ResNet-50 model, but requires very high non-overlapped communication overhead for the large BERT-Base model. Power-SGD and Power-SGD* are useful to improve the training performance in large models, and our ACP-SGD moves a further step to improve the performance significantly.

Next, we study the benefits of system optimizations (WFBP and TF) step-by-step for S-SGD, Power-SGD, and ACP-SGD. In the rest of the paper, Power-SGD refers to the implementation of Power-SGD* if it is not specified. We provide three variants for each method: Naive implementation without WFBP and TF, the one with only WFBP, and the one with WFBP and TF. We conduct experiments on two relatively large models ResNet-152 and BERT-Large, as they involve more communication optimizations than ResNet-50 and BERT-Base, respectively. The results are given in Fig. 9.

For S-SGD and ACP-SGD, it is observed that WFBP achieves about $12\%$ improvement over their Naive implementations. However, Power-SGD is not friendly to WFBP, since overlapping gradient computation and gradient compression causes compute resource competition, leading to an average of $13\%$ slowdown over its Naive implementation. Besides, we find that TF is able to provide a significant speedup for any method with WFBP. Specifically, WFBP with TF achieves an average of $1.28\times$, $2.16\times$ and $1.56\times$ speedups than WFBP (without TF) for S-SGD, Power-SGD, and ACP-SGD, respectively. In particular, TF for Power-SGD achieves the most significant improvement among three methods, due to the collective effect of reduced start-up cost and alleviated performance interference. As for ACP-SGD, with the help of WFBP and TF, it can achieve up to $2.14\times$ speedup over its Naive implementation, which validates the importance of two system optimizations for our gradient compression algorithm.

Sensitivity study of buffer size. When using two system optimizations, buffer size is critical to control the trade-off between WFBP and TF. To study the effect of buffer size, we change the buffer size from $0$ to $1500$MB when training the BERT-Large model (with $1282.6$MB parameters), that is, from optimal WFBP without TF to optimal TF without WFBP. We perform Power-SGD and ACP-SGD with two different ranks: 32 and 256, where rank 256 (with $5.4\times$ compression ratio) is suggested to avoid possible performance degradation in [35].

From Fig. 10, we see that ACP-SGD can consistently outperform Power-SGD under different buffer size and rank settings. And ACP-SGD is very robust to the value of buffer size under different ranks, for example, the default buffer size of $25$MB remains a good candidate for ACP-SGD under different ranks (i.e., compression ratios). In particular, for rank=256, ACP-SGD with a buffer size of 25MB significantly outperforms two extreme cases: buffer size of 0MB (no TF) and buffer size of 1500MB (full TF), providing about $50\%$ improvement than both cases. We contribute it to that ACP-SGD can adaptively adjust the compressed buffer size for TF with different ranks, which can largely reduce the trouble for hyper-parameter tuning.

Effect of batch size. Here we compare ACP-SGD with S-SGD and Power-SGD on different batch sizes for training the ResNet-152 model. We vary the per-GPU batch size from 16 to 32 to fully utilize the GPU memory, and keep other configurations the same. We report results in Fig. 11(a). We find that ACP-SGD consistently outperforms S-SGD and Power-SGD under different batch sizes. For example, when using batch size 16, ACP-SGD provides $2.4\times$ and $1.5\times$ speedups than S-SGD and Power-SGD, respectively. Using batch size 32, it gives $1.6\times$ and $1.3\times$ speedups over S-SGD and Power-SGD, respectively. ACP-SGD achieves the best performance since it always has very small compression and communication costs for different batch sizes. Besides, for the three methods, using large batch sizes often provides performance improvement in terms of throughput. For instance, S-SGD, Power-SGD, and ACP-SGD achieve $3.7\times$, $2.8\times$, $2.4\times$ speedups on throughput by increasing batch size from $16$ to $32$. Especially for S-SGD (with the heaviest communication), we observe that its non-overlapped communication overhead drops as the batch size increases. This is because increasing batch size leads to an increase in the computation-to-communication ratio, which in turns benefits S-SGD to overlap more communications with computations. Using larger batch sizes may further reduce the performance gap between S-SGD and ACP-SGD, but it is impractical due to the limited GPU memory capacity.

Effect of rank. The parameter of rank is used to control the compression ratio of Power-SGD and ACP-SGD. The lower the rank, the stronger the compression. To show the effect of rank, we vary the rank from 32 to 256 by a factor of 2 for training BERT-Large (with model dimension of 1024). We report the performance under different ranks in Fig. 11(b), showing that using large ranks leads to an increase in the compression and communication overheads for both Power-SGD and ACP-SGD. Therefore, Power-SGD and ACP-SGD have $3.4\times$ and $2.4\times$ higher total iteration time, respectively, when varying the rank from 32 to 256. However, compared to Power-SGD, ACP-SGD can overlap more communication overheads with gradient computation and compression overheads as the rank increases. For example, ACP-SGD provides $1.9\times$ speedup than Power-SGD when training using rank 32, and this speedup becomes $2.7\times$ for rank 256. In particular, when using rank 256, ACP-SGD has almost $7.3\times$ reduction in the non-overlapped communication time over Power-SGD. In addition, for training BERT-Large, ACP-SGD with rank 256 (i.e., with $5.4\times$ compression ratio) can still achieve about $3.9\times$ improvement over S-SGD.

Effect of the number of GPUs. To study the scalability of ACP-SGD, we vary number of GPUs from 8 GPUs (2 nodes) to 64 GPUs (16 nodes). The interconnect between nodes is 10GbE. We use ring all-reduce for gradient aggregation.

We report the effect of the number of GPUs in Fig. 12. It shows that all three methods have high scalability as the number of GPUs increases. For instance, it is observed that S-SGD, Power-SGD, and ACP-SGD have only an average of $10\%$, $24\%$, and $8\%$ increase in per iteration time, respectively, when scaling from $8$ GPUs to $64$ GPUs. They can scale well because of using 1) ring all-reduce primitive with optimal bandwidth, whose communication complexity stays almost constant with increase in number of GPUs, and 2) tensor fusion to amortize start-up cost. While our testbed is up to 64 GPUs, the high scalability of ACP-SGD implies that it can always outperform S-SGD and Power-SGD using more GPUs.

Effect of network bandwidth. Here we validate the scalability of ACP-SGD on 32 GPUs, connected by different levels of networks, including inexpensive commodity 1Gb/s Ethernet (1GbE), ubiquitous data-center 10Gb/s Ethernet (10GbE), and expensive high-bandwidth 100Gb/s Infiniband (100GbIB).

The results are given in Fig. 13, showing that ACP-SGD performs better than S-SGD and Power-SGD under different networks. For slow 1GbE network, gradient compression methods Power-SGD and ACP-SGD can largely outperform S-SGD. In ResNet-50, Power-SGD and ACP-SGD achieves $5.7\times$ and $7.1\times$ speedups over S-SGD, and the speedups increase up to $11.2\times$ and $23.9\times$ in BERT-Base. For fast 100GbIB network, while the communication overhead of S-SGD can be well alleviated, ACP-SGD can provide about $40\%$ improvement over S-SGD when training BERT-Base.