추론 최적화: torch.compile, 양자화, 배치 처리

학습이 끝나면 추론 성능을 다듬어야 해

모델을 학습한 뒤에는 실제 서빙 환경에 맞게 최적화해야 해. 다음 세 기법을 자주 함께 사용해:

torch.compile(model): TorchInductor로 모델 그래프를 JIT 컴파일해. 코드 한 줄로 1.5~3배 빨라질 수 있어. 전체 지원 범위는 다음 트랙에서 다룰 거야.
양자화: 모델을 줄이고 행렬곱을 빠르게 하려고 수치 정밀도를 fp32에서 int8이나 int4로 낮춰. 현대적인 PyTorch 경로는 torchao야.
배치 처리: 샘플 32개를 하나씩 추론하는 대신 배치 하나로 묶어 처리해. 장치 활용도가 낮으면 실제 비용이 커져.

적용 순서

즉시 실행 모드와 fp32에서 정확성을 먼저 맞추고 출력을 검증해.
하드웨어가 지원하면 bf16으로 바꾸고 정확도가 유지되는지 검증해.
torch.compile(model)을 추가하고 속도와 정확도를 다시 검증해.
더 줄여야 한다면 양자화를 적용해. Transformer에는 동적 int8, LLM에는 int8 또는 int4 가중치 전용 양자화를 고려해.
요청을 서빙한다면 짧은 대기 시간 동안 들어온 요청을 하나의 배치로 묶어.

Python의 추가 비용도 잊지 마

작은 모델에서는 토큰화, 후처리, 직렬화가 추론 시간의 대부분을 차지할 수 있어. 순전파만 재지 말고 HTTP 요청부터 응답까지 전체 경로를 측정해. 해결책이 모델 변경이 아니라 토크나이저 캐시나 더 빠른 JSON 라이브러리일 때도 많아.

Code

torch.compile: 한 줄 속도 향상·python

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english",
)
model.eval()

# JIT-compile. First call is slow (compilation), subsequent calls are fast.
model = torch.compile(model)

tok = AutoTokenizer.from_pretrained("distilbert-base-uncased")
inputs = tok(["A great movie!", "Terrible."], padding=True, return_tensors="pt")

with torch.inference_mode():
    out = model(**inputs)

print(out.logits.argmax(-1))   # tensor([1, 0])

동적 int8 양자화: Transformer용 한 줄 설정·python

import torch
import torch.nn as nn
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
)
model.eval()

# Quantize Linear layers to int8 — weights stored as int8, computed in int8
quantized = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8,
)

# Compare disk size
torch.save(model.state_dict(), '/tmp/fp32.pt')
torch.save(quantized.state_dict(), '/tmp/int8.pt')

import os
print(f"fp32 : {os.path.getsize('/tmp/fp32.pt')/1e6:.1f} MB")
print(f"int8 : {os.path.getsize('/tmp/int8.pt')/1e6:.1f} MB")
# fp32: 268.4 MB
# int8:  72.1 MB     (~4x smaller)

torchao: 현대적인 int4 / 가중치 전용·python

# pip install torchao
import torchao
from torchao.quantization import int4_weight_only, int8_weight_only
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2")

# int8 weight-only — weights stored as int8, computed in fp16/bf16
torchao.quantize_(model, int8_weight_only())

# int4 weight-only — even smaller, designed for LLMs
# torchao.quantize_(model, int4_weight_only(group_size=128))

# Use the model normally — the int4/int8 layers handle on-the-fly dequant
# in their forward.

처음부터 끝까지 추론 성능 측정·python

import time
import torch

@torch.inference_mode()
def benchmark(model, inputs, n=100, warmup=10):
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    for _ in range(warmup):
        model(**inputs)
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n):
        model(**inputs)
    if next(model.parameters()).is_cuda:
        torch.cuda.synchronize()
    elapsed_ms = (time.perf_counter() - t0) / n * 1000
    print(f"{elapsed_ms:.2f} ms / call")
    return elapsed_ms

Exercise

작은 모델 하나를 골라 세 설정의 추론 지연 시간을 재 봐. 즉시 실행 fp32, 즉시 실행 bf16, 컴파일한 bf16을 비교한 뒤 int8 동적 양자화도 추가해. 각 설정의 지연 시간과 별도 검증 배치에서 실행한 간단한 정확성 검사 결과를 함께 기록해. 실제 배포를 결정할 때 이런 표가 꼭 필요해.