Компиляция модели

Квантование экономит память, но деквантование весов внутри каждого matmul может сделать генерацию медленнее, чем в bf16. torch.compile объединяет эту работу в оптимизированные ядра и восстанавливает значительную часть скорости.

Автоматическая компиляция в generate

generate() автоматически компилирует прямой проход модели, если выполняются все перечисленные ниже условия.

• Метод квантования поддерживает компиляцию (compressed-tensors, Fine-grained FP8, GGUF, NVFP4 и torchao).
• Модель использует кэш фиксированного размера, как статический кэш.
• Модель работает на CUDA, XPU, Neuron или TPU.
• Ни одна часть модели не выгружается в оперативную память или на диск.

Передайте cache_implementation="static" в generate(). В примере ниже используется квантование только весов int8 с помощью 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))

Первые вызовы выполняются медленно, пока модель компилируется. Последующие вызовы с теми же формами повторно используют скомпилированные ядра.

Передайте disable_compile=True в generate(), чтобы отключить автоматическую компиляцию.