2206.01861_ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers

Abstract

  • 提出了一种叫ZeroQuant的后训练量化方法,用来高效压缩大型Transformer模型(比如BERT和GPT-3),让模型在推理时更快、更省内存,同时准确度几乎不受影响。

  • ZeroQuant有三大特点:

    1. 针对权重和激活值设计了细粒度、硬件友好的量化方案;

    2. 开发了一种不依赖原始训练数据的经济高效的逐层知识蒸馏算法(LKD);

    3. 优化了系统后端,减少量化带来的额外开销。

  • 主要效果包括:

    • 将权重和激活降到INT8,速度提升4-5倍左右;

    • 结合LKD,将部分权重量化到INT4,内存占用减少3倍;

    • 在大型开源模型(GPT-J 6B,GPT-NeoX 20B)上,INT8量化模型准确度和FP16相当,但效率提升超过5倍。

  • 简而言之,就是用ZeroQuant能让大语言模型跑得更快、更省资源,还保持高精度。

1. Introduction

  • 大规模自然语言模型(如BERT、GPT)性能强,但体积大,计算和内存成本高,部署难度大。

  • 量化(将模型权重和激活用更低精度表示)是减少内存和加速推理的有效方法,但传统量化需要重新训练(QAT),耗时且需训练数据,现实中难以实现。

  • 新兴的零样本量化和后训练量化(PTQ)能减少对训练数据和计算资源的依赖,但现有方法多针对计算机视觉且规模较小,且实际延迟提升不明显。

  • 极端量化(如INT4)常用知识蒸馏提升精度,但对显存和计算资源需求大,因为要同时加载教师和学生模型。

  • 本文提出ZeroQuant,一种端到端的后训练量化方案,支持INT8和混合INT4/INT8量化,特点包括:

    • 细粒度硬件友好量化策略(权重按组,激活按token)减少误差并支持硬件加速;

    • 创新分层知识蒸馏(LKD),分层逐步量化,迭代少,无需原始训练数据,显存占用低,适合大模型;

    • 高度优化的推理后端,消除量化/反量化的高计算开销,显著提升推理速度。

  • 实验结果:

    • ZeroQuant可无重新训练将BERT和GPT-3类模型量化为INT8,准确率保持,速度提升约4-5倍;

    • ZeroQuant+LKD支持混合INT4/INT8量化,内存减3倍,精度轻微下降,量化时间短,且可用其它数据集替代原训练集;

    • 在大规模模型(GPT-J6B、GPT-NeoX20B)上,ZeroQuant带来显著速度提升和资源节省,推理延迟降低超过2倍,系统效率提升超5倍。

  • 简单来说,ZeroQuant解决了大模型量化难、资源需求高的问题,实现了快速、低资源、高效的后训练量化和推理加速。

2. Relative Work

🌟 总体介绍:

  • 模型压缩有很多方法,其中量化是最有前景的方向之一,因为它可以直接减小内存占用和计算量。

  • 本文聚焦于NLP模型的量化研究。

🧠 主流方法 1:训练感知量化(QAT)

  • 是最常见的做法:在训练阶段就考虑量化。

  • 代表性工作用整数(如 INT2 / INT4)来量化 BERT 的权重和激活值。

  • 有些方法引入Hessian矩阵信息,或采用分组量化来更精细控制。

  • 有些使用知识蒸馏 + 数据增强将权重二值化或三值化。

  • 这些方法大多需要完整重训练或微调整个模型,以恢复精度。

  • 这种方式在处理像 GPT-3 这样的大模型时,成本极高,一般研究者难以承受。

🧠 主流方法 2:训练后量化(PTQ)

  • 不需要重训,只在训练后进行量化,更省计算。

  • 但通常会有明显的精度下降,尤其是在 Transformer 类模型中。

  • 有些方法尝试非均匀量化或将异常值保留为高精度(如 FP32),但这种混合精度难以在 CPU/GPU 上高效运行。

  • 也有人提出对某些激活用 FP16 保持高精度,但仍难以完全解决问题。

❗ 当前挑战:

  1. GPT-3 类大模型的 PTQ 方案仍未解决;

  2. 对十亿级以上模型的 PTQ 研究很少;

  3. 缺乏支持精细量化的高效推理系统,难以实现在通用硬件上低延迟运行。

✅ ZeroQuant 的贡献:

  • 它把系统后端(如硬件运行效率)也纳入算法设计中。

  • 支持 BERT 和最大 200 亿参数的 GPT-3 类模型(如 GPT-NeoX)进行高效 PTQ。

  • 能在多种任务上实现高精度和高效率。

3. Background and Challenges

  • 这段内容主要介绍了在大模型(如 GPT-3 和 BERT)上使用 后训练量化(PTQ) 的挑战。

📌 背景知识

  • PTQ(后训练量化):不需要重新训练模型,只通过原模型加上标定数据即可完成量化,效率高于 QAT(训练中量化)

  • 一般做法是:用原始训练数据运行模型,然后根据平均值来调整量化的比例因子。

📌 存在的问题

1. GPT 和 BERT 上低比特量化难度大

  • 目前研究主要集中在 BERT base 的 INT8 权重量化上,但对于:

    • 更低精度(如 INT4)的探索很少;

    • 大规模模型(如 GPT-3)的 PTQ 也很缺乏。

2. INT8 精度就已明显降低性能

  • 用 INT8 激活(W16A8)量化,会明显损失准确率;

  • 进一步将权重也量化到 INT8(W8A8),虽然对某些任务影响不大,但会显著降低语言生成任务(如 Wikitext-2)的表现;

  • 用更低精度(如 W4/8A16)时,模型基本无法生成有意义的文本。

📌 原因分析

🔸 激活范围变化大

  • 每个 token 的激活范围(数值大小)差异大,比如有的 token 激活范围是 35,有的是 8;

  • 如果只用一个固定的量化范围,会使得小范围的 token 表达不清,精度下降。

🔸 注意力权重行之间差异大

  • GPT 的注意力输出矩阵中,不同行的数值大小差异可能高达 10 倍;

  • INT8 还勉强可以表示,但 INT4 只有 16 个数字,表示能力太弱,影响严重。

📌 小结

  • PTQ 虽然高效,但在大模型上易损失性能,特别是生成类任务

  • 主要挑战来自激活和权重范围波动大,尤其在低精度下更严重;

  • 更复杂的技术(如知识蒸馏)虽然能改善精度,但代价太高,需要轻量替代方案。

4. Methodology

4.1 Fine-grained Hardware-friendly Quantization Scheme

  • 背景问题:INT8量化虽然节省资源,但对BERT/GPT-3这类大模型会造成明显精度损失,原因是INT8无法很好表示每一行权重或每个token激活值的数值范围差异。

  • 改进方案

    • 权重的分组量化:把权重矩阵分成小组,每组单独量化,比整体量化精度更高。

    • 激活值的token级量化:为每个token单独计算量化范围(动态量化),避免激活值范围差异带来的误差。

  • 硬件优化:针对NVIDIA A100等GPU架构设计,结合Warp Matrix Multiply (WMMA) 实现高效计算。

  • 后端优化:通过融合量化操作(如和LayerNorm合并)减少内存访问开销,使token-wise量化既保留精度也提高运行速度。

4.2 Layer-by-layer Knowledge Distillation with Affordable Cost

  • 传统蒸馏问题:需要完整教师模型、原始训练数据,并消耗大量内存/计算资源。

  • 我们的方法(LKD)

    • 一次只量化一个transformer层,用该层的未量化版本做“老师”。

    • 只优化当前一层,降低内存消耗;

    • 不依赖原始标签,只用输入X即可;

    • 不需要完整训练,也不需要存教师模型。

4.3 Quantization-Optimized Transformer Kernels

  • 问题:虽然量化降低了数据量,但频繁的量化/反量化操作会引入额外延迟,抵消性能收益。

  • 优化方案

    • 使用CUTLASS库实现高效INT8矩阵乘法;

    • 通过内核融合(kernel fusion)将激活量化、偏置加法、GeLU、LayerNorm等合并成一个操作,减少数据搬运;

    • 将反量化也融合到GeMM中,进一步减少内存访问。

  • 效果:最终显著提升BERT和GPT-3模型的推理速度。

Figure 2: The illustration of normal (left) and our fused (right) INT8 GeMM.

5. Results

  • 背景与方法概述

    • ZeroQuant 是一种后训练量化(Post-Training Quantization, PTQ)方法,

    • 目标是无需重新训练或只需极少计算,就能对大模型(如 BERT、GPT-3)进行高效的权重量化,保持高精度、降低成本。

  • 实验设置

    • 测试模型包括:BERT (base 和 large)、GPT-3 模型(350M、1.3B)、GPT-J 6B、GPT-NeoX 20B。

    • ZeroQuant-LKD 是在 ZeroQuant 基础上增加了轻量级蒸馏(Layer-wise Knowledge Distillation, LKD)。

  • 量化表示说明

    • WxAy:表示权重使用 x 位量化,激活使用 y 位。

    • 例如 W4/8A16:MHSA 层权重为 INT8,FFC 层为 INT4;激活部分 Attention 计算使用 FP16,其余为 INT8。

🔍 BERT 实验主要结论

  • W8A8 下,普通 PTQ 精度下降 >10 点,但 ZeroQuant 仅下降 0.2,效果接近 QAT。

  • W4/8A16 极限量化下,PTQ 几乎完全失效,但 ZeroQuant 仍能达到 81.65,ZeroQuant-LKD 更可达 82.35。

  • 相比 QAT,ZeroQuant-LKD 仅需几十秒/GPU 就能达到相似精度,训练开销减少 90 倍以上。

🔍 GPT-3 模型实验

  • W8A8 量化对准确率型任务影响小(仅下降 1.1%),对生成任务影响大(下降 4.7);ZeroQuant/LKD 可显著减小损失。

  • W4/8A16 极限量化下,ZeroQuant 仍有效,LKD 提升显著,例如 PPL 从 88.6 降到 30.6,仅花费 3.1 GPU 小时。

推理速度提升

  • 相比 FP16:

    • BERT 模型使用 INT8 速度提升约 2.3–5.2 倍

    • GPT-3 模型生成 50 个 token,速度提升约 4 倍

🚀 大模型扩展性验证

  • GPT-J 6B 在 3 个生成任务中,ZeroQuant 保持精度,速度提升 3.67 倍。

  • GPT-NeoX 20B 中发现 Attention 输入特别敏感,因此保留该部分为 FP16,其余使用 INT8,最终达到 5.2 倍吞吐提升,GPU 使用从 2 块降为 1 块。

🧪 消融实验结果

  • 对 BERTlarge 进行分析:

    • Group-wise weight quantization 由随机预测提升至 66.52。

    • 加上 token-wise activation quantization,再提升 14.54 点。

    • 再加 LKD,提升 0.56 点。

🔐 无原始训练数据也能用 LKD

  • 即使没有原始训练集(如 PILE),使用 Wikipedia 或随机数据也能训练 LKD 并显著提升性能。

    • 用 Wikipedia 数据,效果几乎等同原始数据。

6. Conclusions

  • 随着大模型越来越大,如何高效部署它们变得很重要。

  • 虽然已有研究把量化方法用在BERT上,但目前还没人做过以下三点:

    1. 应用于GPT-3这类10亿参数以上的大模型;

    2. 使用超低精度(比如INT8)的量化;

    3. 提供完整的线上推理部署方案。

  • 本文提出了一种精细的压缩方法,对权重和激活值都进行了INT8量化,适用于最多20B参数的模型(如GPT-NeoX-20B)。

  • 同时,引入了一种便宜有效的逐层知识蒸馏技术,使模型在几乎不损失准确率的前提下,比FP16版本缩小了3倍。

  • 最后,还搭建了系统后端,实测BERT提速可达5.19倍,GPT-NeoX提效达5.2倍。

Appendix A Background

一、Transformer结构(A.1)

Figure A.1: The illustration of a Transformer-block.

  • Transformer通常由三部分组成:嵌入层(embedding)、多个编码器/解码器层(encoder/decoder layers)、最后的分类器(classifier)。

  • 本文关注的是Transformer中最耗计算和内存的部分——encoder/decoder层(也叫Transformer Block)。

  • 每个Transformer Block由两个子模块组成:

    • 多头自注意力(MHSA)

    • 前馈神经网络(FFC)

  • Transformer模型有三种类型:

    • 仅编码器(如BERT)

    • 仅解码器(如GPT)

    • 编码器+解码器(如T5)

  • 本文主要研究前两种,但方法也适用于第三种。

二、量化基础(A.2)

  • 量化是把高精度数值(如FP16/FP32)转成低精度(如INT4/INT8),以减少模型大小和提高运行速度

  • 本文使用对称的均匀标量量化器,即: $xquantize=round(clamp(x/S,2(bit1),2(bit1)1))$

    • bit 是使用的位数(比如 8 表示 INT8)

    • S 是缩放因子

    • clamp 限制范围,round 是四舍五入

  • 权重矩阵的量化,S 是固定的,一般取最大绝对值。

  • 激活值(activations)来说,由于它们在推理时变化,S 需要动态计算。但为了加速,实际中通常会用训练数据预先估计一个静态 S,这样能减少延迟,但也可能影响量化效果。

Appendix D Details about System Optimization

整体优化目标:

通过将模型权重和激活量化为 INT8 格式,使用专门的 Tensor Core(比 FP16/FP32 高效 2-4 倍)来加速矩阵乘法(GeMM),以提升推理效率。

主要优化方法:

  1. 定制高效计算计划(Schedule):

    • 使用 CUTLASS 工具根据不同输入(如 batch size、序列长度等)生成多种执行计划。

    • 在运行时选择填充最少、执行最快的计划。

  2. 高效的激活量化实现:

    • 把激活量化操作与其他操作(如 Bias 加法、GELU、LayerNorm)融合在一起,减少内存访问,提升效率。

    • 使用 FP16 转 INT8 提高内存带宽利用率。

  3. 高效的反量化(Dequantization):

    • INT8 GeMM 输出是 INT32,需要转回 FP16。

    • 将反量化过程与 GeMM 结果一起处理,避免额外开销。

    • 在计算 GeMM 时提前加载缩放因子,进一步减少延迟。

  4. 小模型优化 - CUDA Graph 加速:

    • 对于 BERT 等小模型,INT8 推理速度快,但 CPU-GPU 通信开销变大。

    • 用 CUDA Graph 记录推理过程后复用,减少每次推理的启动开销,特别适合小模型。

总结

核心是讲 通过量化、融合操作、并行优化和 CUDA Graph 重用技术,在 GPU 上高效地实现低精度推理,最大限度提升吞吐量并降低延迟。