优化推理

使用大型语言模型（LLMs）进行推理可能具有挑战性，因为它们必须存储和处理数十亿个参数。为了加载一个 700 亿参数的 Llama 2 模型，全精度权重需要 256GB 内存，半精度权重需要 128GB 内存。当今最强大的 GPU - A100 和 H100 - 也只有 80GB 内存。

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

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

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

静态 kv-cache 和 torch.compile

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

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

静态 kv-cache 通过将 kv-cache 大小预分配为最大值来解决此问题，因此您可以将其与 torch.compile 结合使用，以实现高达 4 倍的加速。您的加速效果可能因模型大小（较大的模型加速效果较小）和硬件而异。

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

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

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

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

1. cache_implementation

2. StaticCache

3. 编译整个 generate 函数

解码策略

解码也可以进行优化以加速生成。您可以使用轻量级的辅助模型来生成候选 token，速度比 LLM 本身更快，或者您可以使用这种解码策略的变体，这种变体对于输入接地的任务尤其有效。

推测性解码

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

对于每个输入 token，在正向传递期间每次都会加载模型权重，当模型具有数十亿个参数时，这既缓慢又繁琐。推测性解码通过使用第二个更小更快的辅助模型来缓解这种减速，该辅助模型生成候选 token，这些 token 由较大的模型在单次正向传递中验证。如果验证的 token 是正确的，则 LLM 实际上“免费”获得了它们，而无需自己生成它们。精度没有下降，因为验证正向传递确保生成的输出与 LLM 自己生成它们时生成的输出相同。

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

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

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

贪婪搜索

采样

Prompt 查找解码

Prompt 查找解码是推测性解码的一种变体，它也与贪婪搜索和采样兼容。Prompt 查找对于输入接地的任务（例如摘要）特别有效，在这些任务中，prompt 和输出之间经常存在重叠的词。这些重叠的 n-gram 被用作 LLM 候选 token。

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

贪婪解码

采样

注意力机制

Transformer 模型的一个已知问题是，自注意力机制的计算和内存需求随着输入 token 数量的增加而呈二次方增长。这种限制在处理更长序列的 LLMs 中被放大。为了解决这个问题，请尝试 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 内存的限制，这可以降低内存使用率，并使加载 LLMs 以进行推理更容易。

如果您不受 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 上更新