Q-Palette: Fractional-Bit Quantizers Toward Optimal Bit Allocation for Efficient LLM Deployment
Q-Palette: Fractional-Bit Quantizers Toward Optimal Bit Allocation for Efficient LLM Deployment
作者/机构: Deokjae Lee1,2, Hyun Oh Song1,2∗, 1首尔大学, 2神经处理研究中心
A1 主要贡献
本文研究仅权重量化中的训练后量化(PTQ),该方法无需重新训练,仅使用少量或无需校准数据即可对大型语言模型(LLM)的权重进行量化。仅权重量化对于降低LLM推理的内存占用和延迟至关重要,尤其是在内存受限的小批量推理场景中,例如在边缘设备上进行个性化推理。尽管其重要性日益凸显,但LLM中带有重尾异常值的不规则权重分布使量化变得复杂。为解决此问题,近期出现了基于旋转的方法,通过将权重转换为更规则、异常值更少的近高斯分布来减少量化误差。
本文首先提出了一个核心问题:如果在任意小数比特宽度下都有理想的高斯量化器,那么在固定的内存预算下,应如何跨层分配比特以最小化性能下降?基于线性定理(linearity theorem),该定理将量化引起的性能下降近似为各层均方误差的加权和,本文推导出了高斯化权重下信息论上的最优比特分配策略。分析揭示,要接近理论预测的量化性能,必须使用能够紧密匹配其理论失真界限的细粒度小数比特量化器。然而,现有复杂的高斯量化器(如网格编码量化)仅在有限的整型比特宽度(如2、3、4比特)下实现了融合内核,且对大于1的批量大小支持不足或没有支持。
为了将理论洞察与实际量化相结合,本文引入了 Q-Palette,这是一套功能多样的小数比特量化器集合,涵盖了从提供近乎最优失真的网格编码量化器(TCQ)到为实现更快推理而优化的更简单的向量和标量量化器,从而覆盖了不同的准确率-延迟权衡。本文为这些量化器提供了优化的CUDA内核,支持广泛的小数比特宽度,并比以往的复杂量化方法(如QTIP)支持更广的批量大小。为了实现更精细的比特宽度控制,本文进一步提出了 half-TCQ,这是一种新颖的TCQ变体,它在单层内混合使用两种不同比特宽度的TCQ量化器(例如2.5比特和3.0比特)以实现中间比特宽度(如2.75比特),并扩展了CUDA内核以支持此变体。
为了展示Q-Palette的实用性,本文将其集成到一个资源受限的 混合方案量化(MSQ) 框架中。为进一步改善准确率-延迟权衡,本文提出了 融合感知MSQ(fusion-aware MSQ),这是首个联合优化量化器选择和层融合的MSQ方法,引入了一个新的优化维度(见图1)。在这种方法中,共享相同输入的线性层(例如Transformer块中的query、key和value投影)可以融合成一个单一的线性层,从而减少内存访问和内核启动次数。通过结合层融合,融合感知MSQ在准确率-延迟权衡上取得了显著的增益。在LLaMA 2、LLaMA 3和Qwen模型上的大量实验表明,本文的MSQ框架与Q-Palette相结合,在内存和延迟受限的设置下,始终优于强大的无数据和数据感知仅权重量化PTQ基准。
A3 背景知识
2.1 用于训练后量化的线性化代理目标
基于Hessian近似的性能退化。先前的研究已经证明,由PTQ引起的神经网络性能下降可以近似为层级代理损失的线性组合,这些代理损失源自二阶Hessian近似【14, 7】。在LLM的背景下,性能通常由困惑度衡量,线性定理表明,由量化引起的困惑度增加可以被精确地近似为每层量化误差的加权和【35】。这个代理目标在无数据量化场景中尤其有价值,因为在这种场景下无法获得基于Hessian的代理损失。形式上,该线性化代理目标表示为:
$$\mathcal{L}(\{Q(W_l)\}_{l=1}^L) - \mathcal{L}(\{W_l\}_{l=1}^L) \approx \sum_{l=1}^L a_l \underbrace{\|Q(W_l) - W_l\|^2 / \|W_l\|^2}_{=: \text{err}(Q; W_l)},$$其中$L(·)$表示困惑度损失,$W_l \in \mathbb{R}^{d_{in_l} \times d_{out_l}}$是第$l$层权重矩阵,$Q(·)$是量化函数,$a_l$是第$l$层的经验估计敏感度系数,$err(Q; W_l)$是第$l$层的归一化量化误差。
内存受限的MSQ问题。利用此代理目标,可以将包含一组候选量化器$\mathcal{Q}$的内存受限MSQ问题公式化为一个多项选择背包问题(MCKP):
$$\begin{aligned} \begin{aligned} \underset{P_{lq} \in \{0,1\}}{\text{minimize}} & \quad \sum_{l=1}^{L} a_{l} \left( \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot \text{err}(Q_{q}; W_{l}) \right) \\ \text{subject to} & \quad \sum_{q=1}^{|\mathcal{Q}|} P_{lq} = 1, \quad \forall 1 \leq l \leq L, \\ & \quad \sum_{l=1}^{L} \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot \text{bit}(Q_{q}; W_{l}) d_{l}^{\text{in}} d_{l}^{\text{out}} \leq M, \end{aligned} \end{aligned}$$其中$Q_q$表示一个候选量化器,$bit(Q_q; W_l)$是权重矩阵$W_l$被$Q_q$量化后每个权重分量的平均比特数,$P_{lq} \in \{0, 1\}$是一个二元指示变量,用于为第$l$层选择量化器$Q_q$,$M$表示分配给量化模型的总内存预算(以比特为单位)【46】。该公式将MSQ明确地转化为一个基于线性化性能代理的组合优化问题,为在严格内存约束下进行最优比特分配提供了一个有原则的框架【35, 7】。
A2 方法细节
3 Q-Palette: 小数比特量化器
3.1 动机与设计目标
信息论最优比特分配。基于第2节介绍的理论基础,我们现在推导信息论上的最优比特分配策略,并讨论其对实用量化器设计的影响。假设权重矩阵通过不连贯处理被高斯化,量化问题可以看作是一个高斯信源编码问题。在这种情况下,经典的率失真理论为期望量化误差建立了一个基本下界:$E[err(Q)] \geq 2^{-2bit(Q)}$【10】。假设存在能够在任意小数比特宽度$b_l \geq \eta$下达到此界限的理想高斯量化器,那么内存受限的MSQ问题(1)可以简化为:
$$\begin{aligned} \begin{aligned} \underset{b_{l} \geq \eta}{\operatorname{minimize}} & \quad \sum_{l=1}^{L} a_{l} 2^{-2 b_{l}} \\ \text { subject to } & \quad \sum_{l=1}^{L} b_{l} d_{l}^{\text {in }} d_{l}^{\text {out }} \leq M, \end{aligned} \end{aligned}$$其中$b_l$是分配给第$l$层的小数比特宽度,$\eta > 0$是为避免将0比特分配给某层等退化情况而引入的最小比特宽度阈值。该公式有一个闭式解,如定理3.1所述。
定理3.1 (理想高斯量化器下的最优比特分配)。如果预算$M$是可行的,即$M \geq \eta \sum_{l=1}^{L} d_{in_l}d_{out_l}$,那么问题(2)的最优小数比特分配$\{b_l^*\}$由下式给出:
$$b_{l}^{*}=\max \left\{\eta, \frac{1}{2 \ln (2)}\left(\ln \frac{a_{l}}{d_{l}^{\text {in }} d_{l}^{\text {out }}}\right)+C\right\}, \quad \forall 1 \leq l \leq L,$$对于满足内存约束$\sum_l b_l^* d_{in_l}d_{out_l} = M$的常数$C$。证明见A节。
理论与实践的差距。然而,在实践中,量化必须使用一组有限的非理想量化器$\mathcal{Q} = \{Q_1, \dots, Q_N\}$来执行,这在理论最优性和实际可达性能之间引入了差距。这个差距的大小主要取决于两个因素:(1)每个量化器接近理想失真界限$2^{-2bit(Q)}$的程度,以及(2)可用的比特宽度能多精细地逼近最优小数比特分配$b_l^*$。进一步的分析请参见B节。
设计目标。这激发了对既准确又能在细粒度小数比特区间上可用的实用量化器的需求。此外,量化器通常在失真和计算效率之间存在权衡:更复杂的量化器可能提供更低的误差,但会带来更高的推理成本。这些考虑促使我们设计一个实用的量化套件,该套件支持细粒度的小数比特宽度,同时考虑到量化误差和计算效率之间的权衡。在这些目标的指导下,我们设计了Q-Palette,它是一个功能多样的量化器集合,旨在平衡不同部署场景下的量化误差和计算效率。
3.2 Q-Palette中的量化方案
Q-Palette量化器系列。Q-Palette包括三个量化器系列:非均匀标量量化(NUQ)、向量量化(VQ)和网格编码量化(TCQ),涵盖了量化误差和推理延迟权衡的范围。本节我们简要介绍这些量化方案。
非均匀标量量化 (NUQ)。与使用等间距区间的均匀标量量化器不同,NUQ通过将每个标量权重分配给一个非均匀间隔的查找表(LUT)中的一个条目来进行量化【27】。我们通过对随机高斯样本进行k-means聚类来构建NUQ码本【34】,从而为高斯化权重优化码本条目。NUQ的反量化开销较低,能够实现高效推理【20, 40】。
向量量化 (VQ)。VQ将权重向量划分为固定维度的组,并将每个组分配给预先计算好的码本中最接近的向量条目【51】。在Q-Palette中,我们特别实现了2D VQ,通过对随机2D高斯样本进行k-means聚类来生成码本。高效的内核实现依赖于大小为2的幂的LUT,这使得紧凑的比特级表示成为可能。因此,0.5比特间隔的小数比特宽度成为高效实现的自然选择。例如,一个形状为$(2^3, 2)$的LUT(八个2D向量)用3个比特编码两个元素,有效地实现了每个标量权重1.5比特。这种构造使VQ能够支持0.5比特间隔的小数比特宽度。
网格编码量化 (TCQ)。TCQ被认为是一种复杂的高斯信源量化器,能实现近乎最优的失真性能【19】。最近,QTIP为TCQ的位移变体引入了优化的CUDA内核【36, 50】,该变体将每个高维实值向量$v$编码为一个二进制位移表示$r$:
$$\hat{\mathbf{v}}[i \cdot V:(i+1) \cdot V]=\operatorname{LUT}(\mathbf{r}[i \cdot s: i \cdot s+L]),$$其中每个$V$维子向量由一个长度为$L$的滑动位窗口表示,每一步移动$s$位。这种编码方式实现了$s/V$的有效小数比特宽度。虽然之前的TCQ内核(QTIP)支持整数比特宽度(例如,对于$V=2$,位移$s = 4, 6, 8$对应2, 3, 4比特)并且仅限于单个批次,我们通过引入小数比特宽度支持(例如,对于$V=2$,位移$s = 3, 5, 7, 9, 10$对应1.5, 2.5, 3.5, 4.5, 5.0比特)和支持高达8的批次大小推理的优化内核,显著扩展了其实际适用性。为了构建LUT,我们遵循QTIP的协议,该协议也基于对高斯样本的k-means聚类。
Half-TCQ。此外,我们引入了half-TCQ,这是一个简单的扩展,能够在中间的小数比特宽度间隔上进行量化。具体来说,给定一个权重矩阵$W \in \mathbb{R}^{d_{in} \times d_{out}}$,half-TCQ按行划分矩阵,并对每个分区应用不同比特宽度的量化。例如,为了实现2.75比特的量化,half-TCQ将$W[: d_{in}/2]$量化为2.5比特,将其余一半$W[d_{in}/2 :]$量化为3比特。为了保持计算效率,我们实现了一个专用的CUDA内核,该内核在单个内核调用中执行half-TCQ的融合反量化和矩阵乘法。如图2所示,基于TCQ的方案,包括half-TCQ,始终能实现接近理论下界的量化误差,优于更简单的量化器。有关这些量化器的更详细解释,包括量化算法,请参阅C节。
3.3 为提高效率的实现细节
降低旋转开销。不连贯处理(Incoherence processing)会沿着输入和输出维度对权重进行旋转($W \rightarrow RWR'$),并进行逐张量缩放,从而产生近似标准的正态分布。然而,每次旋转也需要在推理过程中在线旋转激活值(例如,$X \rightarrow XR$),这会带来显著的计算开销【6, 49, 50】。我们通过仅沿输入维度旋转权重($W \rightarrow RW$)并应用逐输出通道缩放来减少这种开销,从而将每个旋转后的列归一化为近似标准的正态分布。此外,我们在具有相同输入的线性层之间共享旋转矩阵(例如,Transformer块中的query/key/value或gate/up投影)。结合这些技术,每个Transformer块的在线旋转次数从14次减少到4次。
内核实现。我们实现了两种类型的CUDA内核,针对不同的推理场景进行了优化:(i)基于Tensor Core的内核和(ii)基于CUDA Core的内核。我们基于Tensor Core的内核支持TCQ、NUQ和VQ,扩展了QTIP中支持单批次的实现,该实现利用了warp级的mma(矩阵乘法累加)指令【38】。将高效的反量化逻辑集成到这个框架中需要大量的工程努力,包括将量化后的权重精确映射到mma指令片段。总的来说,我们的实现将内核支持从整数比特宽度(2, 3, 4比特)扩展到了小数比特宽度(1.5, 2.5, 3.5, 4.5, 5.0比特)。为了在更大批次大小时进一步最小化开销,我们对每个量化权重只遍历一次,直接加载并执行寄存器级的反量化,无需中间存储。输入激活值缓存在共享内存中,以便在多个权重乘法中高效复用,显著提高了推理效率。
CUDA Core 内核实现。我们基于CUDA Core的内核支持NUQ和VQ,扩展了最初为NUQ设计的Any-Precision LLM内核【40】。具体来说,我们用更简单的位打包编码替换了位平面编码,以简化反量化过程。此外,我们将Anyprecision的表查找合并技术整合到基于Tensor Core和CUDA Core的NUQ内核中,用于2、3和4比特宽度,进一步减少了反量化开销。表1总结了Q-Palette中量化方案支持的内核实现和比特宽度。为保证透明度和可复现性,两种内核类型的完整实现都包含在我们的代码发布中,并在D节提供了额外的内核分析。
端到端解码延迟加速。表2总结了使用Q-Palette中的量化器对LLaMA 3.1 8B模型在不同比特宽度(2.0-4.25比特)下进行量化后的端到端解码延迟加速,并与两个基准进行了比较:使用FLUTE内核的NormalFloat和QTIP【11, 20, 50】。我们的量化器在支持更广泛、更细粒度的小数比特量化的同时,始终提供更优的推理速度。对于较小的批次大小(batch size = 1),基于CUDA Core的内核通常优于基于Tensor Core的内核,而在较大的批次大小(batch size = 8)时,基于Tensor Core的内核通常提供更好的延迟。尽管与NUQ和VQ相比,TCQ的延迟开销略高,但我们的TCQ量化器仍然比基准量化器实现了显著更快的解码速度,清楚地展示了在真实世界推理工作负载中的实际效率提升。请注意,我们的TCQ量化器扩展了QTIP,在批次大小为1时,主要区别在于支持更广泛的比特宽度和减少在线Hadamard变换的次数,从而降低了旋转成本并提高了解码延迟加速,例如在2比特时从2.91倍提升到3.57倍。
4 使用Q-Palette的混合方案量化
4.1 资源约束下的混合方案量化
资源受限的MSQ问题。在第2节中介绍的内存受限MSQ公式可以自然地推广到更广泛的资源约束,例如推理延迟。遵循先前的一次性混合精度量化框架【14, 7, 35】,我们将资源受限的MSQ公式化为MCKP:
$$\begin{aligned} \begin{aligned} \underset{P_{lq} \in \{0, 1\}}{\text{minimize}} & \sum_{l=1}^{L} \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot \ell_{lq} \\ \text{subject to} & \sum_{q=1}^{|\mathcal{Q}|} P_{lq} = 1, \quad \forall 1 \leq l \leq L, \\ & \sum_{l=1}^{L} \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot c_{lq} \leq C, \end{aligned} \end{aligned}$$其中,损失项$\ell_{lq}$表示为第$l$层选择量化器$Q_q$所导致的性能损失估计,成本项$c_{lq} \triangleq \text{cost}(Q_q; W_l)$表示分析得到的资源成本(例如,内存或延迟),总资源约束用$C$表示。损失项$\ell_{lq}$可以根据可用信息以多种方式实例化。例如,在无数据设置中,我们如公式(1)所示,将损失项近似为$\ell_{lq} = a_l \cdot \text{err}(Q_q; W_l)$,其中$a_l$是遵循HIGGS协议计算的逐层敏感度系数【35】。我们通过量化一个随机高斯矩阵得到的$Q_q$的预计算失真来估计$err(Q_q; W_l)$。这使得在无数据场景中能够快速估计损失项,而无需完全实现量化流程。
求解器。虽然存在用于求解MCKP的动态规划算法【43】,但为了灵活性和实现简便,我们直接使用Google OR-Tools提供的SCIP求解器【42】。有关损失项计算和成本分析的更多细节,请参阅E节。
Q-Palette的有效性。我们通过比较Q-Palette内不同量化器子集的内存受限MSQ结果,在图3中展示了Q-Palette的有效性。具体来说,TCQ-ALL包括Q-Palette中所有可用的TCQ量化器,而TCQ-2,3,4和VQ-2,3,4则反映了QTIP和HIGGS中支持的整数比特宽度量化器【50, 35】。虽然QTIP最初是一个单方案基准,但我们构建的整数比特宽度子集TCQ-2,3,4可作为一个合理的参考,以评估更广泛量化器支持所带来的好处。结果清楚地突显了TCQ-ALL的优势,表明Q-Palette提供的更广泛的量化器支持始终能带来更优的性能。
4.2 融合感知的混合方案量化
层融合优化。层融合是加速DNN模型推理速度的一种广泛使用的优化技术【26, 12】。在每个Transformer块内,某些线性层,例如{query, key, value}投影或{up, gate}投影,共享相同的输入,可以融合成一个单一的线性层。例如,我们可以连接权重矩阵并计算$X(W_q \oplus W_k \oplus W_v)$,而不是分别计算$XW_q, XW_k, XW_v$,然后在计算后分割输出(参见图1右侧)。层融合可以减少内核启动和内存访问的次数,从而为推理加速提供更多机会。
融合感知MSQ框架。我们提出了融合感知MSQ,这是一种新颖的MSQ框架,它将量化与层融合这一额外的设计维度联合优化。融合感知MSQ同时决定1)如何对层进行分组融合,以及2)为每个融合组分配哪个量化器。标准MSQ(公式3)为每个(层,量化器)对引入一个二元决策变量,而融合感知MSQ则为每个(可融合层组,量化器)对定义一个二元变量。这里,一个可融合层组是共享相同输入的一组层。对于通用的Transformer模型,我们可以将每个Transformer块b的所有可融合层组的集合写为 $G_b = {{q_b, k_b, v_b}, {q_b}, {k_b}, {v_b}, {u_b, g_b}, {u_b}, {g_b}, {o_b}, {d_b}}$ 。可融合层组的总集合是$\mathcal{G} = \cup_{b=1}^{B} \mathcal{G}b$,其中$B$是Transformer块的数量。对于每个$g \in \mathcal{G}$和量化器$Q_q \in \mathcal{Q}$,我们引入一个二元变量$P \in {0, 1}$,表示$g$中的所有层都被融合并由$Q_q$量化。
融合感知MSQ的ILP公式。融合感知MSQ问题被公式化为:
$$\min_{P_{g q} \in\{0,1\}} \sum_{g \in \mathcal{G}} \sum_{q=1}^{|\mathcal{Q}|} P_{g q} \cdot \sum_{l \in g} \ell_{l q}$$约束条件:
其中$\ell_{lq}$是损失项,$c_{gq}$表示与由$Q_q$量化的组$g$对应的融合层的分析成本(例如,延迟)。
约束说明。为了确保解决方案的有效性,施加了两个约束。排他性分配(C1):每一层必须且只能属于一个活跃的(组,量化器)对;在所有包含给定层$l$的可融合组$g$中,只有一个相关的变量$P_{gq}$可以为1。资源约束(C2):所有激活组的总分析成本不得超过资源预算$C$。
优势与求解。该公式明确地捕捉了由融合带来的延迟改善,因此与忽略层融合的MSQ公式相比,提供了更好的准确率-延迟权衡(参见图1)。由于目标函数和约束条件都是$P_{gq}$的线性函数,公式(4)也是一个整数线性规划(ILP)问题。我们使用OR-Tools中的SCIP求解器来解决这个ILP问题【42】。需要注意的是,与非融合感知的MSQ(公式3)相比,融合感知的MSQ引入了1.71倍的决策变量,同时保持了相同数量的约束。
A4 实验
实验设置。我们评估了我们的方法在LLaMA 3系列(LLaMA 3.1-8B, 70B, 3.2-1B, 3B)、LLaMA 2系列(LLaMA 2-7B, 13B)和Qwen 2.5-7B模型上相对于基准的量化性能【18, 48, 55】。在无数据场景中,我们考虑了单方案量化基准——HQQ(均匀)、NormalFloat(NUQ)、HIGGS-Single(VQ)和无数据QTIP(TCQ)——以及MSQ基准HIGGS-MSQ【2, 11, 50, 35】。在数据感知场景中,我们使用QTIP作为基准。所有实验中,我们的方法和基准都严格在PTQ设置下评估,不进行任何重新训练。性能主要通过WikiText2困惑度和在ARC-easy, ARC-challenge, HellaSwag, PiQA, 和WinoGrande上的平均零样本准确率来衡量【37, 8, 56, 3, 44】。对于延迟评估,我们对较高比特率(4.02, 4.25比特)的HQQ使用Gemlite内核,对NormalFloat(3.25, 4.25比特)和3.25比特的HQQ使用FLUTE内核【22, 20】。由于QTIP仅支持单批次推理,我们通过重复调用内核来模拟更大的批次大小。
方法命名。我们用Ours-TCQ-x表示使用Q-Palette中的TCQ-x进行单方案量化,Ours-MSQ-Mem表示使用Q-Palette的TCQ量化器进行内存受限的MSQ(第4.1节),Ours-MSQ-Lat表示使用所有Q-Palette量化器和Tensor-Core内核进行延迟受限的融合感知MSQ(第4.2节)。关于其他模型、消融研究和详细实验设置的更多结果,请参阅G节。
5.1 无数据量化结果
结果分析。表3总结了在LLaMA 3模型上的无数据量化性能。Ours-TCQ-x在WikiText2困惑度方面持续优于所有单方案基准,并实现了有竞争力的零样本准确率。值得注意的是,Ours-MSQ-Mem超越了所有基准方法,清楚地证明了Q-Palette的有效性。图4a在更广泛的内存范围(2.25-4.25比特)内评估了Ours-MSQ-Mem。我们的方法实现了帕累托最优性能,显著优于基准方法。引人注目的是,我们的2.875比特模型实现了与3.25比特HIGGS-MSQ模型相当的WikiText2困惑度,从而获得了1.13倍的更高压缩率和更优的困惑度。图4b和4c比较了Ours-MSQ-Lat和Ours-TCQ-x与基准方法在吞吐量-困惑度上的权衡。两种变体在批次大小为1和8的情况下都比基准方法实现了显著的吞吐量提升,大幅扩展了帕累托前沿。
5.2 数据感知量化结果
结果分析。我们进一步在数据感知设置下评估我们的方法,将我们的Ours-MSQ-Mem方法与最先进的QTIP基准(无重新训练)在LLaMA 2-7B和13B模型上进行比较。在此设置中,我们在量化过程中使用了与QTIP相同的代理Hessian,并计算我们MSQ的损失项$\ell_{lq}$作为使用量化器$Q_q$量化权重$W_l$所引起的实际验证困惑度下降。如表4所示,与基准相比,我们的方法始终实现了更优的困惑度和零样本准确率。此外,在批次大小为8的情况下,我们优化的内核在2比特和3比特的LLaMA 2模型上都实现了超过4倍的吞吐量提升,展示了我们优化内核在批次大小为8时的实际优势。
A7 补充细节
6 相关工作
不连贯处理(Incoherence processing)。以往处理LLM量化中异常值的方法主要依赖于启发式技术【28, 54, 29】。最近,一种理论上更可靠的方法,即不连贯处理,被引入以系统性地解决权重不规则性问题【6】。不连贯处理在量化前对权重矩阵应用旋转矩阵,显著抑制异常值并将分布转换为近似高斯形式【6, 1, 45】。这种高斯化使得可以使用复杂的高斯量化器,如格向量量化【49】和网格编码量化【50】。然而,目前的实现仅支持有限的整数比特宽度和小的批次大小的高效内核,限制了它们的实用性,而我们提出的Q-Palette通过引入小数比特量化器和支持更广泛批次大小的优化CUDA内核,直接解决了这一限制。
其他近期基于旋转的方法。更新的基于旋转的方法通过应用学习到的矩阵变换(如缩放或仿射变换)来进一步提高量化性能【32, 30, 23, 47】。然而,这些方法主要针对权重-激活值量化,并且需要校准数据来学习这些变换,而我们的工作专注于纯权重量化PTQ,即使在无数据设置下也适用,并且特别适合内存受限的小批量推理。
其他纯权重量化PTQ方法。一些更简单的PTQ方法优先考虑计算和实现的效率。HQQ通过半二次优化采用无数据的均匀量化【2】。NormalFloat使用高斯分位数构建用于非均匀标量量化的查找表【11】。FLUTE为基于LUT的非均匀量化器提供了最先进的CUDA内核,并支持逐组缩放【20】。尽管这些方法效率高,但与TCQ等复杂量化器相比,通常会产生更高的量化误差。
混合精度与混合方案量化。混合精度量化(MPQ)在给定约束下优化逐层的比特分配【53】。对于视觉模型,HAQ和HAWQ-V2引入了基于二阶信息的MPQ代理目标【13, 14】。Chen等人通过明确地纳入如延迟等多种资源约束,并将问题表述为MCKP,推广了这些方法【7】。最近,HIGGS引入了线性定理,这是一种专为LLM量化量身定制的无数据线性代理【35】。在这些工作的基础上,并借鉴编译器优化研究的见解【26, 12】,我们提出了一个新颖的融合感知混合方案量化框架,该框架联合优化量化器选择和层融合决策,从而实现了更优的准确率-延迟权衡。
A5 结论
总结。本文研究了仅权重量化PTQ作为压缩LLMs的解决方案,这对于内存受限、小批量推理任务尤其有益。鉴于LLMs中不规则的权重分布使量化变得复杂,我们利用了近期将权重分布高斯化的旋转方法,从而能够对最优比特分配进行理论分析。基于此视角,我们推导出了在固定比特预算下的信息论最优比特分配策略,证明了细粒度且能紧密逼近高斯失真率界限的小数比特量化器对于实现近乎最优的量化效率至关重要。为将此理论发现转化为实际效益,我们引入了Q-Palette,一个多功能的小数比特量化器套件,从提供近乎最优失真的复杂网格编码量化方案,到为快速推理而优化的更简单的向量和标量量化器,每种都通过优化的CUDA内核在广泛的比特宽度上高效实现。我们进一步将Q-Palette集成到一个MSQ框架中,提出了一个新颖的融合感知MSQ方法,该方法在给定资源约束下联合优化量化器选择和层融合决策,有效改善了推理延迟。实验评估验证了我们的MSQ框架与Q-Palette及融合感知优化相结合,在LLaMA 2和LLaMA 3模型上始终优于现有的基准方法,实现了更优的准确率-内存和准确率-延迟权衡。
A6 附录
A 最优比特宽度证明
本节我们正式推导论文主体部分所述的最优比特分配结果。在权重矩阵已通过不连贯处理被高斯化的假设下,量化问题可以看作是一个高斯信源编码问题。我们回顾内存受限的混合方案量化(MSQ)公式为:
$$\begin{aligned} \begin{aligned} \underset{P_{lq} \in \{0, 1\}}{\text{minimize}} \quad & \sum_{l=1}^{L} a_l \left( \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot \text{err}(Q_q; W_l) \right) \\ \text{subject to} \quad & \sum_{q=1}^{|\mathcal{Q}|} P_{lq} = 1, \quad \forall 1 \leq l \leq L, \\ & \sum_{l=1}^{L} \sum_{q=1}^{|\mathcal{Q}|} P_{lq} \cdot \text{bit}(Q_q; W_l) d_l^{\text{in}} d_l^{\text{out}} \leq M, \end{aligned} \end{aligned}$$其中$a_l$是第$l$层的经验估计敏感度系数,$err(Q; W_l)$是第$l$层的归一化量化误差,$Q_q$表示一个候选量化器,$bit(Q_q; W_l)$是权重矩阵$W_l$被$Q_q$量化后每个权重分量的平均比特数,$P_{lq} \in {0, 1}$是一个二元指示变量,用于为第$l$层选择量化器$Q_q$,$M$表示分配给量化模型的总内存预算(以比特为单位)【46, 35, 7】。
我们回顾经典的率失真理论为高斯信源的期望量化误差提供了一个基本下界:$E[err(Q)] \geq 2^{-2bit(Q)}$【10】。进一步假设我们能够使用理想的高斯量化器,在任意小数比特宽度$b_l \geq \eta$下精确达到此理论失真界限$E[err(Q)] = 2^{-2bit(Q)}$,那么内存受限的MSQ(问题(1))可以再次写成连续优化问题:
$$\begin{aligned} \begin{aligned} \underset{b_l \ge \eta}{\text{minimize}} \quad & \sum_{l=1}^L a_l 2^{-2b_l} \\ \text{subject to} \quad & \sum_{l=1}^L b_l d_l^{\text{in}} d_l^{\text{out}} \le M, \end{aligned} \end{aligned}$$其中$b_l$是分配给第$l$层的小数比特宽度,$\eta > 0$是为避免将0比特分配给某层等退化情况而引入的最小比特宽度阈值。在这里,我们用其期望$E[err(Q)]$替代实际的量化误差$err(Q)$。我们通过实验证明,在LLMs中遇到的典型权重矩阵维度下,量化误差的方差极低(见表5),从而为这一替代提供了经验支持。此外,我们假设敏感度系数$a_l$是非负的($a_l \geq 0$),这是一个合理的假设,因为预训练的权重通常代表了局部最优。基于这个简化的优化问题,我们现在推导最优小数比特分配的闭式解。
定理3.1 (理想高斯量化器下的最优比特分配)。如果预算$M$是可行的,即$M \geq \eta \sum_{l=1}^{L} d_{in_l}d_{out_l}$,那么问题(2)的最优小数比特分配${b_l^*}$由下式给出:
$$b_{l}^{*}=\max \left\{\eta, \frac{1}{2 \ln (2)}\left(\ln \frac{a_{l}}{d_{l}^{\text {in }} d_{l}^{\text {out }}}\right)+C\right\}, \quad \forall 1 \leq l \leq L,$$对于满足内存约束$\sum_l b_l^* d_{in_l}d_{out_l} = M$的常数$C$。
证明。我们首先构建问题(2)的拉格朗日函数,通过拉格朗日乘子$\mu_l \geq 0$明确地包含约束$b_l \geq \eta$,并通过$\lambda \geq 0$包含预算约束:
$$\mathcal{L}\left(\left\{b_{l}\right\}, \lambda,\left\{\mu_{l}\right\}\right):=\sum_{l=1}^{L} a_{l} 2^{-2 b_{l}}+\lambda\left(\sum_{l=1}^{L} b_{l} d_{l}^{\text {in }} d_{l}^{\text {out }}-M\right)-\sum_{l=1}^{L} \mu_{l}\left(b_{l}-\eta\right) .$$对拉格朗日函数关于$b_l$求导并令其等于零以找到驻点,我们有:
$$\frac{\partial \mathcal{L}}{\partial b_{l}}=-2 \ln (2) a_{l} 2^{-2 b_{l}}+\lambda d_{l}^{\text {in }} d_{l}^{\text {out }}-\mu_{l}=0.$$由于对所有$l$,$\mu_l \geq 0$且互补松弛性要求$\mu_l(b_l^ - \eta) = 0$【4】,我们有两种情况:
情况1:如果$b_l^ > \eta$,互补松弛性意味着$\mu_l = 0$,因此:
两边取对数并显式整理项,我们得到:
$$b_{l}^{*}=\frac{1}{2 \ln (2)}\left(\ln \frac{a_{l}}{d_{l}^{\text {in }} d_{l}^{\text {out }}}\right)+\frac{1}{2 \ln (2)}(\ln (2 \ln (2))-\ln (\lambda))=\frac{1}{2 \ln (2)}\left(\ln \frac{a_{l}}{d_{l}^{\text {in }} d_{l}^{\text {out }}}\right)+C(\lambda) .$$其中,常数项$C(\lambda)$定义为$C(\lambda) := \frac{1}{2 \ln(2)} (\ln(2 \ln(2)) - \ln(\lambda))$。
情况2:如果$b_l^* = \eta$,我们直接得到:
$$b_{l}^{*}=\eta.$$结合这些情况,得到最优的小数比特分配:
$$b_{l}^{*}=\max \left\{\eta, \frac{1}{2 \ln (2)}\left(\ln \frac{a_{l}}{d_{l}^{\text {in }} d_{l}^{\text {out }}}\right)+C(\lambda)\right\},$$其中常数$C(\lambda)$被选择以使得内存约束
$$\sum_{l=1}^{L} b_{l}^{*} d_{l}^{\text{in}} d_{l}^{\text{out}} \leq M,$$是紧的(即等式成立)。这个等式条件是自然产生的,因为目标函数(2)在$b_l$上是非增的(由于非负性假设$a_l \geq 0$)。因此,增加$C(\lambda)$直到约束恰好被满足不会使目标恶化,从而完成了证明。□
为了经验性地验证我们用期望值来近似量化误差的做法,我们多次量化随机高斯矩阵,并观察到量化误差(失真)的方差始终很低。具体来说,我们对32个形状为(4096, 4096)的标准高斯随机矩阵进行了量化,这与LLaMA 3.1-8B自注意力查询投影层的维度一致。表5报告了归一化量化误差$(| \hat{W} - W |_2^2 / | W |_2^2)$值的均值和标准差。结果显示方差很低,支持了我们的假设。
B 量化最优性差距分析
在实践中,我们通常只能使用一组有限的量化器$\mathcal{Q} = {Q_1, \dots, Q_N}$,这些量化器可能无法达到理论上的率失真最优性。在此约束下,原始的内存受限MSQ问题(1)仍然可以求解,但得到的解可能偏离在理想高斯量化器假设下推导出的最优解(定理3.1)。本节我们正式分析这种量化最优性差距。
设$Q_l^$表示通过求解问题(1)为每层$l$从量化器集合$\mathcal{Q}$中选出的最优量化器。那么,量化最优性差距,定义为实际最优解与理论最优小数比特宽度解$b_l^$(在定理3.1中推导)的目标值之差,可以表示为:
$$\sum_{l=1}^{L} a_{l}\left(\operatorname{err}\left(Q_{l}^{*} ; W_{l}\right)-2^{-2 b_{l}^{*}}\right).$$我们将总差距分解为两个直观的项,即失真差距和比特分配差距,如下所示:
$$\underbrace{\sum_{l=1}^{L} a_{l}\left(\operatorname{err}\left(Q_{l}^{*}\right)-2^{-2 b_{l}^{*}}\right)}_{\text {Total gap }}=\underbrace{\sum_{l=1}^{L} a_{l}\left(\operatorname{err}\left(Q_{l}^{*}\right)-2^{-2 \operatorname{bit}\left(Q_{l}^{*}\right)}\right)}_{\text {Distortion gap }}+\underbrace{\sum_{l=1}^{L} a_{l}\left(2^{-2 \operatorname{bit}\left(Q_{l}^{*}\right)}-2^{-2 b_{l}^{*}}\right)}_{\text {Bit allocation gap }},$$为了简化,我们用$err(Q_l^)$和$bit(Q_l^)$分别缩写$err(Q_l^; W_l)$和$bit(Q_l^; W_l)$。
由于经典的率失真理论以及$b^$作为连续优化问题(2)解的最优性,该分解中的每一项都是非负的。具体来说,第一项,$err(Q_l^) - 2^{-2bit(Q_l^)}$, 量化了每个实际量化器$Q_l^$接近理论高斯失真界限的程度。第二项,$2^{-2bit(Q_l^)} - 2^{-2b_l^}$, 衡量了可用的比特宽度${bit(Q) | Q \in \mathcal{Q}}$逼近最优比特分配${b_l^*}$的好坏程度。
为了研究量化器集合的设计如何影响差距的每个组成部分,表6报告了LLaMA 3.1-8B在2.5比特和3.25比特约束下的失真和比特分配差距。可以得出两个关键观察:
- 量化器质量的影响。VQ-2,3,4和TCQ-2,3,4显示出可比的比特分配差距,但TCQ-2,3,4产生的失真差距要小得多。因此,它们的性能差异主要源于量化器质量而非比特分配。
- 更宽比特宽度支持的影响。将TCQ-2,3,4与TCQ-ALL进行比较,揭示了比特分配差距的大幅减少,表明更丰富的小数比特宽度支持能够实现更准确的比特分配,更接近理论理想值。
这一分析推动了Q-Palette的设计,它提供了高质量的TCQ量化器和广泛的小数比特宽度覆盖,以减少失真和比特分配差距。分析这些因素为改进实用量化器设计和选择更有效的量化器集以减小量化最优性差距提供了洞见。受此分析的启发,我们专门设计了Q-Palette,它是一套多功能的小数比特量化器,包括能紧密逼近理论失真界限的TCQ,并提供了广泛的小数比特宽度支持。
C Q-Palette中量化器的更多细节
本节我们提供Q-Palette中每个量化器系列的实现细节。对于每个量化器,我们描述:(i) 码本构建,(ii) 反量化,和 (iii) 量化过程。量化步骤依赖于特定于量化器的四舍五入(RTN)算子,其过程根据数据可用性而不同:
-
无数据场景:我们将每个权重矩阵划分为标量元素(NUQ)或向量(VQ, TCQ)。每个分区都使用RTN算子独立量化,其结果的二进制表示被连接起来形成最终的量化权重表示。
-
数据感知场景:我们采用了先前方法中介绍的块状LDLQ框架【17, 49, 50】。具体来说,对于每个权重矩阵,我们在Hessian近似的指导下以块为单位进行量化,块大小因量化器而异:NUQ为1,VQ为2,TCQ为16。此方法从第一行到最后一行顺序处理权重行,根据Hessian和累积的量化误差迭代更新权重,并通过RTN算子量化每个更新后的权重。有关详细的公式和额外的理论背景,请参阅QUIP#和QTIP【49, 50】。
C.1 非均匀标量量化 (NUQ)
码本构建。对于比特宽度为$b$的NUQ,我们使用flash1dkmeans(一种快速的1D k-means算法【25】)构建LUT,该算法应用于$10^8$个随机采样的标准高斯值。我们将聚类数设置为$k = 2^b$,得到一个LUT $\in \mathbb{R}^{2^b}$。
反量化。给定LUT,一个二进制表示$r \in {0, 1}^b$被反量化为:
$$ \mathrm{dq}(\mathbf{r} ; \mathrm{LUT}):=\mathrm{LUT}[\operatorname{int}(\mathbf{r})], $$其中$int(r)$将二进制表示$r \in {0, 1}^b$转换为其在$[0, 2^b - 1]$范围内的相应整数索引。
量化。NUQ的RTN算子 $RTN: \mathbb{R} \rightarrow {0, 1}^b$ 将一个标量输入$v \in \mathbb{R}$映射到最近的LUT条目:
$$\operatorname{RTN}(v ; \mathrm{LUT}):=\underset{\mathbf{r} \in\{0,1\}^{b}}{\operatorname{argmin}}|v-\operatorname{LUT}[\operatorname{int}(\mathbf{r})]| .$$量化遵循上述针对无数据和数据感知场景的通用程序。
C.2 向量量化 (VQ)
码本构建。对于比特宽度为$b$的VQ,我们使用scikit-learn实现的2D k-means算法构建码本,该算法采用Lloyd算法【41, 34】。我们将超参数设置为max_iter=300和tol=1e-6,并将该算法应用于随机标准高斯样本,对于比特宽度$b \leq 5$使用$10^8$个样本,对于比特宽度$b > 5$使用$10^7$个样本。聚类数设置为$k = 2^{2b}$,得到一个LUT $\in \mathbb{R}^{2^{2b} \times 2}$,由$2^{2b}$个2D向量组成。
反量化。给定LUT,一个二进制表示$r \in {0, 1}^{2b}$的反量化方式与NUQ类似,但现在映射到一个向量:
$$\mathrm{dq}(\mathbf{r}; \mathrm{LUT}) := \mathrm{LUT}[\mathrm{int}(\mathbf{r})] \in \mathbb{R}^2,$$其中$int(r)$将二进制表示$r \in {0, 1}^{2b}$转换为其在$[0, 2^{2b} - 1]$范围内的相应整数索引。
量化。特定于VQ的RTN算子,$RTN: \mathbb{R}^2 \rightarrow {0, 1}^{2b}$,将一个2D输入向量$v \in \mathbb{R}^2$映射到最近的LUT条目:
$$\operatorname{RTN}(\mathbf{v} ; \mathrm{LUT}):=\underset{\mathbf{r} \in\{0,1\}^{2 b}}{\operatorname{argmin}}\|\mathbf{v}-\operatorname{LUT}[\operatorname{int}(\mathbf{r})]\|_{2} .$$然后量化遵循上述针对无数据和数据感知场景的通用程序。
C.3 网格编码量化 (TCQ)
C.3.1 通用TCQ
码本构建。我们遵循与QTIP【50】相同的协议,使用基于Lloyd算法的scikit-learn的k-means实现【41, 34】。具体来说,我们将$2^{20}$个随机采样的2D标准高斯向量(经过适当缩放)聚类成$2^{tlut_bits}$个簇,得到簇中心$tlut \in \mathbb{R}^{2^{tlut_bits} \times 2}$。然后我们使用来自QTIP代码库的quantlut_sym函数,通过混合码本构建方法构建最终的码本LUT $\in \mathbb{R}^{2^L \times 2}$:
def quantlut_sym(tlut, L, tlut_bits):
with torch.no_grad():
lut = torch.arange(1 << L, device=tlut.device)
lut = (lut + 1) * lut
sflp = 1 - ((lut >> 15) & 1) * 2
lut = (lut >> (16 - tlut_bits - 1)) & ((1 << tlut_bits) - 1)
lut = tlut[lut]
lut[:, 0] = lut[:, 0] * sflp
return lut
遵循QTIP,我们为所有TCQ量化器设置$L = 16$。对于比特宽度$b \leq 4$,我们将$tlut_bits$设置为9;对于新的小数比特宽度4.5和5.0,分别设置为10和11。
反量化。我们采用QTIP中的带尾咬合(tail-biting)的TCQ位移变体。给定一个二进制表示$r \in {0, 1}^{sT/V}$,我们明确定义参数如下:
- $s$:位移大小,对于比特宽度$b$设为$s = 2b$。
- $V$:向量大小,固定为$V = 2$。
- $L$:码本长度(或滑动窗口大小),固定为$L = 16$。
- $T$:网格大小,设为$T = 256$。
然后反量化通过带尾咬合的滑动窗口LUT索引进行:
$$ \operatorname{dq}(\mathbf{r} ; \mathrm{LUT}):=\operatorname{concat}_{i=0}^{T / V-1} \mathrm{LUT}[\mathbf{r}[i \cdot s: i \cdot s+L]] \in \mathbb{R}^{T} $$其中超过$r$长度的索引会因尾咬合而回绕,从而实现$s/V$的比特宽度。
量化。给定一个目标向量$v \in \mathbb{R}^T$进行量化,我们使用与QTIP中详述的相同的RTN算子,该算子利用维特比算法找到最优的二进制表示$r \in {0, 1}^{sT/V}$,该表示被反量化为最接近向量$v$的向量$\hat{v} = dq(r; LUT)$【16, 50】。量化过程遵循前面描述的通用的无数据和数据感知框架。
C.3.2 Half-TCQ
码本构建。对于half-TCQ,即一半权重使用比特宽度$b$量化,另一半使用$b+0.5$量化,我们遵循与上述TCQ在比特宽度$b+0.5$下完全相同的码本构建过程。
反量化。反量化分别处理权重矩阵的两个分区:前半部分使用比特宽度为$b$的二进制表示,后半部分使用比特宽度为$b+0.5$的表示。然后将每一半得到的向量连接成一个完整的反量化权重向量。
量化。我们将对应于TCQ-$b$的RTN算子应用于权重的前半部分,并将对应于TCQ-($b+0.5$)的RTN算子应用于后半部分。这个过程在无数据和数据感知(块状LDLQ)场景中都一致使用。
D CUDA内核的附加细节和性能分析
我们实现了两种类型的CUDA内核:(i)基于Tensor Core的内核和(ii)基于CUDA Core的内核。在这里,我们首先详细介绍我们基于Tensor Core的内核实现。然后,我们描述我们基于CUDA Core的内核实现。最后,我们对我们的内核进行额外的性能分析。
D.1 基于Tensor Core的内核实现
实现细节。我们基于Tensor Core的内核支持多种量化方案(TCQ、NUQ和VQ),并通过扩展QTIP内核来实现,这些内核最初支持整数比特宽度(2、3、4比特)的TCQ【50】。具体来说,对于TCQ,我们通过仔细扩展QTIP提供的基于warp级mma指令的实现,引入了对细粒度间隔的小数比特宽度(例如1.5、2.5、3.5、4.5、5.0比特)的优化支持。此外,我们借鉴QTIP的内核设计原则,为NUQ和VQ实现了高效的基于Tensor Core的内核。这些扩展涉及重要的工程工作,特别是在将量化权重精确映射到Tensor Core mma指令片段方面。为了在更大批量大小时进一步减少开销,我们对每个量化权重仅遍历一次,在加载时直接执行寄存器级反量化,无需中间存储。输入激活值缓存在共享内存中,以便在多个权重乘法中高效复用,从而大幅提高推理效率。
简化的内核结构。下面,我们提供一个简要的内核结构,突出显示关键函数、其用途和文件位置:
// kernels/tcq-kernels/src/inference.cu
__device__ void load_reg_cs<R>(compressed, idx, laneId, ®s) {
// 将量化的TCQ权重(比特宽度R/2)映射到mma片段
// 支持小数比特宽度 (1.5, 2.0, 2.5, ..., 4.5, 5.0)
}
// kernels/vq-tensor-kernels/src/inference.cu
__device__ void load_reg_cs<R, LUT_TYPE L>(compressed, idx, laneId, ®s) {
if (L == LUT_TYPE::SQ_LUT) {
// 将量化的NUQ权重(比特宽度R)映射到mma片段
} else if (L == LUT_TYPE::VQ_LUT_2) {
// 将量化的VQ权重(比特宽度R/2)映射到mma片段
}
}
// for TCQ fused kernel
// - kernels/tcq-kernels/src/inference.cu: kernel_decompress_gemm_combt
// for TCQ-Half fused kernel
// - kernels/vq-tensor-kernels/src/inference.cu: kernel_decompress_gemm
// for VQ and NUQ fused kernel
__global__ void tensor_core_kernel(...) {
// 在共享内存中管理LUT和输入
// 在寄存器中管理量化权重
// 通过‘load_reg_cs’将量化权重映射到Tensor Core mma片段
// 使用Tensor Core mma指令进行矩阵乘法例程
// 归约后写入最终结果
}
// Dequantization-only kernels
// - kernels/tcq-kernels/src/inference.cu: kernel_decompress
// for TCQ dequantization kernel
// - kernels/tcq-kernels/src/inference.cu: kernel_decompress_combt
// for TCQ-Half dequantization kernel
// - kernels/vq-tensor-kernels/src/inference.cu: kernel_decompress
// for VQ and NUQ dequantization kernel
__global__ void kernel_decompress(...) {
// 独立反量化权重(无矩阵乘法)
}
完整实现。完整的基于Tensor Core的内核实现包含在我们的公开代码发布中(目录kernels/tcq-kernels和kernels/vq-tensor-kernels)。
D.2 基于CUDA Core的内核实现
实现细节。我们基于CUDA Core的内核明确利用CUDA Core指令来支持NUQ和VQ量化方案,扩展了最初为NUQ开发的Any-Precision LLM内核【40】。具体来说,我们用更简单的位打包编码替换了原始的位平面编码,以简化反量化过程。
内核结构与实现。下面,我们提供一个简要的内核结构,突出显示关键函数、其用途和文件位置:
// - kernels/vq-cuda-kernels/src/gemm_routines.cu
__device__ void vq_pack_dequant_routine<nbits, vec_sz>(Bcode, B_row, shC) {
// 使用查找表'shC'将量化的VQ权重'Bcode'(比特宽度nbits/vec_sz)
// 解包到half2数组'B_row'中
// (例如,1.5比特量化: nbits=3, vec_sz=2)
}
// - kernels/sq-cuda-kernels/gemm_routines.cu
__device__ void pack_dequant<nbits>(Bcode_row, B_row, shC) {
// 使用查找表'shC'将量化的NUQ权重'Bcode'(比特宽度nbits)
// 解包到half2数组'B_row'中
}
// General CUDA Core kernel structure
// - kernels/sq-cuda-kernels/gemm_routines.cu: sq_gemm_fp16
// for NUQ fused kernel
// - kernels/vq-cuda-kernels/src/gemm_routines.cu: vq_pack_gemm_fp16
// for VQ fused kernel
__global__ void cuda_core_kernel(...) {
// 在共享内存中管理LUT
// 在寄存器中管理量化权重
// 通过相应的pack_dequant_routine将量化权重解包到half2数组
// 使用CUDA Core半精度FMA (hfma2)指令进行矩阵乘法
// 归约后写入最终结果
}
// Dequantization-only kernels
// - kernels/sq-cuda-kernels/gemm_routines.cu: pack_dequant_kbit_store
// for NUQ dequantization kernel
// - kernels/vq-cuda-kernels/src/gemm_routines.cu: \
// vq_pack_dequant_kbit_store
// for VQ dequantization kernel
__global__ void kernel_decompress(...) {
// 独立反量化权重(无矩阵乘法)
}
完整实现。完整的基于CUDA Core的内核实现包含在我们的公开代码发布中(目录kernels/sq-cuda-kernels和kernels/vq-cuda-kernels)。
D.3 附加性能分析
RTX3090上的性能。本节补充了论文主体中的表2,通过在不同的硬件配置(使用RTX3090 GPU)上提供额外的性能分析。表7总结了在RTX3090 GPU上,量化后的LLaMA 3.1-8B模型相对于FP16基准的解码延迟加速比,补充了论文主体中报告的RTX4090 GPU上的结果。我们的量化器在所评估的两种批次大小(1和8)下都持续优于基准方法。
硬件依赖性。值得注意的是,在RTX3090上,对于批次大小为1的情况,我们基于CUDA Core(‘CC’)和Tensor Core(‘TC’)的内核之间的延迟加速差距比在RTX4090上观察到的要大(表2)。这种差异凸显了内核性能的显著硬件依赖性,证明了提供多种内核实现的重要性。这种灵活性使得可以根据特定的硬件平台和工作负载需求选择最优的内核。
E 混合方案量化的实现细节
E.1 损失项计算
无数据损失项。在无数据场景中,我们通过线性定理来估计损失项【35】,即$\ell_{lq} = a_l \cdot \text{err}(Q_q; W_l)$。具体来说,我们使用通过量化随机标准高斯矩阵预先计算的失真值来近似量化误差$err(Q_q; W_l)$,从而避免对每个权重矩阵$W_l$进行显式量化。为了确定敏感度系数$a_l$,我们采用了HIGGS中介绍的程序【35】。首先,我们从给定的LLM中随机生成128K个token。然后,对于每一层$l$,我们注入按特定范数$n_{li} = \frac{i}{16}$($1 \leq i \leq 16$)缩放的随机高斯噪声,并测量在这128K个token上计算的KL散度损失的增加量:
$$\Delta \mathcal{L}_{l i}=\mathcal{L}\left(\left\{W_{l^{\prime}}+\delta_{l^{\prime} l} \cdot n_{l i} \cdot\left\|W_{l}\right\|_{2} \cdot \epsilon_{l} /\left\|\epsilon_{l}\right\|_{2}\right\}_{l^{\prime}=1}^{L}\right)-\mathcal{L}\left(\left\{W_{l^{\prime}}\right\}_{l^{\prime}=1}^{L}\right), \quad \epsilon_{l} \sim \mathcal{N}(0, I)$$其中$\delta_{l'l}$是克罗内克δ(如果$l' = l$则为1,否则为0),确保噪声仅注入到第$l$层,$\epsilon_l$是与第$l$层维度相匹配的标准高斯噪声矩阵。根据线性定理,损失的增加近似遵循线性关系$\Delta L_{li} \simeq n_{li}^2 a_l$,使我们能够通过对数据点$(n_{li}^2, \Delta L_{li})$进行线性拟合来估计$a_l$。这个过程需要进行$16 \times L$次KL散度损失的计算,其中$L$是总层数,但它可以以一种易于并行的方式执行。此外,计算出的敏感度系数$a_l$可以被所有无数据的MSQ场景重用,只产生一次性的计算成本。
数据感知损失项。对于数据感知场景,我们使用两种不同类型的损失项:线性和实际。
-
基于线性的损失项。与无数据场景类似,我们利用基于线性定理的近似。然而,我们将基于随机生成的128K个token计算的KL散度损失替换为基于RedPajama数据集中的1M个token计算的困惑度损失【52】。
-
实际损失项。在这里,我们显式地计算由量化引起的实际困惑度增加。具体来说,我们使用来自RedPajama数据集的256K个token,并如下计算损失项:
$\ell_{lq} := L({Q_{l'}^{\text{default}}(W_{l'})} \text{, 将第l层替换为} Q_q(W_l)) - L({Q_{l'}^{\text{default}}(W_{l'})})$,
其中$Q_l^{\text{default}}$表示第$l$层的默认量化器(例如,对于2比特量化我们使用TCQ-2,对于3比特量化我们使用TCQ-3)。这明确地测量了用量化器$Q_q$量化第$l$层所引起的实际经验困惑度增加。
对于论文主体表4中的数据感知实验,我们专门使用了‘实际’损失项。此外,在G节提供的表9的消融研究中,我们明确比较了使用两种类型的损失项‘实际’和‘线性’的性能。
E.2 延迟分析
分析方法。对于延迟分析,我们测量与每个量化器对应的CUDA内核的执行时间。由于诸如归一化层、旋转和自注意力操作等计算在不同量化器之间是相同的,我们专门分析每个融合的反量化和矩阵乘法内核的延迟。为了准确估计整体推理延迟,我们首先使用torch.compile测量几个随机选择的量化配置的端到端推理延迟。然后,我们减去内核延迟的总和,以估计由公共计算(例如,归一化层)引起的延迟开销,并将此开销均匀地分配到所有量化器配置文件中。
融合内核的延迟分析。在考虑融合内核时,我们分别测量与每种融合模式对应的内核延迟。为了考虑由于不同融合模式导致的延迟开销变化,我们相应地调整了这种开销偏差。
E.3 优化器
工具。为了解决混合方案量化优化问题,我们使用了Google的OR-Tools优化套件【42】,具体是利用SCIP求解器,时间限制为60秒。
F 图1的设置
实验概述。在图1中,我们展示了基于Q-Palette的量化框架与使用FLUTE内核的NormalFloat基准【11, 20】的定性比较,评估在LLaMA 3.1-8B模型上进行,使用RTX4090 GPU,批次大小为1。具体来说,我们验证了序列长度为8192时的WikiText2困惑度,并测量了与FP16基准相比的推理加速,具体设置如下:
NormalFloat (基准)。对于NormalFloat,我们采用FLUTE,码本大小为$2^3$,组大小为64,从而得到平均比特宽度为3.25。我们利用FLUTE代码库中发布的FLUTE内核【11, 20】,该内核为在RTX4090 GPU上的推理进行了优化,以测量推理延迟并计算加速比。
使用TCQ-3.25的单方案量化 (我们)。我们使用Q-Palette中的TCQ-3.25量化器对LLaMA 3.1-8B模型的所有层进行统一的无数据量化。这对应于half-TCQ方案,即将一半的权重矩阵量化为3.0比特,另一半量化为3.5比特。
使用Q-Palette的MSQ (我们)。我们利用Q-Palette中可用的全套量化器作为我们的搜索空间。对于NUQ和VQ量化器,在优化过程中同时考虑了Tensor-Core和CUDA-Core内核实现。敏感度系数$a_l$遵循HIGGS协议计算【35】,详见E节,并且我们利用了预先计算的失真值,如论文主体4.1节所述。给定这些敏感度系数、失真值和预先分析的延迟测量值,我们求解延迟受限的MSQ优化问题(论文主体中的公式(3)),以在各种延迟约束下识别帕累托最优的量化器选择。从得到的准确率-延迟权衡曲线中,我们选择了一个在延迟和困惑度上都明显优于TCQ-3.25基准的配置。
使用Q-Palette的融合感知MSQ (我们)。我们通过额外分析融合线性层组合的延迟测量值,并将层融合纳入我们的MSQ公式中,求解融合感知的优化问题(论文主体中的公式(4))。这种方法明确捕捉了通过层融合可实现的延迟减少,从而能够联合优化量化方案和层融合决策。我们选择了一个在推理速度和困惑度上都明显优于无层融合的MSQ基准的量化配置。
与这些场景对应的详细量化配置在图1中进行了可视化。
G 实验设置和附加结果
G.1 实验设置
G.1.1 评估指标详情
语言建模性能。为了评估语言建模性能,我们在WikiText2数据集上测量困惑度【37】,对于LLaMA 2模型使用4096个token的序列长度,对于LLaMA 3模型使用8192个token的序列长度。此外,我们报告了在五个下游任务上的零样本准确率:ARC-easy、ARC-challenge、HellaSwag、PiQA和WinoGrande【37, 8, 56, 3, 44】。零样本评估是使用lm_eval库(版本0.4.4)进行的。
G.1.2 设备详情
RTX 4090 GPU实验。我们的RTX 4090 GPU实验使用RunPod提供的云环境进行,硬件和软件规格如下:
* GPU: NVIDIA RTX 4090
* CPU: AMD EPYC 7B13 64核处理器
* 操作系统: Ubuntu 22.04.5
* CUDA版本: 12.4
RTX 3090 GPU实验。我们的RTX 3090 GPU实验使用我们的本地机器进行,详情如下:
* GPU: NVIDIA RTX 3090
* CPU: AMD EPYC 7402 24核处理器
* 操作系统: Ubuntu 22.04.1
* CUDA版本: 12.4
G.1.3 基准配置
HQQ [2]。根据官方文档,推理加速内核(如Gemlite)仅支持axis=1的配置。因此,我们使用以下配置:
* 4.25比特: nbits=4, group_size=64, axis=1 (Gemlite内核)
* 4.02比特: nbits=4, group_size=1024, axis=1 (Gemlite内核)
* 3.25比特: nbits=3, group_size=64, axis=1 (FLUTE内核)
对于4比特的实例,我们遵循HQQ文档中的最佳实践使用Gemlite内核【22, 2】;对于3比特的实例,我们使用FLUTE内核报告推理时间【20】。
NormalFloat [11]。我们利用FLUTE的NormalFloat实现,配置与HQQ类似【20, 2】:
* 4.25比特: nbits=4, group_size=64
* 4.02比特: nbits=4, group_size=1024
* 3.25比特: nbits=3, group_size=64
由于公开可用的优化FLUTE内核不支持1024的组大小,我们只报告了4.25比特和3.25比特配置的推理延迟结果。
QTIP [50]。对于无数据QTIP,我们将Hessian矩阵近似为单位矩阵,并遵循与原始数据感知QTIP相同的算法实现。对于数据感知QTIP,我们使用来自relax-ml HuggingFace仓库的公开可用的Hessian近似,该近似是使用6144 × 4096个token计算的。
HIGGS [35]。由于HIGGS的实现未公开,我们直接引用其论文中的结果。具体来说,对于混合方案基准,HIGGS为每个比特宽度仅提供一个结果,我们在比较中直接使用。对于单方案VQ基准,HIGGS为每个比特宽度报告了多种配置;我们选择了实现最低WikiText2困惑度的配置。尽管这些单方案VQ配置由于码本大小非2的幂次方,在实践中可能无法高效实现,但为完整性和比较目的,我们仍将其包括在内。
G.2 附加结果
G.2.1 在其他模型上的内存受限混合方案量化结果
通用性验证。为了证明我们方法的普适性,我们将我们的MSQ方法应用于Qwen 2.5-7B(非LLaMA模型)和LLaMA 3.1-70B(大规模模型),并在不同的比特宽度约束下与HIGGS-MSQ进行比较【55, 18, 35】。在这里,对于LLaMA 3.1-70B,由于其模型尺寸大得多,我们使用64K个token(即E节中解释的默认设置的一半)来估计敏感度系数,同时保持量化流程的其余部分不变。
结果分析。如图5所示,在相同的比特宽度约束(3.25, 4.00, 4.25)下,我们的方法在两个模型上都持续优于HIGGS-MSQ。此外,与HIGGS-MSQ相比,我们的方法在更低的比特宽度下实现了相当或更好的困惑度。对于Qwen 2.5-7B,我们的3.5比特模型与HIGGS-MSQ在4.00比特时的性能相匹配,而我们的3.75比特结果略优于HIGGS-MSQ在4.25比特时的结果,从而实现了高达12.5%的内存节省。在LLaMA 3.1-70B上也观察到类似的趋势,我们的3.40比特和3.63比特结果分别略优于HIGGS-MSQ在4.00和4.25比特时的结果,在困惑度更优的情况下实现了高达15%的内存节省。这些结果证明了我们MSQ方法的广泛适用性。
G.2.2 关于层融合和CUDA-Core内核集成的消融研究
实验设置与结果。表8展示了一项消融研究,评估了层融合和集成CUDA-Core内核对量化后的LLaMA 3.1-8B模型推理吞吐量和WikiText2困惑度的影响。我们比较了Q-Palette中的单方案基准(VQ-2、TCQ-2和TCQ-3)与几种MSQ变体:无层融合的MSQ(仅Tensor Core内核)、仅使用Tensor Core内核的融合感知MSQ,以及结合Tensor Core和CUDA Core内核的融合感知MSQ。结果显示,与单方案量化和无融合的MSQ相比,应用融合感知MSQ在WikiText2困惑度方面有显著改善,突显了联合优化量化方案和层融合的有效性。例如,无融合的MSQ在223 tokens/sec时达到20.33的困惑度,而我们的融合感知MSQ在224 tokens/sec时困惑度显著降低至7.79。此外,将CUDA Core内核与Tensor Core内核一起集成,可带来进一步的性能提升。
G.2.3 数据感知MSQ中损失项选择的影响
实验对比。表9提供了一项消融研究,比较了在LLaMA 2-7B的数据感知MSQ量化中使用的两种不同损失项定义(‘线性’与‘实际’)。‘线性’损失项通过使用线性定理测量的敏感度系数来有效近似困惑度损失的增加,从而可以在多个量化配置中重复使用,类似于无数据场景(见E节)。相比之下,‘实际’损失项明确计算由量化引起的经验(验证)困惑度增加。我们的结果表明,使用计算效率高的‘线性’损失项所实现的零样本准确率提升与使用‘实际’损失项获得的结果相当,这表明更简单且可重用的基于线性的方法在实践中也是有效的。此外,两种损失项方法都实现了相似的推理吞吐量,进一步证实了计算效率高的‘线性’损失项的实用性。
G.2.4 MSQ与线性定理代理在无不连贯处理情况下的适用性
研究问题。一个自然的问题是,当由于硬件限制而无法使用不连贯处理(IP)时,基于线性定理代理的MSQ框架【35】是否仍然适用。该代理目标需要敏感度系数和逐层量化误差;在使用IP时,权重近似为高斯分布,允许这些误差通过随机高斯矩阵预先计算,如E节所述。
实验设计与结果。为了研究无IP的情况,通常采用逐组量化来处理权重异常值,我们禁用了IP并使用固定的高斯训练码本进行逐组量化(组大小=64)。我们比较了三种比特分配策略:1)统一比特宽度,2)使用高斯假设误差的MSQ,该策略依赖于缓存的高斯失真估计来计算$err(Q_q; W_l)$并求解公式(1),以及3)使用真实量化误差的MSQ,该策略直接测量逐层失真$err(Q_q; W_l)$并求解公式(1)。
结论。如表10所示,即使在无IP的情况下,MSQ也显著优于统一比特宽度。此外,使用高斯假设误差所达到的困惑度与使用真实量化误差的结果几乎相同(例如,在3.25比特时为7.34对7.33),这初步证明了基于缓存的高斯误差估计在无IP的情况下可能仍然可靠。
H 局限性与未来工作
局限性。我们引入了Q-Palette,一个全面的量化器套件,涵盖了内存占用、推理延迟和量化误差之间的广泛权衡,为不同的部署场景提供了多样的选择。为了展示其有效性,我们将Q-Palette集成到一个MSQ框架中,并验证了其在PTQ设置下实现更优性能-效率权衡的能力。然而,我们的框架是围绕一次性MSQ目标设计的,这些目标依赖于端到端损失的逐层二阶近似,并且主要适用于不涉及重新训练的场景【35】。虽然Q-Palette也可以作为基于重新训练的量化工作流程(如量化感知训练)的基础构建块,这在计算预算和数据更充足的情况下可能更可取,但我们并未在该设置下评估其有效性,因为我们的主要焦点是无数据或轻量校准的PTQ。因此,将Q-Palette扩展到基于重新训练的量化工作流程是未来工作的一个有希望的方向。
另一个局限性在于计算敏感度系数的成本。目前,如E节所述,在无数据设置中评估这些系数需要$O(L)$次的KL散度损失计算,随着模型规模的增长,这可能成为一个瓶颈。尽管这种计算可以以一种易于并行的方式执行,并且得到的系数可以在所有MSQ运行中重复使用,从而构成一次性成本,但在目标内存或延迟集合固定且重用有限的情况下,开销仍可能不可忽略。因此,开发进一步降低此成本的方法是未来研究的一个有趣方向。
未来工作。一个有前景的方向是将Q-Palette扩展到权重-激活值量化。在这项工作中,我们专注于纯权重量化PTQ,这在内存受限、小批量推理设置中特别有效,例如在笔记本电脑或移动设备上,其中内存带宽而非计算是主要瓶颈。然而,在某些硬件加速器上,如高通Hexagon NPU,其原生仅支持整数(例如INT8)GEMM,激活值量化对于发挥其全部性能至关重要。因此,将Q-Palette扩展以支持权重-激活值量化是实现更广泛部署的一个自然方向。一个潜在的方法是采用两阶段方案:(1)首先使用均匀的W8A8量化器将权重 量化到INT8,以实现硬件兼容性;(2)然后应用第二步压缩,使用Q-Palette量化器的变体将INT8权重进一步量化为x比特表示,其码本被限制在INT8网格上。在推理过程中,压缩后的权重被反量化回INT8,然后与INT8量化的激活值一起使用整数GEMM进行处理,从而在减少内存使用的同时与仅支持INT8的硬件兼容。类似的想法在Q-Serve中被引入,它分两阶段量化权重,首先到对称的INT8网格,然后到不对称的INT4【31】。我们认为探索此类扩展是未来工作的一个有趣方向。
💬 评论讨论
欢迎在这里分享您的想法和见解!