beautifulsoups.dev
Published on

PyTorch PT2E 양자화로 INT8 모델 배포하기

Authors

서빙 환경에서 모델 최적화의 목표는 단순합니다. 지연시간은 줄이고, 비용은 낮추고, 정확도 하락은 최소화하는 것. 그중 INT8 양자화는 CPU 추론에서 특히 효과가 크고, 최근 PyTorch 2 계열에서는 torch.export를 중심으로 한 PT2E(Post Training 2 export) 양자화가 “내보내기 가능한 그래프”를 기반으로 더 일관된 변환을 제공하는 방향으로 자리 잡고 있습니다.

이 글에서는 PT2E 양자화의 핵심 개념부터, 준비(prepare) - 보정(calibration) - 변환(convert) - 내보내기(export) - 실행(runtime) 흐름을 코드로 묶어 INT8 모델 배포까지 연결합니다. 또한 실제 운영에서 흔히 겪는 실패 케이스(그래프 브레이크, 데이터 분포 불일치, 성능이 안 나오는 경우)도 함께 다룹니다.

운영 배포 관점에서 모델 최적화는 리소스 이슈와도 자주 엮입니다. 추론 파드가 OOMKilled로 죽거나, 재시작 루프에 빠지는 문제를 함께 겪는다면 아래 글도 같이 참고하세요.

PT2E 양자화란 무엇인가

PyTorch의 양자화는 크게 세 가지 축으로 이해하면 정리가 쉽습니다.

  1. Eager mode quantization: 모듈에 옵저버를 붙이고 eager하게 변환. 유연하지만 그래프 기반 최적화와 분리되어 관리가 번거로움.
  2. FX Graph mode quantization: torch.fx로 그래프를 캡처해 변환. 그래프 기반이지만 2.x의 torch.compile 및 export 흐름과 완전히 일치하지는 않음.
  3. PT2E quantization: torch.export로 얻은 exportedProgram(안정적인 그래프 표현)을 기반으로 양자화를 수행. 즉, 배포 가능한 그래프를 먼저 확정하고 그 위에서 양자화를 적용합니다.

PT2E의 장점은 다음과 같습니다.

  • export 친화적: 변환된 모델을 torch.export 기반으로 저장하거나, 이후 런타임(예: ExecuTorch, AOT 흐름 등)으로 넘기기 쉬움
  • 일관된 그래프: eager에서 발생하기 쉬운 “실행 경로마다 다른 그래프” 문제를 줄임
  • 컴파일러 및 백엔드 최적화 연결: CPU에서 x86/fbgemm 계열 최적화, ARM에서 qnnpack 등과 결합이 용이

다만 전제가 있습니다. export 가능한 형태로 모델이 작성되어야 하고, 동적 제어 흐름이나 일부 파이썬 사이드 이펙트는 그래프 캡처를 방해할 수 있습니다.

언제 INT8 양자화를 선택해야 하나

INT8 양자화가 특히 잘 맞는 경우는 다음과 같습니다.

  • CPU 추론 비중이 크고, 배치가 작아 지연시간이 중요한 서비스
  • 모델이 Conv/GEMM 비중이 큰 비전/추천/일부 NLP 계열
  • 정확도 민감도가 극단적으로 높지 않거나, 보정 데이터로 충분히 커버 가능한 경우

반대로 아래 조건에서는 먼저 다른 최적화를 고려하는 편이 낫습니다.

  • GPU 중심 추론: INT8은 GPU에서도 가능하지만, PyTorch 네이티브 양자화가 곧바로 GPU 성능으로 이어지지는 않음
  • 극단적으로 정확도 민감: 특히 분포 변화가 큰 입력에서 PTQ(Post Training Quantization)만으로는 손실이 커질 수 있음

전체 파이프라인 개요

PT2E 기반 INT8 배포는 보통 아래 순서로 진행합니다.

  1. 모델을 eval()로 두고 대표 입력으로 torch.export 가능한지 확인
  2. Quantizer 설정(백엔드, per-tensor vs per-channel 등)
  3. prepare_pt2e로 옵저버 삽입
  4. 보정 데이터로 prepared_model을 몇 백~몇 천 step 실행
  5. convert_pt2e로 INT8 연산으로 변환
  6. 변환된 모델을 저장 및 서빙 코드에 연결

이제 최소 예제로 한 번 끝까지 연결해보겠습니다.

실전 코드: PT2E로 INT8 변환하기

아래 예시는 간단한 MLP를 대상으로 하지만, 실제로는 ConvNet, Transformer의 일부 블록에도 동일한 흐름을 적용할 수 있습니다.

주의: PyTorch 버전에 따라 API 경로가 조금씩 바뀝니다. 글의 코드는 PyTorch 2.x 계열에서 동작하는 PT2E 흐름을 기준으로 작성했습니다. 환경에 따라 torchao 또는 torch.ao 하위 API가 달라질 수 있으니, 실제 적용 시 설치된 버전의 문서를 함께 확인하세요.

1) 준비: 모델과 대표 입력

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim=256, hidden=512, out_dim=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )

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

model = MLP().eval()
example_inputs = (torch.randn(8, 256),)

여기서 중요한 포인트는 두 가지입니다.

  • 반드시 eval() 상태에서 보정 및 변환을 수행(드롭아웃, 배치정규화 등 영향)
  • example_inputs는 실제 서빙 입력의 shape/분포에 가깝게 구성

2) export 그래프 확보

PT2E는 export 가능한 그래프가 핵심이므로 먼저 export부터 확인합니다.

import torch

exported = torch.export.export(model, example_inputs)

export 단계에서 자주 막히는 원인은 다음과 같습니다.

  • 입력 shape가 런타임에 따라 크게 바뀌는데, 이를 동적 shape로 표현하지 않음
  • forward 안에서 파이썬 분기, 리스트 append 같은 사이드 이펙트가 많음
  • 일부 연산이 export에서 지원되지 않음

이 경우에는 (1) 모델 코드를 정리하거나 (2) 지원되는 연산으로 치환하거나 (3) shape 제약을 명확히 하는 식으로 접근합니다.

3) Quantizer 구성 및 prepare

PT2E 양자화는 보통 XNNPACK(모바일/ARM) 또는 FBGEMM(x86 서버) 백엔드를 염두에 둡니다. 서버 CPU라면 대개 fbgemm 계열이 유리합니다.

from torch.ao.quantization.quantizer.x86_inductor_quantizer import X86InductorQuantizer
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e

quantizer = X86InductorQuantizer()
# 기본 설정을 사용하거나, 필요 시 op별/패턴별 설정을 커스터마이즈할 수 있습니다.
quantizer.set_global(torch.ao.quantization.quantizer.x86_inductor_quantizer.get_default_x86_inductor_quantization_config())

prepared = prepare_pt2e(exported, quantizer)

여기서 prepared는 옵저버(통계 수집기)가 삽입된 상태입니다. 아직 INT8로 변환된 것이 아니라, 보정 데이터를 통해 scale/zero_point 같은 양자화 파라미터를 추정하는 단계라고 보면 됩니다.

4) Calibration(보정)

보정은 “학습”이 아니라 “통계 수집”입니다. 따라서 torch.no_grad()로 대표 데이터 몇 백~몇 천 배치를 흘려보내는 식으로 진행합니다.

def calibration_data(num_batches=200, batch_size=8, in_dim=256):
    for _ in range(num_batches):
        # 실제 서비스 입력 분포를 반영하는 것이 가장 중요합니다.
        yield torch.randn(batch_size, in_dim)

prepared_module = prepared.module()
prepared_module.eval()

with torch.no_grad():
    for x in calibration_data():
        _ = prepared_module(x)

보정 데이터 품질이 성능을 좌우합니다.

  • 로그에서 샘플링한 실제 트래픽이 가장 좋음
  • 전처리(정규화, 토크나이즈 등)까지 포함한 “실제 입력”이어야 함
  • 분포가 바뀌는 서비스라면 주기적으로 재보정 또는 모니터링 필요

5) Convert: INT8 모델 생성

quantized = convert_pt2e(prepared)
quantized_module = quantized.module().eval()

with torch.no_grad():
    y = quantized_module(example_inputs[0])
    print(y.shape)

이 시점에서 연산 일부가 INT8로 바뀌며, 디퀀타이즈/퀀타이즈 노드가 적절히 삽입됩니다. 실제로 얼마나 양자화되었는지는 그래프를 프린트하거나, 프로파일링으로 확인하는 것이 좋습니다.

성능을 실제로 끌어올리는 체크리스트

INT8로 “변환”은 했는데 “빨라지지 않는” 경우가 꽤 흔합니다. 아래 항목을 우선 점검하세요.

1) 백엔드와 커널이 맞는가

  • x86 서버에서 fbgemm/inductor 경로가 제대로 타는지
  • 컨테이너 이미지에 필요한 CPU instruction set(예: AVX2/AVX512)이 가능한지
  • PyTorch 빌드가 해당 최적화를 포함하는지

2) 연산이 실제로 양자화되었는가

모델이 다음 패턴을 충분히 포함하지 않으면 양자화 이득이 작을 수 있습니다.

  • 큰 GEMM(Linear) 또는 Conv
  • 양자화 지원 패턴(예: Linear + ReLU)이 fuse 가능한 형태

반대로 LayerNorm, Softmax, 일부 activation은 양자화가 제한적이거나 효과가 작아 FP로 남는 경우가 있습니다.

3) Calibration 데이터가 대표성을 가지는가

보정 데이터가 실제 입력보다 “너무 깨끗”하면(예: 랜덤 정규분포) 실서비스에서 clipping이 늘고 정확도가 크게 떨어질 수 있습니다.

  • 추천/검색: feature 분포가 heavy-tail이면 특히 주의
  • 비전: 조명/노이즈/리사이즈 정책이 보정과 서빙에서 일치해야 함

4) 배치/스레딩 설정

CPU INT8 성능은 런타임 스레딩 영향이 큽니다.

import torch

torch.set_num_threads(8)
torch.set_num_interop_threads(1)

Kubernetes 환경이라면 CPU limit, OMP_NUM_THREADS, KMP_AFFINITY 같은 설정도 함께 봐야 합니다.

배포 관점: 저장, 로딩, 롤백 전략

PT2E 양자화 모델은 “어떤 포맷으로 저장할 것인가”가 중요합니다.

  • 단순 파일 저장: torch.save로 state를 저장하고 동일한 PyTorch 버전에서 로드
  • export 기반 저장: torch.export 결과물을 저장해 런타임에서 재사용

버전 호환성이 민감한 환경이라면, 빌드 파이프라인에서 PyTorch 버전을 고정하고, 모델 아티팩트에 메타데이터(버전, 백엔드, calibration 정보)를 함께 남기는 것을 권장합니다.

또한 운영에서는 “성능 개선”보다 “안전한 롤백”이 더 중요할 때가 많습니다.

  • FP32 모델과 INT8 모델을 동시에 배포하고 트래픽을 점진적으로 전환
  • 정확도 모니터링 지표(클릭률, 에러율, 분류 정확도 샘플링 등)를 함께 설계
  • 문제가 생기면 즉시 FP32로 되돌릴 수 있는 플래그 제공

CI에서 아티팩트 캐시가 꼬여서 이전 빌드가 재사용되지 않거나, 반대로 캐시 때문에 잘못된 바이너리가 배포되는 문제도 빈번합니다. 모델 빌드/패키징을 GitHub Actions로 운영한다면 아래 글의 체크리스트가 도움이 됩니다.

자주 겪는 오류와 디버깅 포인트

1) export 실패

  • 원인: 지원되지 않는 연산, 파이썬 제어 흐름, 동적 shape 처리 미흡
  • 대응: 모델 단순화, 연산 치환, 입력 shape 제약 명확화, export-friendly 코드로 리팩터링

2) 정확도 급락

  • 원인: calibration 데이터 분포 불일치, per-tensor 양자화로 인한 표현력 부족
  • 대응: 보정 데이터 개선, per-channel 적용 검토, 특정 레이어는 FP로 남기기

3) 성능 개선 미미

  • 원인: 양자화 가능한 연산 비중이 낮음, 스레딩/배치 설정 문제, 커널 미적용
  • 대응: 프로파일링으로 병목 확인, 백엔드 재점검, 스레딩 튜닝

마무리

PT2E 양자화는 “그래프를 먼저 고정하고 그 위에서 양자화를 수행”한다는 점에서, 배포 가능한 형태로 모델을 다루기에 훨씬 자연스러운 흐름을 제공합니다. 핵심은 세 가지입니다.

  • export 가능한 모델 구조로 정리
  • 대표성 있는 calibration 데이터로 통계 품질 확보
  • 실제 서빙 환경에서 스레딩/백엔드/프로파일링까지 포함해 검증

이 과정을 제대로 잡아두면, FP32 대비 INT8에서 눈에 띄는 지연시간 개선과 비용 절감 효과를 얻을 수 있고, 무엇보다 “배포 파이프라인” 관점에서 재현성과 롤백 가능성을 함께 확보할 수 있습니다.

Discuss on Twitter • View on GitHub