吞吐量提升近30倍！田渊栋团队最新论文解决大模型部署难题|内存|算法|插件功能-欧洲杯足彩官网

新智元报道

编辑：alan

【新智元导读】大语言模型在实际部署中，存在内存和输入长度限制的问题。最近，田渊栋团队一举解决这两大难题，将推理系统的吞吐量提高了近30倍。

大型语言模型 （llm） 在今年可谓是风光无限。不过惊艳的效果背后是一个巨大的模型以及夸张的硬件资源。

llm在现实中部署时通常会面临两个难题：昂贵的kv缓存成本，以及对长序列的泛化能力差。

近日，田渊栋团队发表了一篇论文，成功解决以上两个难题，并将推理系统的吞吐量提高了近30倍！

论文地址：https://arxiv.org/pdf/2306.14048.pdf

代码地址：https://github.com/fminference/h2o

这个成果也将在neurips'23上展示。

下面，我们来看一下这两个难题的具体情况，以及论文提供的星空体育网站入口官网手机版的解决方案。

首先是缓存，kv缓存用于存储生成过程中的中间注意力键和值，以避免重新计算。

通常，除了模型参数外，还会将大量瞬态信息（kv缓存）存储在gpu内存中，这部分的内存占用，与序列长度和批处理大小线性相关。

例如，一个输入批次大小为128、序列长度为1024的300亿参数模型需要180gb的kv缓存。

其次，由于硬件限制，llm会以固定的序列长度进行预训练（例如llama-2使用固定长度4k的序列）。

然而，这其实也对推理过程中的注意力窗口施加了限制，使得模型在面对更长输入序列时无法发挥作用，阻碍了更广泛的应用。

对此，论文提出了一种实现kv缓存的新方法，显著减少了内存占用，且在长输入序列的任务中表现良好。

方法基于这样一个事实：在计算注意力分数时，一小部分tokens贡献了大部分的价值，——这里称这些tokens为heavy hitters (h2)。

通过综合调查，作者发现h2的出现是自然的，且与文本中词组的频繁共现密切相关，而去除它们会导致显著的性能下降。

基于此，作者提出了heavy hitter oracle（ h2o ），一种kv缓存逐出策略，可动态保持最近的tokens和h2 tokens的平衡。

另外，作者将kv缓存驱逐表述为一个动态的子模块问题，为提出的驱逐算法提供了理论保证。

最后，作者使用opt、llama和gpt-neox在各种任务中验证算法的准确性。

其中，在opt-6.7b和opt-30b上实现的h2o，将deepspeed zero-inference、hugging face accelerate和flexgen这三个推理系统的吞吐量分别提高了29倍、29倍 和3倍，且在相同的批量大小下，h2o最多可以减少1.9倍 的延迟。

论文细节

上图为在 llm 生成中部署不同 kv 缓存策略的符号图；左下为h2o的框架概述；右下为不同策略下的准确性与内存消耗的对比。

我们可以看出，将前几种方法应用于预训练的llm ，会导致高未命中率并降低精度。

解决kv缓存问题，面临着三个技术挑战。

首先，目前尚不清楚是否可以限制kv缓存的大小——原则上，每个解码步骤可能需要访问所有先前的注意力键和值。

其次，确定保持生成准确性的最佳逐出策略是一个组合问题。

最后，即使可以暴力破解最佳策略，在实际应用程序上部署也是不可行的。

幸运的是，作者通过研究发现了一些有趣的结果。

小缓存大小的稀疏性：即使在密集训练时，llm的注意力矩阵在推理时也有超过95% 的稀疏率（图a）。这适用于各种预训练的llm。

因此，在每个生成步骤中，只需要5% 的kv缓存就足以解码相同的输出tokens，这表明kv缓存大小最多可以减少20倍，而不会降低精度。

heavy hitters（ h2 ）：注意力区块中所有tokens的累积注意力分数都遵循幂律分布（图b）。这表明存在一小群有影响力的tokens，这些tokens在生成过程中至关重要，是重量级tokens （ h2 ）。这使我们可以摆脱组合搜索问题，并确定保持准确性的逐出策略。

低成本策略的贪婪算法：在每个解码步骤中保留基于局部统计数据的h2（仅将前面tokens的注意力分数相加）与考虑未来tokens的注意力一样有效（图d）。

基于上述内容，作者定义了在大小受限的kv缓存中， llm的生成过程，并提出了heavy-hitter oracle （ h2o ），该框架利用了上面提到的性质，并使用简单、低成本的驱逐策略。

方法与分析

llm的生成过程包括两个不同的阶段：

提示阶段：使用输入序列来生成kv缓存（由键和值嵌入组成），类似于llm训练期间采用的前向传递；

tokens生成阶段：利用和更新kv缓存以增量方式生成新tokens 。每个生成步骤都依赖于先前生成的tokens。

本文的重点是在tokens生成阶段提高kv缓存的注意力效率，从而加速llm推理。

作者定义了具有有限kv缓存大小的生成过程，包括注意力查询矩阵q和键矩阵k。

驱逐策略：

以及采用了驱逐策略的生成过程：

接下来讨论在不降低精度的情况下，减少kv缓存大小的可能性。

上图中，（a）表示预训练opt模型中的注意力稀疏性；（b）表示累积注意力分数相对于相应单词的分布（红色散点）和数据中单词的共现次数（灰色曲线），x轴表示词汇表中的单词索引；（c）表示具有完整kv缓存的基线模型与本文模型（h2o）的性能比较；（d）表示具有完整kv缓存的基线模型、具有局部统计量的h2o、具有全局统计量的h2o和仅具有最新kv（局部）的模型之间的比较。

给定由查询矩阵q和键矩阵k计算的归一化注意力得分softmax矩阵，将阈值设置为每行最大值的百分之一，并计算相应的稀疏度。

然后在wiki-text-103的验证集上使用预训练的opt模型进行零样本推理，绘制注意力块内的逐层稀疏性，并可视化了归一化的注意力得分矩阵。

结果如下图所示，尽管llm是密集训练的，但由此产生的注意力得分矩阵是高度稀疏的，几乎所有层的稀疏度都超过95%。

注意力块的稀疏性表明，生成下一个tokens时，不需要访问所有先前的键和值嵌入，所以可以逐出不必要的kv嵌入，也就减少了生成过程中对kv缓存的需求。

不过，逐出的策略需要谨慎，因为一旦驱逐了重要的kv，由于llm生成的顺序依赖性，可能会破坏llm的性能。

作者研究发现，注意力区块内所有tokens的累积注意力分数都遵循幂律分布，如下图所示。

这表明存在一小部分在生成过程中至关重要的tokens，也就是前文谈到的heavy-hitters （h2）。

此外，每个单词的累积注意力分数（红点）与它们在数据中的共现（灰色曲线）具有高度相关性。

作者基于以上现象设计了一种贪婪驱逐策略：

在生成过程中，当令tokens数量超过分配的kv缓存预算时，根据其累积的注意力分数统计数据，以及缓存中的本地tokens来保留重量级tokens。

一般而言，需要使用整个生成过程中的统计数据，才能得到最理想的结果，但这在实际部署中显然是不可行的，因为无法访问未来生成的tokens。

于是，作者进行了下图的实验，发现在每个解码步骤中使用局部统计数据计算的局部h2 ，与考虑未来tokens的情况效果差不多（红线和蓝线）。

随后，作者将这种动态注意力分数计算（有空间限制）定义为一种新的动态的子模块问题（dynamic submodular type problem）：

利用上面的形式定义kv缓存驱逐策略：

上图展示了驱逐算法，以及说明性示例。这里假设kv缓存的预算大小为3 ，完成第四个解码步骤后，根据累积的注意力分数逐出与第三个token关联的kv嵌入，被逐出的kv嵌入在后续解码步骤中将不可访问。

另外，作者还提到了实际实现中的细节。比如，为了保证i/o效率，我们在驱逐存储的kv时不会交换内存，而是直接填充新添加的kv。

实验结果

论文的实验选用了三个具有代表性的llm模型系列，包括opt，llama和gpt-neox-20b 。

选取了8个评估任务：copa ， mathqa ， openbookqa ， piqa ， rte ， winogrande ， xsum ， cnn/daily mail 。

实验的硬件采用nvidia a100 80gb gpu。

考虑到h2o所采用的缓存策略，这里除了完整的kv缓存（full），还将本地缓存策略（local）也作为基线方法。

由上面的图和表可知：在不同的kv缓存预算下，本文提出的方法（h2o）在各种不同条件的测试中都优于local策略。

同时，在低于20%的kv缓存预算之下，h2o实现了与全kv嵌入模型（full）相当的性能，且在更具挑战性的长序列生成任务、xsum和cnn/daily mail中表现良好。

参考资料：

https://arxiv.org/abs/2306.14048

https://twitter.com/tydsh/status/1731778956330971507