提高大模型计算速度除了使用推理加速的技术,让模型“算得更快”,还可以使用模型压缩的技术,让模型“算得更少”。本文参照https://zhuanlan.zhihu.com/p/667455383这篇博客分享的内容对模型压缩技术进行一个初步的介绍。
模型压缩技术分类
-
量化(Quantization):使用低精度( $\leq$ 16位)存储模型权重
-
精简Attention:使用一些变种的Attention算法来减少模型计算量
-
投机采样:使用小模型辅助大模型进行推理
模型推理阶段的显存分配
在推理阶段主要有三部分数据会放到显存里
-
KV Cache:前序token序列计算的 $k,v$ 结果,会随着后面token推理过程逐步存到显存里。存储的量随着batch和sequence length长度动态变化
-
模型参数:包括Transformer、Embedding等模型参数会存到显存里。模型大小固定后,这个存储空间是固定的。
- 包括FFN、Norm层、dropout、embeding参数等
-
运行时中间数据:推理过程中产生出的一些中间数据会临时存到显存,即用即释放,一般占用空间比较小。
-
包括在单token推理前向计算时,通过每层产出的计算结果。
-
输入的batch数据也要存到GPU显存上
-
一个KV Cache的计算过程
以Qwen-72B为例,忽略模型GQA的设置,以朴素的MHA方式计算。
模型参数名 | 模型层数 | 注意力头数 | 注意力头向量维度 |
---|---|---|---|
大小 | 80层 | 每层64个头 | 每个头128维 |
$$num_{kv} = 2 * l * n_h = 2 \times80 \times64 = 10240$$
其中,2表示1个 $k$ 和1个 $v$ 。也就是说一个token就需要缓存10240个 $k,v$ 。这些 $k,v$ 如果按照FP32进行计算存储,每个参数占4个字节,那么一个token的参数占用为:
$$token\_mem_{kv} = 2 * num_{kv} * d_{h} = 4 \times 10240 \times 128= 5.24(MB)$$
这些 $k,v$ 如果按照更广泛使用的半精度(bf16)参数进行计算存储,每个参数占2个字节,那么一个token的参数占用为:
$$token\_mem_{kv} = 2 * num_{kv} * d_{h} = 2 \times 10240 \times 128= 2.62(MB)$$
这样也就确定了一个token计算需要的缓存的 $k,v$ 数量和存储量。那么对于一个实际的推理场景,还要考虑批量batch(B)和序列长度sequence length(S)两个维度。以下面的单条文本推理和并发推理场景为例:
单条短文本场景
参数设置 | 参数精度 | batch大小 | 序列长度 |
---|---|---|---|
参数大小 | bf16 | 1 | 2048 |
此时KV Cache的大小计算:
$$mem_{kv} =token\_mem_{kv} * B * S = 2.62(MB) \times 1 \times 2048 = 5.366GB$$
并发长文本场景
参数设置 | 参数精度 | batch大小 | 序列长度 |
---|---|---|---|
参数大小 | bf16 | 32 | 4096 |
此时KV Cache的大小计算:
$$mem_{kv} = token\_mem_{kv} * B * S = 2.62(MB) \times 32 \times 4096 = 343.4GB$$
除了KV Cache需要占用显存之外,模型参数也需要占用存储,推理阶段模型参数占用的存储空间是固定的。假设模型参数量为 $\Phi$ ,以bf16半精度做推理,则参数量约为 $2\Phi$ (字节)。以qwen-72B为例,参数占用存储空间约为:
$$mem_p = 2 * \Phi = 2 \times 72 = 144G$$
重新审视上述两个场景的现存分配:
-
对于场景1:模型存储 $mem_p = 144G$ ,KV Cache存储 $mem_{kv} = 5.366GB$ ,模型的参数存储占主导,使用80G的A100,至少需要2张卡做推理。
-
对于场景2:模型存储 $mem_p = 144G$ ,KV Cache存储 $mem_{kv} = 343.4GB$ ,KV Cache存储占主导,使用80G的A100,至少需要7张卡做推理。
补充:在实际推理过程中,如何选择batch大小是根据实际测试而来的,尽可能在单卡装下完整模型参数和KV Cache,这时候卡内带宽很高,性能更出众。
减少KV Cache的优化方法
-
量化压缩:通过量化的方法,使用更低的bit存储KV,使得KV结果进一步压缩。
-
共享KV:通过共享Attention参数的方式,让多个注意力头共享一个KV,从而压缩KV的存储。
-
窗口KV:通过设计一个计算KV的窗口,KV Cache只保存窗口内的结果,超出窗口的KV会被丢弃,通过这个方法能减少KV的存储,但是会带来一定的长文推理效果损失。
量化(Quantization)
量化(Quantization)是指使用精度更低的单位来表示模型的权重或中间层变量,以节省空间和加速模型推理速度的方法。
-
根据量化的时机,量化可以分为:
-
训练时量化(Quantization-Aware Training, QAT):需要模型重新训练,例如QLora方法
-
训练后量化(Post Training Quantization, PTQ):可以量化预训练好的模型,不需要重新训练模型,例如AWQ方法、GPTQ方法
-
-
根据量化后的数据类型分类,量化可以分为:
-
浮点数量化(Float Quantization):直接用低精度的浮点数单位表示原来的浮点数值,技术相对简单,目前很多模型加载默认会使用FP16了
-
整数量化(Integer Quantization):整数量化是将原来的浮点数值放缩后使用整数近似表示,把矩阵的浮点数范围映射到整数的范围。
-
-
根据量化的范围上,量化可以分为:
-
只量化权重(Weight Only):只量化模型权重,推理时是INT乘Float
-
权重与激活同时量化(Weight and Activation):这里的激活实际就是每一层的输入,对于矩阵乘法的Y=WX,同时量化W和X,推理时是INT乘INT
-
目前Weight and Activation的Weight 8bit Activition 8bit(W8A8)与FP16水平相当,而Weight Only的Weight 4bit Activition 16bit(W4A16)也与FP16相差无几了。
-
根据量化的粒度,量化可以分为:
-
Tensor粒度(per-tensor):整个矩阵一起量化
-
Token粒度(per-token)和Channe粒度(per-channel):每行/每列单独量化,X的每一行代表一个Token,W的每一列代表一个Channel
-
Group粒度(per-group):上述两种方法的折衷,多行/多列分为一组,每组分别量化
-
LLM.int8()
- 原始简单的RTN(Round To Nearest)方法:
在进行整数量化时,先将原始 $float$ 矩阵 $W$ 除以整数 $INT-N$ 能表示的最大值得到缩放因子 $S$ ,进行缩放,得到量化后的 $int$ 矩阵 $W_{int}$ 。在推理时,首先使用 $W_{int}$ 进行推理,然后再根据 $S$ 放缩回 $float$ 结果。这种方式很简单直接,推理精度下降也相对比较明显。
-
LLM.int8()方法的特点:LLM.int8()并不会加速模型推理,反而会让推理变慢,其主要用途是节省空间。
-
LLM.int8()的核心思想:在模型中,有些特征是很重要的,用FP16单独计算,剩下的量化成INT8计算
根据LLM.int8()的分析:
-
在小于3B的模型上,RTN算法效果还可以,但是大于6B的模型上,RTN算法效果就会变差
-
模型的中间层结果,即Activation,
-
6.7B规模以上的模型中包括大量的离群值特征,一个离群值可能降低所有其他值的量化精度。因此,希望每个张量有多个缩放常数,以便将离群值的影响限制在每个块内。
-
因此,LLM.int8()开发了混合精度分解,其中大量离群值特征维度的小部分(约0.1%)以16位精度表示,而其他99.9%的值以8位表示,从而降低了内存消耗。
具体而言,就是通过离群(outlier)检测,把输入 $X$ 中包含异常值的列挑出来直接做16位的矩阵乘法,剩下的使用int8乘法,最后把它们累加起来。
GPTQ
GPTQ方法是对原始的OBQ(Optimal Brain Quantization)方法的改进。
OBQ方法是最早源于LeCun的OBD算法(Optimal Brain Damage),到后来的OBS(Optimal Brain Surgeon),再到2022年的OBC(Optiaml Brain Compression)方法一路流传下来的模型压缩方法。
-
OBD:这是一种经典的模型剪枝方法,其假设不同参数对目标函数的影响是独立的,以此简化目标函数的泰勒展开,并使用Hessian矩阵计算参数的“贡献度”来进行神经网络剪枝。
-
OBS:该方法在ODB的基础上去除了独立性假设,在抹去一个权重 $w_q$ 时同时计算出一个补偿 $\delta_{F}$ 应用于剩余的权重上,从而使得抹去权重 $w_q$ 增加的误差被抵消。具体而言就是通过计算Hessian矩阵的逆得到每个参数的重要性排序,这样就确定了参数剪枝的次序,同时每次剪枝的时候,其他的参数也进行一次更新从而减少误差。
-
OBQ:该方法将OBS的剪枝思想推广到了模型量化中,可以将剪枝看作一种特殊的量化,常用的量化是把数值近似到一个接近的值,而剪枝实际上可以看做直接把数值近似成0。这个方法本身不错,但是其效率比较低。
具体而言,OBQ方法通过在量化过程中保留对损失函数影响最小的权重,来补偿因量化引起的精度损失。该方法将压缩问题转化为了“在一已知要把 $w_{ij}$ 变为 $Quant(w_{ij})$ 的前提下,如何变化 $W_i$ 的其他值使loss的变化最小”。此时,整个量化过程如下:
- 对 $W$ 的每一行分别量化,在量化第 $q$ 行 $w_q$ 时,根据下面这个公式照出对loss影响最小的列,然后更新 $w_q$ ,其中 $\bf H$ 时loss关于 $W$ 的Hessian矩阵
$$w_q = \mathrm{argmin}_{w_{q}} \frac{(quant(w_q)-w_q)^2}{[\mathbf{H}^{-1}]_{qq}}$$
- 对 $w_q$ 剩下的参数进行更新以补偿量化误差
$$\delta_F = - \frac{w_q-quant(w_q)}{[\mathbf{H}^{-1}]_{qq}} · [\mathbf{H}^{-1}]_{:,q}$$
- 通过下面这个公式,将 $\bf H^{-1}$ 剔除
$$\mathbf{H}^{-1} = (\mathbf{H}^{-1} - \frac{1}{[ \mathbf{H}^{-1}]_{qq}} \mathbf{H}^{-1}_{:,q} H^{-1}_{q,:})$$
重复这三个步骤直到全部矩阵都被量化。其中最后一步更新 $\textbf{H}$ 涉及到矩阵乘法,复杂度较高(为 $O(d_{row}*d^3_{col})$ ),所以这个算法对大模型来说很难实现。
为了提高量化效率,GPTQ被提出来了。
-
GPTQ方法的特点:在OBQ的基础上进行改进,在降低量化算法复杂度的同时保留了模型的精度,常用于大模型的高效量化。
-
GPTQ的核心思想:通过最小化量化引入的输出误差,实现高精度低bit量化。具体而言,GPTQ在后量化过程中,针对每一层权重矩阵,利用一小部分校准数据,最小化量化前后模型输出的差异。
-
具体而言,GPTQ在OBQ的基础上做了如下改进:
-
取消贪心算法,选择固定顺序:OBQ采用了贪心策略,先量化对目标影响最小的参数,这就要对每行进行单独计算,因为每行里各元素的量化顺序不一样。但是GPTQ认为,量化的顺序并不重要,对精度影响不大,因此对每一行都使用相同的顺序去量化。这项改进是的参数矩阵每一行的量化可以做并行的矩阵计算(per-channel quantization)。这个改进让量化的复杂度优化到了 $O(max(d_{row}*d^2_{col},d^3_{col}))$ ,大大提升了计算效率。
-
懒惰批量更新:OBQ中对权重的更新是一个个单独进行的,需要很多内存但不怎么需要算力,这导致性能瓶颈在了GPU的内存带宽,但实际上同一个特征矩阵 $W$ 不同列间的权重更新是不会互相影响的。因此GPTQ采用类似于缓存的策略(延迟批处理),把矩阵分块缓存,在更新一列后只存在缓存中,整个block更新完再写会内存并更新整个 $\mathbf{H}$ 。这种方法缓解了带宽的压力,大幅提升了计算速度。
-
Cholesky分解:OBQ方法受浮点数运算影响较大,在 $H^{-1}$ 的计算时容易引入数值误差。因此GPTQ提出了两种方法,一是往 $H^{-1}$ 里加入一个小的常数项。另一个是观察到其实只需要用到 $\mathbf{H}$ 矩阵上三角形的信息,所以可以用数值稳定的Cholesky分解提前计算好所有需要的信息,避免在更新的过程中再计算。
-
-
以伪代码的形式呈现算法可以看做:
-
初始化量化结果 $\textbf{Q}$ 为一个零矩阵,初始化误差矩阵 $\textbf{E}$ 为一个零矩阵
-
对 $\textbf{H}^{-1}$ 进行Cholesky分解,得到 $\textbf{H}^{-1}$ 的逆矩阵信息
-
迭代处理权重列的量化
-
对每个列进行量化,更新Q矩阵中的对应列
-
计算量化误差,更新误差矩阵 $\textbf{E}$
-
更新权重矩阵 $\textbf{W}$ 中的权重,以减少误差
-
重复上述步骤,直到所有的权重列都被量化
-
最终,返回量化后的权重矩阵 $\textbf{Q}$
总体而言,GPTQ有以下两个优点:
-
INT4量化能够节省接近4倍的内存,这是因为反量化操作发生在算子的计算单元附近,而不是在GPU的全局内存中
-
由于用于权重的位宽较低,因此可以节省数据通信的时间,从而潜在地提升推理速度
在INT4精度下,GPTQ可以与FP16效果相当,而INT3下也只低了5%~10%。推理速度方面,GPTQ可以比FP16快3~5倍。
SmoothQuant
当模型规模更大时,单个token的值变化范围也较大(与LLM.int8()的观察一直,存在大量的离群特征),导致token相较weight难以量化。因此,很多量化方法都关注权重量化,在计算过程中主要就是权重参与,而输入token的激活值仍然是FP16的。但是这种只量化weight不量化Activation的想法,在矩阵计算时就需要使用FP16乘以INT8的矩阵乘法,这种计算效率是肯定比INT乘以INT8的计算效率低的。
SmoothQuant意图解决这种激活值幅度变化大难以量化,而权重幅度变化相对较小的情况,平衡两者的量化难度。SmoothQuant通过数学上等价的逐通道缩放变化,将量化的难度从激活值转移到模型权重上。具体而言,SmoothQuant引入了平滑因子 $s$ 对权重 $W$ 和激活值 $X$ 进行缩放:
$$Y = (X diag(s)^{-1}) · ((diag(s))W) = \hat{X}\hat{W}$$
-
为了减少激活的量化难度,可以令 $s_j = max(|X_j|), j=1,2,...,C_j$ ,即第 $j$ 个channel的最大值。但是这种方法权重的量化难度会变得难以量化,因此需要引入另一个超参转移强度 $\alpha$
-
得到平滑因子 $s_j = max(|X_j|)^{\alpha} / max(|W_j|)^{1-\alpha}$ ,这里的 $\alpha$ 可以根据激活值和权重的量化难易程度进行调整。对于大多数模型来说 $\alpha = 0.5$ ( $\alpha$ 取值较小时,激活值量化难度较大,反之则模型权重量化难度大)
得到这样Smooth变换后的激活值和权重矩阵,可以再采用per-token或per-tensor的量化方式。
SmoothQuant不会影响模型的推理效率,因为:
-
weight的放大可以再离线阶段完成
-
而本层Activation的缩小可以被融合至上一层的计算中,不会增加额外的kernel调用
根据SmoothQuant的论文显示,SmoothQuant后推理的Latency可以降低20%-40%,而且得益于需要的内存变少了,批量推理的吞吐量可以提升数倍。
AWQ
AWQ的核心思想:“权重并不同等重要”。AWQ观察到权重里面有0.1%~1%的权重是非常重要的权重,如果对这些权重不量化,保持FP16的精度,那么会大大减少量化误差。
为了寻找这些重要权重,AWQ做了相关的分析实验,与SmoothQuant里的经验结论一致,权重的分布相对平缓,因此依据权重分布选择重要权重是不合理的。与之相对,选择观测激活值的分布,选择对应输出比较大激活值的权重channel,认为这是一个比较重要的权重,依据此构建了一个按照激活值分布量化的算法,成为Activation-aware Weight Quantization,即激活感知权重量化算法。(这个方法有一定的局限性:保留显著权重为FP16可以提升量化后模型效果,但是这种方法并不硬件友好,会使模型效率降低。因此采用类似SmoothQuant的放缩方式进行处理,从而通过通过放大显著channel来减少其相对量化误差)
具体而言,AWQ的量化通过网格搜索找到每个channel的最佳缩放因子scale:
-
原因:虽然使用更大的缩放比,显著权重的量化效果更好,但是如果是用来非常大的缩放比,非显著权重的channel被迫使用更小的动态范围,从而降低了模型性能。因此需要采用网格搜索的方法得到一个最优的缩放比,使得减少显著权重量化损失的同时也不能增加其他权重的量化损失。
-
方法:为了同时考虑显著权重和非显著权重,AWQ通过最小化某层量化前后的差值来寻找最优缩放因子。
$$s^{*} = \mathop{\arg\min}\limits_{s} \cal{L}(s)$$
$${\cal L}(s) = | Q({\mathbf W}·diag(s))(diag(s)^{-1} · {\mathbf X}) - \mathbf W \mathbf X |$$
其中 $s$ 时per-channel scaling的缩放因子, $s^{-1}· {\mathbf X}$ 可以被融合至前面的算子
- 搜索策略:由于上述方法的量化函数 $Q$ 不可微,所以AWQ定义了搜索空间用于寻找最优的 $s$ ,这个 $s$ 只与激活值大小 $s_{\mathbf{X}}$ 有关,论文通过一个超参数 $\alpha$ 来平衡显著通道权重和非显著通道权重。另外使用weight clipping的方法来减少量化误差。
$$s = {s_{\mathbf{X}}}^{\alpha}$$
$$\alpha^{*} = \mathop{\arg\min}\limits_{s} \cal{L}({s_{\mathbf{X}}}^{\alpha})$$
-
其中 $s_{X}$ 是校准数据集的平均激活值,描述了权重的重要性
-
上式将 $s$ 定义为以 $s_{\mathbf{X}}$ 为底数,超参数 $\alpha$ 为指数的固定形式,这个形式将缩放因子 $\alpha$ 的搜索过程,简化为超参数 $\alpha$ 的搜索过程
-
好处:AWQ这个方法不依赖于任何反向传播或重建,只依赖校准集中activation的大小,与activation的实际分布无关(GPTQ与实际分布有关),因此能很好地保持LLM在不同领域和模式上的泛化能力,而不会过拟合到校准集。
精简Attention
因为Attention的计算占据了模型推理的很大一部分计算资源,所以通过精简一些Attention计算的方法可以有效减少模型的运算量。
-
共享Attention:原本在每一层的Attention计算时,都使用多个注意力头来进行计算,即Multi Head Attention,MHA。可以通过注意力头共享一些参数的方式来减少运算量
-
稀疏Attention:因为Attention机制本身是一个 $O(N^2)$ 的计算,即当前token要与前面所有token都计算注意力。但是并不是所有token都一样主要,所以可以跳过某些不重要的token,从而减少计算量。
共享Attention参数-由MHA的思考
- 和先前模型加速技术中提到的一致,大模型的吞吐量很大程度被内存带宽限制住了,而内存中很大部分存的都是KV Cache。在MHA中每个注意力头都使用一组Q、K和V,从而依据多个头关注上下文的不同方面,但是这种多注意力头的形式也带来了计算量和KV Cache占用的增长。自然可以想到通过在一层的多头Attention之间共享KV参数来减少KV Cache的占用,提升模型吞吐量。
共享Attention参数-MQA (Multi Query Attention)
在MHA中,输入分别经过 $W_Q$ 、 $W_K$ 和 $W_V$ 的变换之后,都切成了 $n$ 份( $n$ 即是头数),同时维度也从 $d_{model}$ 降到了 $d_{head}$ ,此后分别进行attention计算然后再进行拼接。
在MQA中,经过线性变换后,只对 $Q$ 进行切分,而 $K$ 、 $V$ 则直接在线性变换的时候把维度降到了 $d_{head}$ (并不进行切分),然后用切分后的 $Q$ 分别和同一份 $K$ 、 $V$ 进行attention计算,之后把结果拼接起来。
简单来说,就是MHA中,每个注意力头的 $K$ 、 $V$ 是不一样的,而MQA这里,每个注意力头的 $K$ 、 $V$ 是一样的,值是共享的。而其他步骤都和MHA一样。如此需要缓存的 $K$ 、 $V$ 值就从所有头的量变成了一个头的量。
这种方法由于共享了多个头的参数,限制了模型的表达能力,虽然能更好地支持推理加速,但是在效果上比MHA差。
共享Attention参数-GQA (Grouped Query Attention)
因为MQA对效果会有影响,但是MHA占据的缓存很多,那么GQA提出了一个折中的办法,既能减少MQA效果的损失,又相比MHA需要更少的缓存。
在GQA中, $Q$ 和MHA/MQA的做法不变,使用多套共享的 $K$ 、 $V$ ,形成多个group,同一个group内的 $Q$ 共享同一套 $K$ 、 $V$ ,不同group的 $Q$ 所用的 $K$ 、 $V$ 不同。
如此一来,MHA可以看做group为n(头数)的GQA,而MQA可以看做group为1的GQA。
共享Attention参数-从MHA得到MQA和GQA
-
从MHA得到MQA:将MHA中H个头的 $K$ 、 $V$ ,分别做mean pooling得到一个 $K$ 和 $V$ ,用得到的 $K$ 和 $V$ 继续训练
-
从MHA得到MQA:将MHA中H个头的 $K$ 、 $V$ ,分别做mean pooling得到H个 $K$ 和 $V$ ,用得到的 $K$ 和 $V$ 继续训练
共享Attention参数- MLA
多头潜在注意力(MLA)将潜在特征表示纳入注意力机制,从而降低计算复杂度并改善上下文表示。其核心思想是对KV进行压缩后,再送入标准的MHA算法中,用一个更短的K,V向量来进行计算,从而减少KV Cache的大小。
- 计算代表KV Cache的潜在向量
$$[c_{t}^{KV}] = W^{DKV}{\mathbf h_t}$$
其中, $c_t^{KV}$ 是在时间步t计算的键值缓存潜在向量。 $W^{DKV}$ 是一个权重矩阵,用于将隐藏状态 $\mathbf h_t$ 映射到键值缓存空间,这一步可以通过神经网络映射得到。总体而言, $c_t^{KV}$ 相对原来的 $\mathbf h_t$ 要小很多。
- 计算key和value的潜在向量
$$[k_{t,1}^C;k_{t,2}^C;...;k_{t,n_h}^C] = k_{t}^C = W^{UK}c_{t}^{KV}$$
$$[v_{t,1}^C;v_{t,2}^C;...;v_{t,n_h}^C] = v_{t}^C = W^{UK}c_{t}^{KV}$$
$k_t^{C}$ 是键的潜在向量,通过将 $c_{t}^{KV}$ 与权重矩阵 $W^{UK}$ 相乘得到,这一步上采样,通过潜向量特征 $c_{t}^{KV}$ 映射得到较大的 $k_t^C$ 用于后续的注意力计算。 $v_t^{C}$ 的计算同理。
- 计算旋转位置编码
$$k_t^{R}={RoPE}\left(W^{K R} \mathbf{h}_t\right)$$
${RoPE}$ 是一种相对位置编码方式,用于在键向量中引入位置信息。
- 组合潜向量 $k$ 和位置编码 $k$ 得到最终的key向量
$$k_{t, i}=\left[k_{t, i}^C ; k_t^R\right]$$
最终的键向量 $k_{t,i}$ 是通过将内容相关的键向量 $k_{t, i}^C$ 和位置编码 $k_t^R$ 连接起来得到的。
- 计算query的潜在向量
$$c_t^{Q} = W^{DQ}h_t$$
$$[q_{t,1}^{C};q_{t,2}^{C};...;q_{t,n_h}^{C}] = q_t^C=W^{UQ}c_t^Q$$
$$[q_{t,1}^{R};q_{t,2}^{R};...;q_{t,n_h}^{R}] = q_t^R={RoPE}({W^{QR}c_t^Q})$$
$$q_{t,i}=[q_{t,i}^C;q_{t,i}^R]$$
与 $k$ 向量的计算类似,通过潜在向量计算得到后续MHA计算的查询向量 $q$ 。
- 注意力计算
最终的注意力输出 $u_t$ 是通过query $\left(q_{t, i}\right)$ ,key $\left(k_{j, i}\right)$ 和value $\left(v_{j, i}^C\right)$ 结合起来计算得到的:
$$o_{t,i} = \sum_{j=1}^{t} {Softmax}_j(\frac{\mathbf{q}_{t,i}^T \mathbf{k}_{j,i}}{\sqrt{d_h+d_h^R}})\mathbf{v}_{j,i}^C$$
$$\mathbf{u}_t = W^o[\mathbf{o}_{t,1};\mathbf{o}_{t,2}; ...;\mathbf{o}_{t,n_h}]$$
$\mathbf{o}_{t, i}$ 是第 $i$ 个注意力头的输出。
通过上述的计算过程可以看到MLA是如何在推理时减少KV Cache的,实际上就是使用了降维的方法,把 $k,v$ 降维到一个隐空间中,然后在推理时只缓存 $c_t^{KV}$ 即可。这个 $c_t^{KV}$ 会在后面使用 $W^{UK},W^{UV}$ 升维,之后的计算与MHA是一致的。实际上,MLA计算时并不需要显式的计算出 $k_{t,i},v_{t,i}$ ,因为降维矩阵 $W^{UK}$ 和 $W^{UV}$ 可以被 $W^{UQ}$ 和 $W^{O}$ 通过预先计算的方式吸收。以不带位置编码时的 $q_t^Tk_t$ 为例:
$$q_t^Tk_t = (W^{UQ}c_t^Q)^TW^{UK}c_t^{KV} = c_t^Q(W^{UQ})^TW^{UK}c_t^{KV}=c_t^{Q}Wc_t^{KV}$$
可以看到,并不需要计算出 $W^{UK}$ ,可以把 $W^{UK}$ 吸收进 $W^{UQ}$ ,这样就避免了重复计算,并且可以让缓存的 $c_t^{KV}$ 直接参与计算,而不需要显式计算出 $k$ 。如此MLA在减少KV Cache的同时保留了KV Cache本身减少重复计算的功能。(实际操作时, $q$ 的输入也改为了低秩投影形式,这与减少KV Cache无关,主要是为了减少训练期间参数量和相应的梯度所占的显存)
但是这种吸收无法直接应用 ${RoPE}$ 编码:
$$q_tk_i^T = (x_tW_qR_t)(c_iW_kR_i)^T = x_t(W_qR_{t-i}W_k^T)c_i^T$$
此时, $W_qR_{t-i}W_k^T$ 与位置差 $t-i$ 相关,就无法合并为一个固定的投影矩阵了。
因此,MLA将 $QK$ 的位置信息部分和内容信息部分解耦开了,专门生成多头的 $q_t^R$ 和共享的 $k_t^R$ (共享的 $k_t^R$ 指的是每个头的 $k$ 都用这同一个 $k_t^R$ )。然后再通过拼接的方式构造完整的 $Q$ 和 $K$ :
$$q_{t,i}^Tk_{j,i} = [c_t^Q(W^{UQ})^T,(q_t^R)^T]\begin{bmatrix} W^{UK}c_t^{KV} \ k_t^R \end{bmatrix}$$
$$[c_t^Q(W^{UQ})^T,(q_t^R)^T]\begin{bmatrix} W^{UK}c_t^{KV} \ k_t^R \end{bmatrix} = c_t^Q(W^{UQ})^TW^{UK}c_t^{KV} +(q_t^R)^Tk_t^R$$
在此设计下,MLA既减少了KV Cache,也保住了 ${RoPE}$ 的应用,同时还确保了没有重复计算。MLA节约了大量缓存并且性能很好,这种提升可以理解为虽然MLA增大了计算了,但是KV Cache的减少降低了显存和带宽的压力,在token的生成阶段,主要的速度瓶颈是带宽瓶颈和显存瓶颈,因此使用MLA可以明显提高生成的速度。
稀疏Attention-由self-attention的思考
对self-attention而言,需要对序列中的任意两个向量都需要计算相关度。上图的左子图显示了注意力矩阵,右子图显示了关联性,这意味着每个元素都跟序列内所有元素有关联。如果想要节省显存,加速计算,一个思路就是减少相关度计算,即认为每个元素只与序列中的一部分元素相关,这便是稀疏Attention的基本原理。
稀疏Attention-Atrous Self Attention
Atrous Self Attention受“膨胀卷积”(Atrous Convolution)的启发,对相关性进行了约束,强行要求每个元素只跟它相对距离为 $k, 2k, 3k,...$ 的元素关联,其中 $k \gt 1$ 是预先设定的超参数。从上左子图的注意力矩阵看,就是强行要求相对距离不是 $k$ 的倍数的位置的注意力为0(图中白色代表0)。
在这种方法下,实际每个元素都只跟大约 $\frac{n}{k}$ 个元素算相关性,这样运行效率和显存占用都变成了 $O(n^2/k)$ ,也就是说能直接降低到原来的 $\frac{1}{k}$ 。
稀疏Attention-Local Self Attention
在计算机视觉领域,自注意力机制一直被统称为“Non Local”,与之相对的Local Self Attention就是放弃了全局的元素关联,重新引入局部关联。在实际处理上,就是约束每个元素只与前后 $k$ 个元素以及自身有关联。从上左子图的注意力矩阵来看,就是相对距离超过 $k$ 的注意力都直接设为0。实际而言,Local Self Attention和普通卷积很像,都保留了 $2k+1$ 大小的窗口,然后在窗口内进行一些运算。区别只是普通卷积是把窗口展平然后接一个全连接层得到输出,现在是窗口内通过注意力来加权平均得到输出。
在这种方法下,实际每个元素都只跟 $2k+1$ 个元素算相关性,这样一来理想情况下运行效率和显存占用都变成了 $O((2k+1)n) \backsim O(kn)$ ,也就是说随着 $n$ 而线性增长,不过这个方法牺牲了长程关联性。
稀疏Attention-Sparse Self Attention
Sparse Self Attention是对Local Self Attention和Atrous Self Attention的交替使用,两者累积起来,理论上也是能学习到全局关联性的,同时还节省了显存。从上左子图的注意力矩阵来看,就是除了相对距离不超过 $k$ 的、相对距离为 $k, 2k, 3k,...$ 的注意力都直接设为0。如此一来,Attention就具有“局部紧密相关和远程稀疏相关”的特性。
稀疏Attention-Longformer
Longformer针对上图(a)的经典self attention进行改进,提出了Sliding Window Attention(滑窗机制,图b,类似Local Attention)、Dilated Sliding Window(空洞滑窗机制,图c)和Global+sliding window(融合全局信息的滑窗机制,图d)。
-
Sliding Window Attention:该机制设定了一个窗口 $w$ ,它规定序列中的每个token只能看到 $w$ 个token,其左右两侧能看到 $\frac{w}{2}$ 个token,因此它的时间复杂度是 $O(n \times w)$ 。同时,因为transformer模型结构本身是层层叠加的结构,模型高层比低层有着更宽广的感受野,从而有能力去建模融合全部序列信息的全局表示。在第 $l$ 层,感受野的范围将是 $l \times w$
-
Dilated sliding Window:在对一个token进行self attention操作时,普通的Sliding Window Attention只能考虑到长度为 $w$ 的上下文,在不增加计算量的情况下,Longformer使用Dilated Sliding Window在进行Self Attemtion的两个相邻token之间放入大小为 $d$ 的间隙,这样序列中的每个token的感受野范围可扩展到 $d \times w$ 。在第 $m$ 层,感受野的范围将会是 $m \times d \times w$ 。研究表明,在进行Multi-Head Self Attention时,通过将一些头不设置Dilated Sliding Window以让模型聚焦在局部上下文,在某些头上设置Dilated Sliding Window以让模型聚焦在更长的上下文序列,这样可以提高模型表现。
-
Global+silding Window:以上两种Attention机制不能完全适应所有task-specific任务,因此提出了Global+silding Window的机制,这种方法设定某些位置的token能够看见全部的token,同时其他的所有token也能看见这些位置的token,相当于将这些位置的token“暴露”在最外面。这些位置与具体的任务有关,例如对于分类任务,这个带有全局视角的token是"CLS";对于QA任务,这些带有全局视角的token是Question对应的这些token。需要注意的是,计算Global Attention的 $k,q,v$ 时使用和计算Sliding Window的 $k,q,v$ 不同的参数。
稀疏Attention-StreamingLLM
根据StreamLLM的实验分析,对于上图(b)的Window Attention而言,其会缓存最近的 $L$ 个token的KV,但是一旦开头的token的KV被驱逐出Cache,模型推理的表现就会急剧下降。为解决这个问题的一个优化是图(c)的重算方法,让Window Attention重新计算,从每个新token的 $L$ 最近token中重建KV Cache。虽然这个方法在长文本上表现良好,但是由于上下文重新计算中的二次注意力计算,导致时间复杂度为 $O(TL^2)$ ,计算速度相当慢。
经过进一步的分析,发现了如下现象:
-
仅仅减少注意力计算的第一个token,模型推理的PPL值就会急剧上升。这意味着开头的第一个token非常关键。
-
通过分析attention每一层的注意力分布发现:
-
第一层和第二层的注意力图中离当前处理token最近的token收到了更多的attention,即attention矩阵对角线位置的值相对更大。
-
除了网络前面的两层,模型在所有的layer和head都会给开头的几个token更多的attention值。
-
基于上图对于Attention的观察,StreamingLLM提出了“attention sink”的概念来解释了Window Attention失败的原因:输入给LLM推理开头的几个initial tokens是非常特殊的,类似水池的排水口,吞噬了大量的attention。同时intial tokens与被预测token的距离如何,语义信息是什么都不重要,重要的只是它的绝对位置。也就是说前几个位置上的token不管是啥,对维持LLM推理的稳定性都很关键。
补充:attention sink出现的原因(一种解释)
在Attention机制中,Softmax的输出代表了key/query的匹配程度的概率。如果softmax在某个位置的值非常大,那么在反向传播时,这个位置的权重就会大幅度地更新。但是有时候attention并不确定哪个位置更值得关注,不过因为softmax需要所有位置的值的和为1,因此必须给某些位置较大的权重,这就可能导致错误的权重更新,这个错误在后续的过程中也很难被纠正。
对此,Evan Miller发表了一篇博客《Attention Is Off by One》,在博客中他解释了这个softmax带来的问题,并提出了一个改进softmax的方法:通过给softmax的分母加1,让所有位置值可以加和不为1,从而给予Attention可以不对任何位置“表态”的权利。
$$(softmax_1(x))_i = \frac{exp(x_i)}{1+\sum_jexp(x_j)}$$
StreamingLLM在这个观点的基础上进一步认为,模型倾向于将不必要的注意力值转嫁给特定的token,而这些token就是初始token。
基于上述的观察研究,StreamingLLM设计了Window Attention的改进版,其在当前滑动窗口方法基础上,重新引入了一些初始token的kv在注意力计算中使用。由此,StreamingLLM中的KV缓存可以分为两个部分:
-
attention_sink是4个初始token,稳定了注意力的计算。
-
rolling KV Cache保留了最近的token,这个窗口是固定的,下图中为3。
此外,还是用了一些小改动来给attention注入位置信息,从而使得StreamingLLM可以无缝地融入任何使用相对位置编码的自回归语言模型,如RoPE和ALiBi。通过上述办法,StreamingLLM不需要做训练,只用把初始token设置为4就可以获得不错的长输入下的推理表现。
如果解除不训练模型的限制,通过Pre-training LLMs with attention sinks可以获得更好的表现。作者提出了两种方法:
-
指定一个全局可训练的attention sink token,称为“sink token”,它将作为不必要的注意力的存储库,从而把初始token作为attention sink的作用转移到sink token上。
-
用类似Miller提出的SoftMax off-by-one的变体替换传统的Softmax函数,StreamingLLM称为Zero Sink方法
投机采样
投机采样的核心思想:并非所有token的生成都很难,前面几个token的生成相对比较容易,可以使用小模型代劳,从而让小模型先生成一部分token再由大模型验证。
投机采样的特点:投机采样并不会改变模型结构,而是通过让decoding阶段可以并发提高计算效率。
投机采样的流程(假设有一个和大模型近似的小模型):
-
小模型对prompt进行 $n$ 次推理,生成 $n$ 个token,记录所有的logits。
-
将prompt和生成的 $n$ 个token组成新的prompt,一起送到大模型推理1次,得到推理结果的logits。(此时 $n$ 个token可以同时参与计算,计算访存比显著提升)
-
将大模型和小模型的logits做对比,如果发现所有推理结果一致(计算相似度),则保留这些token,重新让小模型进行推理。
-
如果发现第 $k$ 个token不一致,则保留第 $1...k$ 个token,大模型重新推理第 $k$ 个token,然后再重新让小模型开始新一轮推理。
这个投机采样的算法流程虽然看起来像是一个近似算法,但是它实际上是数学完备的,可以保证最终输出的结果和直接使用大模型推理的结果严格一致。本质上来说投机采样是利用了大模型推理n个token需要推理n次,而验证结果只需要推理1次的思想。
一种理解方式:投机采样采用了一个原始目标模型和一个比原始模型小的多的近似模型,近似模型用于进行自回归串行采样,而原始目标模型用于评估采样结果。在解码过程中,某些token的解码相对容易,可以交给小模型处理,而有些token的解码很困难,交给大模型处理。这里的小模型可以采用与原始模型相同的结构,但参数更少,从而让计算量更小,而且减少了内存访问的需求。
参考笔记
LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale论文解读https://fancyerii.github.io/2024/01/16/int8/
模型量化 Quantization https://banxian-w.com/article/2024/9/11/2772.html#4.2%20GPTQ%20-%204%20%E4%BD%8D%E9%9D%9E%E5%AF%B9%E7%A7%B0%20PTQ
GPTQ: 模型量化,穷鬼救星 https://zhuanlan.zhihu.com/p/616969812
大模型量化技术原理-LLM.int8()、GPTQ https://blog.csdn.net/scgaliguodong123_/article/details/136176382
LLM(十一):大语言模型的模型量化(INT8/INT4)技术 https://zhuanlan.zhihu.com/p/627436535
大语言模型推理加速技术:模型压缩篇 https://zhuanlan.zhihu.com/p/667455383
大语言模型量化相关技术 https://zhuanlan.zhihu.com/p/664054739
理解Attention:从起源到MHA,MQA和GQA https://zhuanlan.zhihu.com/p/686149289
Attention进阶史(MHA,MQA,GQA,MLA)https://www.gnn.club/?p=2729
一文看懂什么是Longformer https://zhuanlan.zhihu.com/p/405317918
LLM推理技术之StreamingLLM:如何拥有无限生成能力 https://zhuanlan.zhihu.com/p/659875511
缓存与效果的极限拉扯:从MHA、MQA、GQA到MLA https://kexue.fm/archives/10091