Faster Distributed Deep Net Training: Computation and Communication Decoupled Stochastic Gradient Descent

1 Introduction

The training of deep neural networks is resource intensive and time-consuming. With the expansion of data and model scale, it may take a few days or weeks to train a deep model by using mini-batch SGD on a single machine/GPU. To accelerate the training process, distributed optimization provides an effective tool for deep net training by allocating the computation to multiple computing devices (CPUs or GPUs).

When variants of mini-batch SGD are applied to a distributed system, communication between computing devices will be incurred to keep the same convergence rate as minibatch SGD. As a matter of fact, the extra communication cost in a distributed system is the main factor which prevents a distributed optimization algorithm from achieving the linear time speedup, although the computation load is the same as its single machine version. In addition, the communication cost, which is often linearly proportional to the number of workers, can be extremely expensive when the number of workers is huge. Therefore, it is critical to reduce the communication bottleneck to make better use of the hardware resources.

∗The corresponding author.

Given that the total amount of communication bits equals the number of communications multiplied by the number of bits per communication, several works are proposed to accelerate training by reducing the communication frequency [Zhou and Cong, 2018; Yu et al., 2018; Stich, 2018] or communication bits [Lin et al., 2017; Stich et al., 2018; Wen et al., 2017; Alistarh et al., 2017]. However, even when the communication frequency or the number of bits per communication is reduced, hardware resources are not fully exploited in traditional synchronous distributed algorithms because of the following two reasons: (1) only partial resources can be used when workers are communicating with each other and (2) the computation and the communication are interdependent in each iteration. Specifically, the computation of the t-th step relies on the communication of the (t-1)-th step and the communication in the t-th step follows the computation of the t-th step. To address that, another line of research aims to run computation and communication in parallel by using stale gradients [Li et al., 2018], but the communication complexity is still high as all gradients need to be communicated.

To tackle this dilemma, we propose computation and communication decoupled stochastic gradient descent (CoCoDSGD) for distributed training. Instead of waiting for the completion of communication as in traditional synchronous SGD, workers in CoCoD-SGD continue to calculate stochastic gradients and update models locally after the communication started. After the communication finishes, each worker updates its local model with the result of communication and the difference of the local models. As a result, CoCoD-SGD can make full use of resources by running computation and communication in parallel. In the meantime, workers in CoCoDSGD communicate with each other periodically instead of in each iteration. As a benefit, the communication complexity of CoCoD-SGD is lower. By running computation and communication in parallel while reducing the communication complexity, faster distributed training is achieved. In addition, CoCoD-SGD is suitable for both homogeneous and heterogeneous environments. Further, it is theoretically justified with a linear iteration speedup with respect to the total computation capability of hardware devices.

Contributions of this paper are summarized as follows:

• We propose CoCoD-SGD, a computation and communication decoupled distributed algorithm, to make full use of hardware resources and reduce the communication complexity.
• From the theoretical perspective, we prove that CoCoDSGD has a linear iteration speedup with respect to the number of workers in a homogeneous environment. Besides, it has a theoretically justified linear iteration speedup with respect to the total computation capability in a heterogeneous environment.
• When training deep neural networks with multiple GPUs, CoCoD-SGD achieves a better time speedup comparing with existing distributed algorithms.
• In both homogeneous and heterogeneous environments, experimental results verify the effectiveness of the proposed algorithm.

2 RelatedWork

A conventional framework for distributed training is the centralized parameter server architecture [Li et al., 2014], which is supported by most existing systems such as Tensorflow [Abadi et al., 2016], Pytorch [Paszke et al., 2017] and Mxnet [Chen et al., 2015]. In each iteration, the parameter server, which holds the global model, needs to communicate with O(N) workers. This becomes a bottleneck which slows down the convergence when N is large. Therefore, the decentralized ring architecture using the Ring-AllReduce algorithm became popular in recent years. In Ring-AllReduce, each worker only transports O(1) gradients to its neighbors to get the average model over all workers.

The Ring-AllReduce algorithm can be directly combined with synchronous stochastic gradient descent (S-SGD), which has a theoretically justified convergence rate O( 1 √TM ) for both general convex [Dekel et al., 2012] and non-convex problems [Ghadimi and Lan, 2013], where T is the number of iterations and M is the mini-batch size. Such rate shows its linear iteration speedup with respect to the number of workers since increasing the number of workers is equivalent to increasing the mini-batch size. Although only O(1) gradients are needed to be communicated per worker in each iteration for Ring-AllReduce, O(N) handshakes are needed for each worker. Therefore, the communication cost grows as the number of the workers increases.

There have been many attempts to reduce the communication bottleneck of S-SGD while maintaining its linear iteration speedup property. Among them, Local-SGD is a variant of S-SGD with low communication frequency, in which workers update their models locally and communicate with each other every k iterations. It has been proved to have linear iteration speedup for both strongly convex [Stich, 2018] and non-convex [Zhou and Cong, 2018; Yu et al., 2018] problems. QSGD [Alistarh et al., 2017] and TernGrad [Wen et al., 2017] compress the gradients from 32-bit float to lower bit representations. Sparse SGD [Aji and Heafield, 2017] is another method that communicates part of the gradients in each iteration and is proved to have the same convergence rate as SGD [Stich et al., 2018; Alistarh et al., 2018]. Although the communication complexity is reduced in the above methods, the hardware resources are not fully utilized because only a part of them is used for communication.

To fully utilize hardware resources, asynchronous stochastic gradient descent (A-SGD), a distributed variant of SGD based on the parameter server architecture, is proposed. After one communication started, a worker uses a stale model to compute the next stochastic gradient without waiting for the completion of the communication. A linear iteration speedup is proved for both convex [Agarwal and Duchi, 2011] and non-convex [Lian et al., 2015] problems. However, stale gradients may slow down the training and make it converge to a poor solution especially when the number of workers is large, which implies a large delay [Chen et al., 2016]. As a remedy, Pipe-SGD [Li et al., 2018] is proposed to integrate the advantages of A-SGD and Ring-AllReduce. Specifically, it employs Ring-AllReduce to get the average gradient and updates the model with a stale gradient. Pipe-SGD can control the delay as a constant because of the efficiency of the Ring-AllReduce algorithm. Nevertheless, one disadvantage of Pipe-SGD is that its communication complexity is O(T), which is higher than Local-SGD and QSGD. In comparison, the proposed algorithm which makes full use of hardware resources and has a lower communication complexity, achieves the advantages of both Pipe-SGD and Local-SGD.

3 Algorithm

In this section, we introduce the CoCoD-SGD algorithm with the following three techniques: (1) computation and communication decoupling, (2) periodically communicating and (3) proportionally sampling.

3.1 Preliminary Problem Definition

We focus on data-parallel distributed training, in which each worker can access only part of the data. We use D to denote the full training dataset and Di to denote the local dataset stored in the i-th worker. We have D = D1 ∪ · · · ∪ DN and Di ∩ Dj = ∅, ∀i ̸= j. The objective can be written as

where fi(x) := 1 |Di| ∑ξi∈Di f(x, ξi) is the local loss of the i-th worker and pi's define the partition of data among all workers. Specifically, pi is proportional to size of the local dataset on the i-th worker: pi = |Di| |D| and we have ∑Ni=1 pi = 1. We use P to denote the distribution of pi's.

Notations

• ∥ · ∥ indicates the ℓ2 norm of a vector.
• f ∗ represents the optimal value of (1).
• E indicates a full expectation with respect to all the randomness, which includes the random indexes sampled to calculate stochastic gradients in all iterations.
• Ci represents the computation capability, which indicates the computing speed, of the i-th worker.

Algorithm 1 CoCoD-SGD

Input: The number of workers N, the number of iterations T, communication period k, mini-batch sizes M1, · · · ,MN and initial point ˆx0 ∈ Rd .

• 1: while t < T do
• 2: Worker Wi does:
• 3: Run Step I and Step II in parallel.
• 4: Step I : Store xit in the memory and communicate with other workers to get the weighted mean of all local models: ˆxt = ∑Ni=1 Mi ∑Nj=1 Mj xit .

5: Step II :

Compute a mini-batch stochastic gradient

end for

• 8: end while

3.2 Computation and Communication Decoupled Stochastic Gradient Descent

The complete procedure of CoCoD-SGD is summarized in Algorithm 1. On one hand, the goal of CoCoD-SGD is twofold: (1) running computation and communication in parallel and (2) reducing the communication complexity. To achieve that, we propose the following two techniques:

• Computation and communication decoupling: Different from Pipe-SGD which uses stale gradients to decouple the dependency of computation and communication, CoCoD-SGD continues to update models in workers locally after the communication starts. If one communication starts after the t-th update, workers need to communicate with each other to get the average model ˆxt. After the communication starts, the i-th worker continues to update its local model xit by mini-batch SGD (line 4-5). As a result, computation and communication can be executed simultaneously.
• Periodically communicating: In practice, when the number of workers is large, the communication time may exceed the time of one update, which will cause idle time of the computing devices. Therefore, we let workers keep on updating the local models k times instead of waiting for the completion of the communication after only one local update (line 5). In this way, the communication complexity can be also reduced.

On the other hand, CoCoD-SGD is designed to be suitable for heterogeneous distributed system, where workers may have different computation capabilities. Intuitively, when all workers use the same batch size to compute stochastic gradients, the faster workers need to wait for the slower ones after finishing their updates. Therefore, we use proportionally

sampling in CoCoD-SGD:

• Proportionally sampling: Workers in CoCoD-SGD use different mini-batch sizes Mi's to compute stochastic gradients. The batch sizes are proportional to the computation capabilities of workers, i.e., Mi Mj = Ci Cj . Under this setting, all workers will finish k updates at the same time. Meanwhile, to let all workers finish one epoch simultaneously, we proportionally divide the dataset among workers, i.e., |Di| |Dj| = Ci Cj . In addition, we define the average model ˆxt as the weighted mean of all local models: ˆxt = ∑Ni=1 Mi ∑Nj=1 Mj xit (line 4).

Besides, different from QSGD [Alistarh et al., 2017] and Sparse-SGD [Stich et al., 2018], which are based on the parameter server architecture, CoCoD-SGD can communicate with the Ring-AllReduce algorithm by the definition of ˆxt .

After finishing local iterations and the communication, the i-th worker updates its local model with the average model and the difference of the local models by xit+k = ˆxt+(xit+k - xit ) (line 6).

4 Theoretical Analysis

In this section, we provide the theoretical analysis for CoCoD-SGD and show that CoCoD-SGD has the same convergence rate as S-SGD. In addition, we show that CoCoDSGD has lower communication complexity and better time speedup. Due to the space limit, all proofs are deferred to the supplemental material† . In the subsequence analysis, we will use the following definitions.

Definition 1 We denote the total number of iterations and the total time used to converge when using N workers as TN and TN, respectively. Then the iteration speedup (IS) and the time speedup (TS) are respectively defined as

4.1 Main Results

Before establishing our main results, we introduce the following assumptions, all of which are commonly used in the analysis of distributed algorithms [Lian et al., 2015; Aji and Heafield, 2017; Yu et al., 2018].

Assumption 1

1. Lipschitz gradient: All local functions fi's have LLipschitz gradients

(2) Unbiased estimation:

† It can be downloaded from the anonymous link: https://github.com/IJCAI19-CoCoD-SGD/Supplemental-Material

1. Bounded variance: There exist constants σ and ζ such that

1. Dependence of random variables: ξ i,j t 's are independent random variables, where t ∈ {0, 1, · · · , T - 1}, i ∈ {1, 2, · · · ,N}, and j ∈ {1, 2, · · · ,Mi}.

To evaluate the convergence rate, the metric in nonconvex optimization is to bound the weighted average of the ℓ2 norm of all gradients [Ghadimi and Lan, 2013; Lian et al., 2015; Yu et al., 2018].

Theorem 1 Under Assumption 1, if the learning rate satisfies γ ≤ 1L , we have the following convergence result for Algorithm 1:

where k is the communication period and

Choosing the learning rate γ appropriately, we have the following corollary.

Corollary 1 Under Assumption 1, when the learning rate is set as γ = 1 σ√ T ∑Ni=1 Mi and the total number of iterations

satisfies

we have the following convergence result for Algorithm 1:

Corollary 1 indicates that the convergence rate of the weighted average model is O (1/√∑Ni=1 MiT ), which is consistent with S-SGD [Ghadimi and Lan, 2013]. Next, we establish the linear iteration speedup in both homogeneous and heterogeneous environments, and show the communication complexity of CoCoD-SGD.

Remark 1 (linear iteration speedup in the homogeneous environment). For CoCoD-SGD in a homogeneous environment, all workers use the same mini-batch size: M1 = M2 = · · · = MN = M. According to (12), CoCoD-SGD converges at the rate O (1/√NMT ). Consequently, to achieve the ϵapproximation solution, O(1/(NMϵ2 ) iterations are needed, which means CoCoD-SGD has a linear iteration speedup with respect to the number of workers according to the definition of IS in (2).

Remark 2 (linear iteration speedup in the heterogeneous environment). With Proportionally Sampling, we have Mi/Mj = Ci/Cj. According to (12), to achieve the ϵapproximation solution, the number of iterations required is O (1/(∑Ni=1 Mi · ϵ2 )), that is O (C1/ (∑Ni=1 Ci · M1ϵ2)), which means CoCoD-SGD has a linear iteration speedup with respect to the total computation capability of all workers according to (2).

Remark 3 (communication complexity). From (11), we have T ≥ max{O(N),O(Nk2 ),O(N 3k4 )}. Therefore, when the communication period is bounded by O (T 14 /N 34 ), the convergence rate in Corollary 1 is achievable. As a result, the total communication complexity of CoCoD-SGD is O (T/ (T 14 /N 34 )), that is O (T 34N 34 ).

Remark 4 (choice of learning rate). When we run CoCoDSGD for a fixed number of epochs, ∑Ni=1 MiT will be a constant. The learning rate suggested in Corollary 1 can be written as γ = ∑Ni=1 Mi σ√∑Ni=1 MiT . Thus, for CoCoD-SGD with N workers, the learning rate should be set as γN = N · γ1 in the homogeneous environment and γN = ∑Ni=1 Ci/C1 · γ1 in the heterogeneous environment.

4.2 Time Speedup Analysis

Next, we compare and analyse the time speedup of CoCoDSGD and other distributed algorithms when they are applied in practice.

When multiple nodes with GPUs are used to train a deep neural network, GPUDirect communication, where all computation threads and communication threads are executed in GPUs, is widely adopted [Paszke et al., 2017; Chen et al., 2015; Sergeev and Balso, 2018] and verified to be efficient [Potluri et al., 2013]. The computation threads can make full use of GPU resources especially when the size of mini-batch is large, while the communication threads can only use partial of the GPU resources, which is assumed to be a% for ease of analysis. We further denote the time of one computation as Tcomp and the time of communication when using N workers as T N comm.

A comparison of the running processes of CoCoD-SGD and other standard distributed algorithms is shown in Figure 1. All algorithms have provable linear iteration speedup. Thus, the time speedup equals (N · T av 1 )/T av N according to (3), where T av N is denoted as the average time required to finish one iteration among N workers. For S-SGD, computation and communication are sequential. Therefore, its time speedup is

On the other hand, Pipe-SGD decouples the dependence of computation and communication by using a stale gradient, but the communication complexity is still O(T) as it needs to communicate all gradients. Accordingly, its time speedup is

According to Corollary 1 in [Yu et al., 2018] and Remark 3, Local-SGD and CoCoD-SGD have the same convergence rate as S-SGD when the communication period is O(T 1 4 /N 34 ). The difference between Local-SGD and CoCoD-SGD is that CoCoD-SGD runs computation and communication in parallel while Local-SGD runs them sequentially. As a result, their time speedups are

and

respectively, where k = O(T 14 /N 34 ). As (16) is bigger than (13), (14), and (15), we can verify that CoCoD-SGD achieves the best time speedup.

5 Experiments

In this section, we validate the performance of CoCoD-SGD in both homogeneous and heterogeneous environments.

5.1 Experimental Settings

Hardware. We evaluate CoCoD-SGD on a cluster where each node has 3 Nvidia Geforce GTX 1080Ti GPUs, 2 Xeon(R) E5-2620 cores and 64 GB memory. The cluster has 6 nodes, which are connected with a 56Gbps InfiniBand network. Each GPU is viewed as one worker in our experiments.

Software. We use Pytorch 0.4.1 [Paszke et al., 2017] to implement the algorithms in our experiments, and use Horovod 0.15.2 [Sergeev and Balso, 2018], OpenMPI 3.1.2‡ , and NCCL 2.3.7§ to conduct the GPUDirect communication with the Ring-AllReduce algorithm.

Datasets. We use two datasets for image classification.

• CIFAR10 [Krizhevsky and Hinton, 2009]: it consists of a training set of 50, 000 images from 10 classes, and a test set of 10, 000 images.
• CIFAR100 [Krizhevsky and Hinton, 2009]: it is similar to CIFAR10 but has 100 classes.

Tasks. We train ResNet18 [He et al., 2016] and VGG16 [Simonyan and Zisserman, 2014] on the two datasets.

Baselines. We compare CoCoD-SGD with S-SGD, PipeSGD [Li et al., 2018] and Local-SGD [Stich, 2018]. All of them support Ring-AllReduce communication.

Hyper-parameters. We use the following hyperparameters.

• Basic batch size: 32 for both ResNet18 and VGG16.
• Basic learning rate: For both networks we start the learning rate from 0.01 and decay it by a factor of 10 at the beginning of the 81st epoch.
• Momentum: 0.9.
• Weight decay: 10-4 .
• Communication period and gradient staleness: Since the variance of stochastic gradients is higher in the beginning, we set the communication period to 1 for the first 10 epochs and 5 for the subsequential epochs. The staleness of gradients in Pipe-SGD is set to 1 as suggested in [Li et al., 2018].

5.2 Homogeneous Environment

In a homogeneous environment, all workers have the same computation speed. So for N workers, we set γN = Nγ as suggested in Remark 4. And the learning rate warm-up scheme proposed in [Goyal et al., 2017] is adopted.

Comparison of convergence rate. Figure 2 shows the training loss with regard to epochs of ResNet18 and VGG16 on 16 GPUs. All algorithms have similar convergence speed, which validates the theoretical results claimed in Section 4.1.

‡ https://openmp.org

§N Luehr. Fast multi-gpu collectives with nccl, 2016.

Comparison of convergence speed. Figure 3 shows the test accuracy regarding time on 16 GPUs and Table 1 shows the best test accuracies of all algorithms on the two datasets. We evaluate the accuracy on the test set during training, but we only accumulate the time used for training. As shown in Figure 3, CoCoD-SGD achieves almost 2× and 2.5× speedup against S-SGD for ResNet18 and VGG16 respectively. Although Pipe-SGD can run computation and communication in parallel, it is still slower than Local-SGD and CoCoD-SGD since its communication complexity is higher. CoCoD-SGD converges faster than others since it not only runs computation and communication simultaneously but also has a lower communication complexity. The results verify our theoretical results claimed in Section 4.2. In the meanwhile, we can observe from Table 1 that CoCoD-SGD does not sacrifice the test accuracy on both datasets and may get better results than S-SGD.

Comparison of time speedup. Figure 4 shows the time speedup for ResNet18 and VGG16 when the number of workers increases from 1 to 16. We run the experiments for 120 epochs on 1 GPU and multiple GPUs. The speedup for ResNet18 is better due to its smaller model size. On both tasks, CoCoD-SGD achieves the fastest convergence and the best time speedup, which validates our time speedup analysis.

5.3 Heterogeneous Environment

To simulate a heterogeneous environment, we use 16 workers in our experiments and reduce the computation speed of 8 workers by half. For S-SGD and Pipe-SGD, Proportionally Sampling proposed in Section 3.2 can be also applied in the heterogeneous environment, which is equivalent to increasing the mini-batch size. And we denote corresponding algorithms as S-SGD+PS and Pipe-SGD+PS respectively. When one algorithm employs Proportionally Sampling, the batch size is set to 32 for the slower workers and 64 for the faster workers. In the heterogeneous environment with 16 workers, we set γ ps 16 = 0.01∗(8+8∗2) for algorithms with Proportionally Sampling and γ16 = 0.01 ∗ 16 for others as suggested in Remark 4. In addition, the asynchronous version of Local-SGD (AsynLocal-SGD), which is proposed for heterogeneous environments in [Yu et al., 2018], is included in comparison with CoCoD-SGD in our experiments.

Comparison of convergence rate. Figure 5 presents the training loss regarding epochs of ResNet18 and VGG16 on 16 GPUs, which exhibits similar convergence rate for all the algorithms in comparison. Besides, we can see from Table 2 that CoCoD-SGD does not lose the test accuracy.

Comparison of convergence speed. Figure 6 shows the curves of the test accuracy regarding time. As we can see, CoCoD-SGD converges fastest and achieves almost 2.5× and 3× speedup against S-SGD for ResNet18 and VGG16 respectively. Comparing with the results in Figure 3, S-SGD converges slower due to the heterogeneous computation, while CoCoD-SGD is robust. When equipped with Proportionally Sampling, S-SGD and Pipe-SGD converge faster since the hardware utilization is improved.

6 Conclusion

In this paper, we propose a computation and communication decoupled stochastic gradient descent (CoCoD-SGD) for distributed optimization. In comparisonwith existing distributed algorithms, CoCoD-SGD can run computation and communication simultaneously to make full use of hardware resources and has lower communication complexity. CoCoD-SGD is also theoretically justified to have the same convergence rate as S-SGD and obtains linear iteration speedup in both homogeneous and heterogeneous environments. In addition, CoCoD-SGD achieves faster distributed training with superior time speedup when comparing with others. Experimental results demonstrate the efficiency of the proposed algorithm.

Acknowledgments

This research was supported by the National Natural Science Foundation of China (No. 61673364, No. 91746301) and the Fundamental Research Funds for the Central Universities (WK2150110008). We also gratefully acknowledge the support of Cheng Li and Youhui Bai from USTC for providing the experimental environment.