beautifulsoups.dev
Published on

PyTorch 2.2+ PT2E로 INT8 양자화 Edge 최적화

Authors

서버에서 학습한 모델을 Edge 디바이스로 내릴 때 가장 흔한 병목은 연산량보다 메모리 대역폭과 캐시 미스입니다. FP32 모델은 파라미터와 activation이 커서 메모리 트래픽이 폭증하고, 배터리 소모와 지연시간이 함께 악화됩니다. 이때 INT8 양자화는 모델 크기와 메모리 접근량을 줄여 Edge에서 체감 성능을 끌어올리는 가장 현실적인 선택지입니다.

PyTorch 2.2+에서는 PT2E라고 부르는 Export 기반 양자화 파이프라인이 주류가 되고 있습니다. 기존 FX Graph Mode 대비 export 가능한 그래프를 중심으로 설계되어, 컴파일러 친화적이며 Edge 배포에 필요한 형태로 정리하기 쉽습니다. 이 글에서는 PT2E로 INT8 양자화를 적용하는 전체 흐름과, Edge에서 흔히 터지는 함정들을 함께 정리합니다.

또한 운영 관점에서 재현 가능한 실험과 자동화를 위해 CI 설정과 캐시 전략이 중요합니다. 관련해서는 GitHub Actions 캐시 미스 원인 7가지와 해결도 함께 참고하면 좋습니다.

PT2E 양자화가 필요한 이유: Edge 관점에서의 이점

INT8 양자화는 단순히 모델 파일이 작아지는 것에 그치지 않습니다.

  • 메모리 대역폭 절감: weight와 activation이 작아지면 DRAM 왕복이 줄어들어 지연시간이 감소합니다.
  • 캐시 효율 증가: L2, LLC에 더 많은 텐서가 올라가면서 연산 유닛이 덜 놀게 됩니다.
  • 전력 효율 개선: Edge에서 배터리, 발열 제약이 있는 경우 매우 큰 차이가 납니다.
  • 백엔드 최적화 활용: x86에서는 FBGEMM, ARM 계열에서는 QNNPACK 또는 타 백엔드가 INT8 커널을 제공하는 경우가 많습니다.

다만 INT8은 공짜가 아닙니다. 정확도 하락, 연산 그래프 제약, 지원되지 않는 연산의 fallback 같은 비용이 존재합니다. PT2E는 이런 비용을 내보낼 수 있는 그래프로 먼저 정리하고, 그 위에 관측과 변환을 얹는 방식으로 문제를 줄입니다.

큰 그림: PT2E INT8 양자화 파이프라인

PT2E 기반의 전형적인 흐름은 다음과 같습니다.

  1. torch.export.export로 Exported Program 확보
  2. Quantizer와 qconfig 설정으로 관측기 삽입 prepare 단계
  3. 대표 입력으로 calibration 수행 observer 통계 수집
  4. convert로 INT8 연산으로 변환
  5. 실행 백엔드에서 벤치마크 및 정확도 검증

핵심은 1번입니다. Export 단계에서 그래프를 안정적으로 고정해야, 이후 변환이 일관되게 동작합니다.

실전 코드: PT2E로 INT8 Post Training Quantization 적용

아래 예제는 간단한 CNN을 대상으로 PT2E 기반 INT8 PTQ를 수행하는 흐름을 보여줍니다. 실제 프로젝트에서는 모델이 더 복잡하므로, 먼저 작은 모듈 단위로 성공시키고 점진적으로 확장하는 접근이 안전합니다.

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

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)
        return self.fc(x)

model = SmallCNN().eval()
example_inputs = (torch.randn(1, 3, 32, 32),)

2) Export: 변환 가능한 그래프 확보

import torch.export

with torch.no_grad():
    exported = torch.export.export(model, example_inputs)

Export가 실패한다면 모델 내부에 데이터 의존 제어흐름, 특정 Python side-effect, 동적 shape 처리 등이 섞여 있을 수 있습니다. 이 경우 먼저 forward를 단순화하거나, 입력 shape를 고정하는 식으로 Export 가능성을 높여야 합니다.

3) Quantizer 설정과 prepare

PyTorch의 PT2E 양자화는 torch.ao.quantization 계열 API를 사용합니다. 백엔드에 따라 권장 quantizer가 달라질 수 있으며, 일반적으로 x86 서버나 데스크탑은 FBGEMM, ARM 모바일은 QNNPACK 계열을 고려합니다.

from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
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())

prepared = prepare_pt2e(exported, quantizer)

여기서 symmetric 설정은 흔히 weight와 activation을 대칭 양자화로 두는 구성이며, Edge에서 성능과 구현 단순성 측면에서 많이 사용됩니다. 정확도가 민감한 모델이라면 per-channel weight quantization, activation의 scheme 등을 조정해야 합니다.

4) Calibration: 대표 데이터로 통계 수집

PTQ의 성패는 calibration 데이터 품질에 크게 좌우됩니다. Edge 배포라면 실제 디바이스에서 들어오는 입력 분포를 최대한 반영하는 것이 핵심입니다.

def calibrate(prepared_model, num_batches=32, batch_size=8):
    prepared_model.eval()
    with torch.no_grad():
        for _ in range(num_batches):
            x = torch.randn(batch_size, 3, 32, 32)
            _ = prepared_model(x)

calibrate(prepared)

실제 환경에서는 랜덤 텐서가 아니라, 전처리까지 포함한 샘플 데이터를 넣어야 합니다. 특히 이미지 모델은 resize, normalize 방식이 바뀌면 activation 분포가 크게 달라져 양자화 정확도가 급락할 수 있습니다.

5) Convert: INT8 모델로 변환

quantized = convert_pt2e(prepared)
quantized.eval()

with torch.no_grad():
    y = quantized(*example_inputs)
print(y.shape)

이 시점의 모델은 INT8 커널을 최대한 사용하도록 변환되어 있으며, 지원되지 않는 연산은 일부 FP32로 남을 수 있습니다. Edge 최적화에서는 어떤 연산이 INT8로 내려갔는지를 반드시 확인해야 합니다.

Edge 최적화 체크리스트: 성능이 안 나오는 이유 10가지

INT8로 변환했는데도 지연시간이 기대만큼 줄지 않는 경우가 많습니다. 아래 항목을 순서대로 점검하면 원인 파악이 빨라집니다.

1) 진짜 INT8 커널을 타는지 확인

양자화는 되었는데 실행은 FP32로 돌아가는 상황이 흔합니다. 원인은 보통 다음입니다.

  • 백엔드 미지원 연산 포함
  • 연산 패턴이 fuse되지 않아 dequant, quant가 과도하게 삽입
  • 입력이 FP32로 들어와서 매번 변환 비용이 큼

가능하면 모델의 대표 구간을 프로파일링하고, quant, dequant가 어디서 발생하는지 확인하세요.

2) Calibration 데이터 분포 불일치

PTQ는 통계 기반입니다. 실제 입력의 밝기, 노이즈, 해상도, 전처리 차이가 observer 통계를 망가뜨리면 정확도 하락이 커집니다.

  • 최소 수백에서 수천 샘플까지 늘려 비교
  • 클래스 불균형이 심하면 샘플링 전략 조정

3) Conv와 Linear만 INT8이고 나머지가 FP32

Conv, Linear는 잘 양자화되지만, attention 계열이나 특정 activation, resize, normalization이 FP32로 남는 경우가 많습니다. 이때는 경계에서의 dequant 비용이 병목이 됩니다.

4) 채널 정렬과 메모리 레이아웃

Edge CPU에서는 NHWC가 유리한 경우가 있고, 특정 커널은 채널 정렬이 맞지 않으면 느려집니다. 모델이 NCHW로 고정되어 있다면 백엔드별 권장 레이아웃을 검토해야 합니다.

5) Batch 크기 1 최적화가 되어 있는지

Edge는 대부분 batch 1입니다. 서버에서 batch 32 기준으로만 성능을 봤다면, Edge에서 결과가 다르게 나옵니다.

6) 스레딩과 affinity

모바일 및 임베디드 환경에서는 스레드 수를 늘린다고 빨라지지 않습니다. 오히려 context switching과 캐시 경쟁으로 느려질 수 있습니다.

  • torch.set_num_threads 조정
  • big.LITTLE 환경에서는 affinity 고려

7) 정확도 평가는 단일 metric으로 끝내지 않기

Top-1만 보지 말고, task에 맞는 metric을 추가하세요. 예를 들어 detection이면 mAP, segmentation이면 IoU 등으로 봐야 합니다.

8) Quantization friendly 모델 구조인지

ReLU6, hard-swish, layernorm 등은 양자화 난이도가 다릅니다. 가능하다면 학습 단계에서 QAT를 고려하는 편이 낫습니다.

9) 입력과 출력의 dtype 경계 최소화

가능하면 전처리부터 후처리까지 dtype 변환을 줄이고, 파이프라인 전체에서 불필요한 float 변환을 제거하세요.

10) CI에서 재현 가능한 벤치마크를 만들기

양자화는 환경에 민감합니다. 커널, 스레딩, CPU feature에 따라 결과가 달라집니다. GitHub Actions에서 측정할 때도 캐시와 환경 고정이 중요합니다. 캐시 관련 문제는 GitHub Actions 캐시 미스 원인 7가지와 해결에서 자주 겪는 실수를 정리해 두었습니다.

벤치마크 코드: FP32 vs INT8 지연시간 비교

Edge 배포 전 최소한의 마이크로벤치마크로 변환 효과를 확인해야 합니다. 아래는 매우 단순한 측정 예시입니다.

import time
import statistics
import torch

def benchmark(model, x, warmup=20, iters=100):
    model.eval()
    times = []
    with torch.no_grad():
        for _ in range(warmup):
            _ = model(x)
        for _ in range(iters):
            t0 = time.perf_counter()
            _ = model(x)
            t1 = time.perf_counter()
            times.append((t1 - t0) * 1000.0)
    return {
        "mean_ms": statistics.mean(times),
        "p50_ms": statistics.median(times),
        "p95_ms": statistics.quantiles(times, n=20)[18],
    }

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

fp32_stats = benchmark(model, x)
int8_stats = benchmark(quantized, x)

print("fp32", fp32_stats)
print("int8", int8_stats)

주의할 점은 다음과 같습니다.

  • CPU 고정 주파수, 성능 모드 여부에 따라 결과가 흔들립니다.
  • OS 스케줄링 영향이 큰 환경에서는 반복 횟수를 늘려야 합니다.
  • 실제 Edge에서는 전처리, 후처리 시간을 포함해야 체감 성능과 맞습니다.

정확도 방어 전략: PTQ로 부족하면 QAT 고려

PTQ로 정확도가 크게 떨어진다면 QAT가 다음 선택지입니다. QAT는 학습 중에 fake quant를 넣어 양자화 오차에 모델이 적응하게 합니다.

  • PTQ: 빠르고 간단하지만 분포 민감
  • QAT: 학습 비용이 있지만 정확도 보존에 유리

Edge에서 제품 품질을 맞춰야 하는 경우, PTQ로 빠르게 baseline을 만들고, 목표 metric에 못 미치면 QAT로 넘어가는 흐름이 일반적입니다.

배포 관점: Export 이후 포맷과 런타임 선택

PT2E는 export 중심이므로, 이후 어떤 런타임으로 실행할지 결정해야 합니다.

  • PyTorch 런타임으로 그대로 실행
  • ExecuTorch 같은 Edge 런타임 고려
  • ONNX로 변환해 타 런타임 사용

선택 기준은 단순합니다.

  • 디바이스에 PyTorch 런타임을 포함할 수 있는가
  • INT8 커널 최적화가 목표 디바이스에서 충분한가
  • 모델 업데이트와 롤백을 어떻게 할 것인가

운영에서 인증 토큰이나 시간 동기화 문제로 장애가 나면 배포 자체가 막힐 수 있습니다. 배포 파이프라인이 Spring 기반이라면 Spring Security JWT 401 원인 - 시계오차·키롤오버 같은 이슈도 함께 점검하는 편이 안전합니다.

마무리: PT2E INT8 양자화는 Edge 최적화의 출발점

PyTorch 2.2+의 PT2E 양자화는 export 가능한 그래프를 기반으로 하기 때문에, Edge 배포에서 요구하는 재현성과 컴파일러 친화성을 확보하기 좋습니다. 다만 성능과 정확도는 코드 한 줄로 해결되지 않습니다.

  • Export 단계에서 그래프를 안정화하고
  • Calibration 데이터를 실제 분포에 맞추고
  • INT8 커널 적용 범위를 확인하며
  • 벤치마크와 CI로 재현 가능한 측정을 구축

이 4가지를 지키면, INT8 양자화가 단순한 실험을 넘어 제품 성능 개선으로 이어질 가능성이 크게 올라갑니다.

Discuss on Twitter • View on GitHub