torch.compile 모델을 ONNX+TensorRT INT8로 배포하기

서빙 환경에서 GPU 추론 지연시간을 줄이려면 torch.compile 만으로 끝내기 어렵습니다. torch.compile 은 PyTorch 런타임 안에서 그래프를 최적화해주지만, 운영에서는 보통 프레임워크 의존성을 줄이고, 엔진 캐싱, INT8 양자화 같은 배포 친화 기능이 필요합니다. 이때 가장 흔한 조합이 ONNX 변환 후 TensorRT 엔진으로 굳히는 방식입니다.

이 글에서는 다음 목표를 한 번에 다룹니다.

PyTorch 2.0의 torch.compile 모델을 운영 관점에서 다루는 법
ONNX export 시 자주 터지는 함정(동적 shape, opset, unsupported op)
TensorRT에서 INT8(PTQ)로 캘리브레이션하고 엔진을 빌드하는 실전 흐름
성능/정확도/안정성 체크리스트

추론 최적화 과정에서 메모리 압박이 커지는 경우가 많으니, 로컬 LLM이나 대형 모델에서 OOM을 겪고 있다면 Transformers 로컬 LLM OOM? KV 캐시 절감 5가지도 함께 참고하면 좋습니다.

전체 파이프라인 개요

권장 흐름은 아래와 같습니다.

PyTorch에서 모델 준비 및 eval() 고정
(선택) torch.compile 로 개발 단계에서 성능 검증
ONNX로 export(가능하면 torch.onnx.dynamo_export 계열 활용)
ONNX 그래프 정리(상수 folding, shape inference, 간단한 폴리싱)
TensorRT로 FP16 또는 INT8 엔진 빌드
캘리브레이션 데이터셋/프로파일/엔진 캐시를 배포 아티팩트로 고정

핵심 포인트는 torch.compile 결과물을 그대로 ONNX로 내보내는 것이 목적이 아니라, “PyTorch에서 검증한 모델을 ONNX로 안정적으로 내보내고, TensorRT에서 최종 최적화를 수행”하는 것입니다. 즉, torch.compile 은 개발/벤치마크 단계에서 가치가 크고, 배포 엔진은 TensorRT가 최종 책임을 지는 형태가 일반적입니다.

1) PyTorch 2.0 `torch.compile` 의 위치 정리

torch.compile 은 내부적으로 TorchDynamo가 Python 코드를 그래프 형태로 캡처하고, AOTAutograd/Inductor 등이 커널을 최적화합니다. 하지만 다음 이유로 운영 배포의 “최종 포맷”으로는 제약이 있습니다.

환경 종속성: PyTorch 버전, CUDA, 드라이버, 컴파일 캐시 경로 등
워밍업 비용: 첫 호출에서 그래프 캡처 및 컴파일이 발생
그래프 브레이크: 동적 제어 흐름/특정 연산에서 그래프가 끊기면 성능이 불안정

따라서 실무에서는 torch.compile 로 “이 모델이 어느 정도까지 빨라질 수 있는지”를 확인하고, 최종 배포는 ONNX+TensorRT로 고정하는 패턴이 많습니다.

컴파일 예시

import torch

model = ...
model.eval().cuda()

compiled = torch.compile(
    model,
    mode="max-autotune",  # 또는 "reduce-overhead"
    fullgraph=False
)

x = torch.randn(1, 3, 224, 224, device="cuda")
with torch.inference_mode():
    y = compiled(x)

여기서 중요한 점은 ONNX export는 compiled 를 대상으로 하지 않는 편이 안전하다는 것입니다. 일반적으로는 원본 model 을 export 대상으로 두고, ONNX/TensorRT에서 최적화를 수행하는 쪽이 재현성과 디버깅이 쉽습니다.

2) ONNX export: `torch.onnx.export` vs Dynamo Export

PyTorch 2.x에서는 ONNX export 경로가 크게 두 가지입니다.

torch.onnx.export: 전통적인 exporter
Dynamo 기반 exporter: torch.onnx.dynamo_export 또는 torch.onnx.export(..., dynamo=True) 같은 최신 경로(버전에 따라 API 차이 있음)

가능하면 Dynamo 기반을 먼저 검토하세요. 모델이 torch.compile 친화적일수록 Dynamo export도 잘 되는 경향이 있습니다.

ONNX export 기본 예시(동적 shape 포함)

아래 예시는 배치 차원을 동적으로 두는 전형적인 형태입니다.

import torch

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

dummy = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy,
    "model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={
        "input": {0: "batch"},
        "logits": {0: "batch"}
    }
)

opset 선택 팁

TensorRT 버전에 따라 지원 opset 범위가 다릅니다.
보수적으로는 opset_version=17 또는 16이 무난한 편이지만, 실제로는 사용 중인 TensorRT 릴리스 노트를 기준으로 결정해야 합니다.

흔한 실패 원인

aten:: 계열이 ONNX로 변환되지 못함
grid_sample, 일부 einsum, 특정 layernorm 변형 등에서 변환/지원 불가
동적 shape가 과도하게 들어가 TensorRT 최적화가 깨짐

이때는 1) 모델 구조를 약간 바꾸거나, 2) ONNX 그래프를 단순화하거나, 3) TensorRT 플러그인을 쓰는 선택지가 생깁니다.

3) ONNX 그래프 검증과 폴리싱

ONNX로 내보낸 뒤에는 최소한 다음을 수행하는 것이 좋습니다.

onnx.checker 로 모델 무결성 검사
onnxruntime 로 PyTorch와 출력이 유사한지 sanity check
(선택) onnxsim 으로 그래프 단순화

ONNX 체크 + ORT 추론 비교

import numpy as np
import onnx
import onnxruntime as ort
import torch

onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)

sess = ort.InferenceSession("model.onnx", providers=["CUDAExecutionProvider", "CPUExecutionProvider"])

x = torch.randn(4, 3, 224, 224)
with torch.inference_mode():
    torch_out = model(x).cpu().numpy()

ort_out = sess.run(["logits"], {"input": x.numpy()})[0]

max_diff = np.max(np.abs(torch_out - ort_out))
print("max_diff:", max_diff)

여기서 차이가 크게 난다면, TensorRT로 넘어가기 전에 export 자체가 잘못되었을 가능성이 큽니다.

4) TensorRT 엔진 빌드: FP16부터 안정화

INT8로 바로 가기 전에 FP16 엔진을 먼저 성공시키는 전략이 좋습니다.

FP32(ONNXRuntime)와 FP16(TensorRT) 결과 비교
성능 프로파일링으로 병목 파악
그 다음 INT8로 정확도/성능 트레이드오프를 검증

`trtexec` 로 FP16 엔진 빠르게 만들기

아래 커맨드는 환경에 따라 옵션이 달라질 수 있지만, 가장 빠른 검증용입니다.

trtexec --onnx=model.onnx --saveEngine=model_fp16.engine --fp16

동적 배치가 있다면 프로파일을 지정해야 합니다.

trtexec \
  --onnx=model.onnx \
  --saveEngine=model_fp16.engine \
  --fp16 \
  --minShapes=input:1x3x224x224 \
  --optShapes=input:8x3x224x224 \
  --maxShapes=input:32x3x224x224

5) TensorRT INT8(PTQ) 캘리브레이션

INT8은 크게 두 가지 방식이 있습니다.

PTQ: Post-Training Quantization, 캘리브레이션 데이터로 스케일을 추정
QAT: Quantization-Aware Training, 학습 단계에서 양자화 오차를 반영

TensorRT에서 흔히 쓰는 건 PTQ입니다. 이때 성패는 캘리브레이션 데이터 품질과 대표성에 달려 있습니다.

캘리브레이션 데이터 가이드

실제 트래픽 분포를 최대한 반영
전처리(정규화, resize/crop)를 서빙과 100% 동일하게
수량: 보통 수백~수천 샘플로 시작(모델/도메인에 따라 다름)

`trtexec` 로 INT8 빌드(캘리브레이션 캐시 생성)

TensorRT 버전에 따라 옵션이 조금씩 다르지만, 일반적으로는 INT8 플래그와 캘리브레이션 관련 설정을 추가합니다. 예를 들어 이미지 입력이 고정 shape라면 다음처럼 접근합니다.

trtexec \
  --onnx=model.onnx \
  --saveEngine=model_int8.engine \
  --int8 \
  --calib=/data/calib.cache \
  --shapes=input:8x3x224x224

calib.cache 는 재사용 가능한 아티팩트입니다.
동일한 모델/전처리/입력 shape 조건에서 캐시를 고정하면, CI/CD에서 엔진을 반복 빌드할 때 재현성이 좋아집니다.

만약 동적 배치/동적 해상도가 섞여 있다면, INT8은 프로파일링과 캘리브레이션 조건이 더 까다로워집니다. 이 경우 “서빙 입력 규격을 단순화”하는 것이 가장 큰 최적화인 경우가 많습니다.

6) Python에서 TensorRT 엔진 로드 및 실행

운영에서는 trtexec 로 만든 엔진을 그대로 로드해 실행하거나, Python API로 빌드/로딩을 통합합니다. 아래는 “엔진 로드 후 실행”의 최소 골격 예시입니다(버퍼 관리 등은 서비스 코드에서 래핑 권장).

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def load_engine(path: str):
    with open(path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

engine = load_engine("model_int8.engine")
context = engine.create_execution_context()

# 입력/출력 바인딩 인덱스 찾기
input_name = "input"
output_name = "logits"

inp_idx = engine.get_binding_index(input_name)
out_idx = engine.get_binding_index(output_name)

# 예시 shape
batch = 8
context.set_binding_shape(inp_idx, (batch, 3, 224, 224))

inp = np.random.randn(batch, 3, 224, 224).astype(np.float32)
out_shape = tuple(context.get_binding_shape(out_idx))
out = np.empty(out_shape, dtype=np.float32)

# GPU 메모리 할당
inp_d = cuda.mem_alloc(inp.nbytes)
out_d = cuda.mem_alloc(out.nbytes)

# H2D
cuda.memcpy_htod(inp_d, inp)

bindings = [0] * engine.num_bindings
bindings[inp_idx] = int(inp_d)
bindings[out_idx] = int(out_d)

stream = cuda.Stream()
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)

# D2H
cuda.memcpy_dtoh_async(out, out_d, stream)
stream.synchronize()

print(out.shape)

실서비스에서는 다음을 꼭 챙기세요.

엔진/컨텍스트는 프로세스 시작 시 1회 생성, 요청마다 재생성 금지
입력 shape가 동적이면 set_binding_shape 를 요청마다 수행하되, 프로파일 범위 밖으로 벗어나지 않게 제한
멀티스레드 환경에서 컨텍스트 공유 금지(컨텍스트는 스레드 안전하지 않은 경우가 많음)

7) 정확도 검증: INT8에서 무엇이 무너지는가

INT8로 가면 보통 아래 문제가 발생합니다.

확률 분포가 뭉개져 Top-1/Top-k 정확도가 하락
작은 값에 민감한 회귀 문제에서 오차가 급증
특정 레이어에서 saturation이 발생

대응 전략은 다음 순서가 실용적입니다.

FP16 엔진에서 정확도가 충분히 유지되는지 확인
INT8에서만 깨진다면 캘리브레이션 데이터 재구성
특정 레이어만 FP16으로 남기는 mixed precision(가능한 경우)
그래도 안 되면 QAT 검토

8) 운영 체크리스트: 재현성과 장애 대응

ONNX+TensorRT는 성능은 좋지만, 운영에서 “재현성”을 놓치면 디버깅이 어려워집니다.

모델 버전, ONNX opset, TensorRT 버전, CUDA/cuDNN, 드라이버 버전을 모두 고정
엔진 파일을 빌드 환경에서 생성 후 아티팩트로 배포(런타임 빌드 지양)
캘리브레이션 캐시도 함께 버전 관리
입력 전처리 코드를 학습/서빙 동일하게 유지

또한 GPU 서빙은 메모리/스레딩 이슈가 자주 터집니다. 쿠버네티스에서 운영한다면 노드 회수나 리소스 압박으로 파드가 반복 재시작되는 상황도 생길 수 있는데, 이런 케이스는 EKS Pod Eviction Loop - PDB·우선순위·Spot 정리처럼 인프라 레벨에서 함께 점검해야 합니다.

9) 자주 묻는 질문(실무에서 많이 막히는 지점)

`torch.compile` 한 모델을 그대로 ONNX로 내보내면 안 되나요?

대부분의 경우 권장하지 않습니다. torch.compile 은 실행 최적화 결과를 “다른 런타임 포맷으로 변환”하는 기능이 아니라, PyTorch 런타임 내부에서의 최적화에 가깝습니다. ONNX export는 원본 nn.Module 을 기준으로 안정적으로 내보내고, 최종 최적화는 TensorRT에서 수행하는 게 일반적인 성공 경로입니다.

동적 shape를 크게 열어두면 더 유연하지 않나요?

유연성은 늘지만 성능과 안정성이 떨어질 수 있습니다. TensorRT는 프로파일 범위 안에서 최적화를 수행하므로, 입력 규격을 단순화할수록 더 강한 최적화를 적용하기 쉽습니다.

INT8이 항상 더 빠른가요?

항상 그렇지 않습니다.

메모리 바운드 모델은 FP16과 차이가 작을 수 있음
작은 배치에서는 커널 런치 오버헤드가 지배적일 수 있음
INT8에서 레이어 폴백이 발생하면 오히려 느려질 수 있음

그래서 FP16을 기준선으로 잡고, INT8은 “성능 이득이 확실할 때만” 채택하는 전략이 안전합니다.

마무리

torch.compile 은 PyTorch 내부 최적화로 개발 단계에서 빠르게 성능을 끌어올리는 데 강력하고, ONNX+TensorRT는 운영 배포에서 엔진을 고정하고 INT8까지 적용해 지연시간과 비용을 줄이는 데 강력합니다. 둘을 경쟁 구도로 보기보다, 다음처럼 역할을 분리하면 성공 확률이 높습니다.

개발/벤치마크: torch.compile 로 상한 성능 탐색
배포/운영: ONNX로 내보내고 TensorRT에서 FP16 안정화 후 INT8(PTQ) 적용

다음 단계로는 실제 서비스 트래픽에서 p95/p99 지연시간을 관측하고, 엔진 프로파일과 함께 병목을 좁혀가는 작업이 필요합니다. CI에서 엔진 빌드 캐시가 제대로 먹지 않아 배포가 느려진다면 GitHub Actions 캐시 안 먹을 때 key·restore-keys 디버깅처럼 파이프라인 자체도 함께 최적화해보세요.