大模型推理性能优化之KV Cache解读

这篇具有很好参考价值的文章主要介绍了大模型推理性能优化之KV Cache解读。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

0. 引言

做大模型性能优化的一定对KV Cache不陌生,那么我们对这个技术了解到什么程度呢?请尝试回答如下问题:

  • KV Cache节省了Self-Attention层中哪部分的计算?

  • KV Cache对MLP层的计算量有影响吗?

  • KV Cache对block间的数据传输量有影响吗?本文打算剖析该技术并给出上面问题的答案。

1. KV Cache是啥

大模型推理性能优化的一个常用技术是KV Cache,该技术可以在不影响任何计算精度的前提下,通过空间换时间思想,提高推理性能。网上有一些关于该技术的分析博客,但读过后仍然会很迷糊,甚至可能会被带偏,认为这个Cache过程和数据库读取或CPU Cache加速类似的荒谬结论。刚开始我也有类似误解,直到逐行查阅并运行源码,才清楚了解到其Cache了啥,以及如何节省计算的。

2. 背景

生成式generative模型的推理过程很有特点,我们给一个输入文本,模型会输出一个回答(长度为N),其实该过程中执行了N次推理过程。即GPT类模型一次推理只输出一个token,输出token会与输入tokens 拼接在一起,然后作为下一次推理的输入,这样不断反复直到遇到终止符。

如上描述是我们通常认知的GPT推理过程。代码描述如下:

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer


model = GPT2LMHeadModel.from_pretrained("gpt2", torchscript=True).eval()

# tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
in_text = "Lionel Messi is a"
in_tokens = torch.tensor(tokenizer.encode(in_text))

# inference
token_eos = torch.tensor([198]) # line break symbol
out_token = None
i = 0
with torch.no_grad():
    while out_token != token_eos:
        logits, _ = model(in_tokens)
        out_token = torch.argmax(logits[-1, :], dim=0, keepdim=True)
        in_tokens = torch.cat((in_tokens, out_token), 0)
        text = tokenizer.decode(in_tokens)
        print(f'step {i} input: {text}', flush=True)
        i += 1

out_text = tokenizer.decode(in_tokens)
print(f' Input: {in_text}')
print(f'Output: {out_text}')

输出:

step 0 input: Lionel Messi is a player
step 1 input: Lionel Messi is a player who
step 2 input: Lionel Messi is a player who has
step 3 input: Lionel Messi is a player who has been
step 4 input: Lionel Messi is a player who has been a
step 5 input: Lionel Messi is a player who has been a key
step 6 input: Lionel Messi is a player who has been a key part
step 7 input: Lionel Messi is a player who has been a key part of
step 8 input: Lionel Messi is a player who has been a key part of the
step 9 input: Lionel Messi is a player who has been a key part of the team
step 10 input: Lionel Messi is a player who has been a key part of the team's
step 11 input: Lionel Messi is a player who has been a key part of the team's success
step 12 input: Lionel Messi is a player who has been a key part of the team's success.
step 13 input: Lionel Messi is a player who has been a key part of the team's success.

 Input: Lionel Messi is a
Output: Lionel Messi is a player who has been a key part of the team's success.

可以看出如上计算的问题吗?每次推理过程的输入tokens都变长了,导致推理FLOPs随之增大。有方法实现推理过程的FLOPs基本恒定不变或变小吗?(埋个伏笔,注意是基本恒定)。

3. 原理

在上面的推理过程中,每step内,输入一个token序列,经过Embedding层将输入token序列变为一个三维张量[b, s, h],经过一通计算,最后经logits层将计算结果映射至词表空间,输出张量维度为[b, s, vocab_size]

当前轮输出token与输入tokens拼接,并作为下一轮的输入tokens,反复多次。可以看出第i+1轮输入数据只比第 轮i输入数据新增了一个token,其他全部相同!因此第i+1轮推理时必然包含了第i轮的部分计算。KV Cache的出发点就在这里,缓存当前轮可重复利用的计算结果,下一轮计算时直接读取缓存结果,就是这么简单,不存在什么Cache miss问题。

4. 实现细节

目前各大模型推理都实现了KV Cache,下面就看如何使用了。我们可以在上面代码基础上修改,主要改动:

  • 在推理时新增了past_key_values参数,该参数就会以追加方式保存每一轮的K V值。kvcache变量内容为((k,v), (k,v), ..., (k,v)),即有L个 k,v 组成的一个元组,其中 k 和 v 的维度均为[b, n_head, s, head_dims]。这里可以顺带计算出每轮推理对应的 cache 数据量为2bshL,这里s值等于当前轮次值。可以随着输出tokens的增长,cache数据量呈现线性增加特点。以GPT3-175B为例,假设以 float16 来保存 KV cache,senquence长度为100,batchsize=1,则KV cache占用显存为 2×100×12288×96×2 Byte = 472MB。

  • 推理输出的token直接作为下一轮的输入,不再拼接,因为上文信息已经在 kvcache 中。

代码示例:

import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer


model = GPT2LMHeadModel.from_pretrained("gpt2", torchscript=True).eval()

# tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
in_text = "Lionel Messi is a"
in_tokens = torch.tensor(tokenizer.encode(in_text))

# inference
token_eos = torch.tensor([198]) # line break symbol
out_token = None
kvcache = None
out_text = in_text
i = 0
with torch.no_grad():
    while out_token != token_eos:
        logits, kvcache = model(in_tokens, past_key_values=kvcache) # 增加了一个 past_key_values 的参数
        out_token = torch.argmax(logits[-1, :], dim=0, keepdim=True)
        in_tokens = out_token # 输出 token 直接作为下一轮的输入,不再拼接
        text = tokenizer.decode(in_tokens)
        print(f'step {i} input: {text}', flush=True)
        i += 1
        out_text += text

print(f' Input: {in_text}')
print(f'Output: {out_text}')

通过上面代码只能看到调用层面的变化,实现细节还需看各框架的底层实现,例如Hugging Face的transformers库代码实现就比较清爽,在modeling_gpt2.py中Attention部分相关代码如下:

query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

        if layer_past is not None: # 当输出第一个token后,layer_past就是非None了
            past_key, past_value = layer_past # 取出之前计算好的 key, value
            key = torch.cat((past_key, key), dim=-2) # past_key 与当前 token 对应的 key 拼接
            value = torch.cat((past_value, value), dim=-2) # past_value 与当前 token 对应的 value 拼接

        if use_cache is True:
            present = (key, value)
        else:
            present = None

在 block 层面也有相关代码,大家有空细品吧。还是那句话,说一千道一万不如阅读并运行源码一次。

其实,KV Cache 配置开启后,推理过程可以分为2个阶段:

  • 预填充阶段:发生在计算第一个输出token过程中,这时Cache是空的,计算时需要为每个 transformer layer 计算并保存key cache和value cache,在输出token时Cache完成填充;FLOPs同KV Cache关闭一致,存在大量gemm操作,推理速度慢。

  • 使用KV Cache阶段:发生在计算第二个输出token至最后一个token过程中,这时Cache是有值的,每轮推理只需读取Cache,同时将当前轮计算出的新的Key、Value追加写入至Cache;FLOPs降低,gemm变为gemv操作,推理速度相对第一阶段变快,这时属于Memory-bound类型计算。

这里用图可能更有助理解,下图是一个Decoder Block,含有Self-Attention和MLP,标红部分为KV Cache影响到的内容,即KV Cache开启后,标红的序列长度s变为 1,当batch_size=1时,Self-Attention中的2个dense全都变为gemv操作,MLP中的dense也全都变为gemv操作。看懂这个图就可以答对上面的3个问题啦。

大模型推理性能优化之KV Cache解读,性能优化


Decoder Block of GPT

如下链接也有这方面的定量分析,写的很棒,推荐大家看看。

回旋托马斯x:分析transformer模型的参数量、计算量、中间激活、KV cache

5. 总结

KV Cache是Transformer推理性能优化的一项重要工程化技术,各大推理框架都已实现并将其进行了封装(例如 transformers库 generate 函数已经将其封装,用户不需要手动传入past_key_values)并默认开启(config.json文件中use_cache=True)。本文尝试打开封装分析该技术内部实现,希望对大家有所帮助,文中如有纰漏,欢迎指正。

最后给我们团队的开源项目打个广告,Adlik 是中兴通讯贡献的深度学习推理工具,已获得 Linux AI 基金会支持,目前仍在持续完善中,期待大家的支持和关注。另外,我们也在做深度学习编译器和大模型方向的前沿技术研发工作,欢迎感兴趣的小伙伴加入我们 ~~

大模型推理性能优化之KV Cache解读,性能优化文章来源地址https://www.toymoban.com/news/detail-569466.html

到了这里,关于大模型推理性能优化之KV Cache解读的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • HVDC-MMC互连(1000MW,±320KV)使用聚合MMC模型进行优化的SPS模拟

    微 ❤ 关注“电气仔推送”获得资料(专享优惠) 模型概述: 本示例展示了一个SimPowerSystems(SPS)模型,使用基于模块化多电平变换器(MMC)技术的电压源换流器(VSC)实现了高压直流(HVDC)互连。通过使用聚合MMC模型,对SPS模拟进行了优化。 直流电力传输已经成为国际间

    2024年02月07日
    浏览(13)
  • [SSD NAND 7.1] 闪存系统性能优化方向集锦?AC timing? Cache? 多路并发?

    传送门    总目录 主页 : 元存储的博客_CSDN博客 依公开知识及经验整理,如有误请留言。 个人辛苦整理,付费内容,禁止转载。 内容摘要 1. 优化 AC Timing,提升总线频率 1.1 优化 AC Timing 1.2 优化总线频率 2. 使用 Cache Read/Program 3. 多路并发技术 3.1 多平面(Multi Plane)操

    2024年02月04日
    浏览(34)
  • 大语言模型推理提速:TensorRT-LLM 高性能推理实践

    作者:顾静 大型语言模型(Large language models,LLM)是基于大量数据进行预训练的超大型深度学习模型。底层转换器是一组神经网络,这些神经网络由具有 self-attention 的编码器和解码器组成。编码器和解码器从一系列文本中提取含义,并理解其中的单词和短语之间的关系。 当前

    2024年01月25日
    浏览(14)
  • 优化故事: BLOOM 模型推理

    经过“九九八十一难”,大模型终于炼成。下一步就是架设服务,准备开门营业了。真这么简单?恐怕未必!行百里者半九十,推理优化又是新的雄关漫道。如何进行延迟优化?如何进行成本优化 (别忘了 OpenAI 8K 上下文的 GPT-4 模型,提示每 1000 词元只需 0.03 美金,补全每

    2023年04月17日
    浏览(10)
  • [深入理解NAND Flash (指令篇) ] 闪存系统性能优化方向集锦?AC timing? Cache? 多路并发?

    传送门    总目录 主页 : 元存储的博客_CSDN博客 依公开知识及经验整理,如有误请留言。 个人辛苦整理,付费内容,禁止转载。 内容摘要 1. 优化 AC Timing,提升总线频率 1.1 优化 AC Timing 1.2 优化总线频率 2. 使用 Cache Read/Program 3. 多路并发技术 3.1 多平面(Multi Plane)操

    2024年02月14日
    浏览(19)
  • 掌握大语言模型技术: 推理优化

    堆叠 Transformer 层来创建大型模型可以带来更好的准确性、少样本学习能力,甚至在各种语言任务上具有接近人类的涌现能力。 这些基础模型的训练成本很高,并且在推理过程中可能会占用大量内存和计算资源(经常性成本)。 当今最流行的大型语言模型 (LLM) 的参数大小可以

    2024年01月22日
    浏览(18)
  • “一键”模型迁移,性能翻倍,多语言AltDiffusion推理速度超快

    为了推进 AIGC 行业的降本增效,同时也回应用户的热情要求,OneFlow 正在将业内备受欢迎的相关 Diffusion 模型的加速“一网打尽”。 此前,OneFlow 首度 将 Stable Diffusion 模型加速至“一秒出图”时代 ,极大提升了文生图的速度,在 AIGC 领域引发巨大反响,并得到了 Stability.ai 官

    2024年02月08日
    浏览(16)
  • 百度商业AI 技术创新大赛赛道二:AIGC推理性能优化TOP10之经验分享

    朋友们,AIGC性能优化大赛已经结束了,看新闻很多队员已经完成了答辩和领奖环节,我根据内幕人了解到,比赛的最终代码及结果是不会分享出来的,因为办比赛的目的就是吸引最优秀的代码然后给公司节省自己开发的成本,相当于外包出去了,应该是不会公开的。抱着技术

    2024年02月11日
    浏览(16)
  • OpenVINO 2022.3之七:OpenVINO 预处理API提升模型推理性能

    OpenVINO™ 2022.3 提供OpenVINO™ Runtime原生的用于数据预处理的API函数。 如果没有预处理API,那么输入数据的预处理操作只能放在CPU上实现,CPU完成数据预处理后,再将预处理后的数据传给iGPU、VPU等AI加速计算设备进行推理计算。 有了预处理API后,就能将预处理操作集成到在模型

    2024年02月04日
    浏览(179)
  • LLaMa 原理+源码——拆解 (KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLU)

    Vanilla Transformer 与 LLaMa 的对比:LLaMa与普通的Transformer架构不同的地方,包括采用了 前置 了层归一化(Pre-normalization)并使用 RMSNorm 归一化函数(Normalizing Function)、使用了 旋转位置嵌入(RoPE) 、激活函数由ReLU更换为 SwiGLU ,并且将self-attention改进为 使用KV-Cache的Grouped Quer

    2024年01月25日
    浏览(12)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包