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).
编译模型
量化可以节省内存,但在每次矩阵乘法内部对权重进行反量化可能会使生成速度比 bf16 更慢。torch.compile 会将这些操作融合到优化的内核中,从而恢复大部分速度。
在 generate 中自动编译
当以下条件全部满足时,generate() 会自动编译模型的前向传播。
- 量化方法支持编译(compressed-tensors、Fine-grained FP8、GGUF、NVFP4 和 torchao)。
- 该模型使用固定大小的缓存,例如 static cache。
- 该模型在 CUDA、XPU、Neuron 或 TPU 上运行。
- 模型的任何部分都不会被卸载到 CPU 或磁盘上。
将 cache_implementation="static" 传递给 generate()。以下示例使用 torchao int8 仅权重量化。
import torch
from torchao.quantization import Int8WeightOnlyConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
model_id = "Qwen/Qwen3-0.6B"
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=TorchAoConfig(quant_type=Int8WeightOnlyConfig()),
device_map="cuda",
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The capital of France is", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=32, cache_implementation="static")
print(tokenizer.decode(outputs[0], skip_special_tokens=True))在模型编译期间,最初的调用速度较慢。后续具有相同形状的调用会复用已编译的内核。
将 disable_compile=True 传递给 generate() 以关闭自动编译。