Transformer系列：图文详解KV-Cache，解码器推理加速优化

最新推荐文章于 2024-06-26 19:46:34 发布

Cc不爱吃洋葱

最新推荐文章于 2024-06-26 19:46:34 发布

阅读量1.4k

点赞数 33

文章标签：人工智能 chatgpt agi transformer 语言模型自然语言处理学习

本文链接：https://blog.csdn.net/2401_85328934/article/details/139698693

版权

前言

KV-Cache是一种加速Transformer推理的策略，几乎所有自回归模型都内置了KV-Cache，理解KV-Cache有助于更深刻地认识Transformer中注意力机制的工作方式。

内容摘要

自回归推理过程知识准备
KV-Cache原理介绍
KV-Cache推理效率提升统计
KV-Cache工作流程简述
past_key_values存储结构和作用分析
KV-Cache内存占用、FLOPs下降分析

自回归推理过程知识准备

自回归模型采用shift-right的训练方式，用前文预测下一个字/词，并且前文中的最后一个词经过解码器的表征会映射为其下一个待预测词的概率分布。在训练阶段，句子完整输入给网络，所有位置下的token并行计算。
同理，在预测推理阶段也可以将前文prompt完整输入给训练好的模型，取最后一个位置的表征作为下一个token的概率分布，再通过采样策略确认下一个token，最终将token拼接到前文prompt的末尾准备下一次推理。

GPT自回归工作方式

KV-Cache简要介绍

每步推理都将前文整句输入模型是一种效率低下的方式，原因是存在相同结果的重复推理。令前一次待推理的文本长度为S，下一次为S+1，由于网络中的各项参数已经固定，因此两次推理对于前S个token的计算结果是完全相同的，包括Embedding映射，每一层、每一个注意力头下的KQV映射，注意力权重，以及后续的FFN层都在重复计算。
根据shift-right的性质，下一个token是由当前最后一个token的网络输出所决定的，那能不能仅输入最后一个token来进行推理？答案是否定的，虽然在结果层仅由最后一个token来决定，但是中间的注意力过程它依赖于前文所提供的Key、Value向量来携带前文信息，因此也不能抛弃前文不管。

next token计算依赖

结合以上结论，S+1位置token的推理依赖于两个要素，首先是当前第S个token在网络中完整forward一遍，其次是除最后一个token以外，之前所有的S-1位置的token在每一层、每个注意力头下的Key，Value信息。又已知S-1的每个token的Key，Value信息都是在重复计算，每次计算的结果是相同的，在之前的推理中都计算过但在结果层丢弃了，因此完全可以将Key，Value信息在内存中存储起来，使得它们可以在之后的每步推理中进行复用，这种策略就是KV-Cache。这种方式避免了重复计算，大幅减少了参数的计算量，提高了推理效率。

KV-Cache推理效率提升统计

本例采用GPT-2作为实验对象，测试开启/关闭KV-Cache对推理效率的影响。在HuggingFace实现的GPT2LMHeadModel模型类中，推理阶段内置了KV-Cache选项，通过use_cache来开启和关闭KV-Cache，当use_cache为true时，模型在推理过程中会初始化past_key_values来存储Key、Value向量，并且每一步推理会对它进行维护，如果use_cache为false则past_key_values不生效。GPT2LMHeadModel的推理阶段参数如下

class GPT2LMHeadModel(GPT2PreTrainedModel):
    ...
    def forward(
            self,
            input_ids: Optional[torch.LongTensor] = None,
            past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
            ...
            use_cache: Optional[bool] = None,
            ...
        ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:

给到prompt为“明天降温了”，设置最大推理步长从10到1000不等，采用最简单的贪婪搜索Greedy Search方式，分别将use_cache设置为true和false两种模式，查看CPU和GPU推理下耗时，代码如下

import time
import torch
from transformers import BertTokenizer, GPT2LMHeadModel

tokenizer = BertTokenizer.from_pretrained("./gpt2-chinese-cluecorpussmall")
model = GPT2LMHeadModel.from_pretrained("./gpt2-chinese-cluecorpussmall").to("cuda")

text = "明天降温了"
input_ids = torch.LongTensor([tokenizer.convert_tokens_to_ids(list(text))]).to("cuda")

max_length = list(range(10, 1101, 100))
for i in max_length:
    res = model.generate(input_ids=input_ids, max_length=i, do_sample=False, use_cache=True)

其中CPU下开启/关闭KV-Cache的推理耗时差距更加明显，各步长下推理耗时(秒)统计如下表

推理步长	关闭KV-Cache	开启KV-Cache
10	0.17	0.18
110	8.71	3.11
310	40.93	9.31
510	92.43	15.84
710	178.22	21.82
910	332.96	29.22

随着步长的增长，关闭KV-Cache的推理总耗时呈现出指数级增长，而开启KV-Cache的耗时线性增长，当步长达到900时，前者的耗时已经是后者的十倍以上，通过可视化能够直观感受到两者的效率差距

CPU下开启和关闭KV-Cache的推理随着步长的耗时

进一步计算平均每个token的推理速度，用总耗时除以推理步长，统计图如下，当关闭KV-Cache时，随着步长从10增长到1000，推理一个token从17ms增长到426ms，推理步长越大，效率越来越低，而当开启KV-Cache时，推理一个token的耗时基本稳定维持在30ms左右，只呈现出小数点后第三位上的略微增长趋势，推理长度几乎没有对推理效率产生负面影响。

KV-Cache推理一个token的耗时对比

根据以上实验初步得到结论，随着推理步长的增长，关闭KV-Cache推理效率会越来越低，而开启KV-Cache推理效率基本恒等不变。

KV-Cache工作流程简述

KV-Cache会在模型连续推理的过程中持续调用和更新past_key_values，特别的，当模型首次推理时，past_key_values为空，需要对past_key_values进行初始化，首次推理需将全部文本一齐输入，将中间过程的所有Key，Value添加到past_key_values中。
从第二次推理开始，仅需要输入当前最后一个token，单独对该token做Q，K，V映射，将past_key_values中前文所有的K，V和该token的K，V进行拼接得到完成的Key、Value向量，最终和该token的Query计算注意力，拼接后的Key、Value也同步更新到past_key_values。

KV-Cache的代码实现流程图

past_key_values存储结构分析

KV-Cache会将截止当前各个token在每一层、每个头的Key向量和Value向量存储在内存中，在HuggingFace的代码实现中使用past_key_values变量进行存储，past_key_values是一个矩阵，其维度为**[n, 2, b, h, s, d]**，类似一个六维的矩阵，每个维度的含义如下

第一维 num_layers：在外层是以每一个堆叠的Block为单位，例如堆叠12层，则一共有12组Key、Value信息
第二维 2：代表Key和Value这两个信息对象，索引0取到Key向量，索引1取到Value向量
第三维 batch_size：代表batch_size，和输入需要推理的文本条数相等，如果输入是一条文本，则b=1
第四维 num_heads：代表注意力头的数量，例如每层有12个头，则h=12
第五维 seq_len：代表截止到当前token为止的文本长度，在每一个历史token位置上该token在每一层每个头下的Key，Value信息
第六维 d：代表Key、Value向量的映射维度，若token总的映射维度为768，注意力头数为12，则d=768/12=64

past_key_values结构示意图

past_key_values的结构如上图所示，随着模型推理步长的增长，past_key_values在每一步也同步更新，上一个past_key_values和下一个past_key_values的差异仅仅产生在seq_len这个维度上，具体的，seq_len维度大小会加1，它是由新推理的那一个token所对应的Key，Value拼接到上一个past_key_values的seq_len维度中所导致的，如果除开这个加1的因素，上一个past_key_values和下一个past_key_values在seq_len这个维度上的向量完全相同。
用公式可以更清晰的表达出past_key_values前后的变化，令第一次推理Seq_len等于5，12层每层12个头，Key、Value维度为64，则有

past_key_values前后公式对比

KV-Cache内存占用、FLOPs下降分析

KV-Cache本质上是用空间换时间，存储的Key、Value矩阵会额外占用内存，假设以float16精度来存储，每个token的存储占用公式如下

KV-Cache占用内存计算

公式代表每一层、每一个头下的向量维度之和，乘以2代表Key、Value两者只和，再乘以2代表float16占用两个字节。以LLaMa-7B为例，模型加载占用显存14GB，向量维度4096，堆叠32层，最大推理步长4096，若推理一个batch为2，长度为4096的句子，KV-Cache占用的存储空间为2×2×32×4096×2×4096=21474836480字节，约等于4GB，随着推理的batch增大，推理长度变长，KV-Cache占用的存储空间可能超过模型本身。

KV-Cache和模型自身参数的显存占用示意图

另一方面KV-Cache极大地降低了FLOPs（浮点计算量），表面上KV-Cache省去了之前每个token的Key、Value的计算量，每个token在所有层下计算Key、Value的FLOPs公式如下

每个token计算Key、Value的FLOPs

其中d平方代表从token Embedding到Key或者Value向量的过程，乘以2是矩阵相乘中逐位相乘再相加导致有两个操作，再乘以2代表Key、Value各一个。还是以LLaMa-7B为例，推理一个batch为2，长度为4096的句子，光计算KV一共节省了2×2×32×4096×4096×4096×2=17592186044416 FLOPs的计算量，额外的，不仅省去了前文所有token的Key、Value的映射，由此导致后续这些token的注意力权重计算，注意力的MLP层，FFN前馈传播层也都不需要再计算了，相当于推理阶段的计算复杂度永远等于只对一个token进行完整的forward推理，因此计算量大幅降低。

如何系统的去学习大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

在这里插入图片描述

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

目标：了解AI大模型的基本概念、发展历程和核心原理。
内容：
- L1.1 人工智能简述与大模型起源
- L1.2 大模型与通用人工智能
- L1.3 GPT模型的发展历程
- L1.4 模型工程
- L1.4.1 知识大模型
- L1.4.2 生产大模型
- L1.4.3 模型工程方法论
- L1.4.4 模型工程实践
- L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
内容：
- L2.1 API接口
- L2.1.1 OpenAI API接口
- L2.1.2 Python接口接入
- L2.1.3 BOT工具类框架
- L2.1.4 代码示例
- L2.2 Prompt框架
- L2.2.1 什么是Prompt
- L2.2.2 Prompt框架应用现状
- L2.2.3 基于GPTAS的Prompt框架
- L2.2.4 Prompt框架与Thought
- L2.2.5 Prompt框架与提示词
- L2.3 流水线工程
- L2.3.1 流水线工程的概念
- L2.3.2 流水线工程的优点
- L2.3.3 流水线工程的应用
- L2.4 总结与展望

阶段3：AI大模型应用架构实践

目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
内容：
- L3.1 Agent模型框架
- L3.1.1 Agent模型框架的设计理念
- L3.1.2 Agent模型框架的核心组件
- L3.1.3 Agent模型框架的实现细节
- L3.2 MetaGPT
- L3.2.1 MetaGPT的基本概念
- L3.2.2 MetaGPT的工作原理
- L3.2.3 MetaGPT的应用场景
- L3.3 ChatGLM
- L3.3.1 ChatGLM的特点
- L3.3.2 ChatGLM的开发环境
- L3.3.3 ChatGLM的使用示例
- L3.4 LLAMA
- L3.4.1 LLAMA的特点
- L3.4.2 LLAMA的开发环境
- L3.4.3 LLAMA的使用示例
- L3.5 其他大模型介绍

阶段4：AI大模型私有化部署

目标：掌握多种AI大模型的私有化部署，包括多模态和特定领域模型。
内容：
- L4.1 模型私有化部署概述
- L4.2 模型私有化部署的关键技术
- L4.3 模型私有化部署的实施步骤
- L4.4 模型私有化部署的应用场景

学习计划：

阶段1：1-2个月，建立AI大模型的基础知识体系。
阶段2：2-3个月，专注于API应用开发能力的提升。
阶段3：3-4个月，深入实践AI大模型的应用架构和私有化部署。
阶段4：4-5个月，专注于高级模型的应用和部署。

这份完整版的大模型 LLM 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【`保证100%免费`】

😝有需要的小伙伴，可以Vx扫描下方二维码免费领取🆓

Cc不爱吃洋葱

关注

33
点赞
踩
37

收藏

觉得还不错? 一键收藏
0
评论
Transformer系列：图文详解KV-Cache，解码器推理加速优化

自回归推理过程知识准备KV-Cache原理介绍KV-Cache推理效率提升统计KV-Cache工作流程简述past_key_values存储结构和作用分析KV-Cache内存占用、FLOPs下降分析每步推理都将前文整句输入模型是一种效率低下的方式，原因是存在相同结果的重复推理。
复制链接

扫一扫