大型语言模型推理优化¶
大型语言模型（LLMs）通过生成具有高度理解和流畅性的文本，将聊天和代码补全等文本生成应用提升到了新的水平。然而，使 LLMs 如此强大的一个因素——它们的体积——也为推理带来了挑战。
基本的推理速度较慢，因为每次生成下一个标记时都需要反复调用 LLM。随着生成的进行，输入序列会逐渐变长，这使得 LLM 的处理时间越来越长。此外，LLMs 拥有数十亿个参数，这些参数在内存中存储和处理非常困难。
本指南将展示如何使用 Transformers 中的优化技术来加速 LLM 推理。
Hugging Face 还提供了 文本生成推理 (TGI)，这是一个专门用于部署和提供高度优化的 LLM 推理的库。它包括一些 Transformers 不包含的面向部署的优化功能，例如连续批处理以提高吞吐量和多 GPU 推理的张量并行性。
静态键值缓存和 torch.compile¶
在解码过程中，LLM 会计算每个输入标记的键值对（kv 值）。由于它是自回归的，每次生成的输出都会成为新的输入，因此会重复计算相同的 kv 值。这效率不高，因为它每次都在重新计算相同的 kv 值。
为了优化这一点，可以使用 kv 缓存来存储过去的键和值，而不是每次都重新计算。然而，由于 kv 缓存在每次生成步骤中会增长，并且是动态的，这阻碍了你利用 torch.compile，这是一个强大的优化工具，可以将 PyTorch 代码融合成快速且优化的内核。我们有一整篇关于 kv 缓存的指南 在这里。
静态 kv 缓存解决了这个问题，通过预先分配最大值来预分配 kv 缓存的大小，允许你将其与 torch.compile 结合使用，从而实现最高 4 倍的速度提升。具体速度提升取决于模型大小（较大的模型速度提升较小）和硬件。
目前，只有 Llama 和少数其他模型支持静态 kv 缓存和 torch.compile。请查看 这个 issues 获取实时的模型兼容列表。
根据任务的复杂性，静态 kv 缓存有三种使用方式：
- 基本使用：只需在
generation_config中设置一个标志（推荐）； - 高级使用：处理一个多轮生成或自定义生成循环的缓存对象；
- 高级使用：如果单图对你有意义，可以将整个
generate函数编译成单个图。
选择下面的正确标签以获取每种方式的详细说明。
无论使用哪种策略与 torch.compile 结合，如果你将 LLM 输入左填充到有限的几个值，可以避免与形状相关的重新编译。pad_to_multiple_of 分词器标志 是你的朋友！
基本使用：generation_config¶
高级使用：控制静态缓存¶
高级使用：端到端生成编译¶
以下示例使用 Gemma 模型。所有需要做的就是：
- 访问模型的
generation_config属性并将cache_implementation设置为“static”； - 调用
torch.compile对模型进行编译，以便在静态 kv 缓存下运行前向传播。
就是这样！
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false" # 防止长时间警告
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", torch_dtype="auto", device_map="auto")
model.generation_config.cache_implementation = "static"
model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device.type)
outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# 输出：['The theory of special relativity states 1. The speed of light is constant in all inertial reference']
在内部，generate 将尝试重用同一个缓存对象，从而避免每次调用时的重新编译。避免重新编译对于充分利用 torch.compile 至关重要，你应该注意以下几点：
- 如果调用之间批量大小发生变化或最大输出长度增加，则缓存需要重新初始化，触发新的编译；
- 编译函数的前几次调用较慢，因为函数正在被编译。
对于更高级的静态缓存使用，例如多轮对话，我们建议在 generate() 外部实例化和操作缓存对象。请参见高级使用标签。
投机解码¶
有关更深入的解释，请参阅 辅助生成：通向低延迟文本生成的新方向 博客文章！
自回归的一个问题是，对于每个输入标记，你需要在前向传播期间每次加载模型权重。这对拥有数十亿参数的 LLM 来说既慢又繁琐。投机解码通过使用第二个较小且更快的辅助模型来生成候选标记，这些候选标记由较大的 LLM 在单次前向传播中验证来缓解这种缓慢。如果验证的标记是正确的，LLM 实质上可以“免费”获得这些标记，而不需要自己生成。没有准确性下降，因为验证前向传播确保生成的输出与 LLM 自己生成的输出相同。
为了获得最大的速度提升，辅助模型应该比 LLM 小得多，这样它可以快速生成标记。辅助模型和 LLM 模型还必须共享相同的分词器，以避免重新编码和解码标记。
启用投机解码的方法是加载一个辅助模型并将其传递给 generate() 方法。
贪婪搜索¶
采样¶
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from accelerate.test_utils.testing import get_backend
device, _, _ = get_backend() # 自动检测底层设备类型（CUDA, CPU, XPU, MPS 等）
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("Einstein's theory of relativity states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype="auto").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# 输出：["Einstein's theory of relativity states that the speed of light is constant."]
提示查找解码¶
提示查找解码是一种投机解码的变体，也适用于贪婪搜索和采样。提示查找特别适合于输入依赖的任务，例如摘要，其中提示和输出之间经常有重叠的单词。这些重叠的 n-gram 用作 LLM 的候选标记。
要启用提示查找解码，指定 prompt_lookup_num_tokens 参数中的应重叠的标记数量。然后你可以将此参数传递给 generate() 方法。
贪婪解码¶
采样¶
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from accelerate.test_utils.testing import get_backend
device, _, _ = get_backend() # 自动检测底层设备类型（CUDA, CPU, XPU, MPS 等）
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b")
inputs = tokenizer("The second law of thermodynamics states", return_tensors="pt").to(device)
model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype="auto").to(device)
assistant_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m").to(device)
outputs = model.generate(**inputs, prompt_lookup_num_tokens=3)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# 输出：['The second law of thermodynamics states that entropy increases with temperature.']
注意力优化¶
变压器模型的一个已知问题是，自注意力机制的计算和内存需求随着输入标记数量的增加而呈二次增长。这一限制在处理更长序列的 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",
)
使用 torch.compile 和无填充数据收集器进行微调¶
除了优化推理外，还可以通过在微调过程中利用 torch.compile 和无填充数据收集器来提高大型语言模型的训练效率。这种方法可以显著加快训练速度并减少计算开销。
以下是如何使用 TRL 库中的 SFTTrainer 微调 Llama 模型，同时启用 torch_compile 并使用无填充数据收集器：
#################### 导入 ###################
import math
import datasets
import dataclasses
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments
)
from trl import SFTConfig, SFTTrainer, DataCollatorForCompletionOnlyLM
#################### 加载模型并启用 Flash Attention ###################
model_name = "meta-llama/Llama-3.2-1B"
model = AutoModelForCausalLM.from_pretrained(
model_name,
attn_implementation="flash_attention_2" # 启用 FlashAttention-2
)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
#################### 数据预处理（无填充） ###################
response_template = "\n### Label:"
response_template_ids = tokenizer.encode(
response_template, add_special_tokens=False
)[2:] # 排除特殊标记
data_collator = DataCollatorForCompletionOnlyLM(
response_template_ids=response_template_ids,
tokenizer=tokenizer,
ignore_index=-100,
padding_free=True # 启用无填充收集
)
def format_dataset(example):
return {
"output": example["output"] + tokenizer.eos_token
}
data_files = {"train": "path/to/dataset"} # 替换为你的数据集路径
json_dataset = datasets.load_dataset("json", data_files=data_files)
formatted_train_dataset = json_dataset["train"].map(format_dataset)
################# 训练配置 ############################
train_args = TrainingArguments(
num_train_epochs=5,
per_device_train_batch_size=4,
per_device_eval_batch_size=4,
gradient_accumulation_steps=4,
learning_rate=1e-5,
weight_decay=0.0,
warmup_ratio=0.03,
lr_scheduler_type="cosine",
logging_steps=1,
include_tokens_per_second=True,
save_strategy="epoch",
output_dir="output",
torch_compile=True, # 启用 torch.compile
torch_compile_backend="inductor",
torch_compile_mode="default"
)
# 将 TrainingArguments 转换为 SFTConfig
transformer_train_arg_fields = [x.name for x in dataclasses.fields(SFTConfig)]
transformer_kwargs = {
k: v
for k, v in train_args.to_dict().items()
if k in transformer_train_arg_fields
}
training_args = SFTConfig(**transformer_kwargs)
####################### 微调 #####################
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=formatted_train_dataset,
data_collator=data_collator,
dataset_text_field="output",
args=training_args,
)
trainer.train()
PyTorch 缩放点积注意力¶
缩放点积注意力（SDPA）在 PyTorch 2.0 中自动启用，并且支持 FlashAttention、xFormers 和 PyTorch 的 C++ 实现。SDPA 会在使用 CUDA 后端时选择性能最高的注意力算法。对于其他后端，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 权重来减小其大小。这降低了内存使用量，使得在受限于 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"
)