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

编译模型

量化可以节省内存,但在每次矩阵乘法内部对权重进行反量化可能会使生成速度比 bf16 更慢。torch.compile 会将这些操作融合到优化的内核中,从而恢复大部分速度。

在 generate 中自动编译

当以下条件全部满足时,generate() 会自动编译模型的前向传播。

将 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() 以关闭自动编译。

Update on GitHub