PyTorch 2.0 PT2E+XNNPACK int8 양자화 실전

서버에서 GPU로 학습한 모델을 모바일/엣지 혹은 CPU 추론 환경에 배포할 때 가장 먼저 부딪히는 문제는 지연시간(latency) 과 메모리 대역폭입니다. 특히 CPU에서는 float32 연산이 병목이 되기 쉬워서, int8 양자화가 체감 성능을 크게 올리는 경우가 많습니다.

PyTorch 2.0 이후에는 기존 FX Graph Mode 양자화 흐름과 더불어, PT2E(PyTorch 2 Export) 기반 양자화가 점점 표준 경로가 되고 있습니다. 이 글에서는 PT2E + XNNPACK 조합으로 int8 양자화 모델을 만들고, CPU에서 실행하는 end-to-end 과정을 다룹니다.

목표

• PyTorch 2.x에서 torch.export 기반 그래프를 만들고
• 관측기(Observer)를 삽입해 캘리브레이션을 수행한 뒤
• int8로 변환(convert)하고
• XNNPACK 백엔드에서 추론 성능을 확인

PT2E + XNNPACK 개념 정리

PT2E(Export) 기반 양자화란

PT2E는 모델을 torch.export로 정적 그래프 형태로 내보내고, 그 그래프에 양자화 준비/변환 패스를 적용하는 방식입니다. 장점은 다음과 같습니다.

• 그래프가 더 엄격하게 고정되어 최적화/변환이 안정적
• 런타임 의존성이 줄고, 백엔드별 변환 파이프라인을 구성하기 쉬움
• TorchDynamo/Inductor 계열과의 궁합이 좋아지는 방향

다만 Export는 데이터 의존 분기, 동적 shape, 일부 Python side effect 등에 민감합니다. 즉, “학습 때 돌아가던 eager 모델”이 export에서 바로 통과하지 않을 수 있습니다.

XNNPACK이 하는 역할

XNNPACK은 모바일/CPU 환경에서 널리 쓰이는 고성능 커널 라이브러리로, PyTorch에서는 주로 ARM/모바일 뿐 아니라 일부 x86 CPU에서도 효율적인 연산을 제공합니다.

• int8 GEMM/Conv 커널 최적화
• 연산 fusion 및 레이아웃 최적화에 유리
• QNNPACK 계열과 혼동되기 쉬우나, 최근 흐름은 XNNPACK 중심

준비 사항(버전/환경 체크)

권장 버전

• PyTorch: 2.0 이상(가능하면 2.1+ 권장)
• torchvision: 모델을 쓸 경우 PyTorch와 호환 버전

설치는 예시로만 제시합니다.

pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu

CPU 전용 wheel을 쓰면 환경이 단순해지고, XNNPACK 경로 검증에도 도움이 됩니다.

흔한 함정: export가 안 되는 모델

다음 패턴이 있으면 torch.export에서 막히기 쉽습니다.

• 입력 shape에 따라 흐름이 갈리는 if 분기
• 리스트 append 등 Python side effect
• 데이터에 의존하는 for 루프 길이

이 경우에는 모델 구조를 단순화하거나, export-friendly한 형태로 리팩터링해야 합니다.

예제 모델과 데이터 준비

여기서는 이해를 돕기 위해 작은 CNN을 사용합니다. 실제로는 ResNet, MobileNet 계열에서도 동일한 흐름으로 적용합니다.

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

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

    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()
example_inputs = (torch.randn(1, 3, 32, 32),)

양자화는 보통 대표 데이터셋으로 캘리브레이션을 수행합니다. 여기서는 샘플 텐서로 대체하지만, 실제 환경에서는 수십~수백 배치 정도를 흘려보내는 것이 일반적입니다.

PT2E 양자화 파이프라인: Export → Prepare → Calibrate → Convert

PyTorch의 양자화 API는 버전에 따라 모듈 경로와 함수명이 조금씩 달라질 수 있습니다. 아래는 PT2E 흐름을 이해하기 위한 대표적인 형태이며, 실제 적용 시에는 설치된 PyTorch 버전의 문서를 함께 확인하는 것을 권장합니다.

핵심은 다음 4단계입니다.

1. torch.export.export로 ExportedProgram 생성
2. prepare_pt2e로 observer 삽입(양자화 준비)
3. 캘리브레이션 데이터로 forward 수행
4. convert_pt2e로 int8 변환

1) Export

import torch

# Export: 모델을 정적 그래프로 고정
exported = torch.export.export(model, example_inputs)

여기서 에러가 난다면, 대부분 모델이 export-friendly하지 않은 것입니다. 먼저 export를 통과시키는 것이 1차 목표입니다.

2) Quantizer 설정(XNNPACK)

PT2E에서는 백엔드에 맞는 Quantizer를 선택합니다. XNNPACK용 quantizer를 사용해 연산별 quantization spec을 구성합니다.

from torch.ao.quantization.quantizer.xnnpack_quantizer import XNNPACKQuantizer
from torch.ao.quantization.quantizer.xnnpack_quantizer import get_symmetric_quantization_config

quantizer = XNNPACKQuantizer()
quantizer.set_global(get_symmetric_quantization_config())

• 대칭(symmetric) int8는 CPU 커널에서 성능이 잘 나오는 편입니다.
• 모델/연산에 따라 비대칭(asymmetric) 구성이 더 정확할 수 있지만, 우선은 대칭 구성이 출발점으로 무난합니다.

3) Prepare + 캘리브레이션

from torch.ao.quantization.pt2e import prepare_pt2e, convert_pt2e

prepared = prepare_pt2e(exported, quantizer)

# 캘리브레이션: 대표 입력을 여러 번 흘려 observer 통계 수집
for _ in range(50):
    x = torch.randn(1, 3, 32, 32)
    prepared(x)

캘리브레이션 품질은 정확도에 직결됩니다.

• 실제 서비스 입력 분포와 비슷한 데이터로 수행
• 배치 수가 너무 적으면 scale/zero-point 추정이 불안정
• 전처리(정규화 등)가 동일해야 함

4) Convert

quantized = convert_pt2e(prepared)
quantized.eval()

# 추론 확인
with torch.no_grad():
    y = quantized(torch.randn(1, 3, 32, 32))
print(y.shape)

여기까지 오면 int8 양자화 모델이 만들어진 것입니다. 이제 성능/정확도를 체크합니다.

정확도 체크: float vs int8 출력 비교

분류 모델이라면 top-1 정확도를 비교하는 것이 정석이지만, 최소한의 sanity check로는 동일 입력에 대한 출력 분포 차이를 볼 수 있습니다.

import torch

model_fp32 = SmallCNN().eval()

x = torch.randn(8, 3, 32, 32)
with torch.no_grad():
    out_fp32 = model_fp32(x)
    out_int8 = quantized(x)

# 단순 오차 지표
mae = (out_fp32 - out_int8).abs().mean().item()
maxe = (out_fp32 - out_int8).abs().max().item()
print({"mae": mae, "maxe": maxe})

MAE가 낮다고 항상 정확도가 보장되진 않지만, 변환이 비정상적으로 깨졌는지(예: 특정 레이어가 잘못 양자화되어 출력이 폭주) 빠르게 감지할 수 있습니다.

성능 측정: CPU에서 latency 비교

성능 측정은 워밍업과 스레드 설정이 중요합니다.

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

# 스레드 수는 환경에 따라 튜닝 포인트
torch.set_num_threads(1)

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

fp32_t = bench(model, x)
int8_t = bench(quantized, x)
print({"fp32_ms": fp32_t * 1000, "int8_ms": int8_t * 1000})

팁:

• 작은 배치(예: 1)에서 int8 이득이 더 크게 보이는 경우가 많습니다.
• 너무 작은 모델은 오버헤드가 커서 차이가 미미할 수 있습니다.
• 스레드 수(set_num_threads)에 따라 결과가 크게 달라집니다.

자주 겪는 문제와 해결 포인트

1) 특정 연산이 양자화되지 않는다

PT2E 양자화는 모든 연산을 무조건 int8로 바꾸지 않습니다. 아래 상황이 흔합니다.

• 백엔드(XNNPACK)가 해당 연산의 int8 커널을 지원하지 않음
• 패턴 매칭(fusion) 조건을 만족하지 못함
• 연산 앞뒤 텐서 dtype/레이아웃이 기대와 다름

해결 방향:

• 모델을 Conv + ReLU 같은 표준 패턴으로 정리
• 불필요한 view/permute를 줄여 레이아웃 변환을 최소화
• 지원되는 연산 집합에 맞춰 구조를 조정

2) 정확도가 크게 떨어진다

대부분은 캘리브레이션 데이터 문제이거나, 민감한 레이어(예: 첫 Conv, 마지막 FC)가 양자화에 취약한 경우입니다.

대응 전략:

• 캘리브레이션 데이터를 서비스 입력 분포에 맞추기
• 레이어별로 양자화 제외(exclude) 또는 더 정교한 config 적용
• 가능하면 QAT(Quantization Aware Training) 고려

3) Export 단계에서 실패한다

Export 실패는 양자화 이전에 해결해야 합니다.

• 입력 shape를 고정하고, example_inputs를 정확히 맞추기
• 데이터 의존 분기 제거
• 모듈을 TorchScript/Export-friendly 형태로 리팩터링

실무에서는 “모델 자체 문제”보다 “서빙용 래퍼 코드”가 export를 깨는 경우가 많습니다. 전처리/후처리를 모델 밖으로 분리하는 것도 방법입니다.

배포 관점: 체크리스트

• 정확도: float 대비 top-1/top-5 및 경계 케이스 확인
• 성능: 워밍업 포함, 스레드 수/배치 크기별 측정
• 메모리: weight는 줄어도 activation이 병목일 수 있음
• 재현성: PyTorch 버전/옵션에 따라 결과가 달라질 수 있으니 고정

CI에서 성능 회귀를 잡고 싶다면 캐시/빌드 환경도 중요합니다. 예를 들어 Docker 기반으로 측정 파이프라인을 만들 때는 캐시 전략이 성능/재현성에 직접 영향을 줍니다. 관련해서는 Docker BuildKit 캐시 무효화 원인·해결 8가지 같은 글이 실무에 도움이 됩니다.

또한 양자화 모델을 CI에서 반복 빌드/테스트할 때 캐시가 기대대로 동작하지 않으면 측정값이 흔들릴 수 있습니다. 이 경우 GitHub Actions 캐시가 안먹을 때 9가지 원인도 같이 참고해두면 좋습니다.

정리

PyTorch 2.0의 PT2E 기반 양자화는 torch.export를 중심으로 그래프를 고정하고, 백엔드(XNNPACK)에 맞춰 int8 변환을 수행하는 흐름입니다. 핵심은 다음입니다.

• 먼저 torch.export를 통과시키는 것이 1순위
• XNNPACK quantizer로 prepare 후 충분한 캘리브레이션 수행
• convert 이후에는 정확도/성능을 반드시 함께 검증
• 성능은 모델 구조, 스레드 수, 배치 크기, 지원 연산 여부에 따라 달라짐

다음 단계로는 (1) 실제 ResNet/MobileNet에 적용, (2) 레이어별 quantization config 튜닝, (3) 정확도가 민감하면 QAT로 전환, (4) 모바일 런타임에서의 end-to-end 측정까지 확장하는 것을 권장합니다.