Transformers

加入Hugging Face社区

并获得增强型文档体验

协作模型、数据集和空间

使用加速推理的更快的示例

切换文档主题

开始

LLM 推理优化

大型语言模型 (LLM) 通过生成展现出高度理解和流畅性的文本，将聊天和代码补全等文本生成应用推向了新的高度。但 LLM 的强大之处——即其规模——也给推理带来了挑战。

基本的推理很慢，因为 LLM 必须重复调用以生成下一个标记。随着生成过程的进行，输入序列会增加，这会让 LLM 处理的时间越来越长。LLM 还拥有数十亿个参数，这使得在内存中存储和处理所有这些权重成为一项挑战。

本指南将向您展示如何使用 Transformers 中提供的优化技术来加速 LLM 推理。

Hugging Face 还提供了文本生成推理 (TGI)，这是一个专门用于部署和提供针对推理高度优化的 LLM 的库。它包含 Transformers 中没有的部署导向型优化功能，例如用于提高吞吐量的连续批处理以及用于多 GPU 推理的张量并行。

静态 kv 缓存和 torch.compile

在解码过程中，LLM 会为每个输入标记计算键值 (kv) 值，并且由于它是自回归的，因此每次都会计算相同的 kv 值，因为生成的输出现在已成为输入的一部分。这不是很高效，因为您每次都在重新计算相同的 kv 值。

为了优化这一点，您可以使用 kv 缓存来存储过去的键和值，而不是每次都重新计算它们。但是，由于 kv 缓存会随着每个生成步骤而增长并且是动态的，因此它阻止您利用 torch.compile，这是一种将 PyTorch 代码融合到快速优化内核中的强大优化工具。我们专门为 kv 缓存编写了一整套指南在此处。

静态 kv 缓存通过将 kv 缓存大小预先分配到最大值来解决此问题，这使您可以将其与 torch.compile 相结合，从而实现高达 4 倍的速度提升。您的提速可能因模型大小（更大的模型提速更小）和硬件而异。

目前，只有 Llama 和其他一些模型支持静态 kv 缓存和 torch.compile。请查看此问题以获取实时模型兼容性列表。

根据任务的复杂程度，静态 kv 缓存的使用有三种模式

基本用法：只需在 generation_config 中设置一个标志（推荐）；
高级用法：为多轮生成或自定义生成循环处理缓存对象；
高级用法：如果您认为只有一个图很重要，则将整个 generate 函数编译为一个图。

选择下面的正确选项卡，获取有关每种模式的进一步说明。

无论使用 torch.compile 的哪种策略，如果您将 LLM 输入左填充到一组有限的值，就可以避免与形状相关的重新编译。 pad_to_multiple_of 分词器标志非常有用！

基本用法：generation_config

高级用法：控制静态缓存

高级用法：端到端生成编译

推测解码

有关更深入的解释，请查看辅助生成：迈向低延迟文本生成的新方向博客文章！

自回归的另一个问题是，对于每个输入标记，您需要在每次前向传递过程中加载模型权重。对于拥有数十亿个参数的 LLM 来说，这很慢而且很麻烦。推测解码通过使用第二个更小、更快的辅助模型来生成候选标记来缓解这种减速，这些候选标记会在单个前向传递中由更大的 LLM 验证。如果验证的标记是正确的，那么 LLM 实际上可以“免费”获得它们，而无需自行生成它们。精度不会下降，因为验证前向传递确保生成的结果与 LLM 自行生成的结果相同。

为了获得最大的提速，辅助模型应该比 LLM 小得多，以便它能够快速生成标记。辅助模型和 LLM 还必须共享同一个分词器，以避免对标记进行重新编码和解码。

推测解码仅支持贪婪搜索和采样解码策略，并且也不支持批处理输入。

通过加载辅助模型并将其传递给 generate() 方法来启用推测解码。

贪婪搜索

采样

提示查找解码

提示查找解码是推测解码的一种变体，它也与贪婪搜索和采样兼容。提示查找特别适用于输入基础任务（如摘要），这些任务中提示和输出之间经常有重叠的单词。这些重叠的 n 元组用作 LLM 候选标记。

要启用提示查找解码，请在 prompt_lookup_num_tokens 参数中指定应该重叠的标记数量。然后，您可以将此参数传递给 generate() 方法。

贪婪解码

采样

注意力优化

Transformer 模型的一个已知问题是，自注意力机制的计算量和内存使用量随着输入标记数量的平方增长。这一限制在处理更长序列的 LLM 中尤其突出。为了解决这个问题，请尝试 FlashAttention2 或 PyTorch 的缩放点积注意力 (SDPA)，它们是更高内存效率的注意力实现，可以加速推理。

FlashAttention-2

FlashAttention 和 FlashAttention-2 将注意力计算分解成更小的块，并减少对 GPU 内存的中间读/写操作，从而加快推理速度。FlashAttention-2 在原始 FlashAttention 算法的基础上进行了改进，它还并行化了序列长度维度，并更好地将工作划分为硬件，以减少同步和通信开销。

要使用 FlashAttention-2，请在 from_pretrained() 方法中设置 attn_implementation="flash_attention_2"。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    quantization_config=quant_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

PyTorch 缩放点积注意力

缩放点积注意力 (SDPA) 在 PyTorch 2.0 中自动启用，它支持 FlashAttention、xFormers 和 PyTorch 的 C++ 实现。如果您使用 CUDA 后端，SDPA 会选择性能最高的注意力算法。对于其他后端，SDPA 默认使用 PyTorch C++ 实现。

SDPA 支持 FlashAttention-2，只要您安装了最新的 PyTorch 版本。

使用 torch.backends.cuda.sdp_kernel 上下文管理器来明确启用或禁用三种注意力算法中的任何一种。例如，设置 enable_flash=True 以启用 FlashAttention。

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2b",
    torch_dtype=torch.bfloat16,
)

with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    outputs = model.generate(**inputs)

量化

量化通过将 LLM 权重存储在更低精度来减小 LLM 权重的尺寸。这转化为更低的内存使用量，如果您受到 GPU 内存的限制，它将使加载 LLM 进行推断变得更容易。如果您不受 GPU 的限制，您不一定需要量化您的模型，因为它可能会产生少量的延迟成本（除了 AWQ 和融合的 AWQ 模块），这是由于量化和反量化权重所需的额外步骤。

有许多量化库（有关更多详细信息，请参阅量化指南）可用，例如 Quanto、AQLM、AWQ 和 AutoGPTQ。随意尝试一下，看看哪一个最适合您的用例。我们还建议您阅读 🤗 Transformers 中本机支持的量化方案概述博客文章，它比较了 AutoGPTQ 和 bitsandbytes。

使用下面的模型内存计算器来估计和比较加载模型所需的内存量。例如，尝试估计加载 Mistral-7B-v0.1 的内存成本。

要在半精度中加载 Mistral-7B-v0.1，请在 from_pretrained() 方法中将 torch_dtype 参数设置为 torch.bfloat16。这需要 13.74GB 的内存。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", torch_dtype=torch.bfloat16, device_map="auto",
)

要加载量化模型（8 位或 4 位）以进行推断，请尝试使用 bitsandbytes 并将 load_in_4bit 或 load_in_8bit 参数设置为 True。在 8 位中加载模型只需要 6.87 GB 的内存。

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

quant_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", quantization_config=quant_config, device_map="auto"
)

< > 在 GitHub 上更新

←概述单 GPU 上高效训练的方法和工具→