Fine-tuning Language Models over Slow Networks using Activation Quantization with Guarantees
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Communication compression is a crucial technique for modern distributed learningsystems to alleviate their communication bottlenecks over slower networks.
Despiterecent intensive studies of gradient compression for data parallel-style training,compressing the activations for models trained with pipeline parallelism is still anopen problem.
In this paper, we propose AQ-SGD , a novel activation compressionalgorithm for communication-efﬁcient pipeline parallelism training over slownetworks.
Different from previous efforts in activation compression, instead ofcompressing activation values directly, AQ-SGD compresses the changes of theactivations .
This allows us to show, to the best of our knowledge for the ﬁrst time, thatone can still achieve O (1 / p T ) convergence rate for non-convex objectives underactivation compression, without making assumptions on gradient unbiasedness thatdo not hold for deep learning models with non-linear activation functions.
We thenshow that AQ-SGD can be optimized and implemented efﬁciently, without additionalend-to-end runtime overhead.
We evaluated AQ-SGD to ﬁne-tune language modelswith up to 1.5 billion parameters, compressing activations to 2-4 bits.
AQ-SGDprovides up to 4 .
3 ⇥ end-to-end speed-up in slower networks, without sacriﬁcingmodel quality.
Moreover, we also show that AQ-SGD can be combined with stateof-the-art gradient compression algorithms to enable “end-to-end communicationcompression”: All communications between machines, including model gradients,forward activations, and backward gradients are compressed into lower precision .
This provides up to 4 .
9 ⇥ end-to-end speed-up, without sacriﬁcing model quality.
Introduction
Decentralized or open collaborative training has recently attracted intensive interests [1, 2, 3, 4]. Despitetheir great potential in leveraging geo-distributed powerful GPUs, the computation efﬁciency is severelyhindered by low network bandwidth — typically in the range of 10-400Mbps [2, 3, 4]. Recently,efforts in improving communication efﬁciency have signiﬁcantly decreased the dependency on fastdata center networks — the gradient can be compressed to lower precision or sparsiﬁed [5, 6, 7, 8],which speeds up training over low bandwidth networks, whereas the communication topology can bedecentralized [9, 10, 11, 12, 13, 14], which speeds up training over high latency networks.
Indeed,today’s state-of-the-art training systems, such as Pytorch [15, 16], Horovod [17], Bagua [18], and BytePS [19], already support many of these communication-efﬁcient training paradigms.
However, with the rise of large foundation models [20] (e.g., BERT [21], GPT-3 [22], and CLIP[23]),improving communication efﬁciency via compression becomes more challenging.
Current trainingsystems for foundation models such as Megatron [24], Deepspeed [25], and Fairscale [26], allocatedifferent layers of the model onto multiple devices and need to communicate — in addition to the
[Figure]
AC-GC [29] and
TinyScript [30] assume that the returned gradient is unbiased, whereas
AQ-SGD algorithm does not rely on such an assumption.
[Table]
gradients on the models — the activations during the forward pass and the gradients on the activationsduring the backward pass.
Compressing these activations leads to a very different behavior comparedwith compressing the gradient — simply compressing these activations in a stochastically unbiasedway will lead to biases in the gradient that cannot be measured easily or expressed in closed form .
This either breaks the unbiasedness assumptions made by most gradient compression results [5, 6, 7, 8]or makes error compensation over gradient biases [27, 28] difﬁcult to adopt.
Previous efforts on activation compression [32, 33, 34, 35, 36] illustrate, albeit mostly empirically, thatlarge deep learning models can tolerate some compression errors on these activation values.
However,when it comes to the underlying theoretical analysis, these efforts mostly make assumptions that donot apply to neural networks with non-linear activation functions — the only two recent efforts thatclaim theoretical convergence analysis [29, 30] assume that an unbiased compression on activationsleads to an unbiased error on the gradient.
Not surprisingly, these algorithms lead to suboptimal qualityunder relatively aggressive compression, illustrated in Figure 1a — in many cases, using activationcompression to ﬁne-tune a model might be worse than zero-shot learning without any ﬁne-tuning at all.
In this paper, we focus on the problem of activation compression for training language models overslow networks by asking the following:
• Q1.
Can we design an algorithm for activation compression with rigorous theoreticalguarantees on SGD convergence?• Q2.
Can such an algorithm be implemented efﬁciently without additional runtime overheadand outperform today’s
activation compression algorithms without accuracy loss?
Our answers to both questions are Yes .
(Contribution 1)
We propose AQ-SGD , a novel algorithmfor activation compression.
The idea of AQ-SGD is simple — instead of directly compressingthe activations, compress the change of activations for the same training example across epochs .
Intuitively, we expect AQ-SGD to outperform simply compressing the activations because it enablesan interesting “self-enforcing” dynamics: the more training stabilizes !
the smaller the changesof the model across epochs !
the smaller the changes of activations for the same training exampleacross epochs !
the smaller the compression error using the same #bits !
training stabilizes more.(Contribution 2)
The theoretical analysis of AQ-SGD is non-trivial since we have to analyze the abovedynamics and connect it to SGD convergence, which is quite different from most of today’s resultson gradient compression and error compensation.
Under mild technical conditions and quantizationfunctions with bounded error, we show that AQ-SGD converges with a rate of O (1 / p T ) for non-convexobjectives, the same as vanilla SGD [31, 37].
To the best of our knowledge, AQ-SGD is the ﬁrstactivation compression algorithm with rigorous theoretical analysis that shows a convergence rateof O
(1 / p T ) (without relying on assumptions of unbiased gradient).
(Contribution 3)
We then showthat AQ-SGD can be optimized and implemented efﬁciently, without adding additional end-to-endruntime overhead over non-compression and other compression schemes (it does require us to utilizemore memory and SSD for storage of activations).
(Contribution 4)
We then conduct extensiveexperiments on sequence classiﬁcation and language modeling datasets using DeBERTa-1.5B and
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
[Figure]
M ac h
i n e
M ac h
i n e
M ac h
i n e
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
T r a n s f o r m e r
B l o ck
M ac h
i n e
C C
C C
[Figure]
GPT2-1.5B models, respectively.
We show that AQ-SGD can aggressively quantize activations to 2-4bits without sacriﬁcing convergence performance, where direct quantization of activations fails toconverge; in slow networks, AQ-SGD achieves up to 4 .
3 ⇥ end-to-end speedup.
(Contribution 5)
Lastbut not least, we also show that AQ-SGD can be combined with state-of-the-art gradient compressionalgorithms to enable “end-to-end communication compression”: All data exchanges between machines,including model gradients, forward activations, and backward gradients are quantized into lowerprecision .
This provides up to 4 .
9 ⇥ end-to-end speed-up, without sacriﬁcing model quality.
Overview and Problem Formulation
Training large language models over multiple devices is a challenging task.
Because of the vast numberof parameters of the model and data examples, state-of-the-art systems need to combine different formsof parallelism.
Figure 2 illustrates an example in which such a model is distributed over four machines:(Pipeline Model Parallelism)
The model is partitioned onto Machine 1 and Machine 2 (similarly,
Machine 3 and Machine 4), each of which is holding a subset of layers .
To compute the gradient overthe model using backpropagation, these two machines need to communicate the activations during theforward pass and the gradient on activations during the backward pass.
(Data Parallelism) Machine1 and Machine 3 (similarly, Machine 2 and Machine 4) process the same set of layers for differentmacro-batches .
In this case, each of them will hold a replica of the same model.
After the gradientover their model parameters are ready, they need to communicate the model gradient , usually by takingan average [15, 17, 18].
Communication Compression for Forward Activations and Backward Gradients.
In slownetworks, the communication among all machines often becomes the bottleneck [37].
To improvethe speed of training, one can conduct lossy compression of the data before they are communicated,illustrated as the C module in Figure 2.
When the model ﬁts into a single machine, there have beenintensive efforts on compressing the model gradient [5, 6, 7, 8].
However, when it comes to pipelinemodel parallelism, such compression techniques are less studied.
In this paper, we focus on designingefﬁcient communication compression algorithms to compress both forward activations and backwardgradients.
As we will show later, both can be compressed signiﬁcantly with AQ-SGD without hurtingthe model quality.
We also show that it is possible to combine AQ-SGD with state-of-the-art gradientcompression techniques to enable the end-to-end compression scheme illustrated in Figure 2.
Problem Formulation.
In this paper, we focus on the following technical problem.
Note that, for thesimplicity of notations, we present here the case where the model is partitioned onto K =2 machines.
AQ-SGD works for cases with K > 2 : (1) in experiments, we consider K =8 , i.e., a single model is partitioned onto 8 machines; (2) in the supplementary material we provide the theoretical analysis for K > 2 .
1: Initialize: x 0 , learning rate  , sub-network a (  ) weights x ( a ) , sub-network b (  ) weights x ( b ) , quantization function Q , array of previous messages m initialized to2: for t = 1, ..., T do 3:
Randomly sample ⇠ t 4: if ⇠ t not seen before then 5:
Set m ( ⇠ t )
=
a ( ⇠ t ,x ( a ) t ) 6: else 7:
Update m ( ⇠ t ) m ( ⇠ t )+ Q  a ( ⇠ t ,x ( a ) t )  m ( ⇠ t )  8: end if 9: //
Machine a sends Q  a ( ⇠ t ,x ( a ) t )  m ( ⇠ t )  to Machine b , which knows m ( ⇠ t ) through a local version of m 10:
Update x ( b ) t + x ( b ) t   · r x ( b ) ( f  b ) | m 11: //
Machine b sends Q ( r a ( f  b ) | m ) to Machine a 12:
Update x
( a ) t + x ( a ) t   · Q ( r a ( f  b ) | m ) · r x ( a ) a 13: end for 14: Output: x =( x ( a ) T ,x ( b ) T ) Given a distribution of samples D , we consider the following optimization task: [Equation] (2.1) where F is a loss function, a (  ) and b (  ) correspond to two sets of layers of the model — a (  ) hasmodel x ( a ) and b (  ) has model x ( b ) .
In Figure 2, Machine 1 would hold x ( a ) and Machine 2 wouldhold x ( b ) .
In the following, we call the machine that holds x ( a ) Machine a and the machine that holdsx ( b ) Machine b .
In the standard backpropagation algorithm, the communication between these twomachines are as follows:
• Given a data sample ⇠ , Machine a sends to Machine b the forward activation: a ( ⇠ ,x ( a ) ) • Machine b sends to Machine a the backward gradient on a ( ⇠ ,x ( a ) ) .
Difﬁculties in Direct Quantization.
A natural way at compressing forward activations is to send,instead of a ( ⇠ ,x ( a ) ) , a quantized version m ( ⇠ ,x ( a ) )=
Q ( a ( ⇠ ,x ( a ) )) .
This is the quantization schemethat state-of-the-art methods such as AC-GC [29] and TinyScript [30] use.
Both AC-GC [29] and
TinyScript [30] assume that gradient is unbiased when m ( ⇠ ,x ( a ) ) is an unbiased estimator of a ( ⇠ ,x ( a ) ) .
This enables their convergence rates of O (1 / p T ) .
However, because of the non-linearity of F and b ina deep learning model with non-linear activation functions, an unbiased m ( ⇠ ,x ( a ) ) will lead to biaseson the gradient.
In Appendix, we will provide an example showing that such a gradient bias will hurt
SGD convergence even for a very simple optimization problem.
On the theory side, previous effortson understanding gradient bias [38] have also shown that even bounded bias on gradient can impactthe converges of SGD.
As we will show later, empirically, this bias can indeed lead to suboptimalmodels under aggressive compression.
Notation.
Throughout the paper we use the following deﬁnitions:
• f ⇤ is the optimal value of 2.1.
• N is the number of samples.
• x t =( x ( a ) t ,x ( b ) t ) is the full model at iteration t .
• r f ( · ) is the gradient of function f .
• g ⇠ t ( x t )=
r F ( ⇠ t ; x t ) is the stochastic gradient.
• Q ( · ) is the quantization function used to compress activations.
• m ( · ) is the message exchanged between a and b in the feed forward path.
denotes L
AQ-SGD : Theoretical Analysis and System Implementations
In this section we present the AQ-SGD , with the goal to mitigate the above mentioned difﬁculties thatappear in direct quantization of the activation functions.
3.1 AQ-SGD Algorithm
Algorithm 1 illustrates the AQ-SGD algorithm.
The idea behind Algorithm 1 is simple — instead ofquantizing the activations directly, quantize the changes of activations for the same training exampleacross epochs .
As illustrated in Algorithm 1, for iteration t and the data sample ⇠ t , if it is the ﬁrst timethat ⇠ t is sampled, Machine a communicates the full precision activations without any compression:m ( ⇠ t )=
a ( ⇠ t ,x ( a ) t ) (Lines 4-5).
Both machines will save m ( ⇠ t ) in a local buffer.
If ⇠ t has been sampledin previous iterations, Machine a communicates a quantized version:
[Equation]
where m ( ⇠ t ) was the previous message, stored in the local buffer.
Both machines then update thislocal buffer:
[Equation]
Machine b then use m ( ⇠ t ) as the forward activations, compute backward gradients, and communicatea quantized version of the backward gradient to Machine a (Line 11).
We use
[Equation]
to denote the message error in sending the activations.
Update Rules.
The above algorithm corresponds to the following update rules, at iteration t withsample ⇠ t :
[Equation]
where  is the learning rate, r x ( b ) ( f  b ) | ( m ( ⇠ ,x ( a ) t ) ,x ( b ) t ) is the gradient on x ( b ) using the quantizedforward activations ( m ( ⇠ ,x ( a ) t ) , and Q ( r a ( f  b ) | ( m ( ⇠ ,x ( a ) t ) ,x ( b ) t ) ) is the quantized backward gradient.
Setting x t =( x ( a ) t ,x ( b ) t ) , we can rephrase the update rule as
[Equation]
where g ⇠ ( x t ) is the stochastic gradient and  ⇠ ( x t ) is the gradient error introduced by communicationcompression.
We have  ⇠ =(  ( a ) ⇠ +  ( Q ) ⇠ ,  ( b ) ⇠ ) given by:
[Equation]
where  ( Q ) ⇠ ( x t ) is the error introduced by the gradient quantization in the backpropagation part, whilst ( a ) ⇠ ( x t ) and  ( b ) ⇠ ( x t ) are the errors that the gradients of a and b , respectively, inherit from the biasintroduced in the forward pass.
3.2 Theoretical Analysis
We now prove the main theorem which states that, under some standard assumptions that are oftenused in the literature [31, 37], the convergence rate of AQ-SGD algorithm is O (1 / p T ) for non-convexobjectives, same as vanilla SGD.
Assumptions.
We make several assumptions on the networks and the quantization.
It is importantto note is that we put no restrictions on either the message error  ⇠ , nor the gradient error  ⇠ .•
(A1: Lipschitz assumptions)
We assume that r f , r ( f  b ) and a are L f , L f  b -, and ` a -Lipschitz,respectively, recalling that a function g is L g -Lipschitz
if
[Equation]
Furthermore, we assume that a and f  b have gradients bounded by C a and C f  b , respectively, i.e.kr a ( x ) k C a , and kr ( f  b )( x ) k C f  b .
• (A2: SGD assumptions)
We assume that the stochastic gradient g ⇠ is unbiased, i.e.
E ⇠ [ g ⇠ ( x )]= r f ( x ) , for all x , and with bounded variance, i.e. E ⇠ k g ⇠ ( x ) r
f ( x ) k 2   2 , for all x .
Theorem 3.1.
Suppose that Assumptions A1, A2 hold, and consider an unbiased quantization function
Q ( x ) which satisﬁes that there exists c Q < p 1 / 2 such that E k x  Q ( x ) k  c Q k x k , for all x . 1 Let = 1 3(3 L f + C ) p T be the learning rate, where
[Equation]
Then after performing T updates one has
[Equation]
We present the full proof of Theorem 3.1 in Appendix A, whereas here we explain the main intuition.
The usual starting point in examining convergence rates is to use the fact that f has L f -Lipschitzgradient.
It is well known that this implies
[Equation]
After taking the expectation over all ⇠ t and summing over all t = 1 ,...,T , we easily see that the keyquantity to bound is
P
[Equation]
( x t ) k 2 , where ˜  ⇠ t ( x t )=(
( a ) ⇠ t ( x t ) ,  ( b ) ⇠ t ( x t )) .
On the other hand,the main object that we can control is the message error,  ⇠ .
Therefore, we ﬁrst prove an auxiliary resultwhich shows that k ˜
⇠ t ( x t ) k  (1+ C a ) ` a k  ⇠ t ( x t ) k , for all t .
The key arguments for bounding  ⇠ tclosely follow the self-improving loop described in the introduction, and can be summarized as follows.
Since we are compressing the information in such a way that we compare with all the accumulateddifferences, this allows us to unwrap the changes which appeared since the last time that we observedthe current sample, in an iterative way.
However, since these are gradient updates, they are boundedby the learning rate — as long as we have a quantization method that keeps enough information aboutthe signal, we can recursively build enough saving throughout the process.
In particular, the morestability we have in the process, the smaller the changes of the model and the compression error gets,further strengthening the stability.
Tightness.
The bound is tight with respect to quantization — setting c Q =0 (implying C =0 ), i.e.quantization does not incur any loss, recovers the original original SGD convergence (cf.
Regularization and other optimizers. Assuming A1 and A2 for f , and under further assumptions onb and r x ( b ) b , one can prove that the L 2 -regularized loss ˜ f ( x )= f
( x )+  2 k x k 2 , which results in weightdecay, satisﬁes Assumptions A1 and A2 with slightly weaker constants.
We note that a theoreticalanalysis for other regularization methods or optimizers such as Adam [39], could be an independentstudy and represent an interesting line of future research.
3.3 System Implementations and Optimizations.
Additional storage and communication.
AQ-SGD requires us to store, for each data example ⇠ , thequantized activation m ( ⇠ ) in a local buffer in memory or SSD.
For example, in GPT2-XL training,a simple calculation shows that we need an approximately extra 1TB storage.
When using dataparallelism, it reduces to 1TB / # parallel degree, but also incurs communication overhead if datais shufﬂed in every epoch.
In addition, when the example ⇠ is sampled again, m ( ⇠ ) needs to be (1)loaded from this local buffer to the GPU, and (2) updated when a new value for m ( ⇠ ) is ready.
Optimization.
It is easy to implement and optimize the system such that this additional loading andupdating step do not incur signiﬁcant overhead on the end-to-end runtime.
This is because of the fact
This leaves interesting direction for future exploration as we expect a careful analysis of sparsity together with more advanced quantization functions can make this condition much weaker.
that the GPU computation time for a forward pass is usually much longer than the data transfer timeto load the activations — for GPT2-XL with 1.5 billion model parameters, a single forward pass on6 layers require 44 ms, whereas loading m ( ⇠ ) need 0.2 ms from memory and 12 ms from SSD.
Onecan simply pre-fetch m ( ⇠ ) right before the forward pass, and hide it within the forward pass of otherdata examples.
Similarly, updating m ( ⇠ ) can also be hidden in the backward computation.
It is alsosimple to reduce the communication overhead by shufﬂing data only once or less frequently.
Evaluation
We demonstrate that AQ-SGD can signiﬁcantly speed up ﬁne-tuning large language models in slownetworks.
Speciﬁcally, we show: (1) on four standard benchmark tasks, AQ-SGD can tolerateaggressive quantization on the activations (2-4 bits) and backward gradients (4-8 bits), without hurtingconvergence and ﬁnal loss, whereas direct quantization converges to a worse loss or even diverge, (2) inslow networks, AQ-SGD provide an end-to-end speed-up up to 4 .
3 ⇥ , and (3) AQ-SGD can be combinedwith state-of-the-art gradient compression methods and achieve an end-to-end speed-up of up to 4 . 9 ⇥ .
4.1 Experimental Setup
Datasets and Benchmarks.
We consider both sequence classiﬁcation and language modeling taskswith state-of-the-art foundation model.
For the sequence classiﬁcation task, we ﬁne-tune a 1.5 billionparameter DeBERTa model 2 on two datasets: QNLI and CoLA.
For the language modeling task, weﬁne-tune the GPT2 model with 1.5 billion parameters 3 on two datasets: WikiText2 and arXiv abstracts.
All datasets are publicly available and do not contain sensitive or offensive content.
Detailed setupcan be found in Appendix B.
Distributed Cluster.
We conduct our experiments on AWS with 8-32 p3.2xlarge instances, eachcontaining a V100 GPU.
For a single pipeline, we partition a model onto 8 machines.
When combinedwith data parallelism, we use 32 instances — data parallel degree is 4 and pipeline parallel degreeis 8.
By default, instances are interconnected with 10Gbps bandwidth.
We simulate slow networksby controlling the communication bandwidth between instances using Linux trafﬁc control.
Baselines.
We compare with several strong baselines:
FP32 : in which all communications are in 32 bit ﬂoating point without any compression.
DirectQ
[29, 30]: in which activations and backward gradients are directly quantized.
We use a simple, uniform quantization scheme, which ﬁrst normalizes a given vector into [  1 , and quantize each number into a b -bit integers by uniforming partitioning the range [  1 , 1] into 2 bintervals [35].
Additional details of the conﬁguration can be found in Appendix C.
Hyperparameter Tuning.
We conduct careful tuning for all methods on all datasets.
We performgrid search to choose learning rate from { 2.5e-6, 3e-6, 5e-6, 1e-5 } and macro-batch size from { 32,64, 96 } for best model performance.
We train all models using the Adam optimizer with weight decay.
4.2 Results
Convergence.
We ﬁrst compare the convergence behavior of different methods.
For all compressionmethods, we try various settings: fw x bw y means that we use x bits for forward activation and y bitsfor backward gradients.
Figure 3 shows the convergence behavior for the sequence classiﬁcation andlanguage modeling tasks.
FP32 converges fast since it does not introduce any compression errors.
DirectQ , under aggressive quantization, can converge to a signiﬁcantly worse model, or even diverge.
This is not surprising, given the biases on model gradients that direct quantization introduced.
Onthe other hand, AQ-SGD converges almost as fast as FP32 in terms of number of training steps.
End-to-End Runtime.
We show the end-to-end runtime of different methods under slow networks.
As illustrated in Figure 4, AQ-SGD achieves a 4.3 ⇥ end-to-end speed-up comparing with that of {+FP
 [Figure] [Figure] [Table]+}
[Figure] [Table] [Figure] [Table]
(in terms of time to the same loss), illustrating the importance of communication compression in slownetworks.
Table 2 shows the training throughput.
Another interesting observation is that when thenetwork is 100 ⇥ slower (from 10Gbps to 100Mbps), the training is only about 1 .
28 ⇥
1 .
14 ⇥ slower!
This is exciting — if AQ-SGD were to be deployed in a in real-world geo-distributed decentralizednetworks, the training throughput would be almost as fast as the performance inside a data center!
Moreover, AQ-SGD does not introduce signiﬁcant runtime overhead over direct quantization.
From
Table 2, we see that AQ-SGD is essentially as efﬁcient as direct quantization compression in throughput.
.3 End-to-end Communication Compression
: AQ SGD + QuantizedAdam
AQ-SGD can also be combined with existing methods on gradient compression.
This allows us tocompress all communications during training.
We combine AQ-SGD with QuantizedAdam [40], anerror compensation-based gradient compression algorithm for data parallel training.
We quantize the forward activations with 3 bits, the backward gradient with 6 bits, and model gradientwith 4 bits.
Figure 5 illustrates the convergence and the training throughput under different networkconﬁguration.
We see that AQ-SGD converges well when combined with QuantizedAdam (Figure 5(a,b)).
On the other hand, DirectQ , when combined with gradient compression, converges to a much [Figure] worse model.
In terms of training throughput, with both activation and gradient compression, we canachieve up to 8.5 ⇥ throughput improvement compared to the no-compression baseline (Figure 5(c)).
We also see that both activation and gradient compression are important in terms of improvingend-to-end training throughput — as illustrated in Figure 5(c), disabling any of them will lead to amuch lower training throughput.
Related Work
Distributed training of foundation models.
Modern distributed training of deep neural networksgoes beyond data parallelism [24, 41, 42] due to the advance of the large-scale foundation models [20],such as BERT [21], GPT-3 [22], and CLIP[23].
Popular systems to support foundation model traininginclude Megatron [24], Deepspeed [25], and Fairscale [26].
To scale out the training of large-scalemodels, pipeline parallelism (e.g., Gpipe[41], Pipedream[43, 44]) is a popular option, where the modelis partitioned into multiple stages, different stages are allocated on different GPUs and the exchangeof activations and gradients on activations goes through network communication.
Communication compression of distributed learning.
Communication compression is an effective system relaxation for distributed training, especially in data parallelism [40, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55].
Popular techniques include quantization [5, 6, 7, 8], sparsiﬁcation [56, 57, 58, 59], sketching [60, 61] and error compensation [27].
Recently, TinyScript [30] proposes to compress activations and gradients simultaneously.
Sparse Learning for activation compression.
Sparse learning [30, 62, 63, 64, 65, 66, 67, 68, 69]has become increasingly popular for training neural networks, as it can signiﬁcantly reduce the use ofcomputation and memory while preserving the generalization of such models.
In particular, activationcompression methods [32, 33, 34, 35, 36, 29] are proposed to reduce the memory footprint by adoptinglossless [70, 71, 72] and lossy [73, 29, 74] compression techniques in the training of various deepneural networks (e.g., CNN[66, 75, 76], GNN[77]).
These approaches usually compute the activationwith full precision in forward propagation, adopt the compression method over the activation, and storethe compressed version in DRAM for later use in backward propagation.
Thus, compression does notintroduce any error in forward propagation in contrast to the scenario of communicating compressedactivation in the distributed setting.
Delta-based compression.
Delta-based compression [78] is a classic solution to various systemproblems.
Recent research has also used the property of spatial proximity within activation in neuralnetwork training based on empirical observations [79, 80, 74].
However, to our knowledge, no attempthas been made to consider the proximity of activation through training epochs to enable efﬁcientcompression with theoretical guarantee.
Conclusion
In this paper, we discuss how to adopt communication compression for activations in distributedlearning.
We proposed AQ-SGD , a novel activation compression algorithm for communication-efﬁcientpipeline parallelism training over slow networks.
AQ-SGD achieves O (1 / p T ) convergence rate fornon-convex optimization without making assumptions on gradient unbiasedness.
Our empirical studysuggests that AQ-SGD can achieve up to 4 .
3 ⇥ speedup for pipeline parallelism.
When combined withgradient compression in data parallelism, the end-to-end speed-up can be up to 4 .
9 ⇥ .
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for informationon how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] ,or [N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[N/A]
• Did you include the license to the code and datasets?
[N/A]
• Did you include the license to the code and datasets?
[N/A]
Please do not modify the questions and only use the provided macros for your answers.
Note that the
Checklist section does not count towards the page limit.
In your paper, please delete this instructionsblock and only keep the Checklist section heading above along with the questions/answers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] AQ-SGD requires additional
storage to store the activation for each data example.
We describe the limitations from
the perspectives of algorithm and implementation in Section 3.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Language
models increase the inequality on employment as suggested by DeepMind.
The
acceleration of training may lead to more misuse of language models.
Our work also
makes positive social impact in privacy protection because it can also be applied in split
learning as shown in Appendix.
(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]
(b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main
experimental 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 they were
chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments
multiple times)?
[Yes] We repeated some of the experiments for at least three times
and reported the average in the paper.
Details are reported in Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of 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?
[Yes] Our system was built on PyTorch and
CuPy and was evaluated using publicly available datasets (QNLI, CoLA, WikiText and arXiv abstracts).
All of them use open-source licenses and can be used for
non-commercial educational purposes.
(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?
[Yes]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[Yes]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review Board
[Table]
(IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A] [Table]