ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers
Abstract
How to efficiently serve ever-larger trained natural language models in practice has become exceptionally challenging even for powerful cloud servers due to their prohibitive memory/computation requirements.
In this work, we present an efficient and affordable post-training quantization approach to compress large
Transformer-based models, termed as ZeroQuant.
ZeroQuant is an end-to-end quantization and inference pipeline with three main components: (1) a fine-grained hardware-friendly quantization scheme for both weight and activations; (2) a novel affordable layer-by-layer knowledge distillation algorithm (LKD) even without the access to the original training data; (3) a highly-optimized quantization system backend support to remove the quantization/dequantization overhead.
As such, we are able to show that: (1) ZeroQuant can reduce the precision for weights and activations to INT8 in a cost-free way for both BERT and GPT-3-style models with minimal accuracy impact, which leads to up to 5.19x/4.16x speedup on those models compared to FP16 inference; (2) ZeroQuant plus LKD affordably quantize the weights in the fully-connected module to INT4 along with INT8 weights in the attention module and INT8 activations, resulting in 3x memory footprint reduction compared to the FP16 model; (3) ZeroQuant can be directly applied to two of the largest open-sourced language models, including GPT-J 6B and GPT-NeoX 20B , for which our INT8 model achieves similar accuracy as the FP16 model but achieves up to 5.2x better efficiency.
Introduction
Large-scale natural language models have been widely adopted in different applications, e g., natural language understanding using BERT [64] and generation tasks using GPT-style models [49].
Although those models have achieved cutting-edge accuracy results, as the model size keeps increasing dramatically, the requirements of memory footprint and the computational cost to deploy them become a major bottleneck, even on cloud servers with powerful GPU devices.
One promising way to alleviate this challenge is quantization, which can reduce the bit precision for both weight and activations for lower memory footprint and faster compute (e.g., INT8 Tensor cores on T4/A100).
However, quantization usually requires retraining (also known as quantization aware training, or QAT in short) to recover the accuracy degradation from representation loss of weight and activations.
To enable QAT, the full training pipeline is usually required, including the training data and compute resources, to finetune the model.
Access to those components is now oftentimes not available, and QAT is also a time-consuming process, particularly for those large-scale models.
Recently, zero-shot quantization [10, 47] and post-training quantization (PTQ) [46, 39] are proposed to address the training-data access and compute requirement challenges since PTQ generally requires no (or minimal) retraining.
But most of those works primarily focus on computer vision problems on relatively small scales.
More recently, [7] shows promising PTQ results on BERT.
However, (1) its main focus is on high-precision quantization (INT8/FP16) on BERT base , (2) it does not consider other billion-scale generative models (GPT-3-style models [9]).
More importantly, most of these works do not report real latency improvement, putting the usefulness of these methods in improving inference latency into question.
For example, existing work often do not discuss the quantization/dequantization cost associated with different quantization schemes, which in fact has a big impact to the performance benefit of using low precision.
Besides, for extreme quantization (e.g., INT4), knowledge distillation is usually used to boost performance, which adds another source of expensive computation cost as compared to QAT.
Furthermore, in order to achieve better accuracy performance, hidden-states knowledge distillation, e.g., [3, 81], is usually applied for the quantized model.
This would put significant pressure on the GPU memory and the compute resource requirement since both the teacher and student models needed to be loaded into the GPU memory for training.
In this paper, we present ZeroQuant, an end-to-end post-training quantization and inference pipeline, to address those challenges, targeting both INT8 and INT4/INT8 mixed-precision quantization.
Specifically, our contributions are:
• We apply fine-grained hardware-friendly quantization schemes on both weight and activations, i
e., group-wise quantization for weight and token-wise quantization for activations.
Both quantization schemes can significantly reduce the quantization error and retain hardware acceleration properties.
• We propose a novel layer-by-layer knowledge distillation method (LKD) for INT4/INT8 mixedprecision quantization, where the neural network is quantized layer-by-layer through distillation with minimal iterations and even without the access to the original training data.
As such, at any given moment, the device memory is primarily populated only with a single extra layer’s footprint, making billion-scale model distillation feasible with limited training budget and GPU devices.
• We develop a highly optimized inference backend, which eliminates the expensive computation cost of quantization/dequantization operators, enabling latency speedups on INT8 Tensor cores on modern GPU hardware.
• Our empirical results show that: – ZeroQuant enables quantizing BERT and GPT-3-style models into INT8 weight and activations to retain accuracy without incurring any retraining cost.
Compared to FP16 inference, our INT model achieves up to 5.19x/4.16x speedup on BERT base /GPT-3 350M on A100 GPUs.
– ZeroQuant plus LKD can do INT4/INT8 mixed-precision quantization for BERT and GPT3-style models.
This results in a 3x memory footprint reduction with marginal accuracy loss as compared to the FP16 model.
Also, thanks to the lightweight of LKD, we can finish the quantization process in 33s (10 minutes) for BERT base (BERT large ).
We also demonstrate that LKD can use other datasets to achieve similar performance to the original training data.
– We demonstrate the scalability of ZeroQuant on two of the largest open-sourced language models, i.e, GPT-J 6B and GPT-NeoX 20B , with INT8 quantization.
ZeroQuant can achieve 3.67x speedup over the FP16 model for GPT-J 6B and (2) reduce the GPU requirement for inference from 2 toand latency from 65ms to 25ms for GPT-NeoX 20B (i.e., 5.2x better system efficiency in total).
Related Work
Model compression has been explored from different aspects [26, 38, 40, 35, 44, 21, 25, 51, 19, 76, 41, 27, 56, 60, 29, 61, 69, 34, 15, 39, 32].
Among those, quantization is one of the most promising directions as it directly reduces the memory footprint and compute intensity.
Here, we focus on quantization for NLP models and briefly discuss the related work
The majority of quantization works can be categorized into quantization-aware training (QAT). [57, 78] are the first few works to quantize BERT models using integer numbers for both weight and activations.
Particularly, [57] utilizes Hessian information to push the weight bit-precision to even INT2/INT4, and it also proposes group-wise quantization to quantize the weight matrix in a more fine-grained granularity compared to single matrix quantization.
[22] introduces quantization noise to alleviate the variations of QAT.
[81, 3] leverage very expensive knowledge distillation [27] and data augmentation [29] to ternarize/binarize weights.
[30] combines knowledge distillation [29] and learned step size quantization [20] to quantize the weight to 2–8 bits.
Recently, [62] also uses knowledge distillation to compress GPT-2 models on task-specific problems to INT2.
All those works quantize models using the original training datasets.
More importantly they need retraining or finetuning the full model to recover the accuracy, and such compute cost on extra-large models, like [58, 12], can be hardly affordable for most research labs or practitioners.
One solution to overcome the compute cost challenge is post-training quantization (PTQ).
However,
PTQ often induces a significant drop in accuracy because the network can be sensitive to quantization errors.
Along this line, one of the first works applied to Transformer-based [65] models is [77].
The authors introduce centroid-based quantization method, where outlier numbers use FP32 format and the rest numbers are quantized using non-uniform quantization.
As such, it is hard to get the real inference latency benefit on general compute accelerators, e.g., CPU and GPU, because the parallel processing units in these hardware do not support efficient computation of mixed data types.
More recently, [7] introduces high-precision activation quantization (FP16) for part of the model to overcome the high dynamic activation ranges.
However, to the best of our knowledge, (1) How to apply PTQ on GPT-3-style models while achieving high accuracy has not been studied in any of previous work yet; (2) How to apply PTQ on billion (or even a dozen of billions) scale model is still under-explored; (3) Efficient inference system backend is still missing, especially for fine-grained quantization schemes, making it hard to achieve low latency on commodity hardware.
ZeroQuant resolves all those limitations by considering the system backend into the algorithm design and we verify its capability on both BERT and large-scale GPT-3-style (up to 20 billion, i.e., GPT-NeoX 20B ) models for various tasks.
Background and Challenge 3.1 Transformer Architecture
The transformer architecture usually has three components: an embedding layer, a stack of encoder/decoder layers, and a final classifier. In this paper, we focus on quantizing the encoder/decoder layers, i.e., the transformer block, because it is often the most memory and compute intensive components in the entire architecture. With a transformer block, there are two sub-layers, the multi-head self-attention (MHSA) and the feed-forward connection (FFC). We give a short review later and please refer to [65] for more details.
At high level, transformer models can be broadly categorized to three branches: encoder-only models (BERT) [64], decoder-only models (GPT-3-style) [49], and encoder-decoder models (T5) [50]. In this paper, we focus on encoder-only and decoder-only models but our approach can be applied to encoder-decoder models as well. [Figure] Transformer Block Assume the input of an encoder layer is X , the query, key, value, attention output, FFC dense, and FFC output matrices are W q , W k , W v , W o , W h − 4 h , and W 4 h − h , respectively.
Then the forward propagation of a transformer-block is illustrated in Figure 1, where LN is the layer normalization, Softmax is the softmax operator, and GeLU is the activation function.
Quantization maps high-precision numbers, e.g., FP16/FP32, to its low-precision counterpart, e.g., INT4/INT8, to reduce the model footprint and improve the compute performance. In this work, we use uniform symmetric scalar quantizers.
That is to say, if we have a vector/matrix, x , the quantization is applied as
[Equation]
[Table]
where bit is the number of bit we use to represent the quantized value, and S is the scaling factor.
For weight matrix quantization, S is generally computed as S = max ( abs ( x )) , since the weight matrix is static during inference.
On the other hand, activations’ range is dynamic during inference so that an accurate S requires dynamic calculation during inference.
However, to achieve best latency reduction, coarse-grained static quantization is usually applied in practice, where S is calibrated using training data (e.g., momentum based averaging) and fixed during inference [24].
Although static quantization achieves better latency reduction, it also limits the quantization representation for activations, which is discussed below.
3.3 Post Training Quantization
Post-training quantization (PTQ) exhibits great compression efficiency compared to quantizationaware training (QAT) since PTQ is usually applied to quantize the model without retraining.
A common strategy of PTQ is to feed the training data to the network and calibrate the scaling factor, S , using the running mean.
Please see Appendix A.1 for more details.
Some work has been done for BERT base models [7] with INT8 weight and mixed INT8/FP activation quantization.
However, there is no investigation for (1) even lower bit-precision PTQ on
BERT models and (2) large-scale GPT-3-style models.
Here, we briefly discuss the challenge of the application of PTQ on both BERT (in Appendix C) and GPT-3-style models. [Figure]
The results of GPT-3 350M with PTQ are shown in Table 1.
As can be seen, the INT8 activation quantization (i.e., the row of W16A8) causes the primary accuracy loss.
Further pushing the weight to INT8 (i.e., the row of W8A8) does not change the accuracy of zero-shot evaluation tasks but leads the causal language modeling task (Wikitext-2) to worse perplexity score, which demonstrates the sensitivity of generation tasks as compared to other zero-shot evaluation problems.
For W4/8A16, on some accuracy-based tasks, GPT-3 350M still achieves reasonable performance like OpenBookQA but it loses accuracy on the majority of the rest tasks.
Particularly, for Wikitext-2, GPT-3 350M with
W4/8A16 cannot generate any meaningful text anymore.
Please also see Appendix C for the analysis for BERT.
Dynamic Activation Range To investigate why INT8 activation leads to significant accuracy drop
for both BERT and GPT-3-style models, we plot the token-wise (i.e., the hidden state of each token) range of each activation for different transformer layers of GPT-3 350M in Figure 2 (left).
As can be seen, different tokens have dramatically different activation ranges.
For example, the maximum range of the last layer is around 35 but the minimum range is close to 8.
This larger variance in the activation range makes it difficult to use a fixed quantization range (usually the maximum value) for all tokens to retain the prediction accuracy, because the limited representation power for small range tokens is going to hurt the accuracy performance.
Different Ranges of Neurons in Weight Matrices
Similarly, we plot the row-wise (i.e., the output dimension) weight range of the attention output matrix ( W o ) of GPT-3 350M in Figure 2 (right).
There is a 10x difference between the largest magnitudes of different rows and this leads to the worse generation performance of the INT8 weight PTQ.
This also makes it very challenging when INT quantization is applied as the INT4 only has 16 numbers and a 10x smaller range leads to 2 (or 3) numbers for the representations of those smaller-range rows.
This analysis results also indicate why more expensive hidden-states knowledge distillation [3, 37] is used for ultra-low precision quantization to close the accuracy gap.
However, as the training cost of knowledge distillation for large-scale models is too high, a lightweight and efficient method is desirable for PTQ.
Methodology 4.1 Fine-grained Hardware-friendly Quantization Scheme
As shown in Section 3, even applying INT8 PTQ to BERT/GPT-3-style models leads to significant accuracy degradation.
The key challenge is the representation of INT8 cannot fully capture the different numerical ranges of different rows in weight matrices and different activation tokens.
One way to address this is to use group-wise (token-wise) quantization for the weight matrix (activations).
Group-wise Quantization for Weights Group-wise weight matrix quantization has first been proposed in [57], where a weight matrix W ∈ R n × m is partitioned in to g groups, and each group is quantized separately.
However, in [57], the authors only apply this for quantization aware training.
More importantly, they do not consider the hardware efficiency constraint and they do not have a system backend support.
As such, they lack the real latency reduction benefit.
In our design, we consider the hardware constraint from Ampere Architecture of GPUs (e.g, A100), where the compute unit is based on Warp Matrix Multiply and Accumulate (WMMA) tiling size [54] to achieve the best speedup.
Later, we will show that our group-wise quantization leads to much better accuracy as compared to single-matrix quantization due to its finer-granularity quantization while still achieving great latency reduction.
Token-wise Quantization for Activations As mentioned in Section 3 and Appendix 3.2, a common practice for existing PTQ work is to use static quantization for activation, where the min/max range is calculated at an offline calibration phase.
Such a method might be sufficient for small scale models where the variance in the activation range is small.
However, as analyzed in Section 3, there is a huge variance in the activation range for large-scale transformer models such as GPT-3 350M and
BERT base .
As such, a static quantization scheme (often applied to all tokens/samples) would lead to significant accuracy drop.
One natural idea to overcome this issue is to adopt finer-grained token-wise quantization and dynamically calculate the min/max range for each token to reduce the quantization error from activations.
Our evaluation in Section 5 also shows that token-wise quantization for activation significantly improves the accuracy of GPT-3-style and BERT models.
However, directly applying token-wise quantization using existing DL frameworks, such as the
PyTorch quantization suite, would lead to significant quantization and dequantization cost because token-wise quantization introduces additional operations that lead to expensive data movement overhead between the GPU compute units and the main memory.
To address this issue, we build a highly optimized inference backend for token-wise quantization of transformer models.
For example, the inference backend of ZeroQuant employs so called kernel fusion technique to fuse quantization operator with its previous operator, like layer normalization, to alleviate the data movement cost from token-wise quantization.
Similarly, the dequantization cost of the different GeMMs’ output is alleviated by scaling the INT32 accumulation using both the weight and activation quantization scales, before writing the final FP16 result back to the main memory for the next FP16 operator (like
GeLU).
Those optimization will be discussed in more details in Section 4.3. [Figure]
Token-wise quantization can significantly reduce the representation error for quantized activations.
Also, as it does not need to calibrate the activation range, later we will show that there is no quantization-related cost (e.g., activation range calibration) for a moderate quantization scheme (INT8 weight with INT8 activation) for ZeroQuant.
4.2 Layer-by-layer Knowledge Distillation with Affordable Cost
Knowledge distillation (KD) is one of the most powerful methods to alleviate the accuracy degradation after model compression.
However, there are several limitations of KD, especially for hidden-states
KD on large-scale language models: (1) KD needs to hold a teacher and a student model together during the training, which dramatically increases the memory and compute cost; (2) KD usually requires full training of the student model.
Therefore, several copies (gradient, first/second order momentum) of the weight parameters need to be stored in memory to update the model; (3) KD generally requires original training data, which sometimes are not accessible due to privacy/confidential issues.
To address those limitations, we present our layer-by-layer distillation (LKD) algorithm.
Assume the target model for quantization has N transformer blocks, L 1 , ..., L N , the accessible dataset has input ( X , Y ) , which can be the original training data or datasets from other resources.
Our LKD quantizes the network layer-by-layer and uses its original (i.e., unquantized) version as the teacher model.
More specifically, assume layer L k is going to be quantized, and its quantized version is (cid:98) L k .
Then we use the output of the L k − 1 (i.e., by running inference on X over the first k − 1 layers) as the input of L k and (cid:98) L k , measure the difference, and do the model update to L k , i.e.,
[Equation]
where MSE is the mean square loss, and it can be also replaced by other losses (e.g., KL divergence) as well.
As can be seen, (1) our LKD does not need to hold a separate teacher as we use the same Lto L k − 1 for both teacher/student model.
As such, the only extra model cost we have is L k ; (2) the memory overhead of optimizer states are significantly reduced as the only optimizing layer is L k ; (3) as we never optimize the end-to-end model, the training does not depend on the label anymore.
Later, we will show that LKD does not rely on the original training data in Section 5.6.
4.3 Quantization-Optimized Transformer Kernels
Both optimizing the inference latency and model size is crucial for serving large-scale transformer models in practice.
During inference, the batch size is often relatively small, so the inference latency of the model primarily depends on the time of loading inference needed data from the main memory.
By quantizing the weights and activations to lower precision, we reduce the data volume needed to load those data, which allows more effective use of memory bandwidth and higher loading throughput.
However, simply converting weights/activations to INT8 does not guarantee improved latency because there are additional data movement overhead associated with quantization/dequantization operations as shown in Figure 3 (red box).
Such an overhead becomes expensive and in some cases surpasses the performance benefits of using low precision.
To reap the accuracy improvement from token-wise quantization while obtaining improved latency, we now present our optimizations that maximize the memory bandwidth utilization to speed up inference latency for ZeroQuant.
CUTLASS INT8 GeMM To support INT8 computation, we use CUTLASS [6] INT8 GeMM implementation tuned for different batch sizes.
Unlike standard GPU backend library, such as cuDNN, using CUTLASS allows us to more flexibly fuse quantization operation before and after
GeMM to reduce kernel launching and data-movement overhead.
Fusing Token-wise Activation Quantization Token-wise quantization/dequantization introduce many additional operations that lead to extra data movement cost.
To eliminate these cost, we use [Table] kernel fusion [68] to fuse quantization operation for activation with its previous element-wise and/or reduction operations such as bias-add, GeLU, and LayerNorm into a single operator, as illustrated by the green box in Figure 3.
For the dequantization operation (e.g., dequantizing the integer output from the GeMM operator), we similarly fuse it with our custom GeMM schedule to avoid additional read/write accesses to the main memory as illustrated by the blue box in Figure 3.
By doing the above optimizations, we are able to show significant latency reduction for BERT and GPT-3-style models in Section 5.
Please see Appendix D for more details about our system optimization.
Results
Experimental Details To evaluate the proposed ZeroQuant, we test it on both BERT and GPTmodels.
For BERT, we tested both BERT base and BERT large on GLUE benchmark; and for GPT-3-style models, we tested the GPT-3 350M (i.e., GPT-3-style model with 350M parameters) and GPT-3 1.3B
(i.e.,
GPT-3-style model with 1.3B parameters) on 20 zero-shot evaluation tasks, including 19 accuracybased tasks and 1 language modeling generation task.
To illustrate the scalability of the proposed
ZeroQuant, we also directly apply it to two of the largest open-sourced GPT-3-style models, i.e.,
GPT-J 6B [67] and GPT-NeoX 20B [5].
We use a fixed set of hyperparameters for all the LKD-related experiments even though tuning them may benefit our results.
Please see Appendix A.2 for more training details and see Appendix A.3 for the reported metrics for BERT.
To provide a comprehensive study, we also include a tuning result in Appendix E on BERT and an ablation study for different proposed components in Section 5.5.
Notation Explanation We use WxAy to represent using x-bit for weight quantization and y-bit for activation quantization.
Unless specific explanation, for W4/8, we quantize the MHSA’s weight to
INT8 and FFC’s weight to INT4; for A8/16, we use FP16 activation for self-attention calculation (i.e., the GeMM related to W q/k/v ) and use INT8 for the rest calculation.
We use ZeroQuant to represent the method with only fine-grained quantization schemes and use ZeroQuant-LKD to represent the method with both fine-grained quantization schemes and LKD. A summary of results is shown in Appendix B. We also include (1) the results required by reviewers in Appendix G; (2) the speedup explanation .
5.1 Main Results of BERT
BERT base We report the results of BERT base in Table 2.
For W8A8, the average accuracy of PTQ degrades more than 10 points.
However, ZeroQuant can achieve 83.75 scores, which is only 0. lower than baseline.
Particularly, as ZeroQuant has no activation range calibration phase, the cost of
ZeroQuant is 0 which is even cheaper than standard PTQ.
As compared to [7], our method achieves a better average score (1.29 higher).
Meanwhile, as compared to INT8 activation used in ZeroQuant, [7] uses mixed INT8 and FP16 activation.
We also compare our method with our internal trained QAT and other QAT works [57, 78].
As can be seen, with comparable accuracy results as those QAT methods, ZeroQuant can save the retraining cost from 2900s to 0s for INT8 quantization.
[Table] [Table]
For the more aggressive weight quantization with minimal (or no) training quantization, i.e., W4/8A16,
PTQ fully loses all accuracy (pure random prediction).
However, ZeroQuant can still achieve an 81.65 average score.
On top of ZeroQuant, if we add our LKD, the accuracy can be further boosted to 82.35 with a cost of 31s per task using only a single GPU, which is 93.5x cheaper than INT8 QAT quantization.
We also test ZeroQuant and ZeroQuant-LKD under the W4/8A8 quantization scheme and both of them achieve similar accuracy performance as W4/8A16.
If hyper-parameter tuning is applied to LKD, ZeroQuant-LKD can achieve an 83.22 average score under W4/8A8, which is similar to QAT’s W8A8 result.
Please see Appendix E for more details.
BERT large We test our methods on BERT large as well and the results are shown in Table 3.
Similar to BERT base , ZeroQuant achieves much better accuracy than PTQ methods.
As compared to QAT methods, ZeroQuant has comparable results on larger datasets (like MNLI/QQP) and has better performance on small tasks (e.e., CoLA/MRPC/RTE).
We actually tune QAT for multiple learning rates but cannot get even better performance for those small tasks (see Appendix F for more details).
For more aggressive quantization schemes, like W4/8A16 and W4/8A8, ZeroQuant and ZeroQuantLKD still achieve good accuracy except for RTE but the model size is about 3x smaller than FP counterpart.
This is aligned with the INT8 QAT results, which lose significantly more accuracy on
RTE.
Thanks to the lightweight cost of LKD, it only takes about 550s to finish each task even on
BERT large , which is 13x cheaper than QAT.
5.2 Main Results of GPT-3-style Models
GPT-3 350M
We first test ZeroQuant and ZeroQuant-LKD on GPT-3 350M and report the result in Table 4.
The first interesting finding of zero-shot evaluation on GPT-3-stype models is that the accuracy performance of accuracy-based tasks is more tolerant to quantization than generation tasks.
For instance, W8A8 PTQ has a 1.1% average accuracy drop on 19 accuracy-based tasks as compared to 4.7 points loss on Wikitext-2.
Comparing ZeroQuant with PTQ using W8A8, we can reduce the accuracy gap from 1.1% to 0.2% and the perplexity (PPL) gap from 4.7 to 0.2 with no activation range calibration cost.
For W4/8A16 quantization scheme, PTQ can hardly predict reasonable answers for the majority of tasks and its generation performance on Wikitext-2 is fully crashed.
As a comparison, ZeroQuant still achieves non-trivial performance on some tasks but its generation performance significantly degrades on Wikitext-2.
LKD brings a significant performance boost for this W4/8A16 setting.
Note that
ZeroQuant-LKD increases the accuracy from 33.5 to 37.0 and decreases the PPL from 88.6 to 30. compared to ZeroQuant, and the entire cost of this is just 3.1 hours on a single A100 GPU.
Note that this is about 0.027% GPU hours of the full pretraining cost (128 A100 GPUs for 32 hours).
Similar to W4/8A16, ZeroQuant-LKD achieves much better performance than ZeroQuant on W4/8A8 by using the lightweight LKD.
[Table] [Table]
GPT-3 1.3B
The results of {+GPT-
 3B+} are shown in Table 5.
Similar to GPT-3 350M , for W8A8, ZeroQuant has much better performance than PTQ with less no activation calibration cost, particularly for the generation task Wikitext-
(3.2 points lower).
Also, for W4/8 quantization, LKD can bring non-trivial performance gain for ZeroQuant.
The cost of LKD is about 0.02% of the full pre-training cost (128 A100 GPUs for 120 hours)
5.3 Latency Reduction of BERT and GPT-3-style Models
We compare the inference speed of BERT between FP16 and our INT8 versions in Table 6 on a single 40G-A100 GPU.
Using our efficient quantization kernel implementation and operator fusion, the
INT8 model can achieve 2.27–5.19x speedup on BERT base and 2.47–5.01x on BERT large .
We also include the latency comparison of GPT-3-style models between FP16 and our INT8 version.
Particularly, we use the model to generate the first 50 tokens based on a given text and measure the average latency.
Our INT8 model leads to 4.16x/4.06x speedup for GPT-3 350M /GPT-3 1.3B as compared to the FP16 counterpart.
5.4 A Showcase of GPT-J 6B and GPT-NeoX 20B
To demonstrate the scalability of ZeroQuant, we applied it to two of the largest open-sourced models, i.e., GPT-J 6B and GPT-NeoX 20B , which have 6B and 20B parameters separately.
We report the results of GPT-J 6B in Table 7 on three generation datasets, i.e., PTB [42], Wikitext-2, and Wikitext-103 [43].
As can be seen, as compared to FP16 precision, ZeroQuant achieves similar
PPL on all three different tasks.
To compare the latency, we again use the average latency number to generate the first 50 tokens.
Our W8A8 can get up to 3.67x speedup compared to the FP16 version.
To quantize GPT-NeoX 20B to W8A8 for all GeMMs, the accuracy significantly decreases.
We retrieve the quantization of each weight matrix and of each activation, and finally find out that the activation quantization for the attention calculation (i.e., the input of self-attention) causes the accuracy loss.
We conjecture that this is because of the sensitivity of the self-attention module for extra-large models (20B) but cannot verify this for other models due to the lack of open-sourced extra-large models and the full evaluation pipeline.
As such, we leave the input activation for self-attention in FP16 and quantize the rest to INT8.
The results are shown in Table 8.
Our W8A8/16 achieves similar accuracy performance but can reduce both the GPU resource requirement (from 2 A100 GPUs to 1) and the latency from 65ms to 25ms, which together lead to 5.2x better throughput/efficiency.
5.5 Ablation Study of Different Components
To investigate the performance gain of each component we introduced in Section 4, i.e., group-wise weight quantization, token-wise activation quantization, and lightweight layer-by-layer knowledge distillation, we here do an ablation study on BERT large with W4/8A8.
[Table]
Please see Table I.4 for the
[Table]
The quantization scheme used here is W4/8A8.
Here, GP is the abbreviation of group-wise weight quantization, TQ is the abbreviation of token-wise activation quantization.
[Table]
[Table]
We present the results in Table 9.
As can be seen, group-wise weight quantization boosts the accuracy (random-guess prediction) from PTQ to a non-trivial result (66.52).
Further adding token-wise quantization improves 14.54 points accuracy performance.
On top of those (i [Table] e., ZeroQuant), LKD further brings a 0.56 point gain.
5.6 No Access to The Original Training Data
As mentioned in previous sections, the original training data are oftentimes hard to access due to the privacy and/or confidential issues.
Therefore, we here study the performance of our LKD when there is no direct access to the original training data.
As the distillation objective of our LKD does not depend on the label, the training data used for LKD can be very flexible.
We compare the performance of GPT-3 350M on W4/8A8 quantization scheme using three different training data resources, i.e., random data (using random integer number to generate token ids),
Wikipedia (using Huggingface to get the data 2 ), and original PILE dataset.
The results are shown in Table 10.
Compared to ZeroQuant, LKD using random data can boost the accuracy by 1.1% and reduce the PPL from 92.1 to 40.6.
The reason why random data can still significantly improve the performance is that LKD does not optimize the end-to-end pipeline and it only layer-by-layer learns the internal dependency from the teacher model.
Therefore, random data can also provide meaningful information.
Using Wikipedia data from Huggingface can further improve the accuracy to 36.2 and reduce the PPL to 30.4, which is comparable to the results using the original data.
This indicates that a clean text dataset can be used for LKD when we do not have access to the original full dataset.
Conclusions
With the rapid growth of large model sizes, we have reach a point to consider how to serve those models in practice.
Although several works demonstrate that post-training quantization can be applied to BERT models, to the best of our knowledge, there have been no existing works on (1) billion-scale
GPT-3-style models, (2) ultra-low precision post-training quantization, and (3) end-to-end solution of how to efficiently serve the quantized model online.
In this work, we offer fine-grained compression schemes for both weight and activations to enable INT8 quantization for up to 20B-scale models (GPT-NeoX 20B ).
We also offer a novel affordable layer-by-layer knowledge distillation for ultra-low precision quantization, which leads to 3x model size reduction compared to FP16 model while achieving minimal accuracy degradation.
Furthermore, we provide a system backend support and show up to 5.19x speedup on BERT models and 5.2x better efficiency on GPT-NeoX 20B . A Experimental Details A.1 Details of PTQ on BERT and GPT
For BERT, we use a batch size of 32 and sequence length 128 to calibrate the range of activations. In order to capture the dynamic range, we use 0.95 momentum with 100 iterations, i.e., x max = 0 . 95 x max + 0 . 05 max ( x current − iteration ) , x min = 0 .
95 x min + 0 . 05 min ( x current − iteration ) . For GPT-3-style models, we use the same momentum method but change the batch size to 8 with sequence length 2048. A.2 Details of Main Result
BERT BERT models are trained using the code-base from Huggingface [73]. We show our ZeroQuant method on BERT base and BERT large .
We use the same lower-case tokenizer in BERT large instead of the cased tokenizer in the original paper [16]. When fine-tuning on GLUE [66] tasks ((i.e.,
MRPC [18], STS-B [11], SST-2 [59], QNLI [52], QQP [28], MNLI [71], CoLA [70], RTE [14]). 3 ), we follow the instruction from Huggingface Transformer Library [73]. For ZeroQuant and ZeroQuant-LKD, we use 48 groups for group-wise weight quantization on BERT base and 64 groups for group-wise weight quantization on BERT large , for all the weight matrices.
For LKD, we use 100 iterations with batch size 32 and sequence length 128 for BERT base , and we use 400 iterations for BERT large . We fix the learning rate as 5e-6 for both models on all tasks. However, tuning them may favor ZeroQuant. All the models are trained using a single 40G-A100 GPU (Azure ND A100 instances).
GPT-3-style Models All GPT-3-style models used in the paper are trained using DeepSpeed [53] and Megatron-DeepSpeed Library 4 . The pretraining data are from PILE dataset [23], and the training pipeline and hyperparameters are based on the Megatron-DeepSpeed repository. We use 128 A
GPUs (Azure ND A100 instances) to do the pretraining. It takes about 32 hours to finish the training of GPT-3 350M and 120 hours of GPT-3 1.3B . We evaluate our results on 20 zero-shot evaluation tasks, including 19 accuracy evaluation tasks (i.e., HellaSwag [79], LAMBADA [48], TriviaQA [31], WebQS [4], Winogrande [55], PIQA [63], ARC (Challenge/Easy) [8], ANLI (R1/R2/R3) [72],
OpenBookQA [45], RACE-h [33], BoolQ [13], Copa [1], RTE [14], WSC [36], MultiRC [75], ReCoRD [80]) and 1 language modeling generation task (i.e., Wikitext-2 [43]). For ZeroQuant and ZeroQuant-LKD, we use 64/128 groups for group-wise weight quantization on GPT-3 350M /GPT-3 1.3B for all the weight matrices.
For LKD, we use 1600 iterations with batch size 8 and sequence length 2048 for both GPT-3 350M and GPT-3 1.3B . We fix the learning rate as 5e-6 for both models. However, tuning them may favor
ZeroQuant.
All the quantized models are trained using a single 40G-A100 GPU (Azure ND A100 instances).
We report the performance metric for BERT on GLUE based on Table A.1.
For the average score, if the task only has one metric, we use it for the final result; if the task has two metrics, we compute the average of the two metrics first and use it for the final average score.
For instance, the score of
MRPC used to compute the final average is the mean of its accuracy and F1 score. B Summary of Main Results The accuracy drop varies a lot under different benchmarks and scenarios. Here, we give a summary according to those settings:
[Table]
Particularly, for W4/8, we quantize the MHSA’s weight to INT8 and FFC’s weight to INT4.
Please see for the [Table]
{+• For BERT on GLUE benchmark, the accuracy degradation is usually within 0—0.2 for W8A8 and about 1–2 (except RTE for BERT large ) for W4/8A
 (Table , E.1, and E. • For GPT-3-style models (350M and 1.3B) with W8A8, the average accuracy degradation is about 0.1% and the perplexity drop is about 0.3 points.+}
However, for each of individual tasks, due to the nature of zero-shot evaluation (no task-specific fine-tuning, different number of eval cases between tasks, etc), the accuracy change varies.
Particularly, for some cases, W8A8 even outperforms FP16.
• With W4/8A8, (a) for the 350M GPT-3-style model, the average accuracy degradation is about 2.3% and the perplexity drop is about 9.5; (b) for the 1.3B model, the average accuracy degradation is 2.5% and the perplexity drop is about 2.5.
Note that those results are achieved without tuning the layer-by-layer knowledge distillation. • For GPT-J (6B), our INT8 model has about 0.3 average perplexity loss; and for GPT-NeoX (20B), our INT8 model has about 0.08% accuracy loss. C PTQ challenge of BERT base
From Table C.1, we observe similar results as [7], where the accuracy degradation of INT8 quantization is mainly from activation quantization. Specifically, there is a negligible accuracy drop from INT8 weight quantization (i.e., the row of W8A16). However, with sole INT8 activation (i.e., the row of W16A8), the accuracy decreases from 84.06 to 79.61.
Besides, we also push the weight quantization to a mixed-precision setting with INT4 for weights in FFC and INT8 for weights in
MHSA (i.e., the row of W4/8A16).
This ultra-low precision quantization leads the model to be purely random without meaning prediction. [Figure] a x A b s V a l u e o f E a c h
T o k e n [Figure] [Figure] M a x A b s
V a l u e o f E a c h [Figure] D Details about System Optimization By having the weight and activation quantization, we can use the GeMM schedule that exploits the
INT8 Tensor-core units which provide 2x/4x more compute efficiency compared to the FP16/FP Tensor cores. For this purpose, we adapt the CUTLASS library to produce multiple schedules based on the input sizes we are considering in our application, such as the batch size, sequence length, and the Transformer hidden dimension. To achieve the best latency, we also develop our own efficient parallel implementation of the quantization operator on GPU.
During the inference run-time, based on the total batch size ( batch × seq l en ), we choose the schedule that results in the lowest possible padding when performing the Tensor-core matrix-multiplication operations. To find the best schedule for the GeMM operation, we use the CUTLASS profiler tool that explores the tiling dimensions on the thread-blocks, WARPs, and WMMA (Tensor cores), as the three compute hierarchies available within the Ampere GPU architecture. Then, we find the best schedule by sorting the tile-based schedule based on either peak throughput achieved on the large-batch case, or the maximum memory bandwidth taken from the main memory when the batch size is small. However, there are still several challenges we need to address which are discussed below.
Operation Fusion for Token-wise Activation Quantization. One of the main challenges of our quantization scheme is how to efficiently quantize hidden states before the GeMM operation. In order to remove the overhead, we fuse the activation quantization with its associated element-wise and/or reduction-based operations such as bias-addition, GELY, and LayerNorm. This is due to the fact that each SM takes care of one row (token) of the activation and therefore, we can reuse the computation from the thread registers and compute the quantization scale, avoiding the data movement between
GPU kernels and main memory. Moreover, by converting data from FP16 to INT8, we can utilize the memory bandwidth twice, further improving the inference latency and throughput. Dequantization Associated with GeMM Schedule To utilize the output of integer output from GeMM operator in the following operators, one important step is to dequantize the output by using the scaling factor of the weight and activations.
This dequantization step generally introduces extra overhead for quantized network inference due to the data movement.
As such, we add a custom epilogue, which converts the final accumulated result (from INT32 format) of each row and column of the output to the real value (in FP16 format), using corresponding floating-point quantization scales computed from weight and activation group-wise quantization. By fusing the dequantization with GeMM schedule, we ensure that there is no overhead exposed by using the INT8 operations while producing the FP16 results that are used in the following operation. Furthermore, to effectively combine dequantization with the GeMM operation, we read the two groups of quantization scales for the activation and weight matrices in advance prior to completion of the multiplication of the output matrix. Doing so, we overlap the reading of the extra quantization parameters with the GeMM computation and the GeMM-plus-dequantization can seamlessly work together without stalling the inference pipeline.
Cuda Graph Enhanced Small Model Inference.
As the execution time for specific kernels reduce by optimizing the throughput using the INT8 inference pipeline, the overhead of launching the GPU kernels and the CPU-to-GPU communication become a major bottleneck mostly on small-scale models.
To address this issue, we add the CUDA-Graph support to our inference pipeline that reduces the CPU overhead, by storing the trace of the kernels launched during the inference forward computation, and creating the computation-graph to be reused in the next call to the inference pipeline.
Thus, after storing the graph for the first time, we can replay the graph for the following requests, which substantially improves the performance especially on small models, such as BERT base .
For a fair comparison, we also enable Cuda Graph for FP16 baseline.
Other Operator Fusions Besides quantization/dequantization operator fusion, other fusion methods, which fuse element-wise operations, matrix multiplications, transpositions, and reductions all into a single kernel, are also applied.
These can significantly reduce the number of kernel invocations as well as main memory access to reduce the main memory access latency.
Please see [2] for more details.
[Table]
[Table]
E Tuned Results on BERT
As mentioned in the main text and Appendix A.2, we use the same set of hyperparameters for BERT.
However, tuning them can significantly boost the performance for ZeroQuant.
Here, we tune two hyperparameters, i.e., the learning rate and the number of iterations in order to show the best possible performance of ZeroQuant on both BERT base and BERT large .
Particularly, we choose learning rate from the set {1e-6, 2e-6, 5e-6, 1e-5}, and choose number of iterations from the set {0, 50, 100, 200, 400, 800, 1600}.
Thanks to the lightweight of LKD, the total tuning time for BERT base (including all data loading time, evaluation time, tokenization time, all three quantization schemes, etc) is around 4.5 hours on 8 40G-A100 GPUs (i.e., 36 GPU hours), and the tuning time for BERT large is aroundhours on 8 40G-A100 GPUs (i.e., 128 GPU hours).
We summarize the best results in the Table E.1 and E.2.
F QAT on BERT large
We use four different learning rates for QAT on BERT large , {5e-6, 1e-5, 2e-5, 5e-5}.
The final results we reported in the paper are chosen from the best single run among those four different learning rates.
However, even with such tuning, we are not able to get good performance for BERT large on RTE.
Also, note that the time cost we used in the main text is based on a single run.
if we consider the tuning cost, the total time will be 4 × 7181 s
G Other Results G.1 TinyBERT
Our proposed method can also be applied to small Transformer models as well.
One example is
BERT base used in the main text.
To further demonstrate this, we use TinyBERT (4L-312H 5 ) with
ZeroQuant W8A8 (without LKD) as an example (The checkpoint of CoLA is broken so we did not include the result of CoLA).
The results are shown in Table G.1.
As can be seen, the performance of
ZeroQuant (INT8) is comparable with the original FP16 baseline.
[Table]
[Table]
G.2 CNN Results
We can still apply our quantization scheme on convolutional neural networks with some modification.
For instance, we can use per-output-channel quantization for convolutional kernels as well as perimage quantization for hidden states (images).
As such, this quantization is still satisfied with the hardware compute constraint since we do not break the computation granularity of GPUs (i.e., real speedup can still be achieved on GPUs).
To verify this, we applied this scheme for
ResNet18/ResNet50/ResNet152 with Top-1 and Top-5 accuracy.
The code base is from PyTorch
Example 6 .
As can be seen from Table G.2, the accuracy degradation of ZeroQuant is 0.1% as compared to FP32 models.
G.3 ViT Result
Our method can be applied to ViTs as well.
We chose Huggingface’s ViT checkpoint finetuned on ImageNet-1K to test ZeroQuant (INT8) and PTQ.
We provide both Top-1 and Top-5 accuracy in Table G.3.
G.4 Comparison to Per-column Quantization
As requested by review, we add the comparison between ZeroQuant and per-column quantization used in [74].
First of all, we want to clarify that the authors of [74] mainly focus on the algorithm accuracy evaluation, and they did not provide real speedup results.
However, our paper focuses on both algorithm and hardware constraint/implementation, and we provide real speedup solutions.
Here, we include the comparison of ZeroQuant and the work of [74] (use per-column for weight and per-row for tokens) in Table G.4.
The accuracy of ZeroQuant and the method from [74]. is similar to each other.
G.5 Comparison among Different Quantization Schemes
Asymmetric quantization introduces the zero-point (or bias term), which will affect the inference speed performance (this bias-term introduces the extra matrix-vector production).
That’s the reason why in ZeroQuant we choose symmetric quantization as the quantization scheme.
We did not directly provide asymmetric quantization but MP-PTQ in Table 2 applied asymmetric activation quantization.
To further demonstrate that asymmetric quantization cannot fully resolve the outlier issue (and with higher cost due to the extra matrix-vector production), we provided asymmetric quantization results for PTQ (W8A8) on BERT base and BERT large in Table G.5 and G.6.
As can be seen, asymmetric quantization boosts the PTQ accuracy a lot but it is still >1 point lower than FP16 BERT-base.
However, our ZeroQuant can close this gap to 0.25 without involving any extra
[Table]
[Table]
[Table]
G.6 Five-seed Results
First, we want to clarify that all the results we reported in the main text do not use any hyper-parameter tuning. For W8A8, since there is no training involved, there is no tuning at all. Also, for W4/8A16 or W4/8A8, the main text results use a fixed set of hyperparameters, which means there is no hyperparameter tuning.
This part is mentioned with the detailed setting in Appendix A.
The tuning results in Appendix E are used to demonstrate the low-cost of layer-by-layer knowledge distillation as well as for better comparison if someone applied hyperparameter search. For now, the results we used in the paper are based on a single seed. The main reason is that layerby-layer knowledge distillation is a lightweight training procedure. As such, the weight won’t be significantly changed during LKD finetuning.
To further alleviate the concern about variance/noise, we run the results for BERT-base with W4/8A16 and W4/8A8 with 5 seeds and report the results here.
We follow the same setting as we used in the main text. Particularly, for LKD, we use 100 iterations with batch size 32 and sequence length 128, and we fix the learning rate as 5e-6. For group-wise quantization, we choose 48 groups for all weight matrices. We use seed {40, 41, 42, 43, 44} for this experiment.
The results are shown in Table G.7.
As can be seen, the 5-seed average result is similar to what we have reported in the main text.
G.7 INT8 Comparison between cuBLAS and CUTLASS
For the INT8 kernels, we run a profiling on different tiling schedules for thread-blocks, WARPs, and tensor-cores (WMMA instructions) in order to get the best latency based on the model dimension and batch size.
Note that there are some paddings needed to keep the tenor shapes 16, 32, 64, or 128-aligned in order to maximize the performance when configuring the CUTLASS GeMM kernels.
In our inference system, we choose the best scheduling for each application that minimizes the padding and gives the best performance.
In our experiments, we see much better performance on the INT8 kernels for CUTLASS vs cuBLAS.
Furthermore, we can fuse the dequantization step by modifying the CUTLASS kernels whereas this is not an option for cuBLAS library as it is not open-sourced.
We included detailed profiling results in Table G.8 for BERT-scale matrix sizes and GPT-J/GPTNeoX-scale matrix sizes.
The hardware used is the 40GB-A100. Here we use bsz to represent batch size, and seq to represent sequence length. The time is computed based on 1000 runs average time and the TFLOPS here means the Tera-Flops/s the hardware achieved.
As we can see, our highly optimized CUTLASS INT8 GEMM is much faster than FP16/INT8 CUBLAS GEMM.
Please see Appendix A.3 for the reported metrics [Table] [Table]
G.8 Effectiveness of Group Quantization and Token Quantization
We summarize the results in Table G.9.
Besides directly adopting group quantization (GQ) from
Q-BERT, we also include the real GPU kernel implementation.
Without the GPU kernel part, there is no real speedup benefit for group quantization.
Meanwhile, without token-wise quantization (TQ), the average accuracy of GQ is 66.52.
TQ can further bring this number to 81.06 (14.54 accuracy gain, which is non-trivial).
This demonstrates the benefit of TQ.
G.9 Relative Error Analysis
We use the relative error results for BERT base of the quantized models to support the claim that the proposed method can deal with the dynamic range of full-precision activations and weights.
Particularly, the relative error is computed as
[Equation]
where A means the weight or activation, and the subscript quantize (full) means the quantized (FP16) version.
The results are summarized in Table G.10.
The proposed methods in ZeroQuant can significantly reduce the relative error as compared to the standard quantization method.
G.10 Comparison with MP-PTQ and Q-BERT of BERT large
Since the original Q-BERT and MP-PTQ do not include BERT-large study, we implement our own code to produce the result of Q-BERT based on the paper’s description and use the open-sourced code 7 to produce the result of MP-PTQ (we use –quant-dict "’y’: 16, ’h’: 16, ’x’: 16, ’P’: 16, ’C’: 16" and –quant-setup MSE_logits for STS-B task, and we use –quant-dict "’y’: 16, ’h’: 16, ’x’: 16" for other tasks).
The results are shown in Table G.11.
As can be seen, ZeroQuant achieves similar accuracy as Q-BERT but without any training cost.
For MP-PTQ, we can reproduce the paper’s result for BERT-base models.
However, for BERT-large, all tasks give random guess except CoLA (we tried three different kinds of checkpoints, including
Huggingface’s checkpoint, our own checkpoints, and checkpoints trained using the MP-PTQ repo).
As such, we did not include the rest of the results.
After increasing the bit-precision for activation from 8 to 10 bits, the accuracy of MP-PTQ is significantly boosted.
Therefore, the quantization error of activation is the main cause of MP-PTQ on BERT-large and this is primarily from the dynamic range of activations as what we discussed in Section C.
Also, this result (i.e., MP-PTQ works for BERT-base but not BERT-large out of the box) is aligned with our reported PTQ results in Table 2 and 3, i.e., BERT-large’s
W8A8 has lower accuracy than BERT-base’s W8A8, which means
BERT-large is more sensitive to dynamic activation ranges.
[Table]
[Table]
G.11 Different loss choices of LKD
We tried both L2 and L1 losses, and the results from those two losses are very similar.
Below we provide the results of W4/8A8 on BERT-base.
As can be seen from Table G.12, the accuracy of L loss and L2 loss based LKD is similar to each other.
H Limitations and Future Work
We believe it is critical for every work to clearly state its limitations, especially in this area.
One limitation is that in this work we only focused on natural language models, but it would be interesting to see how ZeroQuant would perform for computer vision models.
We leave this as a future work. Another limitation is that we can only verify the scalability of ZeroQuant up to 20B scale models. If there are new releases of larger open-sourced models, it would be great to test ZeroQuant on those larger models as well. Third, in this paper, we found out that the activation input of self-attention is more sensitive for quantization for the extra-large model (GPT-NeoX 20B ).
However, we are unable to verify this on other extra-large models due to the lack of open-sourced models.
We includes all zero-shot evaluation results in this section for all GPT-3-style models, inlcuding GPT-NeoX 20B .
[Table]
Standard Quantization GQ or TQ
[Table]
, INT
[Table]
[Table]
[Table]
[Table]
[Table] Tasks [Table]
HellaSwag
[Table]
LAMBADA
[Table] TriviaQA [Table] Winogrande
[Table]
[Table]
[Table]
[Table]
PIQA
[Table]
[Table]
[Table]
ARC (Challenge) 22.
[Table]
ARC
[Table]
Easy
[Table] [Table] [Table]
ANLI R 33.
ANLI
[Table]
ANLI R 33.
OpenBookQA 29.
RACE-h 27.
BoolQ 47.
Copa 65.
RTE 50.
WSC 36.
MultiRC 1.
ReCoRD 64.
Wikitext40.
Average Acc 34.
[Table]