Loading...

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

智源社区7个月前发布 智源社区
158 0 0
吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题


  新智元报道  

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

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

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

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

吞吐量提升近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的序列)。

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

对此,论文提出了一种实现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倍 的延迟。

论文细节

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

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

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

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

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

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

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

小缓存大小的稀疏性:即使在密集训练时,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。

驱逐策略:

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

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

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

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

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

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

这表明存在一小部分在生成过程中至关重要的tokens,也就是前文谈到的Heavy-Hitters (H2)。

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

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

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

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

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

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

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

利用上面的形式定义KV缓存驱逐策略:

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

上图展示了驱逐算法,以及说明性示例。这里假设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)也作为基线方法。

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

由上面的图和表可知:在不同的KV缓存预算下,本文提出的方法(H2O)在各种不同条件的测试中都优于Local策略。

同时,在低于20%的KV缓存预算之下,H2O实现了与全KV嵌入模型(Full)相当的性能,且在更具挑战性的长序列生成任务、XSUM和CNN/Daily Mail中表现良好。

参考资料:
https://arxiv.org/abs/2306.14048
https://twitter.com/tydsh/status/1731778956330971507

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题


吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题

吞吐量提升近30倍!田渊栋团队最新论文解决大模型部署难题
© 版权声明

相关文章

暂无评论

暂无评论...