1. 首页 > 科技

KVSharer 共享 KV Cache 基于不相似性实现跨层

一、背景

本文中我们介绍一种最新的 KV Cache 共享论文 KVSharer,与之前常见的层内共享不同,KVSharer 主要关注跨层共享,并且是整个层的共享。

对应的论文:[2410.18517] KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing

对应的代码库:​​​ ​​ ​

二、摘要

LLM 推理过程中对 GPU 内存的需求不断增加,而其中 Attention 的 KV Cache 占据了超过 80% 的空间。当前,大多数现有的 KV Cache 压缩算法主要集中在单个 Transformer 层内的压缩,而较少有研究涉及层间压缩(PS:其实我们之前介绍过的 Character.AI 方案和 MixAttention 都是层间共享)。

本文中,作者提供了一种即插即用的方案,称为 KVSharer,通过在层间共享 KV Cache 来实现层间压缩。作者发现一个反直觉的现象:共享相似性较低的 KV Cache 更能保持模型性能(PS:这里实现时其实有个约束,要同时避免对 LM head 之前的 Hidden State 的影响过大)。实验表明,KVSharer 能够在减少 30% 的 KV Cache 计算,以及内存开销,而对模型性能影响不大,同时还能实现至少 1.3x 的生成加速。此外,作者实验表明,KVSharer 与现有的层内 KV Cache 压缩方案兼容,两者结合可进一步节约内存。

三、方案

3.1 概述

如下图 Figure 2 所示为本文方案的示例,其主要包含几个步骤:

3.2 搜索策略

如下图 Algorithm 1 为详细的搜索算法,其大体包含两个阶段:

KV Cache 相似性计算和初始化(1-4):

共享策略搜索(5-18):

将当前候选 r 添加到共享策略 Z。

替换候选 共享策略 Z 中的所有相关层 (保留靠近输入的层,靠近输出的层共享靠近输入的层),然后使用校准集验证最后一层 Hidden Stage 和未替换时的相似度。

如果相似度小于阈值 T,表示替换后影响较大,则从 Z 中删除当前后续 r。

如果相似度大于等于阈值 T,则可以作为候选。

如果候选集容量满足 C,则直接结束。

PS:需要说明的是,以上是基于贪心搜索策略,这种逐步判断并添加的方式往往只能找到局部最优解,并不一定是全局最优的。要想找到全局最优,通常需要进行全局搜索,比如穷举所有可能的共享组合,或者产生更复杂的搜索算法,但是往往会大幅增加计算成本,因此这种贪心搜索是计算效率和效果的折衷。

3.3 推理

如下图 Figure 3 所示,在推理阶段直接使用离线阶段搜索出来的共享策略,可以看出其实是会存在多个 Layer 共享一个 Layer 的。但是图中似乎有问题,论文中并没有讨论共享的传递性,比如Layer 2 和 3 共享 Layer 1,而 Layer 4 共享 Layer 3,Layer 5 共享 Layer 4,那岂不是 Layer 2,3,4,5 都共享 Layer 1?

如下图所示为代码中的注释(KVSharer/llama_real_share/modeling_llama_kvsharer.py#L364-L371):

四、实验&结果

4.1 精度

如下图 Table 1 所示,作者测试几个常见模型 LLaMA2-7B/13B(PS:如果有 LLaMA3 的结果更有说服力)、InternLM2-7B/20B(中英文能力不错)和 Mistral-7B 在一些主要任务上的精度。其中 Layer 表示实际计算的层数,Layer 越小,表示共享的越多。Percent 表示共享后的平均精度相比原始平均精度的比例。可以看出,共享越多效果越差,并且各个模型表现各不相同,比如在 LLaMA2-13B(40->30) 和 Mistral-7B(32->24) 在共享 25% 时,平均精度就不到 90% 了。

其实,从困惑度(PPL)上来看,影响也是挺大的,尤其是与 H2O 和 PyramidInfer 这些层内压缩结合后:

4.2 推理

如下图 Table 2 所示,作者以 LLaMA2-13B-Chat 模型为例,对比了不同方案对 Memory 的开销以及推理的加速比。其中的 KVSharer(25%)表示 25% 的压缩比,也就意味着 40 层只计算 30 层的 KV Cache。(PS:这里的结论很奇怪,25% 的压缩比理论上最多节约 25% 内存,最多吞吐提升 25%?更何况 Q 和 Attention 还要正常计算,而实际上 1024+4096 时节约了 36% 的内存,吞吐提升 1.53x。)

4.3 消融实验

如下图 Figure 6 所示,作者实验表明:共享相似性较低的 KV Cache 更能保持模型性能。

如下图 Table 3 所示,使用 Wikipedia 或 BookCorpus 作为校准集对精度的影响不是特别大:

如下图 Table 4 所示,随机共享相比本文的 KVSharer 会导致精度下降比较多,证明本文方法的有效性:

如下图 Table 5 所示,KVSharer 在几个模型的 Base 模型和 Chat 模型上的精度影响比较类似,证明了方法的通用性。(PS:不过 PPL 确实影响挺大的)

五、参考链接

本文转载自​​,作者:

本网站的文章部分内容可能来源于网络和网友发布,仅供大家学习与参考,如有侵权,请联系站长进行删除处理,不代表本网站立场,转载者并注明出处:https://www.jmbhsh.com/keji/32686.html

联系我们

QQ号:***

微信号:***

工作日:9:30-18:30,节假日休息