Infra Quant III

🦏 AWQ: ACTIVATION-AWARE WEIGHT QUANTIZATION

对于端侧 LLM 推理而言,低比特权重量化能够显著降低模型的内存占用,但实现高质量的低比特量化并非易事

  • 量化感知训练由于训练成本较高而缺乏效率,而后训练量化在低比特设置下通常会产生较大的精度下降
  • 与本文最为相关的方法是 GPTQ,其利用二阶信息进行量化误差补偿。然而,其重构过程可能使模型过拟合于校准集,并在分布外领域中改变模型原本学习到的特征。由于 LLM 通常被作为具有通用能力的模型使用,这一问题尤为值得关注

本文提出 Activation-aware Weight Quantization(AWQ,激活感知权重量化),这是一种面向 LLM 低比特、仅权重量化的硬件友好型方法

  • 该方法基于一个观察:不同权重对于 LLM 性能的重要程度并不相同。模型中存在少量约占 0.1%–1% 的显著权重(salient weights);若不对这些显著权重进行量化,则能够显著降低量化损失
  • 为确定这些显著权重对应的通道,本文的核心思路是:尽管所研究的是仅权重量化,但在判断权重的重要性时,应当参考激活分布而非权重分布。与较大激活幅值相对应的权重通道更加显著,因为这些通道负责处理更为重要的输入特征

直接将部分重要权重保留为高精度会引入硬件效率较低的混合精度实现。为避免这一问题,本文从权重量化误差出发进行分析,并指出:通过对显著通道进行适当放大,可以降低这些通道的相对量化误差

  • 作者设计一种逐通道缩放(per-channel scaling)方法,通过自动搜索最优缩放因子,在对全部权重进行量化的条件下尽可能减小量化误差。AWQ 不依赖反向传播或重构,因此能够在不过拟合校准集的情况下,更好地保持 LLM 在不同领域和不同模态上的泛化能力
  • 作者进一步设计 TinyChat,一个面向端侧 LLM 推理的高效推理框架,用于将 4 比特量化所带来的理论内存节省转化为实际的推理加速
  • TinyChat 通过运行时即时反量化显著加速线性层计算,同时利用高效的 4 比特权重打包以及 kernel fusion,降低中间 DRAM 访问和 kernel 启动等推理开销。由此,即使计算机系统通常采用字节对齐的存储方式,也能够更充分地获得 4 比特权重量化带来的性能收益

🦖 RELATED WORK

模型量化通过降低深度学习模型中的数值精度来减少模型规模并加速推理。现有量化技术通常可以划分为两类:量化感知训练和后训练量化

  • QAT 依赖反向传播来更新量化后的权重;PTQ 通常不需要训练
  • 由于 QAT 方法难以直接扩展到 LLM 这类大规模模型,因此在 LLM 量化中通常采用 PTQ 方法

现有 LLM 量化研究主要关注两类设置。第一类是 W8A8 量化,即将激活和权重均量化为 INT8。第二类是低比特仅权重量化,例如 W4A16,其中仅将权重量化为低比特整数,而激活保持较高精度

  • 本文主要关注第二类设置。这是因为,仅权重量化不仅能够降低硬件部署门槛,即减少模型所需的内存容量,还能够加速 token 生成过程,因为这种方法能够缓解推理过程中由内存带宽限制所导致的瓶颈
  • 除基础的最近邻舍入方法之外,与本文最相关的方法是 GPTQ。然而,GPTQ 的重构过程可能导致模型对校准集产生过拟合,从而无法充分保持 LLM 在其他模态和领域中的通用能力。此外,对于某些模型,例如 LLaMA-7B 和 OPT-66B,GPTQ 还需要使用额外的重排序技巧才能获得较好的量化效果
  • 除面向通用硬件设计的量化方法之外,SpAtten 提出一种渐进式方法,用于逐步增加 Softmax 计算所采用的比特数

低比特量化 LLM 已成为降低推理成本的一种常用方案,因此已有多种系统针对低比特模型提供实际的推理加速支持

  • GPTQ 为 OPT 模型提供 INT3 kernel;GPTQ-for-LLaMA 则借助 Triton 进一步扩展对经过重排序的 INT4 量化模型的 kernel 支持
  • FlexGen、llama.cpp 和 exllama 采用分组式 INT4 量化,以降低 I/O 成本及模型卸载相关开销
  • FasterTransformer 实现用于仅权重量化的 FP16 × INT4 GEMM,但其采用的是逐张量量化,并不支持分组量化
  • LUT-GEMM 则利用查找表,在 GPU CUDA Core 上执行按位计算。与本文同期开展的工作 MLC-LLM 借助 TVM 后端,在多种边缘 CPU 和 GPU 平台上取得良好的性能

🦧 Improving LLM Quantization by Preserving 1% Salient Weights

量化将浮点数映射为低比特整数,是降低 LLM 模型规模和推理成本的一种有效方法。作者首先提出一种仅权重量化方法,通过保护更为重要的权重,在无需训练或回归的情况下提高量化精度;随后进一步提出一种数据驱动方法,用于搜索能够降低量化误差的最优缩放因子

作者观察到,LLM 中不同权重的重要性并不相同:其中存在一小部分显著权重(salient weights),它们相比其他权重对模型性能更加重要。在不进行任何训练或回归的情况下,如果跳过对这些显著权重的量化,就可以缓解由量化损失导致的性能下降

  • 一种常用的权重重要性判断方式是考察权重的幅值或 L_2范数。然而,实验发现,跳过具有较大权重范数的通道,即按照权重本身选择 FP16 通道,并不能显著提升量化后的模型性能,其改善幅度与随机选择通道相近
  • 值得注意的是,如果依据激活幅值来选择权重通道,那么即使仅将0.1\%的通道保留为 FP16,也可以显著改善模型性能
  • 作者推测,幅值更大的输入特征通常更加重要。因此,将与这些输入特征对应的权重保留为 FP16,可以更好地保存这些重要特征,从而改善模型性能

局限性:尽管仅将 0.1\%的权重保留为 FP16,就能够在几乎不增加总体模型大小的情况下改善量化性能,但这种混合精度数据格式会增加系统实现的复杂性。因此,需要设计一种新的方法,在不真正将这些重要权重保留为 FP16 的情况下,对其进行有效保护

🐫 Protecting Salient Weights by Activation-aware Scaling

为降低显著权重的量化误差,同时避免混合精度实现带来的硬件效率问题,作者提出一种基于逐通道缩放的方法。首先考虑仅权重量化所引入的误差。对于某个权重 group/block \mathbf{w},其线性计算可以表示为y=\mathbf{w}x,量化后的计算则为y=Q(\mathbf{w})x.具体而言,量化函数定义为

Q(\mathbf{w})=\Delta\cdot\operatorname{Round}\left(\frac{\mathbf{w}}{\Delta}\right),\qquad \Delta=\frac{\max(|\mathbf{w}|)}{2^{N-1}}

  • N表示量化比特数,\Delta表示由权重绝对值最大值确定的量化尺度。考虑权重组中的一个元素w\in\mathbf{w}。如果将该权重乘以一个缩放因子s>1,同时将对应的输入除以s,那么浮点计算本身保持不变,而量化后的计算变为

Q(\mathbf{w}\cdot s)\cdot\frac{x}{s} = \Delta' \cdot \operatorname{Round}\left(\frac{\mathbf{w}s}{\Delta'}\right) \cdot x\cdot\frac{1}{s}

  • 其中,\Delta'表示应用缩放因子s后得到的新量化尺度。作者基于实验得到以下观察:\operatorname{Round}(\cdot)所产生的期望舍入误差基本保持不变。由于舍入操作将浮点数映射为整数,其误差大致均匀分布于[0,0.5],因此平均误差约为0.25,即\operatorname{RoundErr}(\cdot)\approx0.25。对单个权重元素进行放大,通常不会改变整个权重 group 中的最大值,因此有\Delta'\approx\Delta.因此,原始量化计算产生的误差可以表示为

\operatorname{Err}(Q(\mathbf{w})x) = \Delta\cdot \operatorname{RoundErr}\left(\frac{\mathbf{w}}{\Delta}\right)\cdot x

  • 采用缩放之后,误差变为

\operatorname{Err}\left( Q(\mathbf{w}\cdot s)\frac{x}{s} \right) = \Delta' \cdot \operatorname{RoundErr}\left( \frac{\mathbf{w}s}{\Delta'} \right) \cdot x\cdot\frac{1}{s}

  • 新误差相对于原始误差的比例近似为\frac{\Delta'}{\Delta}\cdot\frac{1}{s}.当\Delta'\approx\Delta且s>1时,该比值小于 1,因此显著权w的相对量化误差将得到降低

同时考虑显著权重和非显著权重的量化误差,作者选择自动搜索一个逐输入通道的最优缩放因子,使某一层量化前后的输出差异最小。其优化目标可以为

\mathbf{s}^{*} = \arg\min_{\mathbf{s}}\mathcal{L}(\mathbf{s})

\mathcal{L}(\mathbf{s}) = \left\| Q\left( W\cdot\operatorname{diag}(\mathbf{s}) \right) \left( \operatorname{diag}(\mathbf{s})^{-1}X \right) - WX \right\|

  • 其中Q表示权重量化函数,例如采用 group size 为 128 的 INT3 或 INT4 量化;W表示原始 FP16 权重;X表示从一个较小校准集上缓存得到的输入特征;\mathbf{s} 为逐输入通道的缩放因子
  • 对于\operatorname{diag}(\mathbf{s})^{-1}X中的缩放操作,通常可以将其融合到前一个算子中,因此不会引入额外的运行时计算开销。由于量化函数本身不可微,因此无法直接使用标准反向传播优化上述目标。虽然可以采用近似梯度方法,但作者发现这类方法仍然存在收敛不稳定的问题

激活感知的搜索空间:为提高搜索过程的稳定性,作者根据影响缩放因子选择的因素,进一步定义一个简化搜索空间。由于权重通道的重要性实际上由对应的激活尺度决定,因此作者将缩放因子定义为\mathbf{s}=\mathbf{s}_{X}^{\alpha}.并搜索最优的\alpha

\alpha^{*} = \arg\min_{\alpha} \mathcal{L}\left( \mathbf{s}_{X}^{\alpha} \right)

  • 其中\mathbf{s}_{X}表示逐通道激活平均幅值,\alpha是一个用于权衡显著通道与非显著通道保护程度的超参数。在区间[0,1]内通过快速网格搜索寻找最优\alpha
  • 当\alpha=0时,不对权重进行缩放;当 \alpha=1时,则对应搜索空间中最激进的缩放方式
  • AWQ 还进一步对权重进行 clipping,以最小化量化产生的 MSE 误差。实验表明,在 OPT 模型的 INT3-g128 量化设置下,AWQ 能够稳定优于简单的最近邻舍入量化,并在保持更高硬件效率的同时,取得与保留 1% FP16 权重的混合精度方法相近的性能

AWQ 不依赖 GPTQ 式的二阶误差补偿,也不需要反向传播。它利用校准数据的逐通道平均激活幅值构造缩放候选,并通过前向计算比较候选的重构误差,再进行量化裁剪。因此,平均激活统计用于构造搜索空间,整个校准过程仍需使用校准样本评估误差。论文实验表明,该方法对校准数据具有较好的泛化能力

🦃 TINYCHAT: MAPPING AWQ ONTO EDGE PLATFORMS

AWQ 能够显著降低 LLM 的模型规模。然而,要将 W4A16(4-bit 权重、16-bit 激活)量化所带来的理论内存节省转化为实际可测的性能提升,并非易事

  • 以 SmoothQuant 等 W8A8 量化方法为例,其存储与计算采用相同的数据精度,因此反量化过程可以较为自然地融合到计算 kernel 的 epilogue 中。相比之下,W4A16 在内存访问与实际计算过程中采用不同的数据类型,因此,为获得更高的执行效率,反量化操作需要被集成到主要计算循环之中,这也带来额外的实现复杂度
  • 作者提出 TinyChat,一套面向 AWQ 模型推理的高效系统。TinyChat 使用 PyTorch 作为前端,并在后端利用设备特定的指令集,例如 CUDA/PTX、NEON 和 AVX

在 GPU 上,TinyChat 仅针对推理过程中最关键的组件实现专用优化,包括 Attention、Layer Normalization 和 Linear Projection kernel。其前端保持较高的灵活性,以便快速适配新的模型结构。采用 4-bit AWQ 后,TinyChat 在多个 LLM 系列上相较 HuggingFace FP16 实现获得超过 3 倍的 GPU 推理加速。在 CPU 上,TinyChat 则将完整计算图下沉至 C++,以降低运行时开销

对于量化层,由于硬件通常不直接提供 INT4 与 FP16 之间的乘法指令,因此在执行矩阵计算前,需要首先将 INT4 权重反量化至 FP16。为避免将反量化后的权重写回 DRAM,TinyChat 将反量化操作直接融合到矩阵乘法 kernel 中。权重在被读取后立即完成解包和反量化,并直接参与后续矩阵计算,而不是先生成完整的 FP16 权重矩阵并写回显存。这种融合方式同时应用于矩阵-矩阵乘法和矩阵-向量乘法 kernel

SIMD 感知的权重打包,本质上是针对目标硬件的 SIMD 寄存器宽度,提前在离线阶段重新排列低比特权重,使运行时能够用少量向量指令一次性完成多个 INT4 权重的解包,而不是逐元素执行标量 shift 和 bitwise AND。其核心目的是让低比特权重的数据布局更适配目标设备的向量化执行方式,从而降低 INT4 解包开销并提升推理效率

TinyChat 还广泛采用 kernel fusion 来优化端侧 LLM 推理。对于 Layer Normalization,作者将乘法、除法以及平方根等操作融合到单个 kernel 中执行,从而减少中间结果读写以及 kernel launch 开销。对于 Attention 层,TinyChat 将 Q、K、V 三个 projection 融合到一个 kernel 中,同时在运行过程中直接计算位置编码。此外,TinyChat 会预先分配 KV Cache,并在 Attention kernel 内部直接完成 Cache 更新。这种 kernel fusion 对于 Falcon 和 StarCoder 等 forward 实现效率较低的模型尤为有效