모델 컴파일

Quantization은 메모리를 절약하지만, 모든 matmul 내부에서 가중치를 dequantize하는 작업은 bf16보다 생성 속도를 늦출 수 있습니다. torch.compile은 해당 작업을 최적화된 커널로 융합하여 상당 부분의 속도를 회복합니다.

generate에서 자동으로 컴파일

generate()는 다음 조건이 모두 충족되면 모델의 forward pass를 자동으로 컴파일합니다.

양자화 방법은 컴파일(compressed-tensors, Fine-grained FP8, GGUF, NVFP4, torchao)을 지원합니다.
모델은 static cache와 같이 고정된 크기의 캐시를 사용합니다.
모델은 CUDA, XPU, Neuron 또는 TPU에서 실행됩니다.
모델의 어떤 부분도 CPU나 디스크로 오프로드되지 않습니다.

cache_implementation="static"를 generate()에 전달합니다. 아래 예제는 torchao int8 weight-only 양자화를 사용합니다.

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

모델이 컴파일되는 동안 첫 번째 호출은 느립니다. 동일한 shape을 가진 이후의 호출은 컴파일된 커널을 재사용합니다.

자동 컴파일을 끄려면 generate()에 disable_compile=True를 전달하세요.

Update on GitHub

Transformers

모델 컴파일

generate에서 자동으로 컴파일