一种在推理过程中更高效计算Transformer注意力的方法
几乎所有大型语言模型（LLM）都依赖于Transformer神经架构。尽管这种架构因其效率而受到赞扬，但它也存在一些众所周知的计算瓶颈。
在解码过程中，其中之一的瓶颈是每个输入标记的键-值张量对的注意力计算。所有这些张量都必须存储在内存中。
注意：我不会在本文中解释这些键-值对的作用。这是Transformer架构中最复杂和最有趣的方面之一。如果您不知道，请强烈建议阅读Jay Alammar的The Illustrated Transformer。
随着LLM接受越来越长的输入，例如LLM Claude接受100k令牌长的输入，这些张量消耗的内存可能会变得非常大。
简单地将所有这些张量存储在内存中会导致内存超额预留和碎片化。特别是对于长序列的令牌，此碎片化可能使内存访问非常低效。至于超额预留，系统会这样做是为了确保它已经为张量分配了足够的内存，即使它没有全部使用。
为了缓解这些问题，加州大学伯克利分校提出了PagedAttention。
PagedAttention实现在vLLM（Apache 2.0许可证）中，由LLMSYS部署，这是一个由加州大学伯克利分校的学生和教师与UCSD和CMU的帮助创立的开放研究组织。
在本文中，我将解释PagedAttention是什么以及为什么它可以显着加快解码速度。我将在文章末尾展示如何开始使用vLLM来利用PagedAttention对您的计算机上进行推理和提供LLMs。
Transformer的PagedAttention
Kwon等人提出了PagedAttention。
目标是更有效地存储键值张量在GPU VRAM的非连续空间中。
简而言之，PagedAttention的想法是创建虚拟连续块，映射到GPU内存中的物理块。
每个块旨在存储预定义数量的令牌的键值对张量。所有块在虚拟上是连续的，并映射到在推理期间根据需要分配的非连续块中，存在于碎片化的GPU内存中。还会在内存中创建一个简单的索引表，用于将虚拟块与物理块关联。
PagedAttention的核心按需获取这些块中的内容。这是有效的，因为由于块的大小有限，系统获取更少的键值张量。
我们以以下提示为例：
猫正在厨房里睡觉，狗正在
我们为每个令牌都有键值对张量。使用PageAttention，我们可以（任意地）将块大小设置为4。每个块包含4个键值对张量，最后一个块仅包含3个键值对张量。这些块在虚拟上是连续的，但在GPU内存中不一定是连续的，如本文开头的图所示。
对于注意力计算，对于每个查询标记，系统按块一个接一个地获取，如下图所示。
通过按块获取键值对张量，而不是整个张量序列，可以大大加快注意力计算速度。
并行采样推理
PagedAttention 的另一个优点是，在推理期间进行采样时，虚拟块可以共享。通过采样或波束搜索并行生成的所有序列都可以使用相同的虚拟块，避免了重复。
LMSYS 在他们的实验中观察到，波束搜索解码的内存使用量减少了 55%。
LMSYS 报告的 PagedAttention 性能
在自己尝试之前，让我们先看看作者（UC Berkely/LMSYS）使用 vLLM 实现的 PagedAttention 与 Hugging Face 开发的文本生成推理库相比的性能。
根据这些结果，vLLM 看起来要快得多，特别是在多个输出完成的情况下。随着模型变得更大，TGI 和 vLLM 之间的差异会增加。这是预期的，因为更大的模型需要更多的内存，因此更容易受到内存碎片化的影响。
总的来说，vLLM 比 Hugging Face Transformers 库快多达 24 倍。
注意：实际上，我也对 HF 到 TGI 的改进印象深刻。我还没有在我的博客上涵盖 TGI，但我可能会写一篇关于它的指南。TGI 在 Hugging Face 的生产中使用。虽然它看起来比 vLLM 慢得多，但 TGI 有其他优点，例如支持更多模型和功能。
如何在计算机上设置 vLLM
注意：vLLM 尚不支持 CUDA 12。使用较低版本，例如 11.8。
在本节中，我只会讲解如何在计算机上设置和运行 vLLM 的基础知识。对于更高级的用法，您可以查看 vLLM 文档。
截至我写这篇文章时，vLLM 仅支持几种类型的模型：
- GPT-2
- GPT-NeoX 和 Pythia based
- LLaMa based
- OPT based
您可以按照这些说明添加对其他模型的支持。
在下面的代码中，我使用了 Dolly V2（MIT 许可证）。它是基于 Pythia 的聊天模型，由 DataBricks 训练。
我选择了具有 30 亿个参数的最小版本。它可以在消费级 GPU 上运行，具有 24 GB 的 VRAM，例如 nVidia RTX 3080/3090。
安装 vLLM 最简单的方法是使用 pip：
pip install vllm注意：这可能需要最多 10 分钟。
但在我的情况下，在我的计算机和 Google Colab 上，pip 无法安装 vllm 库。 vLLM 的作者确认，在某些 nvcc 版本和环境中存在问题。尽管如此，对于大多数配置，pip 应该无问题地安装 vLLM。
如果您和我处于相同的情况，解决方法就是使用 Docker 镜像。这个对我有效：
docker run --gpus all -it --rm --shm-size=8g nvcr.io/nvidia/pytorch:22.12-py3注意：进入 docker 后，作者建议在安装 vLLM 之前删除 Pytorch：pip uninstall torch。然后，“pip install vllm” 应该可以工作。
然后，我们可以开始编写 Python。
我们首先需要导入 vllm，然后使用 vllm 加载模型。llm.generate() 触发推理。
from vllm import LLM
prompts = ["讲讲重力是什么"] #您可以在此列表中放置多个提示
llm = LLM(model="databricks/dolly-v2-3b") # 加载模型
outputs = llm.generate(prompts) # 触发推理您也可以使用vLLM提供LLM服务。它的工作方式类似于TGI。与我在之前的一篇文章中描述的运行NVIDIA Triton推理服务器相比，它更加简单。
您需要先启动服务器：
python -m vllm.entrypoints.openai.api_server --model databricks/dolly-v2-3b注意：服务器将侦听端口8000。确保它可用或在vLLM配置文件中更改它。
然后，您可以使用以下提示查询服务器：
curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ "model": "databricks/dolly-v2-3b", "prompt": "讲讲重力是什么", "max_tokens": 200 }'就这样！您的计算机上运行着一个非常高效的LLM服务器。
结论
PagedAttention显着加快了推理速度。这是更加经济实惠的LLM人工智能的又一步。
在进一步的实验中，我证实vLLM在提示批处理方面特别高效。为了充分利用vLLM，请考虑针对推理进行批处理策略的优化。
虽然使用大型beam进行搜索可能会因标准注意力计算而受到限制，但使用PagedAttention进行beam搜索更快且更节省内存。
我的下一个实验之一将是将PagedAttention与QLoRa结合使用以减少内存使用量。这应该很简单。它将使在消费者硬件上运行LLM更加高效。
如果您喜欢这篇文章并有兴趣阅读下一篇文章，支持我的最佳方式是使用此链接成为小猪AI会员：
使用我的推荐链接加入小猪AI – Benjamin Marie
作为小猪AI会员，您的会员费的一部分将赠送给您阅读的作家，并且您将获得每个故事的完全访问权限…
小猪AI.com
如果您已经是会员并想支持这项工作，请在小猪AI上关注我。