PyTorch 2.x torch.compile로 추론속도↑ 경량화

서빙 환경에서 PyTorch 모델을 최적화할 때 가장 흔한 목표는 두 가지입니다. 첫째는 추론 지연시간을 줄여 QPS를 올리는 것, 둘째는 GPU 메모리와 배포 아티팩트를 줄여 운영비를 낮추는 것입니다. PyTorch 2.x의 torch.compile은 이 두 목표 중 특히 “속도”에 강력한 기본기 옵션이고, 여기에 혼합정밀도, 양자화, (필요 시) ONNX/TensorRT 같은 런타임을 조합하면 “경량화”까지 함께 가져갈 수 있습니다.

이 글에서는 torch.compile의 동작 원리, 어떤 모델에서 효과가 큰지, 성능을 망치는 흔한 함정, 그리고 경량화를 위한 현실적인 조합(정밀도, 양자화, 그래프 안정화)을 코드 중심으로 정리합니다.

torch.compile이 하는 일: 한 줄로 그래프 최적화 파이프라인 연결

PyTorch 2.x의 torch.compile은 대략 다음 흐름으로 동작합니다.

TorchDynamo가 파이썬 바이트코드를 가로채서 그래프 단위로 캡처
AOTAutograd가 (필요 시) forward 및 backward 그래프를 더 잘 최적화할 형태로 분해
Inductor가 커널 퓨전, 연산 재배치, 코드 생성(Triton 등)을 통해 실행을 빠르게 만듦

중요한 점은 “모든 코드가 자동으로 그래프화되는 것은 아니다”라는 것입니다. 데이터 의존 분기, 파이썬 컨트롤 플로우, 동적 shape, 모델 내부의 커스텀 연산 등이 섞이면 그래프가 자주 끊기고, 그러면 기대한 속도 향상을 못 얻을 수 있습니다.

언제 빨라지나: 효과가 큰 패턴과 작은 패턴

torch.compile의 효과가 큰 경우는 보통 다음과 같습니다.

작은 연산이 많이 쪼개져 있어 커널 런치 오버헤드가 큰 모델
연산 퓨전 여지가 큰 MLP, Conv 블록, Attention 주변의 elementwise 연산
동일한 shape로 반복 호출되는 온라인 추론(그래프 재사용 가능)

반대로 효과가 제한적이거나 튜닝이 필요한 경우는 다음이 많습니다.

입력 shape가 요청마다 크게 달라지는 서비스(그래프 캐시가 깨짐)
파이썬 분기/루프가 많고 텐서 연산이 연속적이지 않은 전처리 포함 파이프라인
이미 TensorRT 등으로 최적화된 엔진을 쓰는 경우(중복 투자)

가장 먼저 적용하는 최소 코드: eval + no_grad + compile

추론 최적화의 기본은 model.eval()과 torch.no_grad()입니다. 여기에 torch.compile을 얹는 최소 예시는 아래와 같습니다.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, d=1024):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d, 4 * d),
            nn.GELU(),
            nn.Linear(4 * d, d),
        )

    def forward(self, x):
        return self.net(x)


def benchmark(model, x, iters=200, warmup=50):
    # GPU 타이밍은 synchronize가 필요
    for _ in range(warmup):
        _ = model(x)
    torch.cuda.synchronize()

    import time
    t0 = time.time()
    for _ in range(iters):
        _ = model(x)
    torch.cuda.synchronize()
    return (time.time() - t0) / iters


device = "cuda" if torch.cuda.is_available() else "cpu"
model = MLP().to(device).eval()

x = torch.randn(32, 1024, device=device)

with torch.no_grad():
    t_eager = benchmark(model, x)

    compiled = torch.compile(model, mode="max-autotune")
    t_compiled = benchmark(compiled, x)

print("eager:", t_eager)
print("compiled:", t_compiled)

mode는 대표적으로 "default", "reduce-overhead", "max-autotune"가 있습니다.
온라인 서빙에서 지연시간이 중요한 경우 "reduce-overhead"가 유리한 케이스가 있고, 처리량을 끌어올리고 싶으면 "max-autotune"가 유리한 케이스가 많습니다.
첫 호출은 컴파일 오버헤드가 있으니, 서버 시작 시 워밍업 요청으로 컴파일을 끝내는 전략이 필요합니다.

추론속도 올리는 실전 조합 1: torch.compile + AMP(bfloat16/float16)

실제 운영에서 속도와 메모리 모두를 가장 쉽게 개선하는 조합은 혼합정밀도입니다. 특히 Ampere 이후 GPU에서는 bfloat16이 안정적인 선택인 경우가 많습니다.

import torch

device = "cuda"
model = model.to(device).eval()
compiled = torch.compile(model, mode="max-autotune")

x = torch.randn(32, 1024, device=device)

with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    y = compiled(x)

팁:

float16은 더 빠를 수 있지만 overflow/underflow 민감도가 커서 모델에 따라 품질 이슈가 생길 수 있습니다.
bfloat16은 표현 범위가 넓어 안정성이 좋지만, GPU/드라이버/커널에 따라 성능 차이가 있습니다.
품질 검증은 반드시 해야 합니다. 특히 분류 확률, 회귀 오차, 생성 모델의 샘플 품질 등.

추론속도 올리는 실전 조합 2: dynamic shape를 줄여 그래프 캐시를 살린다

torch.compile은 입력 shape가 바뀔 때 그래프가 다시 컴파일되거나 그래프가 쪼개질 수 있습니다. 서빙에서 입력 길이가 가변인 NLP 모델(예: 토큰 길이)이라면 다음 중 하나를 고려합니다.

길이를 버킷팅해서 몇 개의 대표 shape로만 들어오게 만들기
패딩 전략을 통일해서 shape 변동을 줄이기
정말 동적 shape가 필수라면, 컴파일 모드나 설정을 보수적으로 가져가기

버킷팅 예시:

import math

def bucket_length(n, buckets=(64, 128, 256, 512, 1024)):
    for b in buckets:
        if n <= b:
            return b
    return int(2 ** math.ceil(math.log2(n)))

# 실제 토큰 길이 n을 대표 길이로 올림
n = 137
L = bucket_length(n)
# input_ids를 L로 패딩해서 모델에 넣는다

이 전략은 “컴파일 캐시 히트율”을 높여 지연시간 분산을 줄이는 데도 도움이 됩니다.

모델 경량화 관점: torch.compile은 파일 크기를 줄이지 않는다

많이 오해하는 지점이 있습니다.

torch.compile은 주로 실행 그래프와 커널 실행을 최적화하는 기능입니다.
모델 파라미터 수를 줄이거나, 체크포인트 파일 크기를 줄이는 “경량화” 자체를 직접 해주지는 않습니다.

경량화 목표가 다음 중 무엇인지 먼저 정의해야 합니다.

GPU VRAM 사용량 감소
CPU RAM 사용량 감소
디스크에 저장되는 모델 파일 크기 감소
네트워크 전송량 감소(배포 아티팩트)

대부분의 경우, 경량화는 정밀도 변경(예: FP32에서 BF16/FP16), 양자화(INT8/INT4), 구조적 변경(프루닝, 지식 증류)로 달성합니다.

경량화 실전 조합 1: BF16/FP16 체크포인트로 저장하기

추론에서 BF16/FP16으로 돌릴 거라면, 아예 가중치를 해당 dtype으로 저장해 로딩 메모리와 전송량을 줄일 수 있습니다.

import torch

model = model.eval().to("cuda")
model = model.to(dtype=torch.bfloat16)

state = model.state_dict()

# 저장
torch.save(state, "model_bf16.pt")

# 로드
loaded = MLP().to("cuda").eval().to(dtype=torch.bfloat16)
loaded.load_state_dict(torch.load("model_bf16.pt", map_location="cuda"))

주의:

일부 레이어(예: LayerNorm, Softmax 주변)는 FP32가 더 안정적인 경우가 있습니다. 이 경우 mixed precision 정책을 더 세밀하게 적용해야 합니다.

경량화 실전 조합 2: 동적 양자화(주로 CPU)로 모델 크기와 지연시간 절충

CPU 추론(또는 CPU fallback)이 중요하면 동적 양자화가 빠른 승부처가 됩니다. Linear 중심 모델(Transformer의 FFN 등)에 특히 잘 먹습니다.

import torch
import torch.nn as nn

model = MLP().eval().cpu()

# Linear 레이어를 동적 양자화
qmodel = torch.ao.quantization.quantize_dynamic(
    model,
    {nn.Linear},
    dtype=torch.qint8,
)

x = torch.randn(32, 1024)
with torch.no_grad():
    y = qmodel(x)

장점: 구현이 간단하고 모델 크기 감소 효과가 즉시 나옴
단점: GPU 추론에는 그대로 적용하기 어렵고, 정확도 영향이 있을 수 있음

경량화 실전 조합 3: torch.compile과 양자화는 “순서”가 중요

현장에서 자주 겪는 문제는 “컴파일한 모델에 양자화를 적용”하거나 “양자화한 모델을 컴파일”할 때 예상치 못한 그래프 브레이크나 성능 저하가 생기는 것입니다.

권장 접근:

먼저 기준선을 만든다: eager FP32, eager AMP, compile AMP
그 다음 경량화: (CPU면) dynamic quantization, (GPU면) PTQ/QAT 또는 외부 엔진
마지막으로 조합: 양자화된 모델이 torch.compile 대상이 되는지 확인

모델/백엔드 조합에 따라 정답이 달라, 측정 기반으로 결정하는 것이 안전합니다.

성능 튜닝 체크리스트: 빨라졌는데도 느린 이유 7가지

컴파일 오버헤드를 요청 경로에 포함함

해결: 서버 시작 시 워밍업으로 컴파일 완료

입력 shape 변동이 심해 그래프 캐시 미스

해결: 버킷팅/패딩으로 shape 수를 제한

전처리/후처리가 병목

해결: 전처리를 텐서화하고 GPU로 올리거나, 병렬화/배치 처리

model.train() 상태로 추론

해결: eval() 고정, dropout/bn 동작 확인

작은 배치에서 커널 런치 오버헤드가 지배

해결: 마이크로배칭, 요청 합치기, mode="reduce-overhead" 실험

동기화(torch.cuda.synchronize())를 과도하게 호출

해결: 측정 코드 외에는 동기화 최소화

메모리 OOM이나 파편화로 성능이 흔들림

해결: 배치 조절, 캐시 정책 점검, 필요 시 메모리 최적화 기법 병행

VRAM이 빡빡한 생성 모델 계열이라면, 메모리 최적화는 속도와 직결되기도 합니다. 관련해서는 Stable Diffusion VRAM OOM 해결 - xformers·VAE 타일링 같은 접근이 동일한 맥락에서 도움이 됩니다.

운영 배포 팁: 재시도, 타임아웃, 리소스 고갈까지 같이 본다

추론이 빨라져도 운영 안정성이 떨어지면 의미가 없습니다. 예를 들어 모델 서버가 외부 API를 호출하거나, 큐 기반으로 트래픽을 받는다면 Rate Limit과 백오프 설계가 중요합니다. 이런 관점에서는 OpenAI 429/Rate Limit 재시도·백오프 설계의 패턴을 그대로 적용할 수 있습니다.

또 하나의 현실 이슈는 파일 디스크립터 고갈입니다. 모델 서버가 로깅, 소켓, 파일 IO를 많이 쓰면 Too many open files가 성능 저하나 장애로 이어질 수 있습니다. 이 경우 리눅스 Too many open files - ulimit·fd 해결 체크리스트가 도움이 됩니다.