Published on

PyTorch 2.x PTQ로 int8 양자화해 4배 경량화

Authors

서빙 환경에서 모델이 무겁다는 문제는 대부분 두 가지로 나타납니다. 첫째는 메모리 점유(가중치와 활성값)로 인해 한 프로세스에 올릴 수 있는 동시 모델 수가 제한되는 문제, 둘째는 CPU 추론에서 연산량과 메모리 대역폭 병목으로 지연 시간이 튀는 문제입니다. 이럴 때 가장 비용 대비 효과가 큰 접근이 PTQ(Post-Training Quantization) 입니다. 학습을 다시 하지 않거나(또는 아주 소량의 보정 데이터만으로) FP32 모델을 int8로 변환해 가중치 메모리를 약 4배(32비트에서 8비트) 줄이고, CPU에서는 int8 커널을 통해 지연 시간을 줄일 수 있습니다.

이 글에서는 PyTorch 2.x의 torch.ao.quantization 및 PT2 Export 기반 흐름을 중심으로, int8 PTQ로 실제 경량화와 성능 개선을 얻는 방법을 단계별로 다룹니다.

PTQ int8가 “4배 경량화”로 이어지는 이유

가중치가 FP32일 때 파라미터 하나는 4바이트입니다. int8로 바꾸면 1바이트가 됩니다. 따라서 가중치만 놓고 보면 이론적으로 4배 줄어듭니다.

다만 실제 프로덕션에서는 다음 요소 때문에 체감이 달라질 수 있습니다.

  • 스케일/제로포인트(양자화 파라미터) 저장 오버헤드가 약간 생깁니다.
  • 활성값(activation)은 설정에 따라 int8로 내려가거나, 일부는 FP32로 남을 수 있습니다.
  • 레이어에 따라 int8 커널이 없으면 dequant/quant가 끼면서 이득이 줄 수 있습니다.

그럼에도 MLP, CNN, 일부 Transformer 블록에서 가중치 메모리 절감 효과는 매우 안정적으로 얻을 수 있고, CPU 추론에서는 특히 효과가 큽니다.

PyTorch 2.x에서의 양자화 옵션 한눈에 보기

PyTorch 양자화는 크게 아래로 나뉩니다.

  • Dynamic Quantization:
    • 가중치만 int8로, 활성값은 런타임에서 동적으로 양자화
    • 구현이 단순하고 LSTM/Linear 계열에서 효과적
  • Static PTQ:
    • 가중치와 활성값을 모두 int8로(가능한 구간)
    • Calibration(보정) 단계가 필요
    • CNN/일반 피드포워드에서 성능 이득이 커질 가능성이 높음

PyTorch 2.x에서는 torch.compile과의 조합, 그리고 Export 기반 양자화 흐름이 계속 개선되고 있지만, 실무에서 빠르게 적용하려면 Dynamic Quantization 또는 Eager/FX 기반 PTQ부터 시작하는 것이 안전합니다.

준비: 모델을 “추론 친화적”으로 정리하기

양자화 전에 아래를 점검하면 실패 확률이 크게 줄어듭니다.

  • model.eval() 상태에서만 변환/평가
  • Dropout, BatchNorm이 학습 모드로 남아 있지 않게 확인
  • 입력 텐서 dtype과 shape를 서빙과 동일하게 고정
  • 가능하면 연산을 nn.Linear, nn.Conv2d, nn.ReLU 같은 표준 모듈로 유지

특히 Transformer류는 LayerNorm, GELU 등에서 양자화 경로가 복잡해질 수 있으므로, 처음에는 Linear 위주의 서브모듈만 양자화하는 전략이 좋습니다.

가장 쉬운 시작: Dynamic Quantization으로 Linear int8 적용

Dynamic Quantization은 “보정 데이터 없이” 빠르게 적용할 수 있어, CPU 서빙에서 즉시 효과를 확인하기 좋습니다.

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim=1024, hidden=2048, out_dim=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, out_dim),
        )

    def forward(self, x):
        return self.net(x)

model_fp32 = MLP().eval()

# Dynamic quantization: Linear 계열을 int8 weight로 변환
model_int8 = torch.ao.quantization.quantize_dynamic(
    model_fp32,
    {nn.Linear},
    dtype=torch.qint8,
)

x = torch.randn(32, 1024)
with torch.inference_mode():
    y = model_int8(x)
print(y.shape)

Dynamic Quantization의 장단점

  • 장점
    • Calibration 불필요
    • 적용이 간단하고 실패가 적음
    • Linear 비중이 큰 모델에서 효과적
  • 단점
    • 활성값까지 완전히 int8로 고정되는 방식이 아니라서, Static PTQ 대비 이득이 제한될 수 있음

Static PTQ 핵심: Calibration(보정)으로 활성값 범위 잡기

Static PTQ는 “대표 입력”으로 모델을 몇 번 흘려보내며, 각 레이어의 활성값 분포를 관찰해 scalezero_point를 정합니다. 이 과정을 Calibration이라고 합니다.

Calibration 데이터는 다음 조건을 만족하면 됩니다.

  • 실제 서빙 입력과 분포가 유사
  • 라벨 불필요
  • 수십~수백 배치 정도면 시작 가능(모델과 도메인에 따라 조정)

FX Graph Mode 기반 Static PTQ 예제(실무형)

아래 코드는 FX 기반 PTQ의 전형적인 흐름입니다.

  1. qconfig 설정
  2. prepare_fx로 관측자(observer) 삽입
  3. Calibration 데이터로 forward 몇 회 실행
  4. convert_fx로 int8 모델로 변환
import torch
import torch.nn as nn
from torch.ao.quantization import get_default_qconfig
from torch.ao.quantization.quantize_fx import prepare_fx, convert_fx

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.ReLU(),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 8 * 8, 100),
        )

    def forward(self, x):
        x = self.features(x)
        return self.classifier(x)

model = SmallCNN().eval()

# 백엔드에 따라 qconfig가 달라집니다.
# x86 서버 CPU라면 보통 fbgemm
qconfig = get_default_qconfig("fbgemm")
qconfig_dict = {"": qconfig}

example_inputs = (torch.randn(1, 3, 32, 32),)
prepared = prepare_fx(model, qconfig_dict, example_inputs)

# Calibration: 실제 입력 분포를 대표하는 데이터로 여러 번 forward
with torch.inference_mode():
    for _ in range(100):
        batch = torch.randn(16, 3, 32, 32)
        prepared(batch)

quantized_model = convert_fx(prepared)

with torch.inference_mode():
    out = quantized_model(torch.randn(4, 3, 32, 32))
print(out.shape)

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

  • 정확도 하락이 큰 경우
    • Calibration 데이터가 너무 적거나 분포가 다름
    • per_channel 양자화 지원 여부 확인(특히 Conv/Linear)
    • 일부 민감 레이어는 FP32로 남기고 부분 양자화 시도
  • 성능이 기대보다 안 나오는 경우
    • int8 커널이 적용되는 연산 비중이 낮음
    • quant/dequant 노드가 과도하게 삽입됨
    • CPU 스레딩 설정, 배치 크기, 메모리 레이아웃 문제

“4배 경량화”를 수치로 확인하기: 파라미터 메모리 추정

모델 전체 메모리는 버퍼, 옵저버, 런타임 오버헤드가 섞여 단순 비교가 어렵습니다. 대신 가중치 텐서 크기를 기준으로 비교하면 경향을 빠르게 볼 수 있습니다.

import torch

def param_bytes(model: torch.nn.Module) -> int:
    total = 0
    for p in model.state_dict().values():
        if isinstance(p, torch.Tensor):
            total += p.numel() * p.element_size()
    return total

fp32_bytes = param_bytes(model)
int8_bytes = param_bytes(quantized_model)

print("fp32 params MB:", fp32_bytes / 1024 / 1024)
print("int8 params MB:", int8_bytes / 1024 / 1024)
print("compression x:", fp32_bytes / max(int8_bytes, 1))

주의할 점은, 양자화 모델은 내부 표현이 일반 FP32 텐서와 다르고(예: packed weight), 백엔드에 따라 저장 형태가 달라질 수 있어 단순 비교가 흔들릴 수 있습니다. 그래도 대략적인 경량화 비율을 확인하는 용도로는 충분합니다.

벤치마크: CPU 지연 시간 비교(최소한의 방식)

성능은 환경에 매우 민감하므로, 최소한 아래 원칙은 지키는 것이 좋습니다.

  • torch.inference_mode() 사용
  • 워밍업 반복 후 측정
  • 스레드 수 고정(torch.set_num_threads)
import time
import torch

torch.set_num_threads(8)

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

x = torch.randn(32, 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 x:", t_fp32 / max(t_int8, 1e-9))

실무에서는 여기에 p95/p99, 배치 크기별 곡선, 동시성(멀티프로세스/멀티스레드)까지 포함해 봐야 합니다.

어떤 모델에 PTQ int8가 특히 잘 먹히나

  • 잘 맞는 편
    • CNN 계열(특히 Conv+ReLU 패턴)
    • MLP/Linear 비중이 큰 추천/랭킹 모델
    • 전통적인 비전 백본 일부
  • 까다로운 편
    • Attention 비중이 높고 연산 그래프가 복잡한 LLM/Transformer 전체 양자화
    • LayerNorm, Softmax, GELU 주변에서 정확도 민감

이 경우에는 “전체 int8”에 집착하기보다, 가장 비싼 Linear 블록만 부분 양자화하거나, 더 나아가 GPTQ/AWQ 같은 가중치 양자화 기법 또는 int4로 가는 전략을 병행합니다.

운영 관점 체크리스트: 배포 시 실수 줄이기

양자화는 코드 한 줄로 끝나는 것처럼 보여도, 배포에서는 다음이 문제를 만듭니다.

  • 프로덕션 CPU가 개발 머신과 달라 커널 최적화가 다르게 적용
  • 컨테이너 이미지에 따라 oneDNN, OpenMP 런타임 차이로 성능 흔들림
  • 모델 파일 크기는 줄었는데, 실제 RSS가 기대만큼 안 줄어 “효과가 없었다”로 오해

CI/CD 파이프라인에서 이미지/런타임 차이로 성능이 흔들릴 때는 인프라 이슈를 함께 점검해야 합니다. 예를 들어 Docker 기반 배포에서 권한이나 실행 환경이 꼬이면 성능 측정 자체가 왜곡될 수 있는데, 이런 류의 문제는 GitHub Actions Docker CI/CD 권한(permis…) 오류 해결 같은 글에서 다루는 체크 포인트(러너 권한, 캐시, 레이어 재사용 등)도 함께 참고할 만합니다.

또한 서버에서 프로세스가 반복 재시작되면 워밍업/캐시가 계속 날아가 지연 시간이 악화될 수 있습니다. 운영 중 원인 추적은 systemd 서비스가 계속 재시작될 때 원인 추적법처럼 로그와 재시작 정책을 함께 보는 방식이 도움이 됩니다.

결론: PTQ int8는 “가장 먼저 해볼 경량화”다

PyTorch 2.x에서 PTQ int8는 여전히 가장 실용적인 경량화 옵션 중 하나입니다.

  • 가중치 기준 최대 4배 수준의 메모리 절감을 기대할 수 있고
  • CPU 추론에서는 특히 지연 시간 개선 여지가 크며
  • Dynamic Quantization으로 빠르게 가능성을 확인한 뒤, Static PTQ로 더 큰 이득을 노리는 흐름이 현실적입니다.

다음 단계로는 (1) 어떤 서브모듈이 병목인지 프로파일링, (2) 부분 양자화 전략, (3) 대표 Calibration 세트 구성, (4) 배포 환경에서의 재현 가능한 벤치마크를 갖추는 순서로 고도화하는 것을 권장합니다.