Compiler un modèle

La quantification permet d’économiser de la mémoire, mais la déquantification des poids à l’intérieur de chaque matmul peut rendre la génération plus lente qu’en bf16. torch.compile fusionne ce travail dans des noyaux optimisés et récupère une grande partie de la vitesse.

Compiler automatiquement dans generate

generate() compile automatiquement le passage avant du modèle lorsque toutes ces conditions sont remplies.

• La méthode de quantification prend en charge la compilation (compressed-tensors, Fine-grained FP8, GGUF, NVFP4 et torchao).
• Le modèle utilise un cache de taille fixe, comme le static cache.
• Le modèle s'exécute sur CUDA, XPU, Neuron ou TPU.
• Aucune partie du modèle n'est déchargée sur le CPU ou le disque.

Passez cache_implementation="static" à generate(). L'exemple ci-dessous utilise la quantification int8 weight-only 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))

Les premiers appels sont lents pendant que le modèle est compilé. Les appels ultérieurs avec les mêmes formes réutilisent les noyaux compilés.

Passez disable_compile=True à generate() pour désactiver la compilation automatique.