Fine-tuning Language Models over Slow Networks using Activation Quantization with Guarantees
Abstract
Communication compression is a crucial technique for modern distributed learning systems to alleviate their communication bottlenecks over slower networks.
Despite recent intensive studies of gradient compression for data parallel-style training, compressing the activations for models trained with pipeline parallelism is still an open problem.
In this paper, we propose AQ-SGD , a novel activation compression algorithm for communication-efficient pipeline parallelism training over slow networks.
Different from previous efforts in activation compression, instead of compressing activation values directly, AQ-SGD compresses the changes of the activations .
This allows us to show, to the best of our knowledge for the first time, that one can still achieve O (1 / √ T ) convergence rate for non-convex objectives under activation compression, without making assumptions on gradient unbiasedness that do not hold for deep learning models with non-linear activation functions.
We then show that AQ-SGD can be optimized and implemented efficiently, without additional end-to-end runtime overhead.
We evaluated AQ-SGD to fine-tune language models with up to 1.5 billion parameters, compressing activation to 2-4 bits.
AQ-SGD provides up to 4 .
3 × end-to-end speed-up in slower networks, without sacrificing model quality.
Moreover, we also show that AQ-SGD can be combined with state-of-the-art gradient compression algorithms to enable “end-to-end communication compression”: 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 sacrificing model quality
Introduction
Decentralized or open collaborative training has recently attracted intensive interests [1, 2, 3, 4].
Despite their great potential in leveraging geo-distributed powerful GPUs, the computation efficiency is severely hindered by low network bandwidth — typically in the range of 10-400Mbps [2, 3, 4].
Recently, efforts in improving communication efficiency have significantly decreased the dependency on fast data center networks — the gradient can be compressed to lower precision or sparsified [5, 6, 7, 8], which speeds up training over low bandwidth networks, whereas the communication topology can be decentralized [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-efficient training paradigms.
However, with the rise of large foundation models [20] (e.g., BERT [21], GPT-3 [22], and CLIP[23]), improving communication efficiency via compression becomes more challenging.
Current training
[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]
systems for foundation models such as Megatron [24], Deepspeed [25], and Fairscale [26], allocate different layers of the model onto multiple devices and need to communicate — in addition to the gradients on the models — the activations during the forward pass and the gradients on the activations during the backward pass.
Compressing these activations leads to a very different behavior compared with compressing the gradient — simply compressing these activations in a stochastically unbiased way 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] difficult to adopt.
Previous efforts on activation compression [32, 33, 34, 35, 36] illustrate, albeit mostly empirically, that large 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 do not apply to neural networks with non-linear activation functions — the only two recent efforts that claim theoretical convergence analysis [29, 30] assume that an unbiased compression on activations leads to an unbiased error on the gradient.
Not surprisingly, these algorithms lead to suboptimal quality under relatively aggressive compression, illustrated in Figure 1a — in many cases, using activation compression to fine-tune a model might be worse than zero-shot learning without any fine-tuning at all.
In this paper, we focus on the problem of activation compression for training language models over slow networks by asking the following:
• Q1.
Can we design an algorithm for activation compression with rigorous theoretical guarantees on SGD convergence?
• Q2.
Can such an algorithm be implemented efficiently without additional runtime overhead and outperform today’s activation compression algorithms without accuracy loss?
Our answers to both questions are Yes .
(Contribution 1)
We propose AQ-SGD , a novel algorithm for activation compression.
The idea of AQ-SGD is simple — instead of directly compressing the 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 enables an interesting “self-enforcing” dynamics: the more training stabilizes → the smaller the changes of the model across epochs → the smaller the changes of activations for the same training example across 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 above dynamics and connect it to SGD convergence, which is quite different from most of today’s results on gradient compression and error compensation.
Under mild technical conditions and quantization functions with bounded error, we show that AQ-SGD converges with a rate of O (1 / √ T ) for non-convex objectives, the same as vanilla SGD [31, 37].
To the best of our knowledge, AQ-SGD is the first activation compression algorithm with rigorous theoretical analysis that shows a convergence rate of O (1 / √ T ) (without relying on assumptions of unbiased gradient).
(Contribution 3)
We then show that AQ-SGD can be optimized and implemented efficiently 2 , without adding additional end-to-end runtime overhead
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
Forward
Activation
Backward
Gradient
Model
Gradient
Macrobatch Macrobatch
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
C C
C C
C C
C C
C C
C C
[Figure]
over non-compression and other compression schemes (it does require us to utilize more memory and SSD for storage of activations).
(Contribution 4)
We then conduct extensive experiments on sequence classification and language modeling datasets using DeBERTa-1.5B and GPT2-1.5B models, respectively.
We show that AQ-SGD can aggressively quantize activations to 2-4 bits without sacrificing convergence performance, where direct quantization of activations fails to converge; in slow networks, AQ-SGD achieves up to 4 .
3 × end-to-end speedup.
(Contribution 5)
Last but not least, we also show that AQ-SGD can be combined with state-of-the-art gradient compression algorithms to enable “end-to-end communication compression”: All data exchanges between machines, including model gradients, forward activations, and backward gradients are quantized into lower precision .
This provides up to 4 .
9 × end-to-end speed-up, without sacrificing model quality.
Overview and Problem Formulation
Training large language models over multiple devices is a challenging task.
Because of the vast number of parameters of the model and data examples, state-of-the-art systems need to combine different forms of 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 over the model using backpropagation, these two machines need to communicate the activations during the forward pass and the gradient on activations during the backward pass.
(Data Parallelism) Machine 1 and Machine 3 (similarly, Machine 2 and Machine 4) process the same set of layers for different macro-batches .
In this case, each of them will hold a replica of the same model.
After the gradient over their model parameters are ready, they need to communicate the model gradient , usually by taking an average [15, 17, 18].
Communication Compression for Forward Activations and Backward Gradients.
In slow networks, the communication among all machines often becomes the bottleneck [37].
To improve the 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 fits into a single machine, there have been intensive efforts on compressing the model gradient [5, 6, 7, 8].
However, when it comes to pipeline model parallelism, such compression techniques are less studied.
In this paper, we focus on designing efficient communication compression algorithms to compress both forward activations and backward gradients.
As we will show later, both can be compressed significantly with AQ-SGD without hurting the model quality.
We also show that it is possible to combine AQ-SGD with state-of-the-art gradient compression 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 the simplicity of notations, we present here the case where the model is partitioned onto K =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 .
Given a distribution of samples D , we consider the following optimization task:
[Equation]
where F is a loss function, a ( − ) and b ( − ) correspond to two sets of layers of the model — a ( − ) has model x ( a ) and b ( − ) has model x ( b ) .
In Figure 2, Machine 1 would hold x ( a ) and Machinewould hold x ( b ) .
In the following, we call the machine that holds x ( a ) Machine a and the machine that holds x ( b ) Machine b .
In the standard backpropagation algorithm, the communication between these two machines 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 ) ) .
Difficulties 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 scheme that 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 / √ T ) .
However, because of the non-linearity of F and b in a deep learning model with non-linear activation functions, an unbiased m ( ξ, x ( a ) ) will lead to biases on 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 efforts on understanding gradient bias [38] have also shown that even bounded bias on gradient can impact the converges of SGD.
As we will show later, empirically, this bias can indeed lead to suboptimal models under aggressive compression.
Notation.
Throughout the paper we use the following definitions:
• 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 .
• ∇ f ( · ) is the gradient of function f .
• g ξ t ( x t ) =
∇ 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 the L 2 -norm.
AQ-SGD : Theoretical Analysis and System Implementations
In this section we present the AQ-SGD , with the goal to mitigate the above mentioned difficulties that appear 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 of quantizing the activations directly, quantize the changes of activations for the same training example across epochs .
As illustrated in Algorithm 1, for iteration t and the data sample ξ t , if it is the first time that ξ t is sampled, Machine a communicates the full precision activations without any compression:
Algorithm
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 (cid:0) a ( ξ t , x ( a ) t ) − m ( ξ t ) (cid:1) 8: end if 9: //
Machine a sends Q (cid:0) a ( ξ t , x ( a ) t ) − m ( ξ t ) (cid:1) to Machine b , which knows m ( ξ t ) through a local version of m 10:
Update x ( b ) t +1 ← x ( b ) t − γ · ∇ x ( b ) ( f ◦ b ) | m 11: //
Machine b sends Q ( ∇ a ( f ◦ b ) | m ) to Machine a 12:
Update x
( a ) t +1 ← x ( a ) t − γ · Q ( ∇ a ( f ◦ b ) | m ) · ∇ x ( a ) a 13: end for 14: Output: x = ( x ( a ) T , x ( b ) T ) m ( ξ t ) = a ( ξ t , x ( a ) t ) (Lines 4-5).
Both machines will save m ( ξ t ) in a local buffer.
If ξ t has been sampled in previous iterations, Machine a communicates a quantized version: Q ( a ( ξ t , x ( a ) t ) − m ( ξ t )) , where m ( ξ t ) was the previous message, stored in the local buffer. Both machines then update this local buffer: m ( ξ t ) ← m ( ξ t ) + Q ( a ( ξ t , x ( a ) t ) − m ( ξ t )) .
Machine b then use m ( ξ t ) as the forward activations, compute backward gradients, and communicate a 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 with sample ξ t :
[Equation]
where γ is the learning rate, ∇ x ( b ) ( f ◦ b ) | ( m ( ξ,x ( a ) t ) ,x ( b ) t ) is the gradient on x ( b ) using the quantized forward activations ( m ( ξ, x ( a ) t ) , and Q ( ∇ 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 communication compression.
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 bias introduced in the forward pass.
3.2 Theoretical Analysis
We now prove the main theorem which states that, under some standard assumptions that are often used in the literature [31, 37], the convergence rate of AQ-SGD algorithm is O (1 / √ T ) for non-convex objectives, same as vanilla SGD.
Assumptions.
We make several assumptions on the networks and the quantization.
It is important to note is that we put no restrictions on either the message error δ ξ , nor the gradient error ∆ ξ .
• (A1: Lipschitz assumptions)
We assume that ∇ f , ∇ ( 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.
∥∇ a ( x ) ∥ ≤ C a , and ∥∇ ( f ◦ b )( x ) ∥ ≤ C f ◦ b .
• (A2: SGD assumptions)
We assume that the stochastic gradient g ξ is unbiased,
i
e. E ξ [ g ξ ( x )]
= ∇ f ( x ) , for all x , and with bounded variance, i e. E ξ ∥ g ξ ( x ) − ∇ f ( x ) ∥ 2 ≤ σ 2 , for all x .
Theorem 3.1.
Suppose that Assumptions A1, A2 hold, and consider an unbiased quantization function
Q ( x ) which satisfies that there exists c Q < (cid:112) 1 / 2 such that E ∥ x − Q ( x ) ∥ ≤ c Q ∥ x ∥ , for all x . 3 Let γ
= 1 3(3 L f + C ) √ 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 -Lipschitz gradient.
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 key quantity to bound is (cid:80)
Tt =1 E ∥
˜∆ ξ t ( x t ) ∥ 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 first prove an auxiliary result which shows that ∥ ˜∆ ξ t ( x t ) ∥ ≤ (1 + C a ) ℓ a ∥ δ ξ t ( x t ) ∥ , for all t .
The key arguments for bounding δ ξ t closely 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 accumulated differences, this allows us to unwrap the changes which appeared since the last time that we observed the current sample, in an iterative way.
However, since these are gradient updates, they are bounded by the learning rate — as long as we have a quantization method that keeps enough information about the signal, we can recursively build enough saving throughout the process.
In particular, the more stability 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 on b and ∇ x ( b ) b , one can prove that the L 2 -regularized loss ˜ f
( x ) =
f
( x ) + λ 2 ∥ x ∥ 2 , which results in weight decay, satisfies Assumptions A1 and A2 with slightly weaker constants.
We note that a theoretical analysis for other regularization methods or optimizers such as Adam [39], could be an independent study 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 ξ , the quantized 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 data parallelism, it reduces to 1TB / # parallel degree, but also incurs communication overhead if data is shuffled 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 and updating step do not incur significant overhead on the end-to-end runtime.
This is because of the fact that the GPU computation time for a forward pass is usually much longer than the data transfer time to load the activations — for GPT2-XL with 1.5 billion model parameters, a single forward pass onlayers require 44 ms, whereas loading m ( ξ ) need 0.2 ms from memory and 12 ms from SSD.
One can simply pre-fetch m ( ξ ) right before the forward pass, and hide it within the forward pass of other data examples.
Similarly, updating m ( ξ ) can also be hidden in the backward computation.
It is also simple to reduce the communication overhead by shuffling data only once or less frequently.
Evaluation
We demonstrate that AQ-SGD can significantly speed up fine-tuning large language models in slow networks.
Specifically, we show: (1) on four standard benchmark tasks, AQ-SGD can tolerate aggressive quantization on the activations (2-4 bits) and backward gradients (4-8 bits), without hurting convergence and final loss, whereas direct quantization converges to a worse loss or even diverge, (2) in slow networks, AQ-SGD provide an end-to-end speed-up up to 4 .
3 × , and (3) AQ-SGD can be combined with 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 classification and language modeling tasks with state-of-the-art foundation models.
For sequence classification, we fine-tune a 1.5B parameter
DeBERTa 4 on two datasets: QNLI and CoLA.
For language modeling, we fine-tune the GPT2 model with 1.5B parameters 5 on two datasets: WikiText2 and arXiv abstracts.
All datasets are publicly available and do not contain sensitive or offensive content.
Detailed setup can be found in Appendix
B.
Distributed Cluster.
We conduct our experiments on AWS with 8-32 p3.2xlarge instances, each containing a V100 GPU.
For a single pipeline, we partition a model onto 8 machines.
When combined with data parallelism, we use 32 instances — data parallel degree is 4 and pipeline parallel degree is 8.
By default, instances are interconnected with 10Gbps bandwidth.
We simulate slow networks by controlling the communication bandwidth between instances using Linux traffic control.
Baselines.
We compare with several strong baselines:
FP32 : in which all communications are in 32 bit floating point without any compression.
DirectQ [29, 30]: in which activations and backward gradients are directly quantized.
We use a simple, uniform quantization scheme, which first normalizes a given vector into [ − 1 , 1] and quantize each number into a b -bit integers by uniforming partitioning the range [ − 1 , 1] into 2 b intervals [35].
Additional details of the configuration can be found in Appendix C. [Figure] [Figure]
Hyperparameter Tuning.
We conduct careful tuning for all methods on all datasets.
We perform grid 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 first compare the convergence behavior of different methods.
For all compression methods, we try various settings: fw x bw y means that we use x bits for forward activation and y bits for backward gradients.
Figure 3 shows the convergence behavior for the sequence classification and language modeling tasks.
FP32 converges fast since it does not introduce any compression errors.
DirectQ , under aggressive quantization, can converge to a significantly worse model, or even diverge.
This is not surprising, given the biases on model gradients that direct quantization introduced.
On the 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 (in terms of time to the same loss), illustrating the importance of communication compression in slow networks.
Table 2 shows the training throughput and Table 3 shows the breakdown of our algorithm.
We note that computation and communication can overlap, so the end-to-end time depends on the larger one of the two.
Another interesting observation is that when the network is 100 × slower (from 10Gbps to 100Mbps), the training is only about 1 .
18 × slower!
This is exciting — if AQ-SGD were to be deployed in a in real-world geo-distributed decentralized networks, the training throughput would be almost as fast as the performance inside a data center!
Moreover, AQ-SGD does not introduce significant runtime overhead over direct quantization.
From
Table 2, we see that AQ-SGD is essentially as efficient as direct quantization compression in throughput. [Table]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure]
4.3 End
Compression
AQ-SGD can also be combined with existing methods on gradient compression.
This allows us to compress all communications during training.
We combine AQ-SGD with QuantizedAdam [40], an error 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 gradient with 4 bits.
Figure 5 illustrates the convergence and the training throughput under different network configuration.
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 worse model.
In terms of training throughput, with both activation and gradient compression, we can achieve 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 improving end-to-end training throughput — as illustrated in Figure 5(c), disabling any of them will lead to a much lower training throughput.
Related Work
Distributed training of foundation models.
Modern distributed training of deep neural networks goes 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 training include Megatron [24], Deepspeed [25], and Fairscale [26].
To scale out the training of large-scale models, pipeline parallelism (e.g., Gpipe[41], Pipedream[43, 44]) is a popular option, where the model is partitioned into multiple stages, different stages are allocated on different GPUs and the exchange of 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], sparsification [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 significantly reduce the use of computation and memory while preserving the generalization of such models.
In particular, activation compression methods [32, 33, 34, 35, 36, 29] are proposed to reduce the memory footprint by adopting lossless [70, 71, 72] and lossy [73, 29, 74] compression techniques in the training of various deep neural networks (e.g., CNN[66, 75, 76], GNN[77]).
These approaches usually compute the activation with full precision in forward propagation, adopt the compression method over the activation, and store the compressed version in DRAM for later use in backward propagation.
Thus, compression does not introduce any error in forward propagation in contrast to the scenario of communicating compressed activation in the distributed setting.
Delta-based compression.
Delta-based compression [78] is a classic solution to various system problems.
Recent research has also used the property of spatial proximity within activation in neural network training based on empirical observations [79, 80, 74].
However, to our knowledge, no attempt has been made to consider the proximity of activation through training epochs to enable efficient compression with theoretical guarantee.
Conclusion
In this paper, we discuss how to adopt communication compression for activations in distributed learning.
We proposed AQ-SGD , a novel activation compression algorithm for communication-efficient pipeline parallelism training over slow networks.
AQ-SGD achieves O (1 / √ T ) convergence rate for non-convex optimization without making assumptions on gradient unbiasedness.
Our empirical study suggests that AQ-SGD can achieve up to 4 .
3 × speedup for pipeline parallelism.
When combined with gradient compression in data parallelism, the end-to-end speed-up can be up to 4 .
9 × .
Acknowledgments
We gratefully acknowledge the support from the Swiss State Secretariat for Education, Research and Innovation (SERI) under contract number MB22.00036 (for European Research Council (ERC) Starting Grant TRIDENT 101042665), the Swiss National Science Foundation (Project Number 200021_184628, and 197485), Innosuisse/SNF BRIDGE Discovery (Project Number 40B2-0_187132), European Union Horizon 2020 Research and
Innovation Programme (DAPHNE, 957407), Botnar Research Centre for Child Health, Swiss Data Science
Center, Alibaba, Cisco, eBay, Google Focused Research Awards, Kuaishou Inc., Oracle Labs, Zurich Insurance, the Department of Computer Science at ETH Zurich, NIH under No. U54EB020405 (Mobilize), NSF under
Nos.
{+CCF
 (Beyond Sparsity), CCF1563078 (Volume+} to Velocity), and 1937301 (RTML); ARL under No. W911NF-21-2-0251 (Interactive Human-AI Teaming);
{+ONR under No. N
+}
(Unifying Weak
Supervision)
; ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine
Learning; N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba,
TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce,
Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), and members of the Stanford DAWN project: Facebook, Google, and VMWare.
The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon.
Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government.
This work was done during Jue Wang’s visit to ETH Zürich, which was supported by
Key Research and Development Program of Zhejiang Province of China (No. 2021C01009) and Fundamental
Research Funds for the Central Universities.
The computation required in this work was provided by Together
Computer ( https://together.xyz/ ).
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on 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 justification to your answer , either by referencing the 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 instructions block 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 reflect 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] 2.
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] 3.
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] 4.
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 noncommercial 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 identifiable information or offensive content?
[Yes] 5.
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]
A Proof of the Main Theorem
In this section we prove Theorem 3.1.
The main idea comes from the “self-enforcing” dynamics described in the introduction of this work: the more training stabilizes → the smaller the changes of the model across epochs → the smaller the changes of activations for the same training example across epochs → the smaller the compression error using the same #bits → training stabilizes more.
We start by providing two auxiliary results.
The first one connects the message error and the gradient error.
Lemma A
[Equation]
For every sample ξ , one has
[Equation]
, and
[Equation]
where ˜∆ ξ ( x ) = (∆ ( a ) ξ ( x ) , ∆ ( b ) ξ ) .
P ROOF : Note that
[Equation]
which together with ∥ ˜∆ ξ ( x t ) ∥ =
∥ ∆ ( a ) ξ ( x t ) ∥
+ ∥ ∆ ( b ) ξ ( x t ) ∥ yields the claim.
■
We now prove that the message error can be efficiently bounded by the true gradient.
Lemma A.2.
For C ′ = 18 c 2 Q l 2 a N 2 1 − 2 c 2 Q , one has
[Equation]
P ROOF :
Let ξ be a fixed sample.
To simplify the exposition
[Equation]
we abuse the notation slightly by a ( x ) = a ( ξ, x ) , m ( x ) = m ( ξ, x ) .
Let T ( ξ ) be the number of realizations of ξ before time T .
Using the definition of δ ξ (recalling that δ ξ ( x t 1 ( ξ ) )
=
0 , since in the first iteration we send the correct signal)
[Equation]
we have
[Equation]
[Equation]
β
[Equation]
is ℓ
[Equation] [Equation]
Transferring the δ ξ part to the LHS, and noting that between every two realizations of ξ at times t k − 1 ( ξ ) and t k ( ξ ) , we can follow updates for t
= t k − 1 ( ξ ) , . .
, t k ( ξ ) − 1 , we get (1 − 2 c 2 Q )
T ( ξ ) (cid:88) k
=
[Equation]
a )
[Table]
{ Cauchy-Schwarz } ≤
[Table]
=
[Table]
[Equation]
[Table]
[Equation]
[Table]
where ω : [ T ] → { 0 , 1 , . . .
} is defined by ω ξ ( t ) =
t k ( ξ ) − t k − 1 ( ξ ) , if t ∈ [ t k − 1 ( ξ ) , t k ( ξ ) − 1] , and ω ξ ( t ) = 0 , if t > t T ( ξ ) ( ξ ) .
We note that for two different samples ξ and ξ ′ , the sums on the LHS are disjoint.
Therefore, summing over all samples ξ and taking the expectation over ξ and all ξ t yields (1 −
[Equation]
since E [ ω ξ ( t )] ≤ N .
Applying ∥
g ξ t ( x t )+∆ ξ t ( x t ) ∥ 2 ≤ 3 ∥
g ξ t ( x t ) ∥ 2 +3 ∥
∆ ( Q ) ξ t ( x t )
∥ 2 +3 ∥
˜∆ ξ t ( x t ) ∥ 2 , bounded variance E ∥
g ξ ( x ) − ∇ f ( x ) ∥ 2 ≤ σ 2 , and Lemma A.1, we get (cid:0)
[Equation]
which implies
[Equation]
Recalling the definitions of C and γ , and the fact that γ · C < 1 3 , we can simplify the LHS to get (cid:0)
[Equation]
yielding the claim.
■
We are ready to prove the main result, with a learning rate of γ =
[Equation]
L f +
[Equation]
) √ T .
P ROOF OF T HEOREM 3.1: Since f has L f -Lipschitz gradient, we know that
[Equation]
Since the quantization is unbiased, implying E Q [∆ ( Q ) ξ ( x )]
= 0 , taking the expectation with respect to the quantization yields
[Equation]
, where ˜∆ ξ t ( x t ) = (∆ ( a ) ξ t , ∆ ( b ) ξ t ) , and we used ( α + β + ρ ) 2 ≤ 3 α 2 + 3 β 2 + 3 ρ 2 .
Taking the expectation over ξ t (simplifying the notation of E Q E ξ to simply E ), and recalling that E [ g
[Equation]
x t )]
= ∇ f ( x t ) , we can bound the first two terms of the RHS by
[Equation]
Lemma A.2 }
Reorganizing the terms, summing over all t ∈ [ T ] and dividing by T yields
[Equation]
where
[Equation]
by the definition of γ .
Applying Lemma A.2 and regrouping the terms now yields
[Equation]
Noting that C ′ C ′′ < C 2 and recalling that γC < 1 / 3 and γL f < 1 /
[Equation]
we get
[Equation]
Since γ · (3 L f + C ) = 1 3 √ T ≤ 13 , the LHS coefficent is at least γ/ 6 , so dividing by γ/ 6 now yields the claim by substituting γ with
[Equation]
■
A.1 Theoretical analysis when K >
In this section we sketch how one can generalize the already provided theoretical analysis for the
K = 2 case.
Instead of Machines a and b , we now suppose that we have a stack a 1 , ..., a K of K Machines such that every pair a i , a i +1 communicates a compressed message, denoted by m i .
We simplify the notation of the complete model by x = ( x (1) , . . .
, x ( K ) ) , and, for a sample ξ , further denote a
[Equation]
[Equation] [Equation]
x (
[Equation]
) t )) , for i ∈
[ K ] .
As in Algorithm 1, for iteration t and the data sample ξ t , if it is the first time that ξ t is sampled, Machine a i communicates to Machine a
i +1 the full precision activations without any compression: m i ( ξ t )
= a ( i ) ( ξ t , x t ) .
If ξ t has been sampled in previous iterations, Machine a
i communicates a compressed version:
[Equation]
where m i ( ξ t ) was the previous message, stored in the local buffer.
Both machines then update this local buffer:
[Equation]
Machine a
i +1 then uses m i ( ξ t ) as the forward activations.
Upon receiving backward gradients from Machine a i +2 , it computes backward gradients, and communicates a quantized version of the backward gradient to Machine a i .
We use
[Equation]
to denote the message error of i -th machine in sending the activations (accumulated also through messages in previous pairs), and denote the total message error by δ ξ
= ( δ (1) ξ , . . .
, δ ( K − 1) ξ ) .
Update Rules.
We can now generalize the update rule for a and b to x
[Equation]
for i = 1 , . . .
, K − 1 , where γ is the learning rate.
We can rephrase the update rule as
[Equation]
where g ξ ( x t ) is the stochastic gradient and ∆ ξ ( x t ) is the total gradient error introduced by communication compression.
We have ∆ ξ = (∆ (1) ξ
[Equation]
∆
[Equation]
Q
[Equation]
1) ξ
[Equation]
) , given by:
[Equation]
i
[Equation]
for all i = 1 , . . .
, K − 1 .
Generalized Assumptions.
In order to state the analogue of Theorem 3.1 for K > 2 , we need to define the corresponding assumptions with respect to Lipschitz properties (whereas Assumption GA is here only for completeness, being the same as A2).
• (GA1: Lipschitz assumptions)
We assume that – f had L f -Lipschitz gradient, – f ◦ a K ◦ . . .
◦ a
i +1 has L f ◦ a K ◦ ... ◦ a i +1 -Lipschitz gradient, and has gradient bounded by C f ◦ a K ◦ ... ◦ a i +1 for all i = 1 , . . .
, K − 1 , – a i is ℓ a i -Lipschitz, and has gradient bounded by C a i , for all i = 1 , . . .
, K .
• (GA2: SGD assumptions)
We assume that the stochastic gradient g ξ is unbiased, i.e. E ξ [ g ξ ( x )]
= ∇ f ( x ) , for all x , and with bounded variance, i.e. E ξ ∥
g ξ ( x ) − ∇ f ( x ) ∥ 2 ≤ σ 2 , for all x .
We have the following analogue of Theorem 3.1.
Theorem A.3.
Suppose that Assumptions GA1, GA2 hold, and let
[Equation]
Consider an unbiased quantization function Q ( x ) which satisfies that there exists c Q < (cid:112) 1 /such
that E ∥ x − Q ( x ) ∥ ≤ c Q ∥ x ∥ , for all x
Then there exists a constant C that depends only on the constants defined above and on √ K and N , such that for the learning rate γ proportional to( C + L f ) √ T , after performing T updates one has
[Equation]
Instead of performing a tedious job of rewriting the proof of the K = 2 case with inherently more constant-chasing, we will simply sketch the differences with respect to the proof of Theorem 3.1.
First we note that, having analogous assumptions as in the case when K = 2 , we can easily prove the following analogue of Lemma A.1.
Lemma A.4.
Let
[Equation]
For every sample ξ , one has
[Equation]
, K − 1 , ∥ ∆ ( i ) ξ
( x ) ∥ ≤ C a i L f ◦ a K ◦ ... ◦ a i +1 ∥
δ ( i )
ξ ∥ + C f ◦ a K ◦ ... ◦
a
i +1 L a i ∥
δ ( i − 1) ξ ∥ , i = 1 , . .
, K − 1 , and ∥ ∆ ( K ) ξ ( x ) ∥
≤ L f ◦ a K ∥
δ ( K − 1) ξ ∥ , implying ∥ ∆ ( Q ) ξ ∥ ≤ ˜ Cc Q .
Comparing this with Lemma A.1, we see that for ˜∆ ξ = (∆ (1) ξ , . . .
, ∆ ( K ) ξ ) , we now have an additional term that depends on δ ( i − 1) ξ .
However, since in the proof of Theorem 3.1 we only rely on ∥ ∆ ξ ∥ 2 , we can proceed even with a crude bound
[Equation]
·∥
[Equation]
Mimicking closely the steps of Lemma A.2, separately for each i and the summing over all i via ∥
[Equation]
, one can straightforwardly yield the following analogue of Lemma A.2.
Lemma A.5.
For C ′ =
K · 36 c
[Equation]
N
[Equation]
−
[Equation]
where l a =
( l a 1 , . .
, l a K ) , one has
[Equation]
Theorem A.3 can now be proved by repeating the steps of the proof of Theorem 3.1, carefully substituting Lemma A.1 by Lemma A.4, and Lemma A.2 by Lemma A.5.
[Table]
arXiv -
[Table]
B Benchmark Dataset
The data statistics can be found in Table 4.
We fine-tune DeBERTa on QNLI and CoLA datasets.
The
QNLI task is to determine whether the context sentence contains the answer to the question.
The
CoLA task aims to detect whether a given sequence of tokens is a grammatical English sentence.
In addition, we fine-tune GPT2 on the WikiText2 training set.
We also collect 30K arXiv abstracts in 2021 to fine-tune GPT2.
Neither corpus is included in the GPT2 pre-training data set.
C Training Task Setup and Hyper-parameter Tuning
Sequence Classification.
We use the AdamW optimizer to fine-tune the model for 10 epochs.
Specifically, for QNLI, we set the learning rate to 3.0e-6, learning rate warm-up steps to 1000, max sequence length to 256, the macro-batch size to 64 and micro-batch size to 8; for CoLA, we set the learning rate to 2.5e-6, learning rate warm-up steps to 250, max sequence length to 128, the macro-batch size to 32 and micro-batch size to 8.
After the learning rate warm-up stage, we decay the learning rate linearly over the training epochs.
Language Modeling.
For both WikiText and arXiv datasets, we use AdamW optimizer with a learning rate of 5.0e-6.
We train the model for 10 epochs with a macro-batch size of 32 and a micro-batch size of 1.
The max sequence length is set to 1024 for both datasets.
After the learning rate warm-up stage, we decay the learning rate linearly over the training epochs.
D Distributed View of AQ-SGD lgorithm
Algorithm 2 shows a multi-node view of AQ-SGD .
For brevity, we omit the first warm-up epoch, where we conduct uncompressed training, and thus we update the previous messages by m ( ξ ) ← a ( ξ, x ) .
Initialize: x 0 , learning rate γ , network a ’s weights x ( a ) , network b ’s weights x ( b ) , quantization function Q , the arrays of previous messages m , where networks a and b each maintain a copy of it.
for t = 1, ..., T do ( on network a )
Randomly sample ξ t ∆ m ( ξ t ) ← Q (cid:0) a ( ξ t , x ( a ) t ) − m ( ξ t ) (cid:1)
Update m ( ξ t ) ← m ( ξ t ) + ∆ m ( ξ t )
Send ∆ m ( ξ t ) to network b ( on network b )
Update m ( ξ t ) ← m ( ξ t ) + ∆ m ( ξ t )
Update x ( b ) t +1 ← x ( b ) t − γ · ∇ x ( b ) ( f ◦ b ) | m
Send Q ( ∇ a ( f ◦ b ) | m ) to network a ( on network a )
Update x
( a ) t +1 ← x ( a ) t − γ · Q ( ∇ a ( f ◦ b ) | m ) · ∇ x ( a ) a end for
Output: x = ( x ( a ) T , x ( b ) T )
E Decentralized Training over Slow Network
Decentralized training for large foundation models recently attracted intensive interests.
Example projects include Learning@home [1], DeDLOC [2], and Training Transformers Together [3].
The goal of these projects is to enable a decentralized open-volunteering paradigm for foundation model training.
As many geo-distributed users contribute their GPUs, these GPUs are often connected via slow networks.
For example, [2] investigates a heterogeneous with bandwidths of 200Mbps, 100Mbps, and 50Mbps; [3] advocates to train Transformer over the Internet with 10-100Mbps bandwidth; [4] considers network bandwidth less than 400Mbps.
In this setting, communication compression is key to performance.
However, when compressing activations, existing methods rely on direct quantization.
This inspired our paper, which provides the first activation quantization method with rigorous theoretical guarantee and outperforms direct quantization.
F Discussion on Tensor Parallelism
We here discuss tensor parallelism [24], and the potential adaptation of our algorithm to it.
In tensor parallelism, the activations are computed across different machines, and need to be aggregated.
Therefore we need to compress activations both before and after allreduce to support tensor parallelism.
Specifically, suppose N machines conduct tensor parallelism, then the output activation is:
[Equation]
and we need to compress communication twice:
[Equation]
We believe that delta compensation could be applied to all Q ( − ) , similar to how previous work handles gradient compression (e.g. Eq.
3 and 4 in [27]).
However, this requires careful further studies, both empirically and theoretically.
We leave activation compression for tensor parallelism as future work.
G Limitation and Potential Future Direction
Additional Storage.
Our algorithm trades storage for communication.
Fortunately, we find this is a reasonable trade-off in our settings.
In the following we show that we can offload the activations to
SSD and hide it within the GPU computation of other data examples.
We compare the throughput under the bandwidth of 10Gb/s. FP32 achieves a throughput of 3.8 seqs/s, while AQ-SGD , either offloading activations to host memory or SSDs, achieves 4.0 seqs/s.
Considering the similar training throughput of the above three settings, we show that the overhead of offloading to
SSD can be successfully hided in GPU computation.
In particular, our largest experiment in this paper requires 172GB storage per machine, which even with a 10x larger dataset can be easily offloaded to SSDs.
For much larger datasets (e.g. 100x), we can use data parallelism to reduce the storage requirement for each machine.
For an even larger dataset, our algorithm might not support it well.
This then requires further studies.
Online Learning.
Our proposal relies on iterating over multiple epochs, which is a common setting.
We understand our current algorithm has limitations in the settings such as online learning.
In the following, we provide a potential solution (to both storage requirement and limitation in online learning) – relaxing AQ-SGD by clustering activations and storing only the centers of the clusters.
Recall that AQ-SGD compresses the “delta” (difference between activations from two epochs) and thereby needs to store activations from the previous epoch.
In the relaxed AQ-SGD , we could use algorithms like clustering or locality sensitive hashing to partition the activations and then we only store the “centers” of each partition/cluster.
When computing the “delta”, we can first identify which partition/cluster the current activation belongs to and retrieve the corresponding “center”.
Then “delta” = activation - “center”.
This will potentially help address storage and online learning limitations.
We will explore this in future work.
[Figure] [Figure]
[Table]
H Additional Results
We provide additional experimental results.
Specifically, we show:
• the convergence results with standard deviation; • the training throughput for different dataset settings; • the numerical stability of training from scratch; • the training results under FP16 precision; • the robustness of AQ-SGD under different hyperparameter settings; • the effectiveness of AQ-SGD in the split learning scenario.
H.1 Convergence Results with Standard Deviation
In the main content, we show the convergence performance of different approaches.
We repeated each experiment three times to ensure reproducibility.
We calculate the moving averages of these convergence curves and then average the results of repeated experiments.
We visualize (shaded areas)
the moving standard deviation in all repeated experiments in Figure 6.
Overall, we observe consistent results for all datasets.
H.2 Throughput under Different Dataset Settings
We show the training throughput under different dataset settings in Figure 5.
In general, the observation is similar to that of the main content: our approach maintains similar throughput even when the network is 100 × slower (from 10Gbps to 100Mbps).
WikiText2 and arXiv have essentially the same throughput results, since we use the same training settings for them.
[Figure] [Figure]
H.3 Results of Training from Scratch
We investigate the numerical stability of AQ-SGD by showing the convergence result of training from scratch, where the model parameters are randomly initialized.
We train WikiText and arXiv datasets for 20 epochs and use the first 10% of steps as warm-up, respectively.
As shown in Figure H.3, we can see that AQ-SGD converges almost as fast as FP32 when training from scratch, which indicates our approach is robust enough even when the model is far from the converged state.
In contrast, the curve of DirectQ becomes flatter in the late training stage, showing a clear gap with FP32 .
H.4 Results of Training under FP
To investigate the convergence performance of AQ-SGD under low-precision training, we here show the results of FP16 training, and compare it with FP32 training.
Figure 8 compares the results under
FP32 and FP16.
In general, the convergence curves are consistent with the FP32 case.
This confirms the effectiveness of AQ-SGD when the activation is already in low precision.
H.5 Hyper-parameter Sensitivity
Here we demonstrate the robustness of our method in various settings.
For fast validation, we focus on evaluating DeBERTa-v3-base 6 on QNIL and CoLA datasets.
We by default use K = 4 devices for pipeline parallel training, 2 bits for forward activation, and 4 bits for backward gradients ( fw2 bw4 ).
Number of Pipeline Stages.
We first investigate the influence of the number of pipeline stages on convergence performance.
Intuitively, partitioning into more pipeline stages leads to more rounds of data compression and communication,resulting in a larger accumulated compression error.
The results of Figures 9a and 9b confirm this intuition.
Specifically, the direct quantization method works not bad when K = 2 , but its performance becomes unsatisfied when we further enlarge K .
In comparison, our approach can maintain similar convergence performance to FP32 .
Number of Bits in Communication.
Figures 9c and 9d compare different methods with different numbers of bits in communication.
We observe that using more bits can improve the convergence performance but lead to higher communication overheads.
In general, our approach achieves better accuracy-efficiency trade-offs.
[Figure]
Number of Bits for Previous Messages.
We may find that storing all previous messages is spaceintensive.
To reduce such requirements, we show that previous messages m can be preserved with low precision.
We here perform quantization on m , where m z means that we use z bits for previous messages.
Figures 9e and 9f show the results with different number of bits of the previous messages.
When only 2 bits are used for the previous messages, despite the fact that it is slightly worse than our default setting, our approach is still significantly better than DirectQ .
And there is no significant performance drop when 8 bits are used for the previous messages.
Pre-trained Model Sizes.
Figures 9g and 9h show the results of the base and large version of
DeBERTa.
Surprisingly, larger models seem to be more tolerant of errors from activation compression than smaller models.
One possible reason is that larger models usually use much smaller learning rates.
So the error of each iteration can be restricted to a smaller range.
Here, we use 2.0e-5 for the base model and 7e-6 for the large model, as suggested in the official repository of DeBERTa.
H.6 Split Learning
Split learning is a scenario of federated learning, where the client trains a shallow part of deep network (known as the cut layer) that accesses the training data, while the rest of the model is trained in a data center.
Clients and server need to exchange the activation and its gradients in the cut, where
AQ-SGD can be adopted.
We evaluated AQ-SGD on a split learning scenario where neither the input data nor its labels are shared with the server—the model is cut twice, one after the first resnet block and one before the last block to generate the prediction.
We evaluate AQ-SGD for split learning over
Cifar10 and Cifar100 with the ResNet34 model.
We set 16 clients and use a Dirichlet distribution with concentration parameter 0.5 to synthesize non-identical datasets.
Following the previously established work of split learning, in each communication round, we conduct local training for each client sequentially.
and each client will train 3 epochs with its local data.
We utilize SGD optimizer with momentum of 0.9, a batch size of 64, and a learning rate of 0.01.
We decay the learning rate to its 10% for every 20 communication rounds.
The datasets are augmented with random cropping and flipping.
To adapt to random cropping, we do the same cropping operation on the retrieved previous message, and only update its non-cropped part.
To adapt to random flipping, we maintain another previous message copy for flipped images, and retrieve and update only the corresponding copy during training.
Figure 10 presents the results of split learning, where fw2 bw8[0.2] means that, for forward pass, we perform 2-bit quantization, and for backward pass, we keep only the top 20% gradients and then perform 8-bit quantization.
We can see that AQ-SGD transfers the activations in 2 bits while maintaining a performance similar to FP32 , which indicates the effectiveness of AQ-SGD in improving the communication efficiency in the split learning scenario.
Furthermore, compared to DirectQ ,
AQ-SGD shows advantages in terms of both the convergence and generalization of the trained model.
[Figure]
I Case Study of Generation Results
We here conduct case study to better understand the generation quality across different methods.
All methods are fine-tuned on WikiText with the same seed.
We use greedy decoding to generate texts.
For DirectQ and AQ-SGD , we use 4-bit in forward pass and 8-bit in backward pass.
We randomly pick documents from the test set of WikiText to construct the prompts and let the model complete the rest.
The generation results after fine-tuning are presented in Tables 6 and 7.
For instance, in Cases 1 and 2, AQ-SGD produces the same results as FP32 ; in Case 3, although
AQ-SGD does not generate the exact same output as FP32 , it still expresses a similar meaning.
This indicates our proposed compression scheme has relatively little impact on the learning process.
In contrast, DirectQ may generate text that is not relevant to the context.
For example, in Case 1, we know that “Matsuo Bash ¯ o was strongly influenced by Du Fu”, so we should expect the text below to show the fact how he was influenced, just like FP32 and AQ-SGD did.
But DirectQ fails to generate reasonable text based on the context.
Other inference results presented in Tables 6 and 7 lead to the same observation.
[Table]
Prompt: ...
Matsuo Bash¯o, the greatest haiku poet, was also strongly influenced by Du Fu; in Oku no Hosomichi, his masterpiece, he cites the first two lines of A Spring
View before a haiku as its introduction and also many of his other haiku have similar wording and themes.
It is said that when he died
Gold:
in Osaka during a long travel, a
copy
of Du Fu’s poetry
was found
with
him
as
one of a few precious items which he was
able to carry around.
FP32 : , his body was placed in a coffin with a Du Fu-like inscription.
AQ
SGD
: , his body was placed in a coffin with a Du Fu-like inscription.
DirectQ : , his body was covered with a red cloth and his head was shaved
[Table]
Prompt : ...
Peppiatt notes that the exhibition showed Bacon no longer needed sensationalist material to make an impact, and was now
Gold : capable
of
creating an intense emotional response through more subtle means ...
FP32 : able to "make a serious contribution to the art of painting" AQ-
SGD
: able to "make a serious contribution to the art of painting"
DirectQ
: able to "make a serious statement
about
the state of the art"
[Table]
Prompt : ...
At this point, Dokhturov’s column appeared behind the French line and joined the battle.
The French were outnumbered more than three to
Gold : one, assaulted in the front by Miloradovich’s column, in the middle by Strik’s and in the rear by Dokhturov.
FP32 : one, but they held their ground and forced the Russians to retreat.
AQ-SGD : one, but they were able to hold their ground.
DirectQ : one, and the Russians were able to hold their ground.
[Table]
Prompt : ...
Between Jupiter and Fort Pierce, the storm knocked down power and telegraph lines.
In the latter city, high waves washed out a portion of the causeway.
In the 1980s, an elderly resident recalled that the storm was the
Gold : most severe on record
in Fort Pierce.
FP32 : worst storm he had
ever seen
AQ-SGD :
worst storm he had ever seen.
DirectQ : worst he had ever experienced.
Case
Prompt : ...
Lesnar was born in Webster, South Dakota, on July 12, 1977.
He was raised on a Webster dairy farm owned by his parents, Stephanie and Richard Lesnar.
He is of German descent.
He has two older brothers named Troy and Chad, and a younger sister named Brandi.
At age 17, he joined the National Guard, where he
Gold : was assigned to an office job after his red-green colorblindness was deemed hazardous to his desire to work with explosives.
FP32 : served in the South Dakota National Guard.
He was a member of the South Dakota National Guard for two years.
AQ-SGD : served in the South Dakota National Guard.
He was a member of the South Dakota National Guard for two years.
DirectQ : served in the South Dakota National Guard for two years.
He was discharged in May 1991.
[Table]
placed the game at # 7 in their "Top Mega Drive Games of All Time" list.
[Table]
game’s graphics, but
[Table]
the game’s easy difficulty level,
[Table]
explaining simply that "the graphics are excellent, but the game is easy to complete."
[Table]
stating that "the game is too easy to be enjoyable".
[Table]
stating that it was "too easy".
[Table]
which was criticized for being too easy.
[Table]
... "Obviously, he has great rapport and communication ability with his team."
[Table]
Sports compared Stevens to legendary coach John Wooden writing
[Table]
: Brad Stevens is winning at Butler the Wooden way – calm and composed on the sideline.
" FP32 : He’s a coach who has been around the block, and he’s not afraid to tell you what he thinks."
[Table]
: Stevens is a coach who has been around the game for a long time, and he has a knack for getting the most out of his players."
[Table]
: Stevens is a coach who has been able to mold a team into a championship contender.
He has a knack for getting the most out of his players, and he has a knack for getting the most out of his players’ personalities."
[Table]
Prompt : ...
The Gore’s social circle in England revolved around his good friend Rufus King, who Gold : was appointed Ambassador to Great Britain in 1796, along with other Massachusetts expatriates.
FP32 : was a member of the Royal Society.
AQ-SGD : was a member of the Royal Society.
DirectQ : was a member of the Royal Society.
[Table]
Prompt : ...
The city has 896 barangays that are known by sequential numbers instead of names.
These barangays are further group into Gold : 100 zones for administrative and municipal purposes.
FP32 : districts, which are further subdivided into wards.
AQ-SGD : districts, which are further subdivided into wards.
DirectQ : districts, which are numbered according to the number of barangays in the district
[Table]
Prompt : ...
The National Spiritual Assembly of the Bahá ’ís of the Philippines
the governing body of the Filipino Bahá’í community, is Gold : headquartered near Manila’s eastern border with Makati.
FP32 : headquartered in the city.
AQ-SGD : headquartered in the city.
DirectQ : headquartered in the city.