概述
Transformers 提供了多种推理优化技术,使模型运行更快、成本更低且更易获取。可选方案包括用于降低内存访问量的替代注意力机制、用于加速执行的代码编译,以及用于提升吞吐量的优化内核。将这些技术叠加使用,可获得最佳性能。
有关设备设置和特定于硬件的配置,请参阅 硬件 指南。在应用以下优化之前,请检查您的硬件支持哪些技术。
内存和速度密切相关,但并不相同。缩小内存占用会使模型“更快”,因为需要移动的数据更少。纯粹的速度优化并不总能减少内存,有时甚至会增加内存使用量。请根据您的使用场景和硬件选择合适的优化方式。
使用下表选择一种优化技术。
| 技术 | 速度 | 内存 |
|---|---|---|
| 编译 | ✅ | |
| 注意力后端 | ✅ | ✅ |
| 内核 | ✅ | ✅ |
| 量化 | ✅ | ✅ |
| 缓存 | ✅ | ✅ |
| 并行 | ✅ | |
| 连续批处理 | ✅ |
本指南将带你快速入门 Transformers 优化。
编译
torch.compile 能够减少 Python 开销、融合算子,并针对您的张量形状和硬件生成调优后的内核。首次运行会进行预热,后续运行则会使用更快的编译路径。
将 固定大小的缓存 传递给 generate() 以自动触发 torch.compile。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", dtype=torch.float16, device_map="auto")
input = tokenizer("The French Bread Law states", return_tensors="pt").to(model.device)
output = model.generate(**input, do_sample=False, max_new_tokens=20, cache_implementation="static")
tokenizer.batch_decode(output, skip_special_tokens=True)[0]避免在 generate() 之外调用
torch.compile(model),以防止模型在每一步都重新编译。
注意力后端
替代的 注意力后端 可以降低内存流量。例如,FlashAttention 通过分块处理注意力计算并避免生成大型中间张量,从而减少内存占用。
在 from_pretrained() 中设置 attn_implementation 以加载优化的注意力后端。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B", attn_implementation="flash_attention_2")内核
Kernels 通过融合操作来提升吞吐量并降低内存占用。Kernels 库以灵活且版本安全的方式从 Hub 加载优化后的计算内核。
以下示例在不安装该包的情况下加载了优化后的 FlashAttention-2 内核。
import torch
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)量化
量化会缩小每个参数的尺寸,从而降低内存占用并提升速度,因为可以执行更多操作。
将量化配置传递给 from_pretrained() 中的 quantization_config 参数。每个量化后端都有不同的配置和不同的参数。下面的示例使用 bitsandbytes 后端将模型量化为 4 位,并配置计算的数据类型。
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
model = AutoModelForCausalLM.from_pretrained(
"allenai/Olmo-3-7B-Think", quantization_config=bnb_config
)缓存
缓存通过复用过去的键和值,而不是为每个 token 重新计算,从而加速生成过程。为了抵消并降低存储过去键和值所带来的内存开销,Transformers 支持将缓存卸载到 CPU。只有当前层保留在 GPU 上。
使用 generate() 中的 cache_implementation 参数来设置缓存策略。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B", attn_implementation="kernels-community/flash-attn2"
)
inputs = tokenizer("The Le Décret Pain states that a baguette must,", return_tensors="pt")
outputs = model.generate(**inputs, do_sample=False, max_new_tokens=50, cache_implementation="offloaded")并行
并行将模型分布到多个设备上,使超出单个设备容量的模型也能快速运行。由于分片开销以及同步结果所需的通信,这种方法会占用更多内存。
张量并行将模型层拆分到多个设备上。使用 DistributedConfig(tp_size=N) 设置设备数量,并将其传递给 from_pretrained() 以启用该功能。
from transformers import AutoModelForCausalLM, DistributedConfig
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Meta-Llama-3-8B-Instruct",
distributed_config=distributed_config,
)
print(model._tp_plan)连续批处理
连续批处理通过动态调度和分块预填充让 GPU 保持忙碌,从而最大化吞吐量。服务应用使用它来并发处理多个传入请求。
使用 generate_batch() 以启用连续批处理。
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-0.6B",
attn_implementation="sdpa",
device_map="auto",
torch_dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
prompts = [
"The Le Décret Pain states that a baguette must",
"Explain gravity in one sentence.",
"Name the capital of France.",
]
inputs = [tokenizer.encode(p) for p in prompts]
generation_config = GenerationConfig(
max_new_tokens=32,
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.pad_token_id,
do_sample=False,
)
continuous_batching_config = ContinuousBatchingConfig(max_batch_tokens=512)
outputs = model.generate_batch(
inputs=inputs,
generation_config=generation_config,
continuous_batching_config=continuous_batching_config,
)
for request_id, output in outputs.items():
text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
print(f"[{request_id}] {text}")