Biên dịch một mô hình

Quantization giúp tiết kiệm bộ nhớ, nhưng việc khử lượng hóa trọng số bên trong mỗi phép nhân ma trận có thể khiến quá trình sinh chậm hơn so với bf16. torch.compile hợp nhất công việc đó vào các kernel được tối ưu hóa và khôi phục phần lớn tốc độ.

Tự động biên dịch trong generate

generate() tự động biên dịch bước truyền thẳng của mô hình khi tất cả các điều kiện sau được thỏa mãn.

• Phương pháp lượng tử hóa hỗ trợ biên dịch (compressed-tensors, Fine-grained FP8, GGUF, NVFP4 và torchao).
• Mô hình sử dụng bộ nhớ đệm có kích thước cố định, giống như static cache.
• Mô hình chạy trên CUDA, XPU, Neuron hoặc TPU.
• Không có phần nào của mô hình được chuyển xuống CPU hoặc ổ đĩa.

Truyền cache_implementation="static" vào generate(). Ví dụ bên dưới sử dụng lượng tử hóa chỉ trọng số int8 của torchao.

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))

Các lần gọi đầu tiên sẽ chậm trong khi mô hình đang được biên dịch. Các lần gọi sau đó với cùng các hình dạng sẽ tái sử dụng các kernel đã được biên dịch.

Truyền disable_compile=True vào generate() để tắt biên dịch tự động.