优化推理

使用大型语言模型（LLM）进行推理可能具有挑战性，因为它们需要存储和处理数十亿个参数。加载一个 70B 参数的 Llama 2 模型，需要 256GB 内存用于全精度权重，128GB 内存用于半精度权重。而目前最强大的 GPU（A100 和 H100）只有 80GB 内存。

除了内存需求之外，推理速度也很慢，因为 LLM 需要反复调用以生成下一个 token。随着生成过程的进行，输入序列会不断增长，处理时间也越来越长。

本指南将向您展示如何优化 LLM 推理以加速生成并减少内存使用。

尝试使用 Text Generation Inference (TGI)，这是一个 Hugging Face 库，专门用于部署和提供高度优化的 LLM 进行推理。

静态 kv-缓存和 torch.compile

LLM 会为每个输入 token 计算键值（kv）对，并且每次都会执行相同的 kv 计算，因为生成的输出成为输入的一部分。然而，每次都执行相同的 kv 计算效率不高。

kv-缓存会存储过去的键和值，而不是每次都重新计算它们。因此，kv-缓存是动态的，并且随着每个生成步骤而增长，这使得您无法利用 torch.compile，这是一种将 PyTorch 代码融合到优化内核中的强大优化方法。

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

关注此问题以跟踪哪些模型（Llama、Gemma、Mistral 等）支持静态 kv-缓存和 torch.compile。

根据您的任务，有几种方法可以使用静态 kv-缓存。

对于基本用例，将 cache_implementation 设置为 "static"（推荐）。
对于多轮生成或自定义生成循环，请直接初始化和处理 StaticCache。
对于更独特的硬件或用例，将整个 generate() 函数编译成单个图可能更好。

无论您如何使用静态 kv-缓存和 torch.compile，请使用 pad_to_multiple_of 将输入左填充到有限的一组值，以避免与形状相关的重新编译。

1. 缓存实现

2. StaticCache

3. 编译整个生成函数

解码策略

解码也可以优化以加速生成。您可以使用一个轻量级的辅助模型来比 LLM 本身更快地生成候选 token，或者您可以使用这种解码策略的一个变体，它在以输入为基础的任务中特别有效。

推测解码

要了解更深入的解释，请查阅辅助生成：低延迟文本生成的新方向博客文章！

对于每个输入 token，模型权重在前向传播期间每次都会加载，这在模型拥有数十亿参数时既慢又麻烦。推测解码通过使用第二个更小、更快的辅助模型来生成候选 token，这些 token 在单个前向传播中由更大的模型进行验证，从而缓解了这种减速。如果验证的 token 是正确的，LLM 基本可以“免费”获得它们，而无需自己生成它们。准确性没有下降，因为验证前向传播确保生成与 LLM 自己生成时相同的输出。

为了获得最大的加速，辅助模型应该比 LLM 小得多，以便它可以快速生成 token。辅助模型和 LLM 模型还必须共享相同的 tokenizer，以避免重新编码和解码 token。

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

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

贪婪搜索

采样

提示查找解码

提示查找解码是推测解码的一种变体，也兼容贪婪搜索和采样。提示查找在输入接地任务（例如摘要）中特别有效，因为提示和输出之间通常存在重叠词。这些重叠的 N-gram 用作 LLM 候选 token。

要启用提示查找解码，请在 prompt_lookup_num_tokens 参数中指定应重叠的 token 数量。然后将此参数传递给 generate()。

贪婪解码

采样

注意力机制

Transformer 模型的一个已知问题是，自注意力机制在计算和内存方面随输入 token 的数量呈二次增长。这一限制在处理更长序列的 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++ 实现。

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

使用 torch.nn.attention.sdpa_kernel 上下文管理器来显式启用或禁用四种注意力算法中的任何一种。例如，使用 SDPBackend.FLASH_ATTENTION 启用 FlashAttention。

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM

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

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)

量化

量化通过以较低精度存储模型权重来减小模型大小。这会降低内存使用量，并在 GPU 内存受限时使 LLM 推理更易于访问。

如果您不受 GPU 限制，则不必对模型进行量化，因为它可能会略微增加延迟（AWQ 和融合 AWQ 模块除外），因为需要额外的步骤来量化和反量化权重。

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

使用下面的模型内存计算器估算和比较加载模型所需的内存。例如，尝试估算加载 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 上更新