大模型量化技术原理-ZeroQuant系列

近年来，随着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规模参数，从而导致模型变得越来越大，因此，我们需要一些大模型压缩技术来降低模型部署的成本，并提升模型的推理性能。
模型压缩主要分为如下几类：

剪枝（Pruning）
知识蒸馏（Knowledge Distillation）
量化Quantization）

本系列将针对一些常见大模型量化方案（GPTQ、LLM.int8()、SmoothQuant、AWQ等）进行讲述。

大模型量化概述
量化感知训练：
- 大模型量化感知训练技术原理：LLM-QAT
- 大模型量化感知微调技术原理：QLoRA
训练后量化：
- 大模型量化技术原理：GPTQ、LLM.int8()
- 大模型量化技术原理：SmoothQuant
- 大模型量化技术原理：AWQ、AutoAWQ
- 大模型量化技术原理：SpQR
- 大模型量化技术原理：ZeroQuant系列
大模型量化技术原理：总结

而本文主要针对大模型量化技术 ZeroQuant 系列进行讲述。

ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers
ZeroQuant-V2: Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation
ZeroQuant-FP: A Leap Forward in LLMs Post-Training W4A8 Quantization Using Floating-Point Formats
ZeroQuant-HERO: Hardware-Enhanced Robust Optimized Post-Training Quantization Framework for W8A8 Transformers

文章较长，建议先点赞收藏，后续再慢慢观看。另外，我撰写的大模型相关的博客及配套代码均整理放置在Github：llm-action，有需要的朋友自取。

ZeroQuant

背景

由于内存/计算要求过高，即使对于强大的云服务器来说，如何在实践中有效地服务越来越大的模型也变得异常具有挑战性。在这项工作中，我们提出了一种高效且经济实惠的训练后量化方法来压缩基于 Transformer 的大模型，称为 ZeroQuant。

ZeroQuant量化方案

将 INT8 PTQ 应用于 BERT/GPT-3 模型会导致准确性显著下降。关键的挑战是 INT8 的表示无法完全捕获权重矩阵中不同行和不同激活Token的不同数值范围。解决这个问题的一种方法是对权重矩阵（激活）使用group-wise（token-wise）量化。

用于权重矩阵的分组量化：

分组量化最先在Q-BERT中提出，其中权重矩阵 $\in R^{n \times m}$ 被划分为 g 个组，每个组单独量化。然而，在Q-BERT中，作者仅将其应用于量化感知训练。更重要的是，他们没有考虑硬件效率约束，也没有系统后端支持。因此，它们缺乏真正的降低延迟。

在我们的设计中，考虑了 GPU Ampere 架构（例如: A100）的硬件约束，其中计算单元基于 Warp Matrix Multiply and Accumulate (WMMA) 的分片（Tiling）大小，以实现最佳加速。与单矩阵量化相比，我们的分组量化由于其更细粒度的量化而具有更好的精度，同时极大的降低了延迟。

用于激活的按token量化：

现有 PTQ 中对激活进行量化常见做法是使用静态量化，其中最小/最大范围是在离线校准阶段计算的。

对于激活范围方差较小的小模型来说，这种方法可能就足够了。然而，GPT-3 和 BERT 等大 Transformer 模型的激活范围存在巨大差异。因此，静态量化方案（通常应用于所有tokens/样本）将导致准确度显著下降。

克服这个问题的一个办法就是采用更细粒度的token-wise量化并动态计算每个token的最小/最大范围，以减少激活引起的量化误差。评估结果表明针对激活的 token-wise 量化显著提高了类GPT-3 和 BERT 模型的准确性。

然而，使用现有的深度学习框架（例如：PyTorch 量化套件）直接应用 token-wise 量化会导致显著的量化和反量化成本，因为 token-wise 量化引入了额外的操作，导致 GPU 计算单元和主存之间产生昂贵的数据移动开销。

为了解决这个问题，ZeroQuant 构建了一个高度优化的推理后端，用于Transformer模型 token-wise 量化。例如，ZeroQuant 的推理后端采用所谓的核融合（kernel fusion）技术将量化算子与其先前的算子（如：层归一化）融合，以减轻 token-wise 量化的数据移动成本。类似地，在将最终 FP16 结果写回到下一个 FP16 算子（如：GeLU）的主存之前，使用权重和激活量化 scales 缩放 INT32 accumulation，可以减轻不同 GeMM 输出的反量化成本。

Token-wise 量化可以显著减少量化激活的表示误差，它不需要校准激活范围，对于ZeroQuant 的量化方案（INT8 权重和 INT8 激活）不存在与量化相关的成本（例如，激活范围校准）。

逐层知识蒸馏

知识蒸馏（KD）是缓解模型压缩后精度下降的最有力方法之一。然而，KD 存在一些局限性，特别是对于大语言模型上的隐藏状态 KD：

1）KD 需要在训练过程中将教师和学生模型放在一起，这大大增加了内存和计算成本；
1. KD 通常需要对学生模型进行充分训练。因此，需要在内存中存储权重参数的多个副本（梯度、一阶/二阶动量）来更新模型；
1. KD 通常需要原始训练数据，有时由于隐私/机密问题而无法访问。

为了解决这些限制，我们提出了逐层蒸馏（LKD）算法缓解精度损失，原网络做老师，量化后的网络做学生。

量化优化的 Transformer kernel

优化推理延迟和模型大小对于在实践中服务大Transformer 模型至关重要。在推理过程中，batch size往往比较小，因此模型的推理延迟主要取决于从主存加载推理所需数据的时间。通过将权重和激活量化到较低的精度，我们减少了加载这些数据所需的数据量，从而可以更有效地使用内存带宽和更高的加载吞吐量。

然而，简单地将权重/激活转换为 INT8 并不能保证延迟的改善，因为存在与量化/反量化操作相关的额外数据移动开销，如图 2（红色框）所示。这样的开销很昂贵，并且在某些情况下超过了使用低精度的性能优势。

为了从 token-wise 量化中获得准确性提高，同时获得更低的延迟，作者针对性提做了优化，可以最大限度地提高内存带宽利用率，以加快 ZeroQuant 的推理延迟。

CUTLASS INT8 GeMM：为了支持 INT8 计算，我们使用针对不同批量大小进行调整的 CUTLASS INT8 GeMM 实现。与标准 GPU 后端库（例如：cuDNN）不同，使用 CUTLASS 允许我们在 GeMM 之前和之后更灵活地融合量化操作，以减少kernel启动和数据移动开销。
融合 Token-wise 激活量化：Token-wise量化/反量化引入了许多额外的操作，从而导致额外的数据移动成本。为了消除这些成本，我们使用核融合将激活的量化操作与其之前的 element-wise 归约操作（例如: bias-add、GeLU 和 LayerNorm）融合到单个算子中，如图 2（绿色框）所示。对于反量化操作（例如，对 GeMM 算子的整数输出进行反量化），我们同样将其与自定义 GeMM 融合，以避免对主存进行额外的读/写访问，如图 2（蓝色框）所示。

ZeroQuant 小结

ZeroQuant 是一个端到端量化和推理管道，具有三个主要组件：

1）针对权重和激活的细粒度硬件友好量化方案；
2）一种新的、经济实惠的逐层知识蒸馏算法（LKD），即无需访问原始训练数据；
1. 提供了一个高度优化的量化系统后端，以消除量化/反量化开销。

实验结果显示：

ZeroQuant 可以将 BERT 和类GPT-3 等模型的权重和激活精度降低到 INT8，对模型准确率的影响最小，同时，与 FP16 推理相比，这些模型的推理速度提高了 5.19 倍/4.16 倍；
ZeroQuant 加上 LKD 可将全连接模块中的权重量化为 INT4，以及注意力模块中的INT8权重和INT8激活，与FP16模型相比，内存占用减少了3倍；
ZeroQuant可以直接应用于GPT-J和GPT-NeoX等，其中我们的INT8模型达到了与FP16模型相似的精度，但效率提高了5.2倍。

ZeroQuant-V2

背景

训练后量化 (PTQ) 已成为一种有前途的技术，可减少大语言模型中的内存消耗和计算成本 (LLMs)。然而，目前缺乏对各种量化方案、模型族和量化位精度的系统检查。

本文将舍入到最近（RTN）、GPTQ、 ZeroQuant 及其变体方法应用于参数范围从 125M 到 176B 的两个不同的模型系列进行了系统性分析。

文本的贡献如下：

1）敏感性分析表明，相比于权重量化，激活量化精度更容易受到的影响，较小的模型在激活量化方面通常优于较大的模型；
(2) 对现有 PTQ 方法进行评估和比较，使模型的尺寸减小，同时最大限度地减少对精度的影响，揭示当前的方法使用 INT4 权重或 INT4 权重和INT8激活进行量化都无法达到原始模型质量；
（3）基于以上见解，本文提出了一种称为低秩补偿（LoRC）的优化方法，该方法采用低秩矩阵以最小的模型参数大小的增加来提升模型质量的恢复。

不同的模型系列在量化方面的表现是否相似？

针对不同的模型系列进行敏感性分析发现：

INT8纯权重量化可以作为标准一种降低LLM内存成本的方法，其准确性的下降可以忽略不计。
小模型的 INT4 仅权重量化会导致精度大幅下降(Class-3)，但这种影响会随着模型大小的增加而减弱 (Class-2)。
与2相反，INT8 激活导致小型模型的准确度下降最小(Class-1)，但较大的模型表现出更大的下降(Class-3)。
对于 INT8 激活，BLOOM 在模型大小达到 176B 时没有出现差异，而 OPT 从 6.7B 模型大小起表现不佳。

现有的量化方法是否以最佳方式利用了最小化LLM大小的潜力？

人们已经提出了许多轻量级优化的方法，这些方法在量化期间更新模型权重。这些方法（如：zeroquant、gptq等），与量化感知训练不同，只需要一小部分训练数据和有限的训练时间。并且 GPTQ 和 ZeroQuant 已被证明在 GPU 资源、时间成本等方面是有效和高效的。本文重点关注 GPTQ 和 ZeroQuant 的变体以及最直接的基线舍入到近邻（RTN）。

RTN 直接对训练数据应用 PTQ，即 $\text{INT}\big({(x-Z)}/{S}\big)-Z$ 。具体来说，对于对称量化，我们设置 $S = ma x (ab s (x))$ 和 $Z = 0$ ；对于非对称量化，我们设置 $S = ma x (x) - min (x)$ 和 $Z = min (x)$ 。
GPTQ 扩展了 OBQ，它尝试优化以下非线性最小二乘问题，$ \min_{\hat W} |Wx - \hat Wx|_2^2 $ ，其中 W 是权重，x 是激活值，$ \hat W $ 是量化权重。 GPTQ 采用二阶方法来求封闭式解。此外，每个权重矩阵的量化是按列/行方式执行的，并且先前列的量化误差将传递到尚未量化的那些列。
ZQ-Global 是ZeroQuant中提出的原始方法，作者将每一层视为一个小型神经网络（也称为子网络），并使用 FP16 子网络作为教师模型，通过数百次迭代来提取量化的模型，即$ \min_{\hat \theta} |f_{\theta}(x) - f_{\hat\theta}(x)|2^2,$ ，其中 $\theta$ 是一组权重， $\hat \theta$ 是量化版本， $f{\theta}$ 是具有参数 $\theta$ 的子网络，x 是输入。因此，它可以显著降低GPU资源需求和时间成本。
ZQ-Local是ZQ-Global的扩展模式，用于进一步降低GPU需求和降低训练成本。我们不使用每个transformer层作为子网络，而是将每个线性层视为子网络。该方法可以看作是一种迭代一阶优化方法（例如：SGD）用于求解 $\min_{\hat W} \|Wx - \hat Wx\|_2^2$ 。

通过实验发现：

GPTQ 通常在仅权重量化方面表现更好，而 ZeroQuant（包括 ZQ-Global 和 ZQ-Local）在权重和激活量化方面产生优异的结果。
除了 OPT-30B 上使用仅权重量化的 GPTQ 之外，测试的其他的量化方法（仅 INT4 权重或 W4A8）均无法实现 Class-1 量化误差。

低秩补偿（LoRC）的优化方法

LoRC 的灵感来自于对量化误差矩阵 $\hat{W}$ 进行低秩矩阵分解，其中 $W$ 表示原始权重， $\hat{W}$ 是量化权重。

LoRC 通过使用两个低秩矩阵 $\hat{U}$ 和 $\hat{V}$ 来近似误差 $\hat{E}=\hat{U}\hat{V}$ 。这样可以通过 $\hat{W}_{\text{lorc}} = \hat{W} + \hat{E}$ 更准确地近似原始权重矩阵 W，从而减少量化误差： $\|W-\hat{W}\|\geq \|W-\hat{W}_{\text{lorc}}\|$ 。

LoRC 包含两个步骤：

步骤一：在误差矩阵 $\Sigma V$ 上实现奇异值分解（SVD），其中 $\in\mathbb{R}^{d_\text{in}\times d_\text{in}}$ 和 $\in \mathbb{R}^{d_\text{out} \times d_\text{out}}$ 是酉矩阵， $\Sigma \in\mathbb{R}^{d_\text{in}\times d_\text{out}}$ 是一个对角矩阵，其对角元素按降序方式排列。

步骤二：公式表示矩阵 $\hat{E} = \hat{U} \hat{V}$ ，其中， $\hat{U}= U_m(\Sigma_m)^{\frac{1}{2}}$ ， $\hat{V}= (\Sigma_m)^{\frac{1}{2}} V_m$

$U_m = U_{:, 1:m} \in\mathbb{R}^{d_\text{in}\times m}$
$V_m = V_{1:m, :} \in\mathbb{R}^{ m \times d_\text{out}}$
$\Sigma_m = \Sigma_{1:m, 1:m} \in\mathbb{R}^{ m \times m}$

LoRC 的目标是使用低秩矩阵实现误差矩阵 $E$ 的近似，同时对模型参数大小的增加影响最小。

例如，对于标准transformer模型，其中每一层都由多头注意力（MHA）模块和 MLP 模块组成。令 h 表示隐藏维度，l 表示层数。参数总数为 $12lh^2$ ，每层包含 $4h^2$ 用于 MHA（用于K、Q、V和投影矩阵）和 $8h^2$ 用于 MLP（两个大小为 h × 4h 和 4h × h 的矩阵）。在每层6个矩阵中添加低秩LoRC后，l层参数总数将达到 18hml（6 * 3hm * l）。因此，向现有模型添加参数的比例为3m/2h。需要注意的是，低秩维度 m 可以小到 4 或 8，而标准隐藏维度 h ≥ 768，使得数字 3m/2h ≤ 0.016。

评估结果

LoRC 可以被视为现有量化方法（例如：RTN、GPTQ 和 ZeroQuant-Local/Global）的补充功能，并且可以与 FGQ 无缝集成。我们进行了实验来评估 LoRC 在 OPT 和 BLOOM 上的性能，通过将激活设置为 FP16 来应用 4 比特、3 比特和 2 比特权重量化。

本文利用 GPTQ量化策略，为了全面了解 LoRC，包含了应用和不应用 FGQ 的结果。结果如表 7 所示，分为两部分：粗粒度权重量化（每行）和细粒度量化（块大小 256）。值得注意的是，我们观察到两个低秩矩阵 $\hat{U}$ 和 $\hat{V}$ 可以量化为 8 比特，而不会出现任何性能差异（表 8）。

得出几个关键的观察结果。

首先，LoRC 始终如一地提高所有比特大小和块大小的性能。
其次，LoRC带来的增强随着比特大小的减小而变得更加显著，尤其是对于W2A16来说更为明显，在大多数场景下与W4A16和W3A16相比，其影响明显更大。
最后，细粒度量化与 LoRC 的结合产生了最令人印象深刻的结果，强调了 LoRC 与 FGQ 集成时的效果。

总体而言，结果显示出使用 LoRC 增强了权重量化性能及其与 FGQ 的兼容性。

ZeroQuant-V2 小结

本文使用不同的 PTQ 方法（例如：RTN、GPTQ、ZeroQuant）和不同的量化覆盖范围（仅权重、权重和激活）对大语言模型的训练后量化 (PTQ) 进行了全面的研究。发现PTQ方法对于提高量化模型质量至关重要，而细粒度量化（FGQ）可以带来可接受的精度和模型大小的权衡。

最后，作者引入了一种称为低秩补偿（LoRC）的优化技术，它与 PTQ 和 FGQ 协同工作，以最小的模型参数大小的增加来改善整个模型质量的恢复。

本文存在的局限性。尽管量化了 10,000 多个实验，但我们的研究仍受到计算资源的限制。这种限制使我们在模型大小多样化和任务多样化之间做出选择。我们战略性地将数据集限制为 WikiText、PTB 和 C4，以专注于更广泛的量化方法的对比。因此，对于本文中检查的两个模型系列和三个数据集，总体来说都更加稳健。然而，将这些发现推广到与本研究所涵盖的任务不同的任务上应谨慎。

ZeroQuant-FP

背景

在大语言模型领域，在计算效率和保持模型质量之间取得平衡是一项艰巨的挑战。为了克服多数量化固有的局限，特别是在处理异常值时，并受到 NVIDIA H100 硬件推出的推动，本研究深入探讨了浮点 (FP) 量化的可行性，特别关注 FP8 和 FP4 作为潜在解决方案。

通过调查显示：

对于 LLMs，FP8 激活始终优于其INT8激活，并且在参数超过 10 亿的模型中，性能优势变得更加明显。
对于权重量化，FP4 表现出与 INT4 相当（即使不是更优）的性能，从而简化了在 H100 等支持 FP 的硬件上的部署。

为了减轻权重和激活之间的差异引起的精确对齐的开销，本文提出了两个权重量化的缩放建议，与标准 W4A8 模型相比，它们对模型性能的影响可以忽略不计。同时还通过集成低秩补偿（LoRC）策略来增强我们的量化方法，并且在较小的模型中也有提升。

ZeroQuant-FP 优化方案

本文选择使用的方法与 GPTQ 保持一致。根据 ZeroQuant-V2 ，本文应用了细粒度权重量化（FGQ），并对激活进行 token-wise 量化。此外，我们还将研究ZeroQuant-V2中提出的附加特征 LoRC（低秩补偿），其目的是通过采用低秩矩阵分解来减少权重的量化误差。

LoRC涉及两个主要步骤：

首先，它对误差矩阵进行奇异值分解（SVD），误差矩阵是原始权重与量化权重之间的差值。因此，误差矩阵被分解为两个酉矩阵和一个对角矩阵。
其次，该方法使用从第一步中的矩阵导出的两个低秩矩阵来制定新的误差近似。然后将该近似值添加到量化权重中，以产生对原始权重的更准确的估计，从而减少量化误差。

将 FP4 转换为 FP8 方案：

本文探索了 FP4 权重和 FP8 激活量化的潜力。但由于对权重 (W) 和激活 (A) 使用不同的精度级别，出现了一个独特的挑战。 W4A8 在 H100 NVIDIA 硬件中的实际软件实现是需要转换 W 的 FP4 以匹配 A 中使用的 FP8 精度。直接反量化然后再次量化的方法可能会对推理效率产生不利影响，因此不是一个可行的解决方案。

为了解决这个问题，我们提出了位移方法。这意味着，我们不让等式（Q(x) = INT(x − Z)/S − Z）中定义的 S 为任何实值比例因子，而是将 S 限制为 2 的幂，即 $S = 2^n$ ，n ∈ N（当n为负数时，S仍然可以表示分数；当n不为负数时，S仍然可以表示整数。）。

本文实现了两种方法：

(M1) 映射到由 2 的幂表示的最接近的值，即让新的scale为 $\hat{S} = 2^{\lceil \log_2(S)\rceil}$
(M2) 首先收集scales形成向量 $\mathbf{S} = [S_1, S_2, \ldots, S_n]$ 。然后取组（group）中的最大值（通常，该集合由矩阵的（多）行组成），记为 $S_{\max}$ ，将这些元素 $S_{\max}/S_i$ 调整为2的幂表示，然后定义 $\hat{S}_i = S_{\max}/ 2^{\lceil \log_2(S_{\max}/S_i)\rceil}$ 。与 (M1) 相比，这提供了更好的近似值。

注意：这种使用 2 的幂的限制，无论是使用 (M1) 还是 (M2)，都可以简化计算，特别是在基于二进制逻辑操作的数字系统中。这是我们优化计算效率和保持模型性能的方法的关键要素。

实验结果

本文基于GPTQ（不带或带LoRC），对激活量化（FP8、INT8）和权重量化（FP8、FP4）到进行全面比较。表2展示了应用权重和激活的不同整型 (INT) 和浮点型 (FP) 量化方法在 LLaMA（上）和 OPT（下）模型上的评估结果。性能以困惑度来衡量（分数越低越好），使用了三个数据集：WikiText-2 (WIKI)、PTB 和 C4。对于每个模型的结果，首先显示整个数据集的平均性能，然后是每个数据集的详细分。

FP8 激活优于 INT8 激活：

对于 LLaMA 和 OPT 模型系列，FP8 激活通常优于 INT8 激活。这一观察结果证实了 FP8 捕获更细致信息的卓越能力。并且对于参数大于 67 亿的较大模型，例如: LLaMA-7b/13b 和 OPT-6.7B/13B，FP8 相对于 INT8 的优势变得更加明显。

FP8 权重可与 INT8 权重相媲美，而 FP4 权重则可能优于 INT4 权重：

当保持 FP8 激活时，各种模型和数据集上 INT8 权重和 FP8 权重量化之间的性能相当。这可能是由于我们在权重量化上使用了 FGQ。有趣的是，当降低权重量化比特时，FP4 表现出优于 INT4 的某些优势，在 LLaMA-7b（15.14 至 16.09）和 LLaMA-13b 模型（11.08 至 11.31）中尤其明显。具体来说，在 LLaMA-7b 的 W4A8 配置下，我们看到 FP4 比 INT4 提高了 0.95，这是一个显著的增益。 FP4 优于 INT4 的性能对于 H100 等已支持 FP8 的硬件设计尤其有利。因此，一个简单修改来适应 FP4 将比实现支持 INT4 权重和 FP8 激活的系统更容易。

LoRC 改进了 W4A8：

低阶补偿 (LoRC) 方法增强了 W4A8 量化方案，减少了量化误差。这种改进在较小的模型中尤其明显，突显了 LoRC 在优化这些计算过程的性能方面的有效性，同时对模型大小的影响很小。

将 FP4 转换为 FP8 的方案对比：

为了最大限度地提高 NVIDIA H100 硬件上的实际延迟加速，我们建议将权重量化的比例因子 S 表示为 2 的幂。为了实现这一目标，我们使用 FP4 进行权重量化和FP8 用于激活量化进行了一系列实验。表 3 列出了使用和不使用 LoRC 的实验结果。

数据显示，限制缩放因子偶尔会导致 LLaMA-7b 和 LLaMA-13b 等模型出现提升，但无论我们使用方法 M1 还是 M2，观察到模型性能通常会出现轻微的退化，W4A8 浮点模型中的质量都会受到影响。M2 通常优于 M1。当我们实施 LoRC 时，可以缓解质量下降的情况，特别是在 OPT-1.3b、LLaMA-7b 和 LLaMA-13b 模型中。因此，提倡使用 LoRC。

ZeroQuant-FP 小结

通过 FP8 激活量化和权重量化（W8A8）导致较小的模型退化，特别是在较大的模型中，FP8 激活和权重量化导致的模型退化可以忽略不计，其性能与原始 FP16 模型相当。

在W4A8浮点模型中，即使对缩放因子施加了约束，也能保持模型质量。

为了在 W4A8 模型中实现真正的效率，从 FP4 到 FP8 的权重转换至关重要。为了减轻这种转换开销，本文提出了权重量化的两种可能的缩放约束建议：

1）将所有缩放因子限制为 2 的幂；
2）将缩放因子放在一个计算组中（例如，权重矩阵的几行可以通过简单的位移位来转换）。

通过分析表明，与传统的 W4A8 配置相比，这两个限制提高了计算效率，同时对模型性能的影响可以忽略不计。

ZeroQuant-HERO

背景

由于机器学习算法和硬件之间的跨学科差距，该领域仍然很大程度上缺少硬件感知的 PTQ 方法，特别是对于基于 Transformer 的模型。

例如，ZeroQuant 为 BERT 和 GPT 模型提出了对激活进行每token动态量化和对权重进行每列量化，以实现良好的准确性。然而，它没有考虑如下两点：

内存限制（memory bounded）算子，例如： LayerNorm 和 attention，并将这些部分留在 FP16/BF16 中。
当没有算子融合机会时，调用额外kernel进行每个token量化的成本，例如：注意输出线性层的 INT8 GeMM 算子。

为了解决这些限制，我们引入了 ZeroQuant-HERO，这是一个完全硬件感知且实用的训练后 W8A8 量化框架。具体贡献总结如下：

ZeroQuant-HERO 在设计时考虑了内存带宽限制和计算密集型运算。因此，该框架可以（有可能）实现最佳硬件性能。
为了进一步提高 ZeroQuant-HERO 的可用性，可以执行 ZeroQuant-HERO 的不同量化级别，即 INT8 运算与 FP16/BF16 对应运算的比率，以实现所需的精度和延迟权衡。

ZeroQuant-HERO 量化方案

ZeroQuant-HERO中，除非特别的注释说明，否则我们均使用 INT8 对称量化。然而，ZeroQuant-HERO方法也适用于其他 8 位精度格式，例如 FP8。

下面我们使用列主序权重矩阵格式来执行 GeMM：

$Y = X W$

其中， $X\in\R^{n\times d}$ 是激活值， $W\in\R^{d\times m}$ 是权重矩阵。

对于权重量化，我们执行按列量化（Zeroquant），即权重的每一列都有自己的缩放因子，

$W = W_{int8}S_w$

其中， $W$ 是重建的权重矩阵， $W_{int8}$ 是 INT8 对应矩阵， $S_w\in\R^{1\times m}$ 是缩放向量。

对于激活量化，我们应用三种不同的量化方案。

Token-wise 量化（TWQ）:

作者用于token量化的第一个量化方案是 TWQ（Zeroquant）。

$X = S_xX_{int8}$

其中， $X$ 是重建的激活， $X_{int8}$ 是 INT8 对应项， $S_x\in\R^{n\times 1}$ 是缩放向量。

这种方法需要动态计算缩放向量 $S_x$ ，这更适合带宽限制算子融合，例如：层归一化 (LN)。

事实上，量化是以零内存开销成本完成的，使用额外的寄存器级操作来计算最小值和最大值，以降低将在以下 GeMM 算子中使用的 LN 的输出精度。另一方面，如果与 GeMM 等计算密集型算子融合，这种扩展方法会损害 Tensor-core 效率，因为会增加寄存器压力并为每个矩阵乘法累加 (MMA) 操作添加更多计算。

Feature-wise 量化 (FWQ):

作者用于token量化的第二种量化方案是 FWQ （Outlier suppression、Smoothquant）。

$X = X_{int8}S_x$

其中， $S_x\in\R^{1\times d}$ 是缩放向量。 $S_x$ 需要在预处理阶段进行校准，即在网络中输入多批数据以获得缩放因子。由于它是预先确定的，因此可以简单地与大多数其他算子融合。与 TWQ 量化方案相比，TWQ 量化方案涉及读取token宽度（token-width）的长度来量化数据并且只能与某些算子融合，而 FWQ 缩放可以与内存限制或计算限制算子融合。

静态量化（SQ）：

$X = X_{int8}S_x=S_xX_{int8}$

其中， $S_x\in\R$ 只是单个实值。与FWQ类似，在预处理阶段也需要进行校准。

ZeroQuant-HERO 核心组件

ZeroQuant-HERO 的三个主要组件：

嵌入层量化
注意力模块量化
MLP 模块量化

嵌入层量化：

Transformer 模型的第一个主要算子是查找表，也称为嵌入。通常，嵌入有三种类型，即标记嵌入（ $X_t$ ）、位置嵌入（ $X_p$ ）和句子类型嵌入（ $X_s$ ）。当批量大小足够大时，与 $X_t$ 相比， $X_p$ 和 $X_s$ 相对较小。应用 layer norm：

$X_{emb} = LN(X_t,~X_p,~X_s)$

其中， $X_{emb}$ 是 layer norm 的输出，LN 算子是与输入 $X_t$ 和输出 $X_{emb}$ 相关的内存带宽限制算子。

为了减少内存带宽开销，我们对 $X_t$ 和 $X_{emb}$ 都执行 TWQ，即

$S_{emb}X_{emb,int8} = LN^{quant}(S_{x_t}X_{t,int8},~X_p,~X_s)$

其中， $LN^{Quant}$ 是量化感知算子。通过利用上述embedding格式，我们将后续操作的通信数据量大致减少了 2 倍。

注意力模块量化：

$X_{q/k/v} = X_{in}W_{q/k/v}$
$X_q X_k^T / \sqrt{d}$
$P = S o f t ma x (A)$
$X_{attn} = PX_v$
$X_o = X_{attn}W_O$
$X_{out} = LN(X_{in},~X_o)$

先对所有激活使用的量化方案进行分类：

TWQ 应用于 $X_{in}$ 和 $X_{out}$ ，以最小的性能开销保持Transformer层输入和输出的高精度，因为缩放逻辑可以融合在预先发生的 LN 算子中。

SQ（静态量化) 应用于 $X_q$ , $X_k$ , $X_v$ , and $P$ ，以提高涉及这些张量的GeMM运算的效率。我们将此逻辑添加到 flash-attention kernel 实现中，并且可以配置每个 GeMM 的 dtype 以保持模型准确性。

FWQ 应用于 $X_{attn}$ 和 $X_o$ ，以降低用于GeMM操作的缩放激活的复杂性，同时保持精度。与 SQ 相比，我们对每个输出元素使用一个scale，因此，此操作的性能成本类似于在线性层添加偏差。

对于 $A$ ，不应用量化。这是由于注意力分数值对其精度的敏感性，量化可能会损害下游任务的模型准确性。

在应用权重量化之前，我们有

$X_{q/k/v,int8}S_{q/k/v} = S_{in}X_{in,int8}W_{q/k/v}$

$S_qS_kX_{q,int8} X_{k,int8}^T / \sqrt{d}$

S_{p}P_{int8} = Softmax^{Quant}(A)

X_{attn, int8}S_{attn} = S_pS_vP_{int8}X_{v,~int8}

X_{o, int8}S_o = X_{attn, int8}S_{attn}W_O
$$

$S_{out}X_{out, int8} = LN^{Quant}(S_{in}X_{in, int8},~X_{o, int8}S_{o})$

这里， $⋅quant \cdot^{quant}$ 是量化感知运算，并且输出 $Softmax^{quant}$ 是非对称 INT8 数，因为 softmax 的输出中没有负值。现在，让我们更深入地研究权重量化和 GeMM 算子。

首先，我们可以应用与Zeroquant相同的核融合来将反量化算子与INT8 GeMM融合。为了进一步减少量化开销，我们可以将 FWQ 和 SQ 的缩放因子融合到 INT8 GeMM 中，因为缩放因子是预先确定的。更重要的是，可以简化 FWQ/SQ 量化作为一个简单的舍入到整数的运算，无需任何除法/乘法，因为缩放因子可以合并到权重矩阵中。以 $X_{q, int8}$ 为例，我们可以定义
$\tilde W_q = W_q / S_q$
$\tilde W_{q,int8}S_q = Quant(\tilde W_q)$

这里 $Q u an t$ 是量化转换运算。之后，post-GeMM 量化算子被简化为

$Round\left( GeMM^{quant}(X_{in, int8},~W_{q, int8},~S_{in},~S_{q})\right)$

其中， $Round(\cdot)$ 是整数运算的舍入。同样，我们也不需要将除以 $\sqrt{d}$ 后的 A 反量化计算。我们可以用 $\tilde d = S_qS_k/\sqrt{d}$ 和 $GeMM^{quant}(X_{q,int8}, X_{k,int8}^T, \tilde d)$ 进行简化。

$S_{attn}$ 和 $S_{o}$ 的缩放因子可以通过以下方式合并为 $W_o$ ：

$\tilde W_o = S_{attn}W_o / S_o$

这样可以显著简化整个kernel实现。之后，LN 算子将两个 INT8 数作为输入，并为后续 MLP 模块输出最终的 INT8 激活值。

MLP 模块量化：

$X_1 = X_{in}W_1$

$A = GELU(X_1)$

$X_2 = AW_2$

$X_{out} = LN(X_{in}, X_2)$

与之前类似，我们首先对所有激活采用的量化方案进行分类。对于 $X_{in}$ 和 $X_{out}$ ，应用 TWQ。对于 $A$ 和 $X_2$ ，应用 FWQ。对于 $X_1$ ，不使用量化。在应用权重量化之前，我们有

$X_1 = S_{in}X_{in,~int8}W_1$

$A_{int8}S_a = GELU^{quant}(X_1)$

$X_{2, int8}S_{x_2} = A_{int8}S_aW_2$

$X_{out} = LN^{quant}(S_{in}X_{in},~X_2S_{x_2})$

与之前类似，可以将缩放因子 $S_a$ 和 $S_{x_2}$ 合并到 W 中以简化计算：

$\tilde W_2 = S_aW_2 / S_{x_2}.$

混合精度推理

针对不同的模型和任务对量化的容忍度不同，并且对准确性和系统效率的权衡也有不同的期望。为了满足各种模型/任务的要求，混合精度推理是量化的解决方案之一。

由于 ZeroQuant-HERO 的模块化设计，我们可以为最终模型设置各种量化级别。

实验结果

ZeroQuant-HERO 不同量化级别的结果如下所示。随着量化级别的增加，整体精度会下降。然而，除了 CoLA（这是一个超级敏感的任务）之外，对于所有其余任务，即使是 ZeroQuant-HERO-M3，与 FP16 相比，仅有较少的精度下降。

总结

本文讲述了微软提出ZeroQuant相关量化优化技术。

ZeroQuant 针对权重和激活提出了一个细粒度硬件友好量化方案，对权重矩阵使用group-wise量化，对激活使用token-wise 量化。同时逐层知识蒸馏算法（LKD）算法缓解精度损失，原网络做老师，量化后的网络做学生。
ZeroQuant-V2 引入了一种称为低秩补偿（LoRC）的优化技术，它与 PTQ 和 FGQ 协同工作，以最小的模型参数大小的增加来改善整个模型质量的恢复。
ZeroQuant-FP 探讨了浮点 (FP) 量化的可行性，同时，为了最大限度地提高 NVIDIA H100 硬件上的实际延迟加速，建议将权重量化的比例因子 S 表示为 2 的幂，并将缩放因子放在一个计算组中，同时可以通过 LoRC 缓解模型质量下降。
ZeroQuant-HERO 是一种新的硬件增强型训练后 W8A8 量化框架，为了解决算法与硬件协调的挑战。设计时考虑了内存带宽限制和计算密集型运算可以（有可能）实现最佳硬件性能。同时为了进一步提高 ZeroQuant-HERO 的可用性，可以执行 ZeroQuant-HERO 的不同量化级别，即 INT8 运算与 FP16/BF16 对应运算的比率，以实现所需的精度和延迟权衡。

码字不易，如果觉得我的文章能够能够给您带来帮助，期待您的点赞收藏加关注~~

参考文档：