PyTorch 2.1+ PT2E로 LLM 4bit 양자화 실전

서빙 환경에서 LLM 비용의 대부분은 결국 메모리 대역폭과 VRAM 점유로 귀결됩니다. 파라미터가 큰 모델은 FP16/BF16로도 충분히 무겁고, KV 캐시까지 얹히면 단일 GPU에서 동시 처리량이 급격히 떨어집니다. 그래서 4bit 양자화는 “품질을 크게 망치지 않으면서” VRAM을 확 줄이는 가장 현실적인 선택지가 됩니다.

PyTorch 2.1+에서는 기존 FX Graph Mode Quantization 흐름을 정리하고, PT2E(Prepare/Convert-to-Export) 기반으로 내보내기(Export) 친화적인 양자화 파이프라인을 강화했습니다. 이 글은 PT2E 관점에서 LLM 4bit 경량화를 시도할 때 필요한 개념, 코드, 그리고 실무에서 자주 부딪히는 함정을 정리합니다.

목표
PyTorch 2.1+에서 PT2E 양자화 파이프라인의 큰 그림 이해
LLM에 적용할 때의 현실적인 전략(가중치 4bit 중심, 활성값은 보수적으로)
성능/정확도/배포 이슈 체크리스트

PT2E 양자화란 무엇인가

PT2E는 이름 그대로 Export를 기준으로 양자화 준비(prepare)와 변환(convert)을 수행하는 흐름입니다. 핵심은 다음과 같습니다.

torch.export로 정적이고 안정적인 그래프를 얻는다
그 그래프를 대상으로 관측(Observer) 또는 가짜 양자화(FakeQuant) 삽입 등 양자화 준비를 한다
캘리브레이션(또는 QAT 학습) 이후, 실제 양자화 연산/패킹으로 변환한다
변환된 모델을 torch.compile/AOTInductor/백엔드로 보내거나, 다른 런타임으로 내보내기 쉬워진다

LLM에서 체감이 큰 지점은, “모델이 크고 복잡할수록” 그래프 안정성이 중요해지고, Export를 중심으로 한 파이프라인이 디버깅/재현성 면에서 유리해진다는 점입니다.

LLM 4bit에서 현실적인 타협점

LLM을 4bit로 만든다고 할 때, 모든 텐서를 무작정 4bit로 내리는 건 대개 실패합니다. 실무적으로는 아래 전략이 흔합니다.

1) 가중치(weight) 4bit + 활성값(activation) 고정밀

가중치만 4bit로 줄여도 VRAM이 크게 절약됩니다.
활성값까지 4bit로 내리면 정확도 하락과 커널/백엔드 제약이 급격히 커집니다.
특히 Attention의 Softmax 주변, LayerNorm, RMSNorm, 로지트 헤드 등은 민감합니다.

2) per-channel(또는 per-group) 스케일

LLM의 Linear 계층은 출력 채널별 스케일(per-channel) 또는 그룹 단위 스케일(per-group)에서 품질이 잘 나옵니다.
4bit에서는 per-tensor 스케일이 품질을 많이 깎을 수 있습니다.

3) outlier 처리(정확도 급락 방지)

일부 채널/행이 큰 값(outlier)을 가지면 4bit 스케일이 그것에 끌려가 전체 정밀도가 무너집니다.
실무에서는 outlier 채널만 고정밀로 남기거나, 그룹 크기를 조정하거나, SmoothQuant류의 스케일 재분배를 씁니다.

PT2E 기반 양자화 파이프라인 개요

PT2E 양자화는 보통 아래 단계로 이해하면 편합니다.

모델 로드 및 추론용 설정(eval, no_grad)
torch.export.export로 Export 그래프 획득
Quantizer 설정(무슨 연산을 어떤 스킴으로 양자화할지)
prepare_pt2e로 관측/가짜양자화 삽입
캘리브레이션(대표 입력을 여러 번 흘려 통계 수집)
convert_pt2e로 실제 양자화 모델로 변환
성능/정확도 검증 및 배포

아래 코드는 “구조를 이해하기 위한” 예시이며, 실제 LLM 4bit는 백엔드(예: Inductor, 커스텀 커널, GPU 4bit GEMM 지원 여부)에 따라 적용 가능한 경로가 달라질 수 있습니다.

코드: PT2E prepare/convert 뼈대

다음 예시는 PT2E 흐름에서 Export 그래프를 만들고, prepare/convert를 수행하는 최소 뼈대입니다. (환경에 따라 API 이름이 조금씩 바뀔 수 있으니, 사용 중인 PyTorch 버전 문서를 함께 확인하세요.)

import torch

# PyTorch 버전에 따라 import 경로가 달라질 수 있습니다.
# 아래는 PT2E 개념을 보여주는 예시입니다.
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


def export_model(model, example_inputs):
    model.eval()
    with torch.no_grad():
        exported = torch.export.export(model, example_inputs)
    return exported


def pt2e_quantize(exported, calibrate_fn):
    # 예시: XNNPACK quantizer (주로 CPU 백엔드)
    # LLM 4bit에서는 GPU/커스텀 백엔드가 필요할 수 있습니다.
    quantizer = XNNPACKQuantizer()

    qconfig = get_symmetric_quantization_config(
        is_per_channel=True,
        is_qat=False,
    )
    quantizer.set_global(qconfig)

    prepared = prepare_pt2e(exported, quantizer)

    # 캘리브레이션: 대표 입력을 준비된 모델에 여러 번 통과
    calibrate_fn(prepared)

    converted = convert_pt2e(prepared)
    return converted

캘리브레이션 함수 예시

LLM은 입력 분포가 다양하므로, 캘리브레이션 데이터가 부실하면 품질이 크게 흔들립니다. 최소한 실제 트래픽과 유사한 프롬프트 길이/도메인을 반영하는 것이 좋습니다.

def calibrate_fn(prepared_model):
    prepared_model.eval()
    with torch.no_grad():
        for _ in range(64):
            # 예시 입력(실제로는 토크나이저 결과 텐서)
            input_ids = torch.randint(0, 1000, (1, 128))
            attention_mask = torch.ones_like(input_ids)
            _ = prepared_model(input_ids=input_ids, attention_mask=attention_mask)

LLM 4bit에서 “PT2E만으로 끝”이 어려운 이유

PyTorch의 양자화 스택은 계속 발전 중이지만, LLM 4bit는 다음 이유로 단순히 prepare/convert만으로 끝나지 않는 경우가 많습니다.

1) 4bit GEMM 커널과 패킹 포맷

8bit/16bit는 상대적으로 표준화된 경로가 많지만, 4bit는 커널/패킹/스케일 적용 방식이 백엔드마다 다릅니다.
“가중치 4bit + 활성값 FP16” 형태라도, 실제로는 4bit 가중치를 어떤 레이아웃으로 패킹하고, GEMM에서 어떻게 언패킹/스케일링할지가 성능을 좌우합니다.

2) 연산자 커버리지

LLM 그래프에는 Linear 외에도 Rotary Embedding, RMSNorm, Attention, KV cache 업데이트 등 다양한 연산이 있습니다.
양자화가 가능한 연산과 아닌 연산이 섞이면, 양자화/비양자화 경계에서 캐스팅이 늘어 성능이 악화될 수 있습니다.

3) 정확도 민감 구간

LayerNorm/RMSNorm, Softmax, 최종 LM head는 작은 오차가 누적되어 품질이 흔들릴 수 있습니다.
실무에서는 “모든 Linear을 4bit”가 아니라, 특정 블록/레이어만 4bit 또는 일부 채널만 고정밀을 남기는 식으로 타협합니다.

실전 전략: 어떤 레이어를 먼저 4bit로 내릴까

LLM에서 VRAM을 많이 잡아먹는 건 대체로 큰 Linear 계층입니다.

MLP의 up/down projection
Attention의 QKV projection, output projection

반면 아래는 보수적으로 접근하는 편이 안전합니다.

Embedding/LM head(특히 vocab가 클 때)
Norm 계열
Softmax 주변

즉, “가장 큰 Linear부터” 4bit를 시도하고, 품질이 무너지면 민감 레이어를 제외하거나 그룹 크기/스케일 방식을 바꾸는 순서가 실패 비용이 적습니다.

검증: 품질과 성능을 같이 봐야 한다

4bit는 VRAM이 줄어드는 만큼 속도가 무조건 빨라질 것이라고 기대하기 쉽지만, 실제로는 다음이 병목이 됩니다.

커널이 최적화되어 있지 않으면 언패킹/스케일 비용이 커져 느려질 수 있음
양자화/비양자화 경계 캐스팅 증가
KV 캐시가 여전히 BF16/FP16이면 전체 VRAM 절감이 제한적

그래서 아래처럼 최소한의 지표를 함께 봐야 합니다.

Perplexity 또는 태스크별 정확도(가능하면 내부 eval)
토큰/초(tokens/sec)
VRAM peak(모델 로드 + 워밍업 + 실제 길이 프롬프트)
배치 크기별 지연(p50/p95)

트러블슈팅 체크리스트

OOM이 계속 난다

4bit로 가중치를 줄여도 KV 캐시가 크면 OOM이 날 수 있습니다.
프롬프트 길이/동시 요청 수가 늘면 KV 캐시가 선형으로 증가합니다.

이럴 때는 양자화 자체보다 서빙 파라미터(최대 토큰, 배치, 동시성)와 캐시 정책을 먼저 점검하는 게 빠릅니다. OOM 원인 추적은 아래 글의 체크리스트가 도움이 됩니다.

리눅스 OOM Killer로 프로세스 죽을 때 원인 추적

CrashLoopBackOff로 서빙이 계속 재시작된다

컨테이너 메모리 제한과 실제 피크 사용량 불일치
모델 워밍업 중 일시적 피크
드라이버/런타임 호환성 문제

쿠버네티스에서 원인을 빠르게 좁히는 방법은 아래 글을 참고하세요.

K8s CrashLoopBackOff 원인별 빠른 진단·복구

GPU 오토스케일이 필요하다

4bit로 줄여도 트래픽 변동이 크면 GPU를 유휴로 두는 시간이 늘어 비용이 새어 나갑니다. 모델 서빙을 자동 확장하는 아키텍처를 함께 고려하는 게 좋습니다.

KServe+Knative로 GPU 모델 자동스케일 배포

배포 관점에서의 권장 운영 방식

버전 고정: PyTorch, CUDA, 드라이버, 양자화 백엔드는 조합에 따라 성능/정확도가 크게 흔들립니다. 최소 단위로 락을 걸고(컨테이너 이미지), 재현 가능한 벤치를 남기세요.
캘리브레이션 데이터 관리: “어떤 프롬프트 분포로 캘리브레이션했는지”가 곧 품질의 일부입니다. 데이터 샘플링 규칙을 코드로 고정하세요.
롤백 플랜: 4bit는 코너 케이스에서 품질 이슈가 늦게 드러날 수 있습니다. 동일 모델의 BF16/FP16 경로를 함께 운영하며 점진 전환하는 편이 안전합니다.

결론

PyTorch 2.1+의 PT2E 양자화는 Export 중심으로 그래프를 안정화하고, 양자화 준비/변환을 구조적으로 관리할 수 있게 해줍니다. 다만 LLM 4bit는 단순 기능 호출만으로 끝나는 문제가 아니라, 4bit 커널/패킹/연산자 커버리지/정확도 민감 구간이 맞물린 시스템 문제에 가깝습니다.

실무에서는 다음 순서가 가장 성공 확률이 높습니다.

큰 Linear 계층부터 “가중치 4bit + 활성값 고정밀”로 제한적으로 시작
per-channel 또는 per-group 스케일로 품질 확보
캘리브레이션을 트래픽 분포에 맞게 구성
tokens/sec, VRAM peak, p95 지연을 함께 측정
OOM/재시작/오토스케일 등 운영 이슈까지 포함해 최종 의사결정

이 흐름으로 접근하면, 4bit는 단순한 최적화가 아니라 서빙 비용 구조를 바꾸는 레버가 됩니다.