モデルのコンパイル

Quantizationはメモリを節約しますが、すべてのmatmul内で重みをdequantizeすると、bf16での生成よりも遅くなる可能性があります。torch.compileはこの処理を最適化されたカーネルに融合し、速度の大部分を回復します。

generate内で自動的にコンパイル

generate() は、これらすべての条件が満たされている場合、モデルのフォワードパスを自動的にコンパイルします。

量子化メソッドはコンパイルをサポートしています(compressed-tensors、Fine-grained FP8、GGUF、NVFP4、およびtorchao)。
モデルは、static cacheのような固定サイズのキャッシュを使用します。
モデルは CUDA、XPU、Neuron、または TPU で実行されます。
モデルのどの部分もCPUやディスクにオフロードされません。

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

モデルのコンパイル中は最初の呼び出しが遅くなります。その後、同じ形状での呼び出しではコンパイル済みカーネルが再利用されます。

自動コンパイルを無効にするには、generate() に disable_compile=True を渡します。

Update on GitHub

Transformers

モデルのコンパイル

generate内で自動的にコンパイル