beautifulsoups.dev
Published on

PyTorch 2 Export+PT2E로 INT8 양자화 실전

Authors

서빙 비용을 줄이려면 보통 FP32 -> FP16을 먼저 떠올리지만, CPU 추론이거나 메모리/캐시가 병목인 환경에서는 INT8이 체감 효과가 더 큰 경우가 많습니다. 문제는 “어떻게 안정적으로 INT8을 만들고, 실제로 빨라졌는지 검증하며, 배포 파이프라인까지 연결하느냐”입니다.

PyTorch 2에서는 torch.export를 중심으로 그래프를 고정하고, 그 위에서 PT2E(Post Training 2 Export) 양자화를 적용하는 흐름이 자리 잡았습니다. 이 글은 Export+PT2E로 INT8 양자화를 끝까지(캘리브레이션부터 속도/정확도 검증, 실전 함정) 다룹니다.

또한 모델 최적화는 결국 배포 환경에서 빛을 봅니다. 컨테이너 빌드/캐시를 같이 손보면 반복 실험 속도가 크게 올라가니, 필요하면 Docker 빌드 70% 단축 - 멀티스테이지·캐시 전략도 함께 참고하세요.

왜 PyTorch 2 Export+PT2E인가

기존의 PTQ(Post Training Quantization)나 QAT(Quantization Aware Training)는 torch.fx 기반 흐름이 많았고, 모델에 따라 그래프 캡처가 흔들리거나(특히 control flow, 동적 shape) 변환 포인트가 애매한 경우가 있었습니다.

PyTorch 2의 방향은 다음처럼 정리할 수 있습니다.

  • torch.export로 입력 스펙에 대해 실행 그래프를 고정한다.
  • 고정된 그래프 위에서 PT2E가 관측(Observer) 삽입 -> 캘리브레이션 -> 변환(Convert) 을 수행한다.
  • 결과적으로 “어떤 연산이 양자화되었는지” 추적이 쉬워지고, 백엔드(x86 fbgemm, arm qnnpack)에 맞춘 설정도 명확해진다.

다만 Export는 “한 번 정해진 입력 제약”이 생깁니다. 즉, 배치 크기나 시퀀스 길이 같은 동적 차원을 어떻게 둘지(또는 고정할지) 설계를 먼저 해야 합니다.

준비: 환경과 백엔드 선택

INT8 PTQ는 CPU에서 특히 효과적입니다. 일반적으로 다음 중 하나를 선택합니다.

  • x86 서버: fbgemm
  • ARM(모바일/엣지): qnnpack

PyTorch에서 엔진 설정은 전역으로 잡는 경우가 많습니다.

import torch

# x86 서버라면 보통 fbgemm
torch.backends.quantized.engine = "fbgemm"
print("quant engine:", torch.backends.quantized.engine)

의존성은 PyTorch 버전에 따라 API가 조금씩 변합니다. 아래 코드는 개념적으로 “Export+PT2E 흐름”을 보여주며, 실제 적용 시에는 사용 중인 PyTorch 버전의 torch.ao.quantization 문서를 함께 확인하는 것을 권장합니다.

대상 모델 예시: CNN 분류기

실전에서는 ResNet, MobileNet, ViT 등 다양한 모델이 대상이지만, 글에서는 흐름을 명확히 하기 위해 간단한 CNN을 예로 들겠습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, stride=2, padding=1)
        self.fc = nn.Linear(64 * 8 * 8, num_classes)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

model = SmallCNN().eval()

입력은 예를 들어 NCHW 이미지 텐서로 가정합니다.

example_inputs = (torch.randn(1, 3, 32, 32),)

1) Export: 그래프를 고정해 양자화의 기반 만들기

torch.export.export는 “이 입력 형태로 실행되는 연산 그래프”를 추출합니다. 이 단계에서 control flow가 많거나 입력 shape가 다양하면 Export가 실패하거나, 특정 입력에만 맞는 그래프가 만들어질 수 있습니다.

import torch

# PyTorch 2 Export
exported = torch.export.export(model, example_inputs)
print(type(exported))

여기서 중요한 포인트는 다음입니다.

  • Export는 모델을 “일종의 IR”로 고정합니다.
  • 이후 양자화 준비(Observer 삽입)는 이 IR 위에서 진행됩니다.

실전 팁:

  • 데이터 전처리(정규화, 리사이즈)를 모델 밖에서 하거나, 모델 안에 넣되 Export 가능한 연산인지 확인하세요.
  • 분기 로직이 있으면 “항상 같은 분기”만 Export되는 문제가 생길 수 있습니다. 이 경우 모델 구조를 단순화하거나, 입력 제약을 명시적으로 관리해야 합니다.

2) Quantizer와 qconfig: 무엇을 INT8로 만들 것인가

PTQ에서 핵심은 “어떤 범위로 quantize 할지”입니다.

  • Activation(활성화): 보통 uint8(0~255)로 양자화, 관측 후 scale/zero-point 결정
  • Weight(가중치): 보통 int8(-128~127) per-channel 양자화가 성능/정확도에 유리

PyTorch의 양자화 설정은 보통 qconfig(또는 backend config)로 표현됩니다. PT2E에서는 “Quantizer” 개념을 통해 어떤 패턴(conv+relu 등)을 어떻게 변환할지 결정합니다.

아래는 흐름을 보여주는 예시 코드입니다(버전에 따라 import 경로/함수명이 다를 수 있습니다).

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

quantizer = X86InductorQuantizer()
quantizer.set_global(get_default_x86_inductor_quantization_config())

prepared = prepare_pt2e(exported, quantizer)

여기서 prepared는 관측 모듈(Observer 또는 FakeQuant 유사 구조)이 삽입된 상태입니다. 이제 캘리브레이션을 통해 activation 범위를 수집해야 합니다.

3) 캘리브레이션: 정확도를 좌우하는 단계

PTQ는 학습 없이 변환하므로, 캘리브레이션 데이터 품질이 곧 결과 품질입니다.

원칙:

  • 실제 추론 분포를 대표하는 샘플을 사용
  • 최소 수백~수천 샘플 권장(모델/도메인에 따라 다름)
  • 전처리까지 포함해 “실제 서빙 입력과 동일한 분포”를 넣기

예시로 더미 데이터 로더를 만들겠습니다.

def calibration_loader(num_batches=50, batch_size=32):
    for _ in range(num_batches):
        yield torch.randn(batch_size, 3, 32, 32)

# 캘리브레이션 실행: prepared 모델로 forward만 수행
prepared.eval()
with torch.no_grad():
    for x in calibration_loader():
        _ = prepared(x)

캘리브레이션에서 흔히 하는 실수:

  • 배치 크기를 지나치게 크게 해서 activation 분포가 달라짐
  • 전처리 누락(정규화/클리핑/리사이즈)
  • 대표성이 낮은 데이터(특정 클래스만 과다)

4) Convert: 실제 INT8 그래프로 변환

캘리브레이션이 끝나면 변환을 수행합니다.

quantized_model = convert_pt2e(prepared)
quantized_model.eval()

이 시점부터는 내부적으로 quantize_per_tensor, quantize_per_channel 등의 파라미터가 고정되고, 가능한 연산은 INT8 커널로 내려갑니다.

체크 포인트:

  • 모든 레이어가 양자화되지는 않습니다. 예를 들어 softmax, 일부 reduction, 특정 reshape 패턴 등은 FP로 남을 수 있습니다.
  • “양자화가 적용된 비율”이 낮으면 성능 향상이 제한적입니다.

5) 정확도 검증: FP32 대비 손실을 수치로 확인

양자화는 손실이 0이 아닙니다. 따라서 최소한 다음을 확인해야 합니다.

  • 단일 배치에서 FP32 vs INT8 출력의 상대 오차
  • 검증 데이터셋에서 top-1, top-5 같은 지표 변화

간단한 수치 비교 예시입니다.

import torch

x = torch.randn(16, 3, 32, 32)
with torch.no_grad():
    y_fp32 = model(x)
    y_int8 = quantized_model(x)

# 평균 절대 오차
mae = (y_fp32 - y_int8).abs().mean().item()
print("MAE:", mae)

실전에서는 MAE만으로 부족하고, 실제 metric을 봐야 합니다. 특히 분류는 logit의 작은 차이가 argmax를 바꾸기도 합니다.

정확도 손실이 큰 경우의 대응 순서:

  1. 캘리브레이션 데이터 개선(가장 효과적인 경우가 많음)
  2. per-channel weight quantization 사용 여부 확인
  3. 특정 레이어를 quantize 제외(민감 레이어)
  4. 그래도 안 되면 QAT 고려

6) 성능 검증: “진짜 빨라졌는지” 측정하는 법

INT8은 다음 조건에서 효과가 큽니다.

  • CPU 추론
  • 메모리 bandwidth/캐시 미스가 병목
  • 배치가 너무 크지 않음(서빙은 보통 작은 배치)

간단 벤치마크(워밍업 포함) 예시:

import time
import torch

def bench(m, x, iters=200, warmup=50):
    m.eval()
    with torch.no_grad():
        for _ in range(warmup):
            _ = m(x)
        t0 = time.perf_counter()
        for _ in range(iters):
            _ = m(x)
        t1 = time.perf_counter()
    return (t1 - t0) / iters

x = torch.randn(1, 3, 32, 32)

t_fp32 = bench(model, x)
t_int8 = bench(quantized_model, x)

print("fp32 ms:", t_fp32 * 1000)
print("int8 ms:", t_int8 * 1000)
print("speedup:", t_fp32 / t_int8)

성능 측정 시 주의:

  • 스레드 수 고정: OMP_NUM_THREADS, MKL_NUM_THREADS 등을 통제하지 않으면 결과가 흔들립니다.
  • 동일한 affinity/NUMA 조건에서 비교하세요.
  • 실제 서빙 입력 shape(배치, 해상도)로 측정하세요.

7) Export 기반에서 자주 만나는 함정들

함정 1: 동적 shape와 Export 제약

Export는 “이 입력 스펙”을 기준으로 그래프를 고정합니다. 입력 shape가 다양하면 다음 중 하나가 필요합니다.

  • 입력 shape를 몇 가지 케이스로 나눠 여러 그래프를 만든다
  • 가능한 범위의 dynamic dimension을 명시적으로 허용한다(버전/기능 지원 확인 필요)

서빙이 batch=1 위주라면, 과감히 batch=1 고정이 성능/안정성에 유리한 경우도 많습니다.

함정 2: 양자화가 기대만큼 적용되지 않음

원인은 보통 다음입니다.

  • 패턴 매칭 실패(예: conv -> relu가 모듈 경계/함수 호출 형태 때문에 분리됨)
  • 지원되지 않는 연산 포함
  • 레이아웃/메모리 포맷 문제

대응:

  • 모델을 nn.Sequential로 단순화하거나, fuse 가능한 구조로 정리
  • 지원되는 op set을 기준으로 전처리/후처리 재배치

함정 3: 정확도 손실이 특정 클래스에서만 급증

이는 캘리브레이션 분포 편향일 가능성이 큽니다.

  • 희귀 클래스 샘플을 캘리브레이션에 포함
  • 실제 트래픽 로그 기반 샘플링

함정 4: “INT8인데도 느림”

가능한 이유:

  • 양자화된 구간이 짧고, dequantize/quantize 오버헤드가 큼
  • 작은 텐서에서는 INT8 커널 이점이 작음
  • 스레드/배치/메모리 병목이 다른 곳에 있음

이 경우는 “양자화 자체”보다 “전체 파이프라인”을 봐야 합니다. 예를 들어 전처리/후처리가 Python에서 병목이면 모델만 빨라져도 체감이 작습니다.

8) 배포 관점 체크리스트

INT8 모델을 만들었으면, 배포 단계에서 다음을 확인하세요.

  • 런타임이 같은 엔진을 쓰는지(fbgemm vs qnnpack)
  • 컨테이너 이미지에 필요한 라이브러리(예: oneDNN, OpenMP 런타임)가 포함되어 있는지
  • 성능 측정이 CI에서 재현 가능한지

실험 반복이 잦다면 Docker 캐시 전략이 생산성을 크게 좌우합니다. 모델 파일/의존성 설치 레이어를 분리하는 방식은 Docker 빌드 70% 단축 - 멀티스테이지·캐시 전략에서 자세히 다룹니다.

또한 “양자화 vs ONNX Runtime”을 비교 검토하는 팀도 많습니다. ONNX로 내보내서 런타임 최적화를 타는 전략은 파이썬 ONNX Runtime로 CNN 10배 경량화·가속 글이 좋은 출발점입니다.

9) 실전 권장 워크플로우(요약)

  1. 베이스라인 측정: FP32 정확도/지연/CPU 사용률
  2. Export: 입력 스펙(배치/shape) 결정 후 torch.export
  3. Prepare(PT2E): quantizer 설정 후 observer 삽입
  4. Calibration: 대표 데이터로 forward-only 수행
  5. Convert: INT8 그래프 생성
  6. 검증: metric 하락폭과 분포별 실패 케이스 확인
  7. 벤치마크: 동일 조건에서 지연/처리량 비교
  8. 배포: 엔진/스레드/컨테이너 의존성까지 고정

마무리

PyTorch 2의 Export+PT2E 흐름은 “그래프를 먼저 고정하고, 그 위에서 양자화를 안정적으로 적용한다”는 점에서 실전성이 높습니다. 핵심은 코드 몇 줄이 아니라, 입력 스펙 설계, 캘리브레이션 데이터 품질, 그리고 성능 측정의 재현성입니다.

다음 단계로는 (1) 특정 레이어만 선택적으로 양자화/제외하는 전략, (2) 정확도 민감 모델에 대한 QAT, (3) Inductor 기반 컴파일과의 결합을 함께 다루면 좋습니다. 원하시면 사용 중인 모델 아키텍처와 목표 환경(x86 서버, ARM 엣지, 배치 크기, 지연 목표)을 기준으로 PT2E 설정을 더 구체화해 드릴게요.

Discuss on Twitter • View on GitHub