Transformers documentation

概述

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

概述

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}")
Update on GitHub