Compilar un modelo

La cuantización ahorra memoria, pero descuantizar los pesos dentro de cada matmul puede hacer que la generación sea más lenta que en bf16. torch.compile fusiona ese trabajo en kernels optimizados y recupera gran parte de la velocidad.

Compilar automáticamente en generate

generate() compila automáticamente el paso hacia adelante del modelo cuando se cumplen todas estas condiciones.

• El método de cuantización admite la compilación (compressed-tensors, Fine-grained FP8, GGUF, NVFP4 y torchao).
• El modelo utiliza una caché de tamaño fijo, como la static cache.
• El modelo se ejecuta en CUDA, XPU, Neuron o TPU.
• Ninguna parte del modelo se descarga en la CPU o en el disco.

Pase cache_implementation="static" a generate(). El ejemplo de abajo usa la cuantización int8 solo de pesos de 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))

Las primeras llamadas son lentas mientras el modelo se compila. Las llamadas posteriores con las mismas formas reutilizan los núcleos compilados.

Pasa disable_compile=True a generate() para desactivar la compilación automática.