SAPipe: Staleness-Aware Pipeline for Data-Parallel DNN Training
Anonymous Author(s)
AffiliationAddress email
Abstract
Data parallelism across multiple machines is widely adopted for accelerating distributed deep learning, but it is hard to achieve linear speedup due to the heavycommunication.
In this paper, we propose SAPipe, a performant system thatpushes the training speed of data parallelism to its fullest extent.
By introducingpartial staleness, the communication overlaps the computation with minimal staleness in SAPipe.
To mitigate additional problems incurred by staleness, SAPipeadopts staleness compensation techniques including weight prediction and delaycompensation with provably lower error bounds.
Additionally, SAPipe presentsan algorithm-system co-design with runtime optimization to minimize systemoverhead for the staleness training pipeline and staleness compensation.
We haveimplemented SAPipe in the BytePS framework, compatible to both TensorFlowand PyTorch.
Our experiments show that SAPipe achieves up to 157% speedupsover BytePS (non-stale), and outperforms PipeSGD in accuracy by up to 13.7%.
Introduction
Deep Neural Networks (DNNs) have achieved ground-breaking performance on a wide range ofdomains, such as computer vision (CV) [10, 17] and natural language processing (NLP)
Meanwhile, the model sizes and data volumes have grown exponentially, making DNN trainingtime-consuming and resource-intensive.
The most common approach to accelerate DNN training is touse data parallelism, scaling DNN training across multiple devices.
Despite the substantial speedup,distributed machine learning systems with data parallelism often cannot fully utilize the computationresources and achieve linear scaling ( i.e. , GPU number times single-GPU training speed), due tonon-negligible communication overhead [31, 2, 23, 13].
Many recent studies have been devoted to developing communication acceleration techniques.
Someworks reduce communication traffic using gradient compression [2] or mixed-precision training [21],while others schedule communication to overlap it with computation.
For example, ByteScheduler [23] and PACE [3] propose preemptive communication scheduling to hide the communicationoverhead within forward computation time.
These communication scheduling approaches reduce thecommunication overhead without affecting the convergence of training, but still cannot fully hidecommunication when the communication-to-computation ratio is high.
A new direction to accelerate distributed DNN training has been explored, which intentionallyintroduces staleness to the training pipeline in order to further increase the overlap between communication and computation.
For example, PipeSGD [19] uses the gradients from the previous iterationin stochastic gradient descent (SGD), resulting in fixed 1-step staleness.
Though introducing fixed staleness may fully overlap communication with computation for some
DNNs, the staleness may also incur severe problems to the convergence of training.
We observesignificant accuracy degradation or even training divergence with only 1-step staleness of gradients in our experiments.
We identify this key limitation and propose 2 solutions: 1) partial staleness, whichintroduces staleness to a limited number of layers; 2) staleness compensation, which compensatesthe 1-step staleness by predicting the gradient to be produced in the next iteration, with optimizedimplementation to reduce the overhead of prediction.
We design a performant and Staleness-Aware communication Pipeline (SAPipe) system for accelerating distributed DNN training, which reduces the communication overhead in distributed training byoverlapping communication with computation, and approaches the linear scaling.
The main contributions of this paper are as follows:
• We propose a partial staleness algorithm, which finds the minimal number of layers to introducestaleness to, so as to keep the training pipeline running without stall (§ 3.1).• We adopt multiple staleness compensation techniques, including delay compensation, weightprediction and their combinations (§ 3.2).• We propose an algorithm-system co-design, kernel fusion, and other runtime optimizations inSAPipe, that are especially designed and implemented to minimize the system overhead of partialstaleness and staleness compensation (§ 3.3).• We provide theoretical guarantees to show that SAPipe achieves the same convergence rate asvanilla SGD, and conditionally better error bounds compared to PipeSGD (§ 4).• We demonstrate that SAPipe outperforms existing frameworks.
SAPipe achieves up to 157%speedups over BytePS (non-stale), and outperforms PipeSGD in accuracy by up to 13.7% (§ 5).
BackgroundPreliminaries.
In distributed deep learning, we solve the following optimization problem withn workers: min x ∈ R d F ( x ) , where F ( x ) = 1 n (cid:80) i ∈ [ n ] F i ( x ) = 1 n (cid:80) i ∈ [ n ] E z i ∼D i f ( x ; z i ) , ∀ i ∈ [ n ] ,is the objective function, x ∈ R d is the set of model parameters ( d is the total number of modelparameters), z i is a mini-batch of data randomly sampled from the local data D i on device i , and f ( · )is loss function.
A typical DNN is composed of m layers, which are concatenated into a flattenedvector x for simplicity.
Distributed training with data parallelism.
Data parallelism partitions the training data ontomultiple devices, i.e. , workers.
Each worker propagates its local data through the model and calculatethe loss ( forward propagation ).
It uses the loss value to compute the gradients of each parameter( backward propagation ), and aggregates them from all workers, before updated to the global model.
To facilitate distributed training, the parameter server [18] and all-reduce [26] are the two mostpopular communication architectures for gradient aggregation.
The detailed process is shown in Algorithm 1, highlighted in blue.
Staleness pipeline.
To hide communication time, previous works ( e.g. , PipeSGD [19]) introduce1-step staleness to the training pipeline, where the forward computation can progress without waitingfor the gradient synchronization.
To be more specific, the gradient aggregation of each layer isexecuted once its backward computation is finished, followed by the optimizer update using theaggregated gradient from the previous iteration.
Thus, the gradient aggregation overlaps with notonly the backward computation and optimizer update of the current iteration, but also the forwardcomputation of the next iteration.
The gradient aggregation initiated in the current iteration will befinished in the next iteration.
We outline the conceptual distributed staleness training pipeline in Algorithm 1, highlighted in orange.
Staleness pipeline results in delayed gradients, which may causesignificant accuracy degradation of the converged model, as shown in our experiments (§ 5).
SAPipe Design
To address the convergence issue of staleness pipeline, we design a staleness-aware system, SAPipe,as shown in Figure 1(a), based on the following key designs. Partial Staleness.
Not all layers in a DNN model require a staleness training pipeline to fully hidecommunication within computation.
Our partial staleness algorithm finds a minimal number of layersto be updated by stale gradients (§ 3.1), while updating the other layers without staleness.
1: Initialize x2: for all iteration t ∈ [ T ] do 3: for all workers i ∈ [ n ] in parallel do 4:
Compute ∇ f ( x t − 1 ; z i,t ) , z i,t ∼ D i 5: if t = 1 then 6:
Same as t > 1 Pass 7: else 8: g i,t ← ∇ f ( x t − 1 ; z i,t ) 9: g i,t ← ∇ f ( x t − 2 ; z i,t − 1 ) 10: g t = 1 n (cid:80) i ∈ [ n ] g i,t 11: x t ← optimizer ( x t − 1 , g t , η t )
1: Initialize x2: for all iteration t ∈ [ T ] do 3: for all workers i ∈ [ n ] in parallel do 4:
Compute ∇ f ( x t − 1 ; z i,t ) , z i,t ∼ D i 5: if t > 1 then 6: g i,t = ∇ f ( x t − 2 ; z i,t − 1 ) 7: g t = 1 n (cid:80) i ∈ [ n ] g i,t 8: ∆ x t
= x t − 1 − x t −9: g DCt ← DC λ ( g t , ∆ x t ) 10: ▷ DC is a func.
defined in Equ.
x t ← optimizer (cid:0) x t − 1 , g DCt , η t (cid:1)
[Figure]
Staleness Compensation.
To further mitigate the problem caused by staleness of the layers chosen bythe partial staleness algorithm, we compensate the staleness of the gradients by multiple approachesincluding delay compensation and weight prediction (§ 3.2).
Optimized Runtime.
Both staleness training pipeline and staleness compensation incur additionaloverhead of computation and memory copy.
To minimize such overhead, we adopt several systemoptimizations to achieve the high-performance runtime (§ 3.3).
3.1 Partial Staleness
Due to the layer-wise structure of DNN models and the reverse order of forward and backward passes,it only requires parts of the gradients to be stale to hide the communication overhead.
Reducing thenumber of stale gradients mitigates the problems caused by staleness, improving convergence andkeeping communication overhead hidden at the same time.
We seek to find the minimal number ofstale gradients such that there is no training pipeline stall, i.e., no waiting time for computing devices.
A naive way is to enumerate all possible combinations of gradients and check if they lead to trainingpipeline stall, which involves prohibitive complexity.
Fortunately, the execution orders of the layersin sequential models are fixed, which provides the opportunity for efficient searching.
Theorem 1.
In a training pipeline without stall, if some forward layer is stale, then all its precedingforward layers are stale.
Detailed proof of Theorem 1 is in the Appendix.
Then we consider the case that the first k layersneed to be trained on staled gradients for fully overlapping its computation and communication, andformulate the following convex program for finding the minimal number of stale layers.
Here thevariables are defined in Table 1.
Notation Description
[Table]
Duration of
[Table] operator i T, t
Total number
[Table]
u i
Duration of forward operator i g t Stochastic gradient g t = 1 n (cid:80) i ∈ [ n ] g i,t
[Table]
Duration of comm.
operator i ( g t )
j
The j th coordinate of g t , j ∈
[ d ]
[Table]
Completion time of operators ( g i,t ) j
The j th coordinate of g i,t , on worker
i
[Table]
( ∇ F t )
j
The j th coordinate of ∇ F ( x t ) , j ∈ [ d ] m
[Table]
◦ Hadamard (coordinate-wise) product x Model parameter x ∈ R d d The number of model parameters minimize k subject to (cid:88) i v i ≤ (cid:88)
i b i + (cid:88)
i u i , m (cid:88) i = k + v i ≤ m − 1 (cid:88) i
= b i + k (cid:88) i = u i , i = 1 , . . .
, m ; 0 ≤ k ≤ m, v i > 0 , b i > 0 , u i > 0 , i = 1 , . .
, m. The first constraint aims to keep communication and computation fully overlapped.
The second constraint ensures that the non-delayed gradients should be synchronized before the start of computationof the corresponding layers.
We assume that the computation time of each layer remains the sameunder different solutions and the execution order follows FIFO.
The above problem can be solvedin O ( m 2 ) time by enumerating k .
The optimal solution gives the minimal delayed gradients andensures maximal training throughput.
In SAPipe, the first k layers are updated with 1-step staleness, while the remaining layers are updated by vanilladistributed SGD without staleness.
The basic version of
SAPipe (referred to as vanila SAPipe) is a mixture of distributed SGD and PipeSGD as given in Algorithm 1.
Fortheoretical analysis, we add the following definition.
Definition 1.
Let κ = d stale d ∈ [0 , 1] denote the portionof model parameters updated by stale gradients, whered stale is the total number of model parameters containedin the k layers that are updated with 1-step staleness.
3.2 Staleness Compensation
We introduce two staleness compensation methods in our distributed training pipeline: delay compensation and weight prediction.
Note that in SAPipe, they are only applied to the first k layers of DNN,chosen by the partial staleness algorithm above, while the remaining layers are updated by normaldistributed training procedure as shown in the blue part of Algorithm 1.
Delay compensation (DC).
Inspired by DC-ASGD [33], we mitigate the effects of delayed gradientson the model with delay compensation.
This method leverages Taylor expansion of the gradientfunction and efficient approximation of the Hessian matrix of the loss function.
In the original designof DC-ASGD, the following delay-compensated gradient is used: g t ≈ g t ′ + λg t ′ ◦
g t ′ ◦
( x t − 1 − x t ′ − 1 ) ,where t > t ′ , g t is the gradient evaluated on x t − 1 , and λ > 0 is a hyperparameter.
Note that the delaycompensation above is a diagonal approximation of the full matrix form: g t ′ ◦
g t ′ ◦
( x t − 1 − x t ′ − 1 ) ≈g t ′ g ⊤ t ′ ( x t − 1 − x t ′ − 1 ) .
In this paper, we prefer to use the full-matrix form of DC to avoid the additionalerror caused by the diagonal approximation:
[Equation]
The detailed algorithm of SAPipe with DC is shown in Algorithm 2.
We will show that by compensating the 1-step staleness, SAPipe-DC achieves a lower error bound than vanilla SAPipe, andconverges as fast as the non-stale baseline.
Weight prediction (WP).
In a 1-step staleness pipeline, model weights are always updated usingthe synchronized gradients in the last iteration, resulting in convergence problems with inconsistentweight.
The goal of WP is to estimate the 1-step-ahead weights for forward pass and backward pass,
1: Initialize x2: for all iteration t ∈ [ T ] do 3: for all workers i ∈ [ n ] in parallel do 4: if t = 1 then 5:
Compute ∇ f ( x 0 ; z i, 0 ) , z i, 0 ∼ D i 6: else 7: g i,t ← ∇ f (˜ x i,t − 1 ; z i,t − 1 ) ▷ Computed in the prev.
iteration 8:
Option 1: ˜
g
i,t ← g i,t ▷ Local gradient 9:
Option 2: ˜
g i,t ← g t −▷ Latest sync.
gradient 10:
Option 3: ˜
g i,t ← DC ( g t − 1 − 1 n g i,t − 1 , ∆ x t ) + 1 n g i,t , where ∆ x t = x
t − 1 − x t −11:
˜ x i,t ← optimizer ( x t − 1 , ˜ g
i,t , η t )
▷ 1-step-ahead weight prediction 12:
Compute ∇ f
(˜ x i,t ; z i,t ) , z i,t ∼ D i
▷ Synchronization is finished in the next iteration 13: g t ← 1 n (cid:80) i ∈
[ n ] g i,t 14: x t ← optimizer ( x t − 1 , g t , η t ) and to obtain 1-step-ahead gradients, in order to counteract the 1-step staleness. We provide threeoptions for weight prediction: 1) WP with local gradient in the current step, 2) WP with the latestsynchronized gradient, and 3) WP with all the above combined and DC.
The algorithm of SAPipe with WP is given in Algorithm 3.
Line 5 computes the normal gradientsin the first step.
In the following steps, Line 7 retrieves the local gradient cached in each worker,and Line 8 to Line 10 provide the options for weight prediction.
We apply an extra optimizer updateusing predicted gradients (Line 11) in each worker to predict 1-step-ahead weights, and conductforward and backward propagation to obtain the 1-step-ahead gradients (Line 12).
Line 13 finishesthe synchronization of the gradients from the last step, and Line 14 uses them to update the model.
Note that the delayed gradients in the last step are computed based on 1-step predicted weights.
Therefore, the 1-step staleness is compensated.
3.3 Optimized Runtime
Figure 1(b) shows the workflow of the staleness pipeline system.
SAPipe’s staleness pipeline incursextra overhead, namely: (1) the computation of staleness compensation; (2) data transfer betweenthe staleness buffer and other modules.
To minimize the system overhead, we present an algorithmsystem co-design.
We fuse the computation operations in DC and WP methods and other smallweight update functions into several batched kernels.
This saves a significant amount of kernellaunching time.
We further optimize the usage of staleness buffer.
We use a double buffer system [28]to handle stale gradients, reducing the transferring overhead, and minimize the memory usage bysharing the staleness buffer between computation and communication.
We also switch the locationof the staleness buffer between CPU memory and GPU memory according to the execution timeof communication and computation pipeline stages.
This overlaps the transferring overhead of thestaleness buffer between communication and computation devices.
Theoretical Results
We establish theoretical guarantees of the convergence of SAPipe for smooth but non-convex problems, using SGD as the optimizer: optimizer SGD ( x t − 1 , g t , η t ) =
x t − 1 − η t g t , with a constantlearning rate η t = η .
4.1 Assumptions
Assumption 1.
(Smoothness)
We assume that f ( x ; z ) , ∀ z , are L -smooth: ∥∇ f ( x ; z ) − ∇ f
( y ; z ) ∥ ≤
L ∥ x − y ∥ , ∀ x, y , which implies f ( y ; z ) − f ( x ; z ) ≤ ⟨∇ f ( x ; z ) , y − x ⟩ + L 2 ∥ y − x ∥ 2 .
Assumption 2.
For any stochastic gradient g i,t = ∇ f
( x i,t − 1 ; z i,t ) , z i,t ∼ D i , where D i is the localdataset on worker i , we assume bounded variance and ℓ 2 -norm: E [ ∥ g i,t − ∇ F i ( x i,t − 1 ) ∥ 2 ] ≤ V 1 ,∥ g i,t ∥ 2 ≤ V 2 , ∀ i ∈ [ n ] , t ∈
[ T ] .
Furthermore, gradients from different workers are independent ofeach other.
Assumption 3.
Introduced in [30], we assume bounded gradient diversity: (cid:80) i ∈ [ n ] ∥∇ F i ( x ) ∥ 2 ∥ (cid:80) i ∈ [ n ] ∇ F i ( x ) ∥ 2 ≤ ρ, ∀ x .
The gradient diversity shows to what extent the local gradients on different workers are distinguishedfrom each other.
It is easy to check that Assumption 3 implies the bounded difference between thelocal and global gradients: 1 n (cid:80) i ∈ [ n ] ∥∇ F i ( x ) − ∇ F ( x ) ∥ 2 ≤ (cid:0) ρn − 1 (cid:1) ∥∇ F ( x ) ∥ 2 .
Assumption 4.
For DC in Eqn.
(1), we assume that the objective function f ( x ) is twicedifferentiable.
Thus, we have the Taylor’s approximation on x ′ ∈ R d with bounded remainder:∥∇ f ( x ) − ( ∇ f ( x ′ )
+ ∇ 2 f ( x ′ )
( x − x ′ ))
∥
≤ M ∥ x − x ′ ∥ 2 , where ∇ 2 f ( x ′ ) is the Hessian matrixevaluated on x ′ .
We also assume that the Hessian approximation error is upper-bounded by ∆ ,∀ x ∈ R d , i.e., ∥∇ f ( x )( ∇ f ( x ))
⊤ − ∇ 2 f ( x ) ∥ 2 ≤ ∆ , with vector-induced matrix norm ∥ · ∥ 2 .
Note that the original paper of DC-ASGD proves that ∆ → 0 when t → ∞ under certain assumptions.
Thus, it is reasonable to assume a bounded approximation error.
Assumption 5.
For simplicity, we assume that for partial staleness, the elements of every gradientare randomly and independently chosen to have 1-step staleness.
Assumption 6.
There exists at least one global minimum x ∗ , where F ( x ∗ ) ≤ F ( x ) , ∀ x. And wedefine the initial gap as R 0 =
F ( x 0 ) − F ( x ∗ ) .
4.2 Convergence Analysis
We derive the following error bounds on the convergence of SAPipe under the above assumptions.
All proofs can be found in Appendix A.
Theorem 2.
Under Assumptions 1, 2, 5 and 6, taking η ≤ 1 L , after T iterations, for vanilla
SAPipe without DC or WP, we have the following error bound: 1 T (cid:80)
Tt =1 E (cid:2) ∥∇ F ( x t − 1 ) ∥ 2 (cid:3) ≤2 R 0 ηT + κErr
+ V ar 0 , where Err
=
L 2 η 2 V 2 and V ar 0 = LηV 1 n .
Remark 1.
Theorem 2 shows that the overall convergence error bound of vanilla SAPipe includes2 main components: the gradient estimation error Err 0 and the variance error V ar 0 .
Err 0 isrooted in the 1-step staleness.
V ar 0 is incurred by the random sampling in stochastic gradients.
Remark 2.
When no layer is updated with staleness, SAPipe is reduced to vanilla distributed SGD.
In this case, the gradient estimation error κErr 0 vanishes with κ = 0 , which results in the exact errorbound of vanilla distributed SGD.
On the contrary, when all the layers are updated with staleness,
SAPipe is reduced to PipeSGD with the gradient estimation error Err 0 and κ = 1 .
Theorem 3.
Under Assumptions 1, 2, 4, 5 and 6, taking a small enough λ so that ληV 2 ≤ 1 − 1 √ 2 ,and η ≤ 1 L , after T iterations, for Algorithm 2 (SAPipe-DC), we have the following error bound:
[Equation]
, where Err DC = 128 η
[Equation] [Equation]
Remark 3.
Regardless of the variance error, the main difference between the error bounds of vanilla
SAPipe and SAPipe-DC is the gradient estimation error.
Note that if we have a small enough Hessianestimation error yielding very small M and ∆ , by taking small enough η and making λ → 1 , theoverall gradient estimation error of SAPipe-DC is smaller than that of PipeSGD, i.e., Err DC ≤ Err 0under certain conditions.
In other words, SAPipe-DC has better convergence compared to PipeSGDwhen the Hessian approximation error is small enough .
Theorem 4.
Under Assumptions 1, 2, 3, 5 and 6, taking η ≤ 1 L , after T iterations, for Algorithm 3 (SAPipe-WP) with Option 1, we have the following errorbound:
[Equation]
0 , where Err WP −
[Equation] [Equation]
Remark 4.
The main difference between Theorem 2 and Theorem 4 is the gradient estimation error,
i.e., V 2 versus (cid:0) ρn − 1 (cid:1) V 2 + (cid:0) 3 − ρn (cid:1) V 1 .
When V 1 and ρ are small enough (i.e., V 1 ≤ 2 − ρ/n 3 − ρ/n V 2 ),
SAPipe-WP with Option 1 produces a smaller error compared to the vanilla SAPipe, i.e., Err WP − 1 ≤
Err 0 under certain conditions.
In other words, SAPipe-WP with Option 1 has better convergencecompared to vanilla SAPipe when the difference between the local datasets on different workers issmall enough .
Theorem 5.
Under Assumptions 1, 2, 5 and 6, taking η ≤ 1 L , after T iterations,for Algorithm 3 (PipeSGD-WP) with Option 2, we have the following error bound:
[Equation]
ar 0 , where Err WP − 2 =
[Equation] [Equation]
Remark 5.
The main difference between Theorem 2 and Theorem 5 is the gradient estimation error,i.e., V 2 versus 1 1 − 2 L 2 η 2 (2 V 1 +2 L 2 η 2 V 2 ) .
When L is small enough (i.e., L 2 ≤ 1 − 2 V 1 /V 2 4 η), SAPipe-WPwith Option 2 produces a smaller error compared to vanilla SAPipe, i.e., Err WP − 2 ≤ Err 0 undercertain conditions.
In other words, SAPipe-WP with Option 2 has better convergence compared tovanilla SAPipe when the objective function is “smooth” enough .
Theorem 6.
Under Assumptions 1, 2, 4, 5 and 6, taking η ≤ 1 L , after T iterations, for Algorithm 3 (SAPipe-WP) with Option 3, we have the following error bound:
[Equation]
0 , where Err WP − 3 =
η 2 L 2 [
[Equation]
Remark 6.
The main difference between Theorem 2 and Theorem 6 is the gradient estimation error,i.e., V 2 versus [8 V 1 n + 4 η 4 M 2 V 22 + 2 η 2 V 2 ( L 2 (1 − λ ) 2 + λ 2 ∆ 2 )] .
When η , M , ∆ are small enoughand n is large enough, SAPipe-WP with Option 3 produces a smaller error compared to vanilla
SAPipe, i.e., Err WP − 3 ≤ Err 0 under certain conditions .
Furthermore, λ provides a trade-offbetween L 2 and ∆ 2 .
If the objective function is relatively “smooth” yielding L < ∆ , then λ → 0is preferred, and otherwise λ → 1 is preferred.
In practice, since both L and ∆ are unknown anddepend on the model architecture and datasets, tuning λ ∈ [0 , 1] is required for better performance.
Corollary 1.
For vanilla SAPipe, SAPipe-DC and SAPipe-WP with all 3 options, taking η ∝min (cid:16) 1 √ T , 1 L (cid:17) , we have 1 T (cid:80)
Tt =1 E (cid:2) ∥∇ F ( x t − 1 ) ∥ 2 (cid:3)
≤ O (cid:16) 1 √ T (cid:17) + O (cid:16) V 1 n √ T (cid:17) , same as vanilla SGD.
Remark 7.
The corollary above shows that vanilla SAPipe, SAPipe-DC and SAPipe-WP with all3 options converge to a critical point where ∥∇ F ( x t − 1 ) ∥ → 0 , when T → + ∞ , and the errordecreases when there are more workers.
In general, vanilla SAPipe, SAPipe-DC and SAPipe-WPhave the same overall convergence rate O (cid:16) 1 √ T (cid:17) .
However, the detailed gradient estimation errorvaries for different algorithms, as shown in Theorems 2, 3, 4, 5, and 6.
Under certain conditionssuch as low gradient variance, low gradient diversity, good smoothness, low Hessian approximationerror, and specific choices of η and λ , one of these algorithms achieves the lowest convergence error.
Hence, hyperparameter tuning and algorithm selection are required in practice .
Evaluation
5.1 Set-up
Testbed.
We evaluate SAPipe on 8 physical machines, each equipped with 90 CPU cores, 320GBmemory, 8 Tesla V100 GPUs with NVLinks, and 100Gbps bandwidth between any two machines.
Models and datasets.
We choose two CV models, VGG16 [27], ResNet50 [10], and two NLPmodels, pretrained GPT-2 [25], Transformer [29], as our benchmark models.
The batch sizes per
GPU are 128 images, 128 images, 80 tokens and 3200 tokens, respectively.
We adopt SGD optimizerwith 0.9 Polyak’s momentum [24] and 5e-5 weight decay when training VGG16 and ResNet50models, and Adam [14] optimizer with (0.9, 0.98) betas for NLP models.
The global learning ratesfor VGG16, ResNet50 and GPT-2 are 0.1, 0.1, and 5e-5, respectively, and we follow the learningrate setting in [29] when training Transformer.
SAPipe uses Option 3 in Algorithm 3 as the defaultstaleness compensation method, with λ empirically set as 0.2.
We train CV models on two datasets: (i) CIFAR-10 [16] and (ii) ImageNet [17].
We fine-tune thepretrained GPT-2 model on (iii) WikiText-2 language modeling dataset [20].
The Transformer modelis trained on (iv) Multi30K [8] for WMT16 English-to-German Multimodal Translation task.
Baselines.
We compare SAPipe with three state-of-the-art communication frameworks: (1)
Horovod [26], a high-performance all-reduce paradigm; (2) BytePS [13], an optimized parameterserver architecture; (3) PipeSGD [19], a pipelined training framework with 1-step staleness.
Allbaselines and SAPipe are run on PyTorch computation framework.
[Table] [Table]
[Table] [Table]
[Figure] [Figure]
perplexity (lower is better) for GPT-2, and test accuracy (higher is better) for others.
[Figure]
5.2 Results and Analysis
Convergence.
Figure 4 shows the convergence curves when training DNN models with 16 GPUs.
We observe that SAPipe converges much faster than baselines.
It not only achieves the highesttraining throughput, but also incurs very little influence on model convergence.
For example, SAPipeis roughly 35%, 67% and 21% faster for VGG16 than BytePS, Horovod and PipeSGD, respectively.
Though PipeSGD has comparable training speed, it converges much slower across models, anddecreases final accuracy by 13% with Transformer.
This reveals the obvious error incurred by 1-stepstaleness.
Staleness mitigation.
Table 2 lists the final model performance of different frameworks.
3 SAPipeachieves comparable test accuracy and perplexity with non-stale baseline, BytePS.
PipeSGD incurssignificant performance drop on four models.
With staleness compensation, SAPipe solves theconvergence problem with staleness pipeline, while greatly boosting the training speed.
Overall,
SAPipe achieves significant improvement in accuracy/perplexity over PipeSGD of 2.6%, 3.5%, 13.7%and 2.13, for VGG16, ResNet50, Transformer and GPT-2, respectively.
We also compare different staleness compensation options of SAPipe with baselines.
The best optionvaries among DNN training jobs.
For example, “SAPipe-WP-OPT1” ( i.e. , predicting weights usinglocal gradients) achieves the highest test accuracy when training CV models on CIFAR-10, and“SAPipe-WP-OPT3” ( i.e. , combined mitigation with WP and DC) is the best option for ImageNet, it.
WikiText-2 and Multi30K datasets.
Other options improve the convergence to some extent ascompared to PipeSGD, but are suboptimal in these cases.
The method to find the best stalenessmitigation option for a given DNN task is our future work.
Scalability.
Figure 3 shows the throughput of baselines and SAPipe when training with differentnumbers of workers.
SAPipe achieves up to 57% and 203% speedups over BytePS and Horovod inall settings.
Horovod does not apply communication scheduling, which leads to the worst throughput.
BytePS enables preemptive communication scheduling, but it does not overlap gradient synchronization of the first few layers with computation.
Thus, the communication overhead in BytePS is stilllarge for GPT-2 model.
PipeSGD performs the best among baselines, while SAPipe still achieves upto 6% speedup over PipeSGD, thanks to our high-performance runtime optimization.
Ablation study.
We evaluate the effectiveness of each component of SAPipe.
Figure 5(a) showsthe throughput speedup and staleness mitigation with partial staleness when training VGG16 on
CIFAR-10 dataset with 8 GPUs.
We observe that with our partial staleness, SAPipe improves thefinal accuracy without slowing down training, by involving less delayed gradients.
We also compare SAPipe with BytePS, PipeSGD and SAPipe without runtime optimization whentraining DNN models with 32 GPUs, as shown in Figure 5(c).
We use the speedup over the BytePSbaseline as the throughput metric.
“SAPipe w/o optimized execution” refers to a naive implementationof staleness-aware pipeline, which incurs non-negligible overhead for staleness mitigation.
SAPipeachieves the highest throughput across all models, and the overhead of staleness mitigation decreasesthe training speed by 7% to 24% over SAPipe.
Without runtime optimization, the staleness pipelinewith mitigation methods could be slower than BytePS baseline by up to 6%.
Sensitivity.
Figure 5(d) varies the value of hyperparameter λ in “SAPipe-WP-OPT3” methodwhen training ResNet50 and Transformer.
We observe that the staleness compensation method is notsensitive to the hyperparameter, and achieves the highest accuracy when it reaches 0.5, 0.2 and 0.2when training on CIFAR-10, ImageNet and Multi30K datasets, respectively.
Related Work
Communication optimization.
Many popular ML frameworks, such as TensorFlow [1] and PyTorch [22], enable overlapping communication with backward propagation by default.
Recentworks [12, 9, 23, 3] further overlap gradient synchronization with forward computation via tensorpartitioning, at the cost of extra overhead.
ASP [6] advocates asynchronous training to improve communication efficiency; without controlling the staleness, it leads to unstable convergence.
SSP [11] isanother asynchronous training protocol with bounded staleness.
PipeSGD [19] restricts the stalenessto one step in training.
However, these methods do not mitigate the extra error caused by staleness,which degrades the model performance.
We focus on synchronous pipelined training, without gradientcompression.
We minimize the number of stale layers, and propose multiple staleness compensationmethods, achieving both high training throughput and comparable accuracy.
Staleness mitigation.
Staleness mitigation is important for asynchronous SGD.
Widely used approaches include staleness-aware rescaling of learning rate [32] and gradients [4], and delay compensation [33].
[15, 5] use linear weight prediction to narrow down the difference between models usedin forward and backward passes for model parallelism.
In our experiments, we find that simply usingdelay compensation or linear weight prediction in PipeSGD results in poor convergence, which callsfor better strategies.
Furthermore, the additional overhead incurred by staleness mitigation motivatesour co-design of training algorithms and system optimization.
However, the best staleness mitigationmethod varies in different DNN tasks, and this is the limitation of SAPipe that requires empiricalexperiments to find the optimal solution.
Conclusion
We present SAPipe, a performant and staleness-aware pipelined system to accelerate distributed DNNtraining without model performance loss.
To fully overlap gradient synchronization communicationwith computation with minimal staleness, we introduce partial staleness, which restricts the numberof layers learned with stale gradients.
To further mitigate convergence issues caused by staleness,
SAPipe adopts weight prediction and delay compensation.
With an algorithm-system co-design,
SAPipe achieves both better error bounds in theory, and high-performance runtime in practice.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes (b) Did you describe the limitations of your work
[Yes] Section 6, related work
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] (d) Have you read the
ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] Section 4. (b) Did you include complete proofs of all theoretical results?
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes] Section 4.1 and source code
(c) Did you report error bars (e
g., with respect to the random seed after running experi- ments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A] Publicly available benchmarks
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]