beautifulsoups.dev
Published on

PyTorch 2.1 PT2E INT8 양자화 실패 해결 가이드

Authors

PyTorch 2.1에서 torch.export 기반의 PT2E(Post Training Quantization 2 Export)로 INT8 양자화를 적용하려다 보면, 기존 FX Graph Mode 대비 디버깅 포인트가 달라져서 “왜 실패하는지”를 파악하기가 더 어렵습니다. 특히 torch.compile/Inductor, 연산자 분해(decomposition), 관측자(observer) 삽입, 백엔드(x86/fbgemm/qnnpack) 선택이 한 번에 얽히면서 에러 메시지가 난해해지는 경우가 많습니다.

이 글은 PyTorch 2.1 기준으로 PT2E INT8 양자화가 실패하는 대표 패턴을 증상별로 분류하고, 원인 확인 → 최소 수정 → 재검증 흐름으로 해결하는 실전 체크리스트를 제공합니다.

문제 진단을 체계화하는 접근은 인프라 트러블슈팅과도 유사합니다. 예를 들어 제한된 시간 안에 원인을 좁혀가는 방식은 systemd 서비스 무한 재시작 10분 진단 체크리스트에서 소개한 “증상에서 가설로 수렴” 패턴과 거의 같습니다.

PT2E INT8 양자화 파이프라인 한 장 요약

PT2E의 핵심 흐름은 아래처럼 이해하면 디버깅이 쉬워집니다.

  1. torch.export.export로 ExportedProgram 생성
  2. prepare_pt2e로 관측자 삽입(캘리브레이션 준비)
  3. 대표 입력으로 캘리브레이션 실행
  4. convert_pt2e로 정수 연산(quantize/dequantize, int8 kernel)로 변환
  5. (선택) torch.compile로 최적화

여기서 실패는 대개 다음 레이어 중 하나에서 발생합니다.

  • Export 단계: 동적 제어 흐름, unsupported op, shape constraint 불일치
  • Prepare 단계: 관측자 삽입 불가 패턴(예: 일부 fused op), qconfig 충돌
  • Convert 단계: int8 커널 미지원, per-channel 요구 불일치, dtype/scale 제약
  • Compile 단계: Inductor가 quantized op를 못 먹거나, decomposition이 달라져 그래프가 변형

아래부터는 “실패 메시지”를 기준으로 가장 흔한 원인과 해결책을 정리합니다.

실패 1: Export 자체가 안 됨 (dynamic shape, control flow)

대표 증상

  • torch.export가 예외를 던지거나, ConstraintViolationError 류 메시지
  • 입력 shape가 바뀌면 export가 깨짐

원인

PT2E는 ExportedProgram을 기반으로 하므로, export 시점에 shape 제약이 명확해야 합니다. 기존 eager/FX에서는 “일단 트레이스”가 되던 코드도 export에서는 더 엄격하게 막힙니다.

해결

  1. 먼저 정적 shape로 export를 성공시키고, 그 다음 dynamic shape를 확장합니다.
import torch
import torch.nn as nn

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, stride=1, padding=1)
        self.relu = nn.ReLU()
    def forward(self, x):
        return self.relu(self.conv(x))

m = M().eval()
example = (torch.randn(1, 3, 224, 224),)

ep = torch.export.export(m, example)
print(ep)
  1. dynamic shape가 필요하면 torch.export.Dim로 제약을 명시합니다. 이때 부등호가 들어간 표현은 MDX에서 문제를 일으킬 수 있으니, 문서/주석에서는 반드시 인라인 코드로 감싸는 습관을 권합니다.
import torch

B = torch.export.Dim("B", min=1, max=8)
example = (torch.randn(1, 3, 224, 224),)

ep = torch.export.export(
    m,
    example,
    dynamic_shapes={
        "x": {0: B}
    },
)
  1. 제어 흐름(if, for)이 입력 데이터에 따라 갈리는 모델은 export 호환 형태로 바꾸거나, 해당 부분을 torch.cond 등으로 표현해야 합니다.

실패 2: prepare_pt2e는 되는데 convert_pt2e에서 터짐

대표 증상

  • convert_pt2e 호출 시 “지원되지 않는 quantized op” 또는 “pattern not found”
  • 변환 후 실행 시 dtype mismatch

원인

  • 선택한 quantization backend가 해당 op 조합을 지원하지 않음
  • per-tensor/per-channel 설정이 연산자 요구와 충돌
  • 관측자 삽입 지점이 예상과 달라 scale/zero_point가 어긋남

해결 전략

1) 백엔드 지정과 qconfig를 명확히 하기

PyTorch 2.1의 PT2E는 torch.ao.quantization.quantizer 및 XNNPACKQuantizer/X86InductorQuantizer 등 경로가 섞여 혼동되기 쉽습니다. 가장 먼저 “내가 어느 백엔드로 가는지”를 고정하세요.

  • 서버 x86 CPU: 보통 x86 또는 fbgemm 계열
  • 모바일/ARM: qnnpack 또는 xnnpack

아래는 x86 경로 예시(개념적으로 가장 많이 쓰는 형태)입니다.

import torch
import torch.nn as nn

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

class M(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.relu(self.conv(x))
        x = self.pool(x).flatten(1)
        return self.fc(x)

m = M().eval()
example = (torch.randn(1, 3, 224, 224),)

ep = torch.export.export(m, example)

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

prepared = prepare_pt2e(ep, quantizer)

# calibration
for _ in range(10):
    prepared.module()(torch.randn(1, 3, 224, 224))

converted = convert_pt2e(prepared)

out = converted.module()(torch.randn(1, 3, 224, 224))
print(out.shape)

핵심은 다음입니다.

  • export 결과인 ep를 기준으로 prepare/convert를 수행
  • 캘리브레이션은 prepared.module()을 통해 실행

2) 캘리브레이션 입력 분포가 “너무 다르면” 실패/품질 악화

엄밀히 말해 캘리브레이션 분포가 달라서 convert가 직접 실패하는 경우는 흔하지 않지만, 실제로는 scale이 극단값으로 잡히면서 후속 단계에서 overflow나 정확도 붕괴로 “실패처럼 보이는” 현상이 생깁니다.

  • 최소한 실제 서비스 입력의 통계 범위를 반영하세요.
  • 배치 정규화가 있는 모델은 eval() 고정이 필수입니다.

3) relu/add/cat 등에서 dtype이 섞이는 문제

양자화 그래프에서 가장 자주 터지는 지점이 “합류 지점”입니다.

  • add 양쪽의 quantization parameter가 달라서 requantize가 필요
  • cat이 채널 축이 아닌 축으로 결합되어 커널이 없거나 느림

해결책은 보통 두 가지입니다.

  • 해당 연산 주변을 FP16/FP32로 남기고 부분 양자화
  • 모델 구조를 조금 바꿔 quantization-friendly 패턴으로 유도

부분 양자화는 “지원되는 블록만” 먼저 성공시키는 데 유리합니다.

실패 3: “지원되지 않는 연산” 때문에 변환이 막힘

대표 증상

  • aten.* 특정 op가 quantized 커널로 내려가지 않음
  • “no quantization pattern for op” 류

원인

PT2E는 모든 연산을 자동으로 INT8로 만들지 않습니다. 특히 아래는 자주 문제를 만듭니다.

  • layer_norm, gelu, 일부 softmax 패턴
  • einsum, 복잡한 reshape/view 조합
  • 커스텀 op 또는 composite op

해결

  1. 문제 op를 식별하고, 그 구간을 FP로 남깁니다.
  2. 가능하면 op를 더 단순한 조합으로 바꿉니다(예: gelu를 근사로 대체).

실전에서는 “한 번에 전체 모델 INT8”보다 “핵심 블록만 INT8”이 성공률이 높고, 성능 이득도 충분한 경우가 많습니다.

실패 4: torch.compile을 켰더니만 실패하거나 느려짐

대표 증상

  • eager로는 되는데 torch.compile(converted.module())에서 실패
  • 컴파일은 되지만 성능이 기대보다 안 나옴

원인

  • Inductor가 특정 quantized op를 아직 최적으로 처리 못함
  • decomposition이 바뀌어 quantization pattern 매칭이 깨짐

해결

  • 양자화 파이프라인을 먼저 안정화한 뒤 torch.compile을 마지막에 붙이세요.
  • 컴파일 실패 시에는 torch._dynamo.config.suppress_errors = True로 우회하기보다, 실패 op를 좁혀서 부분 컴파일/부분 양자화를 선택하는 편이 장기적으로 유지보수에 좋습니다.

간단한 패턴은 다음입니다.

import torch

mod = converted.module()

# 1) 먼저 eager로 기능 검증
x = torch.randn(1, 3, 224, 224)
ref = mod(x)

# 2) 그 다음 compile
compiled = torch.compile(mod)
out = compiled(x)

print((ref - out).abs().max())

실패 5: 정확도 폭락(“성공했는데 실패한 것 같은” 케이스)

대표 증상

  • 변환은 성공했으나 top-1 정확도가 크게 하락
  • 출력이 특정 클래스에 쏠림

원인

  • 캘리브레이션 데이터 부족
  • activation outlier가 큰데 per-tensor로 잡혀서 정보 손실
  • conv/linear weight의 per-channel 적용이 필요

해결

  • 캘리브레이션 샘플 수를 늘립니다(최소 수십~수백, 가능하면 1천 단위)
  • 가능하면 weight는 per-channel을 사용하도록 설정을 확인합니다
  • outlier가 큰 레이어는 FP로 남겨 하이브리드로 타협합니다

정확도 문제는 DB에서 느린 쿼리를 “측정 없이 감으로” 튜닝하면 실패하는 것과 비슷합니다. 캘리브레이션/평가 지표를 먼저 계측해 두는 습관이 중요합니다. 계측 중심 접근은 PostgreSQL 느린 쿼리 튜닝 - auto_explain+pg_stat_statements에서의 관점과도 통합니다.

디버깅을 빠르게 만드는 체크리스트

1) 그래프를 눈으로 확인하기

PT2E에서는 “내가 생각한 그래프”와 “export된 그래프”가 다른 경우가 많습니다.

print(ep.graph_module.graph)

prepared/converted에서도 확인해 어떤 노드에 관측자와 quantize_per_tensor 류가 붙는지 보세요.

2) 최소 재현(MRE)로 줄이기

실패가 나는 모델에서 해당 블록만 떼어내 작은 모듈로 만들면 원인 파악이 빨라집니다. CI에서 캐시/동시성 문제를 최소 재현으로 줄여 원인을 찾는 것과 같은 방식입니다. 필요하면 GitHub Actions 캐시가 안 먹을 때 - key·dir 충돌 디버깅처럼 “환경 변수, 입력, 경로”를 고정해 재현성을 확보하세요.

3) 부분 양자화로 성공 지점부터 확보

  • backbone만 INT8
  • head는 FP16/FP32
  • 또는 conv/linear만 INT8

이렇게 “되는 조합”을 먼저 만들고, 점진적으로 범위를 넓히는 편이 전체 시간을 줄입니다.

실전 권장 워크플로우

  1. eval() 고정, 대표 입력으로 torch.export.export 성공
  2. 목표 백엔드에 맞는 quantizer/qconfig를 명시
  3. prepare_pt2e 후 캘리브레이션을 충분히 수행
  4. convert_pt2e 성공 여부 및 eager 실행 검증
  5. 정확도 평가(변환 전후 비교)로 품질 확인
  6. 마지막으로 torch.compile 적용 및 성능 측정
  7. 실패 op가 있으면 부분 양자화로 우회하고, 구조 변경은 마지막에

마무리

PyTorch 2.1의 PT2E INT8 양자화 실패는 “한 가지 원인”보다 Export 제약, 백엔드 지원 범위, 관측자/스케일, compile 단계의 상호작용에서 발생하는 경우가 대부분입니다. 따라서 해결의 핵심은 한 번에 모든 걸 켜는 게 아니라, 단계별로 성공을 쌓아 올리며 실패 지점을 좁히는 것입니다.

다음 단계로는 (1) 서비스 입력 분포 기반 캘리브레이션 파이프라인 자동화, (2) 실패 op 목록을 팀 내 위키로 축적, (3) 부분 양자화 정책(성능 대비 정확도)을 정해두면 PT2E를 안정적으로 운영할 수 있습니다.

Discuss on Twitter • View on GitHub