PyTorch 2.x PT2E INT8 양자화 실패 원인 7가지

PyTorch 2.x에서 INT8 양자화를 시도할 때, 예전 FX Graph Mode(예: prepare_fx/convert_fx) 경험만으로는 디버깅이 잘 안 되는 경우가 많습니다. PT2E(PyTorch 2 Export) 기반 양자화는 torch.export로 정적 그래프(ExportedProgram) 를 만든 다음, 그 그래프에 대해 관측기 삽입/캘리브레이션/변환을 수행합니다. 이 과정에서 그래프가 “조금만” 동적이거나, 연산이 분해(decomposition)되거나, 백엔드가 기대하는 패턴과 달라지면 변환이 실패하거나 성능이 급락합니다.

이 글은 PT2E INT8 양자화가 실패하는 흔한 원인을 7가지로 정리하고, 각 원인별로 증상, 확인 방법, 해결/우회를 제공합니다.

운영 환경에서 장애를 빠르게 줄이는 관점은 K8s의 CrashLoopBackOff를 진단할 때와 유사합니다. 원인을 “범주화”해 체크리스트로 쪼개면 복구 시간이 급격히 줄어듭니다. 비슷한 접근은 K8s Pod CrashLoopBackOff 원인 7가지와 해결 글의 사고방식과도 통합니다.

PT2E INT8 양자화 파이프라인 빠른 개요

PT2E 흐름은 보통 아래 단계로 요약됩니다.

torch.export.export로 ExportedProgram 생성
Quantizer(예: XNNPACKQuantizer)와 QuantizationConfig 설정
prepare_pt2e로 관측기/가짜양자화 삽입
캘리브레이션(대표 입력 몇 배치 실행)
convert_pt2e로 INT8 연산으로 변환

아래는 “형태”를 보여주는 최소 예시입니다(환경/버전에 따라 import 경로가 다를 수 있으니, 실패 시 설치된 PyTorch 버전 문서를 우선 확인하세요).

import torch
import torch.nn as nn

# PT2E 관련 API는 버전에 따라 위치가 바뀔 수 있습니다.
# 아래는 개념 예시이며, 실제로는 torch.ao.quantization.pt2e 또는
# torch.ao.quantization.quantize_pt2e 경로를 사용합니다.
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)

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

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

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

# 1) export
ep = torch.export.export(model, example)

# 2) quantizer
quantizer = XNNPACKQuantizer()
qconfig = get_symmetric_quantization_config(is_per_channel=True)
quantizer.set_global(qconfig)

# 3) prepare
prepared = prepare_pt2e(ep, quantizer)

# 4) calibrate
for _ in range(8):
    prepared.module()(*example)

# 5) convert
converted = convert_pt2e(prepared)

# 실행 확인
out = converted.module()(*example)
print(out.shape)

이제부터는, 위 파이프라인이 어디서/왜 깨지는지 “실패 원인 7가지”로 들어가겠습니다.

실패 원인 1) Export 단계에서 그래프가 정적으로 고정되지 않음

대표 증상

torch.export.export에서 예외 발생
“data-dependent control flow”, “dynamic shape”, “guard” 관련 메시지
모델은 eager에서는 되는데 export에서만 실패

왜 발생하나

PT2E는 export된 정적 그래프를 전제로 합니다. 아래 패턴이 있으면 export가 흔히 실패합니다.

입력 텐서 값에 따라 분기하는 if x.sum() > 0: 같은 코드
파이썬 리스트/딕셔너리를 텐서처럼 조작
런타임에 레이어를 생성하거나 모듈 속성을 바꾸는 코드
동적 shape에 강하게 의존하는 reshape/view

확인 방법

torch.export.export를 통과하는지부터 확인하고, 실패 메시지에 “control flow”, “data dependent”, “guard”가 있는지 봅니다.

try:
    ep = torch.export.export(model.eval(), example)
except Exception as e:
    print("export failed:")
    print(type(e), e)
    raise

해결/우회

분기 로직을 텐서 연산으로 바꾸거나, 분기 자체를 모델 밖으로 이동
입력 shape를 고정(또는 export가 허용하는 범위 내에서 dynamic shape 명시)
forward에서 파이썬 컨테이너 조작을 최소화

실패 원인 2) Quantizer 백엔드가 지원하지 않는 연산/패턴이 그래프에 남아 있음

대표 증상

prepare_pt2e 또는 convert_pt2e에서 “pattern not found”, “unsupported op”
변환은 되는데 INT8로 안 바뀌고 대부분 FP32로 남음(사실상 실패)

왜 발생하나

PT2E는 “백엔드가 인식 가능한 패턴”을 찾아 INT8로 치환합니다. 예를 들어 XNNPACK은 특정 conv/linear 패턴을 기대합니다.

다음이 흔한 방해물입니다.

커스텀 연산, 일부 activation, 특수한 normalization
einsum, 복잡한 indexing, scatter/gather가 많은 그래프
모델이 export 과정에서 예상치 못한 decomposition으로 바뀌어 패턴이 깨짐

확인 방법

변환 후 그래프에서 INT8 관련 연산이 실제로 생겼는지 확인합니다.

print(converted)
# 또는
print(converted.graph_module.graph)

또는 변환 전/후로 연산 종류를 카운팅해 “양자화가 적용된 비율”을 대략 봅니다.

해결/우회

가능한 한 “양자화 친화적인 블록”으로 재구성(Conv+ReLU, Linear+ReLU 등)
지원되지 않는 부분은 FP32로 남기는 것을 목표로 부분 양자화 전략 채택
백엔드를 바꾸거나(예: 서버면 FBGEMM 계열), 설정을 per-tensor로 단순화해 패턴 매칭을 쉽게 함

실패 원인 3) 관측기(Observer) 캘리브레이션이 부정확하거나 아예 실행되지 않음

대표 증상

변환은 성공했는데 정확도가 크게 하락
일부 텐서의 스케일/제로포인트가 비정상(예: 스케일이 0에 가깝거나 극단적으로 큼)

왜 발생하나

PTQ(Post-Training Quantization)에서 캘리브레이션은 사실상 “학습”에 해당합니다. 다음이 흔한 원인입니다.

캘리브레이션 데이터를 너무 적게 사용
대표성이 없는 데이터(분포가 실제 입력과 다름)
캘리브레이션 실행을 prepared.module()(... )로 안 하고, 실수로 원본 모델을 실행

확인 방법

캘리브레이션 루프가 실제로 prepared 모델을 실행하는지 확인하고, 배치 수를 늘려 accuracy 변화를 봅니다.

prepared_mod = prepared.module()
prepared_mod.eval()

with torch.no_grad():
    for i in range(64):
        x = torch.randn(1, 3, 224, 224)  # 실제론 validation 샘플 권장
        prepared_mod(x)

해결/우회

최소 수십~수백 배치로 캘리브레이션(모델/도메인에 따라 더 필요)
실제 서비스 입력과 동일한 전처리/정규화를 캘리브레이션에도 적용
정확도 민감 구간은 QAT(Quantization Aware Training) 고려

실패 원인 4) dtype, device, eval 모드 불일치로 관측/변환이 꼬임

대표 증상

“expected CPU tensor” 또는 “quantized kernel not available for CUDA” 류 오류
model.train() 상태에서 캘리브레이션을 돌려 통계가 흔들림
AMP/autocast가 켜져서 관측값이 의도치 않게 FP16 기반으로 수집됨

왜 발생하나

많은 INT8 백엔드는 CPU 중심이거나 특정 디바이스만 지원합니다. 또한 BatchNorm/Dropout 등은 eval() 여부가 통계에 큰 영향을 줍니다.

확인 방법

model.training 값 확인
입력/모델 파라미터 device 확인
캘리브레이션 구간에서 autocast 비활성화

model = model.eval().cpu()
example = (example[0].cpu(),)

with torch.no_grad():
    with torch.autocast(device_type="cpu", enabled=False):
        prepared = prepare_pt2e(torch.export.export(model, example), quantizer)
        for _ in range(32):
            prepared.module()(*example)
        converted = convert_pt2e(prepared)

해결/우회

PTQ는 일단 CPU + eval() + no_grad로 고정하고 성공부터 시킨 뒤 확장
AMP는 캘리브레이션/변환 구간에서 끄기
디바이스 지원이 필요한 경우 백엔드/타깃 런타임에 맞는 quantizer 선택

실패 원인 5) 연산이 fusion되지 않아 스케일 전파가 깨지고, 성능/정확도가 동반 하락

대표 증상

변환은 되었는데 기대한 속도 향상이 거의 없음
Conv/Linear 주변에 quantize/dequantize가 과도하게 삽입됨
정확도도 같이 떨어짐

왜 발생하나

INT8의 이점은 “양자화된 연산을 길게 유지”할 때 커집니다. 그런데 그래프가 다음처럼 끊기면 Q/DQ가 남발됩니다.

Conv 다음에 지원 안 되는 activation이 끼어 있음
reshape/permute가 과도하게 들어가 연산 패턴이 분절됨
residual/add 경로에서 스케일이 맞지 않아 dequantize가 필요해짐

확인 방법

변환된 그래프에서 quantize_per_tensor, dequantize 류 노드가 어디에 얼마나 있는지 봅니다.

gm = converted.graph_module
ops = [n.target for n in gm.graph.nodes if n.op in ("call_function", "call_method")]
print("dequantize" in " ".join([str(o) for o in ops]))

해결/우회

블록 구조를 “양자화 우호적”으로 재배치(예: Conv+ReLU를 붙이고 불필요한 view 최소화)
residual 구간은 per-channel/per-tensor 설정을 조정하거나, 일부 구간 FP32 유지
성능 목표라면 “정확도 1% 손실 허용” 같은 기준을 정하고, 그 안에서 Q/DQ 삽입을 최소화하는 방향으로 타협

실패 원인 6) Per-channel 설정/채널 축(axis) 불일치로 변환이 막힘

대표 증상

특정 Conv/Linear에서만 변환 실패
에러 메시지에 “axis”, “channel”, “per_channel” 관련 문구

왜 발생하나

Per-channel weight quantization은 일반적으로 정확도에 유리하지만, 연산 종류/레이아웃/백엔드에 따라 요구하는 채널 축이 다르거나 제한이 있습니다. 특히 export 과정에서 weight 레이아웃이 바뀌거나, 그룹 컨볼루션/Depthwise에서 제약이 생길 수 있습니다.

확인 방법

문제 레이어를 최소 재현으로 떼어내 export/prepare/convert를 돌려봅니다.
per-channel을 per-tensor로 바꿔서 통과하는지 확인합니다.

qconfig = get_symmetric_quantization_config(is_per_channel=False)
quantizer.set_global(qconfig)

해결/우회

우선 per-tensor로 성공 경로를 만든 뒤, 병목 레이어만 per-channel로 점진 적용
그룹/Depthwise Conv는 백엔드 지원 여부를 확인하고, 지원이 약하면 해당 블록만 FP32로 남기기

실패 원인 7) 버전/백엔드 조합 이슈: API 경로 변경, 커널 미지원, 기대 동작 차이

대표 증상

튜토리얼 코드 그대로인데 import 에러
동일 코드가 로컬에서는 되는데 CI/서버에서는 실패
변환은 되는데 특정 CPU에서만 크래시 또는 성능 역전

왜 발생하나

PyTorch 2.x에서 PT2E/quantization 관련 API는 비교적 빠르게 변하고, 백엔드 커널 지원은 OS/CPU feature(예: AVX2/AVX512)와도 연관됩니다. 또한 nightly/stable 간 동작 차이도 종종 있습니다.

확인 방법

실행 환경에서 버전/빌드 플래그를 출력해 기록합니다.

import torch
print(torch.__version__)
print(torch.__config__.show())

해결/우회

동일한 PyTorch 버전으로 고정하고(프로덕션은 특히), 재현 가능한 환경을 먼저 확보
CPU feature가 다른 머신(노트북 vs 서버)에서 성능/동작이 다를 수 있으니, 타깃 머신에서 측정
문서/예제 코드를 가져올 때는 “작성 시점 버전”을 같이 확인

실전 디버깅 체크리스트(한 번에 훑기)

torch.export.export가 통과하는가
변환 후 그래프에 INT8 연산이 실제로 들어갔는가(대부분 FP32면 사실상 실패)
캘리브레이션을 prepared 모델로 충분히 수행했는가
eval()/no_grad/AMP off/CPU 고정으로 단순화했는가
Q/DQ가 과도하게 삽입되어 있지 않은가(성능 실패 신호)
per-channel 옵션을 per-tensor로 낮추면 통과하는가
버전/백엔드/CPU feature를 고정했는가

이런 식의 “원인 범주화”는 LLM 에이전트의 무한 루프를 끊을 때도 유사합니다. 문제를 한 덩어리로 보지 않고, 재현 가능한 관측 지점을 만들어 단계별로 차단합니다. 관심 있다면 LangChain 에이전트 무한루프·툴난사 차단법도 같은 디버깅 사고방식을 다룹니다.

보너스: 실패를 빠르게 재현하는 최소 템플릿

아래 템플릿은 “어디서 실패하는지”를 단계별로 분리해 로그를 남깁니다.

import torch
import torch.nn as nn
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
    XNNPACKQuantizer,
    get_symmetric_quantization_config,
)

def pt2e_int8_debug(model: nn.Module, example_inputs):
    model = model.eval().cpu()
    example_inputs = tuple(x.cpu() for x in example_inputs)

    print("[1] export")
    ep = torch.export.export(model, example_inputs)

    print("[2] quantizer")
    quantizer = XNNPACKQuantizer()
    qconfig = get_symmetric_quantization_config(is_per_channel=True)
    quantizer.set_global(qconfig)

    print("[3] prepare")
    prepared = prepare_pt2e(ep, quantizer)

    print("[4] calibrate")
    prepared_mod = prepared.module()
    with torch.no_grad():
        with torch.autocast(device_type="cpu", enabled=False):
            for _ in range(32):
                prepared_mod(*example_inputs)

    print("[5] convert")
    converted = convert_pt2e(prepared)

    print("[6] run")
    with torch.no_grad():
        y = converted.module()(*example_inputs)
    return converted, y

이 템플릿으로 “export에서 실패인지”, “prepare/convert에서 패턴 문제인지”, “성공했지만 정확도/성능이 실패인지”를 빠르게 분리할 수 있습니다.

마무리

PT2E INT8 양자화 실패는 대체로 export 정적화 실패, 백엔드 패턴 미지원, 캘리브레이션 부실, 모드/디바이스/dtype 불일치, fusion 붕괴, per-channel 축 제약, 버전/커널 이슈로 수렴합니다.

가장 효율적인 접근은 “일단 CPU, eval, per-tensor, 단순 모델”로 성공 경로를 만든 다음, per-channel/부분 양자화/모델 구조 개선을 단계적으로 적용하는 것입니다. 이렇게 하면 실패 원인이 겹쳐도 하나씩 제거하면서 안정적으로 INT8 경로를 확보할 수 있습니다.