论文阅读:《Sequence can Secretly Tell You What to Discard》,减少推理阶段的 kv cache

目前各类大模型都支持长文本,例如 kimi chat 以及 gemini pro,都支持 100K 以及更高的上下文长度。但越长的上下文,在推理过程中需要存储的 kv cache 也越多。假设,数据的批次用 b 表示,输入序列的长度仍然用 s 表示,输出序列的长度用 n 表示,隐藏层维度用 h 表示,层数用 l 表示。kv cache 的峰值显存占用大小 = b ∗ ( s + n ) ∗ h ∗ l ∗ 2 ∗ 2 = 4 b l h ( s + n ) b * (s + n) * h * l * 2 * 2 = 4blh(s + n) b(s+n)hl22=4blh(s+n),这里的第一个 2 表示 k 和 v cache,第二个 2 表示 float16 数据格式存储 kv cache,每个元素占 2 bytes。

然而,目前的大多数 LLM 会使用 GQA 而非 MHA,因此 kv cache 的占用量会更少,以 transformers 的 modeling_llama.py 脚本中的实现为例:

class LlamaAttention(nn.Module):def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None):super().__init__()# ...self.hidden_size = config.hidden_sizeself.num_heads = config.num_attention_headsself.head_dim = self.hidden_size // self.num_headsself.num_key_value_heads = config.num_key_value_headsself.num_key_value_groups = self.num_heads // self.num_key_value_heads# ...def forward(#...) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:# ...key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)# ...

其中,q_len = s + n,bsz = b,self.hidden_size = h,然而,self.num_key_value_heads 会小于 self.num_heads,以 Llama3-8B 为例:

"hidden_size": 4096,
"num_attention_heads": 32,
"num_hidden_layers": 32,
"num_key_value_heads": 8,

k 和 v 的注意力头是 q 的 1/4,因此 kv cache 的峰值显存占用大小还可以继续除以 4,在这里暂时表示为 b l h ( s + n ) blh(s + n) blh(s+n)(注意,不同模型的比例不同,需要根据情况调整计算公式)。

示例:我们继续以 Llama3-8B 为例,来计算不同长度时的 kv cache 显存占用。令 b = 1,n = 32。

  • s = 512: 32 × 4096 × ( 512 + 32 ) = 71 , 303 , 168 ≈ 0.066 G B 32 \times 4096 \times (512 + 32) = 71,303,168 \approx 0.066GB 32×4096×(512+32)=71,303,1680.066GB
  • s = 16,384: 32 × 4096 × ( 1024 + 32 ) = 71 , 303 , 168 ≈ 2.004 G B 32 \times 4096 \times (1024+ 32) = 71,303,168 \approx 2.004GB 32×4096×(1024+32)=71,303,1682.004GB
  • s = 327,680: 32 × 4096 × ( 1024 + 32 ) = 71 , 303 , 168 ≈ 40.004 G B 32 \times 4096 \times (1024+ 32) = 71,303,168 \approx 40.004GB 32×4096×(1024+32)=71,303,16840.004GB

可以看到,随着 context 长度的增加,kv cache 的显存占用量也随之呈线性增长,成为推理的主要瓶颈。在论文《Sequence can Secretly Tell You What to Discard》中,作者介绍了一种优化 KV 缓存的新方法,它能显著减少 KV 缓存的内存占用。通过综合研究,发现在 LLaMA2 系列模型上:

  • 相邻 token 的 query 向量之间的相似度非常高;
  • 当前 query 的注意力计算可以完全依赖于一小部分前面 query 的注意力信息。

基于这些观察结果,作者提出了一种 KV 缓存驱逐策略 CORM,它能在不对模型进行微调的情况下动态保留重要的 kv 对进行推理。

观察实验与结果

LLMs 中的注意力稀疏性

首先探讨 LLM 注意力层的稀疏性,这是减少 KV 缓存大小的有效前提和依据。具体来说,用重要 key 的比例来表示注意力稀疏性。让 q t ∈ R 1 × d q_t \in \R^{1 \times d} qtR1×d 表示第 t 步的 query state 向量, k i ∈ R 1 × d k_i \in \R^{1 \times d} kiR1×d 表示第 i ( 1 < i < t) 步的 key state 向量,其中 d 是隐藏层维度(为简单起见,这里只考虑单个注意力头)。 q t q_t qt k i k_i ki 的归一化注意力分数的计算公式:
α t , i = e x p ( q t k i T / d ) ∑ j = 1 t e x p ( q t k j T / d ) . (2) \alpha_{t, i} = \frac{exp(q_t k_i^T / \sqrt{d})}{\sum_{j=1}^t exp(q_t k_j^T / \sqrt{d})}. \tag{2} αt,i=j=1texp(qtkjT/d )exp(qtkiT/d ).(2)

定义 1(重要 key):当且仅当 α t , i ≥ 1 t \alpha_{t, i} \geq \frac{1}{t} αt,it1,key k i k_i ki 在步骤 t 中被视为重要 key,否则被视为 minor(次要)key。

作者在 PG-19 测试集上使用 LLaMA2-7B 模型进行 zero-shot 推断。绘制了注意力区块内的 layer-wise 稀疏度和head-wise 稀疏度,结果如图 1 所示。

图 1:LLaMA2-7B 的注意力稀疏度。(a) layer-wise 注意力稀疏度。(b) 第 0 层和第 1 层的 head-wise 注意力稀疏度。

结果显示,底层相对密集,而其他层高度稀疏,稀疏度超过 90%。这说明在生成过程中只需要使用一小部分 kv cache 就可以尽量维持原始的生成结果。

相似的 query 对 key 有相似的关注

考虑 token 长度为 T 的序列(i < j ≤ T)中的第 i 和第 j 个 query state 向量 q i q_i qi q j q_j qj。它们的余弦相似度可计算为:

c o s i n e _ s i m i l a r i t y ( q i , q j ) = q i q j T ∣ ∣ q i ∣ ∣ ⋅ ∣ ∣ q j ∣ ∣ . (3) cosine\_similarity(q_i, q_j) = \frac{q_iq_j^T}{||q_i|| \cdot ||q_j||}. \tag{3} cosine_similarity(qi,qj)=∣∣qi∣∣∣∣qj∣∣qiqjT.(3)

考虑第 i 个 key 之前的所有 key states k 1 , k 2 , … , k i − 1 k_1, k_2, \ldots, k_{i-1} k1,k2,,ki1。假设 c o n s i n e _ s i m i l a r i t y ( q i , q j ) = 1 consine\_similarity(q_i, q_j) = 1 consine_similarity(qi,qj)=1,则 q i = m ⋅ q j , m ∈ R + q_i = m \cdot q_j, m \in \R^+ qi=mqj,mR+ q i q_i qi 对前 i - 1 个 key 的关注权重(attention weight)可以表示为:

a t t e n t i o n _ w e i g h t = 1 d ( q i k 1 T , q i k 2 T , … , q i k i − 1 T ) = m d ⋅ ( q j k 1 T , q j k 2 T , … , q j k i − 1 T ) ⋅ (4) attention\_weight = \frac{1}{\sqrt{d}}(q_ik_1^T, q_ik_2^T, \ldots, q_ik_{i-1}^T) = \frac{m}{\sqrt{d}} \cdot (q_jk_1^T, q_jk_2^T, \ldots, q_jk_{i-1}^T)\cdot \tag{4} attention_weight=d 1(qik1T,qik2T,,qiki1T)=d m(qjk1T,qjk2T,,qjki1T)(4)

注意,m 是一个正数,不会影响注意力权重的相对顺序。例如,对于 q i q_i qi 而言,如果 q i k 1 T ≥ q i k 2 T q_ik_1^T \ge q_ik_2^T qik1Tqik2T,那么 q j q_j qj 一定是 q j k 1 T ≥ q j k 2 T q_jk_1^T \ge q_jk_2^T qjk1Tqjk2T。这意味着,如果某个 key 对 q i q_i qi 很重要,那么它对 q j q_j qj 也很重要,尽管重要程度可能会因 softmax 函数而不同。

虽然在实际情况中, c o s i n e _ s i m i l a r i t y ( q i , q j ) = 1 cosine\_similarity(q_i, q_j) = 1 cosine_similarity(qi,qj)=1 几乎是不可能的,但可以提出这样的假设:两个相似的 query 可能对 key 有相似的关注。为了验证这一假设,使用 LLaMA2-7B 提供了两个从 PG-19 中随机抽取的句子的注意力图,如图 2 所示。

图 2:相似的 query 对 key 有相似的关注。作者绘制了一个句子中两个不同层的注意力图。将注意力分数离散化,重要的 key 显示为亮绿色。每个注意力图都有两条红色边框,底边显示当前 query 实际关注的重要 key,另一条边框显示最相似 query 关注的重要 key。

重要 key 以亮绿色标出,观察到,假设是正确的,相似的 query 对重要 key 的关注也是相似的。同时,重要 key 只占很小一部分,尤其是在较深(离输出近的层,下图)的注意力层中,这与上一节中发现的较深层较稀疏的情况是一致的。

另外,下图中序列开头位置处的 key 都是亮绿色,表示这些 key 是重要的 key,这与先前的一些研究,例如 LM-infinite 和 Attention Sink 有着相同的结论。

问题:什么是 Attention Sink?
回答:在标准的注意力机制中,模型会计算输入序列中每个元素(如单词或token)对其他元素的重要性,并据此分配不同的注意力权重。然而,在某些情况下,模型可能会对输入序列中的某些特定元素赋予异常高的注意力权重,这种现象有时被称为 “Attention Sink” —— 摘自 kimi chat 的回答。

问题:LM-infinite 提出了什么观点?
回答:开头的⼏个 Token 是绝对位置的“锚点”:顾名思义,相对位置编码原则上只能识别相对位置,但有些任务可能⽐较依赖绝对位置,通过开头⼏个绝对位置约等于 0 的 Token 作为“锚点”,每个 Token 就能够测出⾃⼰的绝对位置,⽽去掉开头⼏个 Token 后则缺失了这⼀环,从⽽完全打乱了注意⼒模式,导致 PPL 爆炸 —— 摘自前同事龙哥的调研结果;

query 向量的相似性探索

上一节已经验证了两个相似的 query 对 key 有相似的关注,因此还需要验证,在每一步中,是否能找到与当前 query state 足够相似的同层同头的前一个 query state。为了验证这一点,作者对同一序列中 query 向量的余弦相似度进行了可视化处理,如图 3 所示。

图 3:一个长度为 1024 的句子的 query 向量余弦相似度可视化图。图中第 i 行表示第 i 个 query 与之前所有 query 的余弦相似度。该图显示,在大多数情况下,当前 query 与最近的 query 最为相似。

观察到一个有趣的现象,许多图像都显示出明显的斜向颜色分割,最上面的斜向块最接近暗红色,这意味着当前 query 与最近的 query 最为相似。

CORM 方法

在本节中,介绍 CORM,一种无需任何微调过程就能根据最近的 query 注意力信息来减少 KV Cache 的方法。

直观地说,可以直接存储所有 query 及其注意力信息,以备将来参考。在每个生成步骤中,使用当前 query 从之前 queries 中找出最相似的 query’,并使用其保存的注意力信息单独计算重要的 key。然而,这种方法会产生很大的成本:

  • 首先,存储所有 query 会大幅增加内存开销。
  • 其次,每一步都需要进行相似性计算,这也增加了计算开销。

由于大多数情况下当前 query 与最近 query 最为相似,因此可以只使用最近 query 的注意力信息。从图 2 中还可以观察到,只有一小部分 key 被最近的 query 认为是重要的。因此,即使保存了所有在前几步中被认为重要的 key,也能节省大量内存。

定义 2(Long-term Minor Key):只有在最近的所有 r 步(从 t - r + 1 到 t)中都被认为是 minor key 的情况下, k i k_i ki才被认为是 long-term minor key。

CORM 将有一个大小为 w 的最近窗口来记录最近 w 次的 query 信息,并将始终保持最近 r 个 key 不被删除,以防止因观察不足而过早丢弃。在生成过程中,一旦 k i k_i ki 被视为 long-term minor key, k i , v i k_i, v_i ki,vi 将被丢弃。直观地说,w 越大,保存的 key 和 value 越多,压缩率越小,性能越好;反之,w 越小,保存的 key 和 value 越少,压缩率越大,性能越差。因此,性能和压缩率之间需要权衡。


判断某一位置的 key 是否与未来(滑动窗口大小)query 的注意力分数都低于定义 1 的阈值,如果是则说明该 key 是个 minor key,它对于生成提供的信息有限,那么它就可以被“驱逐”。被“驱逐”的 key(value 也一同被“驱逐”)越多,需要参与到注意力计算的 key 也就越少,需要缓存的 kv 对也就越少。

实验结果

对 4K 长文本的 LLaMA2-7B-Chat 进行了评估。表 1 和表 2 总结了 LLaMA2-7B-Chat 的结果。

表 1: 单文档 QA、多文档 QA 和摘要任务的结果(%)。“Full”指使用完整 KV Cache 的 LLaMA2-7B-Chat,“StreamLLM”配置为 4+1020,“Scissorhands”配置为 768+256(窗口大小 = 256),“H2O”配置为 768+256,“CORM”配置为 256+256,以便进行公平比较。为简洁起见,在此使用 ID 表示数据集,ID 与数据集的映射关系见附录 B。

表 2: few-shot learning、合成和代码任务的结果(%)。“Overall”按主要任务类别的宏观平均值计算。这是按英文(EN)任务、中文(ZH)任务和所有(All)任务计算的,两种语言都包括代码任务。

可以得出以下结论:

  • 在各种任务中,CORM 在相同的压缩率下始终优于以前的方法。
  • 在减少 70% 以上 KV Cache 的情况下,CORM 的性能与使用全 KV Cache 的模型相当,甚至在某些任务上超过了它。作者推测这是因为全 KV Cache 中存在一些影响模型输出的噪声,而 CORM 通过丢弃部分 KV Cache 可以在一定程度上消除这些噪声。

看实验的结果很棒,但是否存在某个 key 与未来一段时间(滑动窗口)内的 query 的相似度都很低,在这之后才出现较高的相似度?而此时,如果将该 key 认为是不重要而进行丢弃的话,未来的 query 可能获取不到这部分信息。

看到这篇论文,联想起去年的 EMNLP 的最佳论文《Label Words are Anchors:An Information Flow Perspective for Understanding In-Context Learning》,大家有兴趣的可以去阅读下。同时,似乎也能和 prompt 压缩相关的方法产生关联(具体如何联系,我得重温下相关的论文,看看它们是如何过滤 token),例如 selective context 和微软的 LLMLingua。它们更多地是减少 prompt 的长度,降低调用 API 的费用,在过滤低信息的 token 时,也会改变保留后 token 的相对位置。在实际使用过程中,性能会有所损失。

这篇论文讲述推理阶段,最近另一篇论文《RHO-1:Not All Tokens Are What You Need》,则从训练角度论证了“并非语料库中的所有 token 对语言模型训练都同样重要”,初步分析深入研究了语言模型的 token 级训练动态,揭示了不同 token 的不同损失模式。利用这些见解,推出了一种名为 RHO-1 的新语言模型,采用了选择性语言建模 (SLM),即有选择地对符合预期分布的有用 token 进行训练。这与最近的数据子集挑选方法(论文阅读:A Survey on Data Selection for LLM Instruction Tuning,以对话、文本为粒度挑选高质量训练子集)和 selective context 也能关联上。

内存开销分析

为了减少 KV Cache 的内存开销,最近信息缓存引入了额外的内存开销。我们需要存储最近的 query 信息,这增加了内存开销。不过,这些开销远小于压缩 KV Cache,我们可以使用一小部分内存来避免维持完整的 KV Cache 内存,而不会明显降低性能。另一方面,如图 4 所示,压缩率会随着序列长度的增加而增加,因此相比之下,该组件的内存开销较低。

图 4:压缩率与序列长度之间的关系。从图中可以看出,LLaMA2-7B-Chat 在 CORM 为“256+256”和预算 = 1024 时的压缩率比较接近。

总结

本文研究了 LLM 部署中的一个关键内存瓶颈——KV Cache。受相似 query 对 key 有相似关注以及最近 query 足够相似的启发,作者提出了一种无预算的 KV Cache 驱逐策略 CORM,通过重复使用最近的 query 注意力信息来显著减少显存占用。通过广泛的评估,作者证明 CORM 可以将 KV Cache 的推理显存使用量减少多达 70%,而在各种任务中不会出现明显的性能下降。

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/pingmian/7271.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

【typescript测试 - Jest 配置与使用】

安装 npm install --save-dev types/jestnpm install --save-dev ts-jest配置 tsconfig.json {"compilerOptions": {"types": ["jest"]} }jest.config.js module.exports {preset: ts-jest,testEnvironment: node, };使用 // add.js funct…

C++学习笔记——对仿函数的理解

文章目录 思维导图仿函数出现的逻辑仿函数使用上的巧妙 仿函数的本质仿函数的优势仿函数语法的巧妙 思维导图 仿函数出现的逻辑 我们在学习stack时会遇到一些新的问题&#xff0c;这些问题需要我们使用非类型模板参数去解决&#xff0c;即我们需要在设计类时需要有一个途径去快…

Android硬件加速hardwareAccelerated支持/不支持的绘图接口

Android硬件加速hardwareAccelerated支持/不支持的绘图接口 Android硬件加速也即在Androidmanifest.xml配置开启GPU渲染&#xff1a; <application android:hardwareAccelerated"true" > 配置后&#xff0c;Android将启用GPU渲染&#xff0c;在trace里面看会…

clang:在 Win10 上编译 MIDI 音乐程序(一)

先从 Microsoft C Build Tools - Visual Studio 下载 1.73GB 安装 "Microsoft C Build Tools“ 访问 Swift.org - Download Swift 找到 Windows 10&#xff1a;x86_64 下载 swift-5.10-RELEASE-windows10.exe 大约490MB 建议安装在 D:\Swift\ &#xff0c;安装后大约占…

Docker新建容器 修改运行容器端口

一、修改容器的映射端口 项目需求修改容器的映射端口 二、解决方案 停止需要修改的容器 修改hostconfig.json文件 重启docker 服务 启动修改容器 三、方案 目前正在运行的容器 宿主机的89 端口 映射 容器端口80 3.1测试环境中新建nginx服务 docker run -itd --n…

打破 AI 算力天花板,Meta超大规模AI基础设施架构解读

Meta超大规模AI智算基础设施架构设计 摘要 双重 GPU 集群&#xff0c;每群配备 2.4 万个 H100 芯片&#xff0c;分别采用 RoCE 和 InfiniBand 网络连接。LLaMA3 就是在这两个集群上训练出来的&#xff1b;Meta AI 将部署庞大算力集群&#xff0c;拥有 35 万张 H100 GPU&#x…

【C语言】第一个C程序:hello world

printf简介 printf是C语言提供的库函数&#xff0c;可以在屏幕上打印格式化数据。这里不作展开&#xff0c;只需要知道&#xff0c;如果要打印hello world&#xff0c;就把双引号引起来的"hello world"作为参数传给printf就行了。如果想要在打印后换行&#xff0c;要…

数据分析之Tebleau可视化:树状图、日历图、气泡图

树状图&#xff08;适合子分类比较多的&#xff09; 1.基本树状图的绘制 同时选择产品子分类和销售金额----选择智能推荐----选择树状图 2.双层树状图的绘制 将第二个维度地区拖到产品分类的下面---大的划分区域是上面的维度&#xff08;产品分类&#xff09;&#xff0c;看着…

牛客热题:单链表排序

&#x1f4df;作者主页&#xff1a;慢热的陕西人 &#x1f334;专栏链接&#xff1a;力扣刷题日记 &#x1f4e3;欢迎各位大佬&#x1f44d;点赞&#x1f525;关注&#x1f693;收藏&#xff0c;&#x1f349;留言 文章目录 牛客热题&#xff1a;单链表排序题目链接方法一&…

如何压缩word文档的大小?6个方法教你方便的压缩word文档

如何压缩word文档的大小&#xff1f;6个方法教你方便的压缩word文档 以下是六个常用的软件和方法&#xff0c;可以帮助您方便地压缩Word文档大小&#xff1a; 使用Microsoft Word内置功能&#xff1a; 在Microsoft Word中&#xff0c;您可以使用内置的压缩功能来减小文档的大…

Celery(分布式任务队列)入门学习笔记

Celery 的简单介绍 用 Celery 官方的介绍&#xff1a;它是一个分布式任务队列; 简单&#xff0c;灵活&#xff0c;可靠的处理大量消息的分布式系统; 它专注于实时处理&#xff0c;并支持任务调度。 Celery 如果使用 RabbitMQ 作为消息系统的话&#xff0c;整个应用体系就是下…

25.哀家要长脑子了---哈希表

1.525. 连续数组 - 力扣&#xff08;LeetCode&#xff09; 在我对通义千问的一番折磨下&#xff0c;终于弄清楚一点点了。哈希表存储前缀和数组值 用一个counter来记录nums中0、1数量差值的变化。 哈希表map存储某个特定的counter值首次出现的位置。counter的计算&#xff1a;…

【深度学习】序列模型

深度学习&#xff08;Deep Learning&#xff09;是机器学习的一个分支领域&#xff1a;它是从数据中学习表示的一种新方法&#xff0c;强调从连续的层中进行学习&#xff0c;这些层对应于越来越有意义的表示。 1. 为什么选择序列模型&#xff1f; 循环神经网络&#xff08;RNN…

(二)JSP教程——taglib指令

创建标签文件 首先创建一个Web项目&#xff0c;在webapp/WEB-INF目录下创建一个tags文件夹 在tags文件夹中创建一个oddNumberSum.tag文件&#xff0c;Tag文件时扩展名为.tag的文本文件&#xff0c;其结构和JSP文件非常相似&#xff0c;该文件的目录结构如图所示 创建Tag文件的…

Linux基础之makefile/make

目录 一、背景 二、makefile和make的讲解 2.1 使用方法 2.2 伪目标文件 2.3 文件的属性以及属性的更新 2.4 makefile的自动推导 一、背景 这里会提及为什么要使用makefile和make&#xff0c;以及他们是什么和作用。 会不会写makefile&#xff0c;从一个侧面说明了一个人是…

怎么口语外教一对一课程?这篇文章告诉你答案!

怎么口语外教一对一课程&#xff1f;在当今全球化的时代&#xff0c;英语口语能力已经成为许多人追求的重要技能。为了满足这一需求&#xff0c;市场上涌现出了许多提供一对一口语外教课程的软件。这些软件不仅提供了与母语为英语的外教进行实时交流的机会&#xff0c;还通过互…

遭遇.halo勒索病毒怎么办?如何识别和应对.halo勒索病毒

导言&#xff1a; 近年来&#xff0c;网络安全问题愈发严峻&#xff0c;其中勒索病毒成为了威胁企业和个人数据安全的重要隐患。在2023年初&#xff0c;一种新的勒索病毒——.halo勒索病毒开始在网络上肆虐&#xff0c;给广大用户带来了极大的困扰。本文91数据恢复将对.halo勒…

c3 笔记6 认识css样式表

<link>与import应该如何选择?事实上&#xff0c;使用link与import链接外部样式文件的效果看起来是一样的&#xff0c;区别在于<link>是HTML标记而import属于CSS语法。<link>标记有rel、type与href属性&#xff0c;可以指定CSS样式表的名称&#xff0c;这样就…

控制台调试 hover 后才出现的元素

调试 hover后才出现的元素 打开开发者工具&#xff0c;鼠标放在hover时才出现的元素上&#xff0c;然后点击右键&#xff1b; 不要选中任何选项&#xff0c;将鼠标移动到开发者工具的调试面板中&#xff1b; 按下N键&#xff0c;此时悬浮的元素不会消失&#xff0c;定位成功。…

知到java笔记(4.1--继承的用法以及this和super的用法)

格式&#xff1a; 例子&#xff1a; get set获取父类的私有变量 private属性 this和super区别&#xff1a; this用法 super用法 例子