PyTorch 2.0 PTQ로 LLM 4bit 양자화 실전

서빙 비용이 LLM 운영의 병목이 되는 순간이 옵니다. GPU 메모리와 대역폭이 한정된 환경에서 모델을 더 많이 띄우거나, 더 긴 컨텍스트를 처리하거나, 더 낮은 지연시간을 만들려면 결국 가중치/활성값을 줄이는 양자화를 고민하게 됩니다.

이 글에서는 PyTorch 2.0 생태계에서 PTQ(Post-Training Quantization) 접근으로 LLM을 4bit로 줄이는 실전 흐름을 다룹니다. 핵심은 “학습을 다시 하지 않고(또는 최소화하고)”, 대표 데이터로 캘리브레이션한 뒤, 4bit 가중치 중심의 저정밀 연산 커널을 활용해 추론 효율을 얻는 것입니다.

참고: 배포 파이프라인에서 체크리스트 성격의 글을 좋아한다면, 운영 장애를 빠르게 진단하는 스타일의 글도 함께 참고해보세요. 예: Cloud Run 504 Timeout 원인·해결 9가지

1) PTQ 4bit의 기본: 무엇을 4bit로 줄이나

LLM 양자화에서 보통 “4bit”라고 말할 때는 다음 중 하나(또는 조합)를 의미합니다.

• 가중치 4bit(W4): Linear 레이어의 weight를 4bit로 저장하고, matmul 시에 dequantize 또는 int 연산 커널로 처리합니다.
• 활성값 8bit 또는 16bit(A8/A16): 활성값까지 4bit로 내리면 정확도 손실이 급격히 커지는 경우가 많아, 실무에서는 W4A16 또는 W4A8 조합이 흔합니다.
• 그룹 단위 스케일링(group-wise quantization): 채널별/그룹별로 scale, zero-point를 두어 4bit의 표현력을 보완합니다.

PTQ는 학습을 하지 않으니, “어떤 분포를 기준으로 스케일을 잡을지”가 성능을 좌우합니다. 그래서 캘리브레이션 데이터가 중요합니다.

2) PyTorch 2.0에서의 현실적인 선택지

PyTorch 2.0 자체는 torch.ao.quantization(FX Graph Mode 등)을 제공하지만, LLM 4bit는 다음 이유로 전통적인 int8 PTQ 워크플로우와 다르게 접근하는 경우가 많습니다.

• Transformer 블록의 연산 패턴(특히 Linear 비중)이 크고, 4bit 커널 최적화가 중요합니다.
• 4bit는 하드웨어/커널 의존성이 강해서, “PyTorch 기본 양자화 API만으로 끝”이 잘 안 됩니다.

그래서 실무에서는 보통 다음 조합을 씁니다.

• PyTorch 2.x + Hugging Face Transformers + bitsandbytes 4bit(NF4/FP4)
• 또는 PyTorch 2.x + GPTQ/AWQ 계열 툴링

이 글의 코드는 “PyTorch 2.0 환경에서 PTQ 4bit를 실전처럼” 진행하기 위해, 가장 접근성이 좋은 bitsandbytes 4bit 로딩을 중심으로 설명합니다. (엄밀히 말하면 bitsandbytes는 “양자화된 가중치 포맷+커널”을 제공하고, 캘리브레이션/스케일링은 구현 방식에 따라 다릅니다.)

3) 환경 준비: 버전과 GPU 체크

4bit는 커널 호환성 이슈가 잦습니다. 아래 항목을 먼저 확인하세요.

• CUDA 버전, 드라이버 버전
• GPU 아키텍처(예: Ampere, Ada 등)
• torch 버전과 transformers, bitsandbytes 버전

예시 설치(환경마다 다르므로 개념만 참고):

pip install -U torch transformers accelerate bitsandbytes

버전 확인:

import torch
import transformers

print("torch:", torch.__version__)
print("cuda:", torch.version.cuda)
print("gpu:", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu")
print("transformers:", transformers.__version__)

4) 4bit PTQ 실전: 모델 로딩과 추론 파이프라인

4.1 4bit 로딩(가장 빠른 체감 포인트)

transformers에서 BitsAndBytesConfig로 4bit 로딩을 활성화할 수 있습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import BitsAndBytesConfig

model_id = "meta-llama/Llama-2-7b-hf"  # 예시

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # nf4가 실무에서 무난
    bnb_4bit_use_double_quant=True,      # 2단계 양자화로 오차 완화
    bnb_4bit_compute_dtype=torch.bfloat16 # 연산 dtype
)

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto"
)

model.eval()

여기서 중요한 포인트:

• bnb_4bit_compute_dtype는 정확도/성능에 영향이 큽니다. GPU가 bf16을 잘 지원하면 bfloat16이 안정적인 편입니다.
• nf4는 분포를 고려한 4bit 포맷으로, 단순 int4 대비 정확도 손실이 적은 경우가 많습니다.

4.2 간단한 생성 테스트

prompt = "Explain PTQ 4-bit quantization for LLMs in Korean."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        temperature=0.0,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

이 단계에서 이미 VRAM 사용량이 크게 줄어든 것을 확인할 수 있습니다.

5) “PTQ답게” 만들기: 캘리브레이션 데이터와 품질 평가

bitsandbytes 4bit 로딩은 편하지만, 실무에서 중요한 건 “이 설정이 내 도메인에서 얼마나 버티나”입니다. PTQ의 핵심은 대표 데이터로 품질을 검증하고, 필요하면 설정을 조정하는 것입니다.

5.1 캘리브레이션 데이터 준비

• 길이 분포: 짧은 질의만 있으면 긴 컨텍스트에서 품질이 무너질 수 있습니다.
• 도메인 분포: 법률/의료/CS 등 전문 도메인은 일반 코퍼스와 분포가 다릅니다.
• 토큰화 특성: 숫자, 코드, 표 등은 민감합니다.

간단한 형태의 캘리브레이션/평가 세트 예시:

calib_prompts = [
    "Summarize the following incident report...",
    "Given this JSON, extract fields...",
    "Write a SQL query for...",
    "Explain this Python code...",
]

5.2 품질 측정: 정량(Perplexity) + 정성(샘플)

LLM 품질은 단일 지표로 끝나지 않습니다. 그래도 PTQ 비교에는 Perplexity가 1차 지표로 유용합니다.

아래는 매우 단순화한 perplexity 계산 예시입니다.

import math
import torch

def ppl_on_texts(model, tokenizer, texts, max_length=512):
    model.eval()
    losses = []
    for t in texts:
        enc = tokenizer(t, return_tensors="pt", truncation=True, max_length=max_length)
        enc = {k: v.to(model.device) for k, v in enc.items()}
        with torch.no_grad():
            out = model(**enc, labels=enc["input_ids"])
            losses.append(out.loss.detach().float().cpu().item())
    mean_loss = sum(losses) / len(losses)
    return math.exp(mean_loss)

texts = [
    "Today we discuss post-training quantization.",
    "Large language models require significant memory bandwidth.",
]

print("ppl:", ppl_on_texts(model, tokenizer, texts))

주의할 점:

• 절대값보다 FP16/BF16 원본 대비 상대 비교가 중요합니다.
• truncation, padding 정책이 비교 조건을 바꾸니 고정하세요.

6) 성능 측정: VRAM, 토큰/초, 지연시간

양자화는 정확도만큼이나 “돈이 되는지”를 봐야 합니다.

6.1 대략적인 VRAM 확인

import torch

torch.cuda.reset_peak_memory_stats()
# 간단한 워밍업
prompt = "Hello"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    _ = model.generate(**inputs, max_new_tokens=32)

peak = torch.cuda.max_memory_allocated() / (1024**3)
print("peak_allocated_gb:", peak)

6.2 토큰/초 측정(간단 벤치)

import time
import torch

def bench_generate(model, tokenizer, prompt, new_tokens=128, n=5):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    # warmup
    with torch.no_grad():
        _ = model.generate(**inputs, max_new_tokens=16)

    times = []
    for _ in range(n):
        torch.cuda.synchronize()
        t0 = time.time()
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=new_tokens, do_sample=False)
        torch.cuda.synchronize()
        t1 = time.time()
        times.append(t1 - t0)

    # 생성된 토큰 수 대략 계산
    gen_tokens = out.shape[-1] - inputs["input_ids"].shape[-1]
    avg_t = sum(times) / len(times)
    return gen_tokens / avg_t, avg_t

prompt = "Write a short guide about PTQ for LLMs."
tps, latency = bench_generate(model, tokenizer, prompt)
print("tokens_per_sec:", tps)
print("latency_sec:", latency)

실무 팁:

• 첫 요청 지연시간은 커널 로딩/캐시 영향이 커서, 워밍업을 분리하세요.
• 배치 크기, 입력 길이, max_new_tokens에 따라 병목이 바뀝니다.

7) 정확도가 무너질 때 체크할 것(4bit 함정)

7.1 도메인 프롬프트에서만 급격히 나빠짐

• 캘리브레이션/평가 데이터가 일반 문장 위주면, 코드/숫자/표/전문용어에서 오차가 커질 수 있습니다.
• 해결: 도메인 대표 샘플을 반드시 포함해 상대 비교를 하세요.

7.2 특정 레이어 민감도

LLM은 레이어별 민감도가 다릅니다. 일부 방법(AWQ 등)은 “중요 채널을 더 잘 보존”하는 식으로 이를 보완합니다.

• 단순 4bit 로딩에서 품질이 부족하면, GPTQ/AWQ 계열을 검토할 가치가 있습니다.

7.3 compute dtype 선택 실수

• fp16이 더 빠르지만 오차가 커지는 케이스가 있습니다.
• bf16이 가능한 GPU라면 bf16부터 시도하는 편이 안전합니다.

8) PyTorch 2.x 컴파일과의 관계: torch.compile은 만능이 아니다

PyTorch 2.x의 torch.compile은 모델 형태/커널에 따라 이득이 달라집니다. 4bit 커널(bitsandbytes)은 이미 커스텀 커널을 쓰기 때문에, torch.compile이 큰 이득을 못 보거나 오히려 제약이 생길 수 있습니다.

그래도 시도는 가능합니다.

import torch

# 일부 환경에서만 유효. 문제 발생 시 즉시 원복 권장.
compiled_model = torch.compile(model)

실전 권장 순서:

1. 4bit 적용만으로 VRAM/속도 목표 달성 여부 확인
2. 그 다음에 KV 캐시 최적화, 배치/시퀀스 튜닝
3. 마지막으로 컴파일/서빙 런타임 최적화

운영에서 병목을 추적하는 습관은 다른 영역에도 그대로 통합니다. 예를 들어 “원인 후보를 빠르게 좁히는” 방식은 Chrome INP 나쁨 - Long Task 쪼개기 실전 가이드 같은 성능 디버깅 글과도 결이 같습니다.

9) 운영 관점 체크리스트: 재현 가능하게 만들기

양자화는 로컬에서 성공해도, 배포 환경에서 흔히 흔들립니다.

• 모델 아티팩트 버전 고정: transformers, bitsandbytes, torch 버전 고정
• GPU 타입 고정 또는 최소 보장: 서로 다른 아키텍처에서 성능/정확도 차이 발생 가능
• 프롬프트/샘플 기반 회귀 테스트: “양자화 적용 전/후”를 자동 비교

CI에서 권한/인증으로 배포가 막히는 케이스도 많습니다. 모델 서빙 이미지를 배포할 때 GitHub Actions를 쓴다면 GitHub Actions OIDC에서 AWS AssumeRoleAccessDenied 해결 같은 이슈가 실제로 발목을 잡습니다.

10) 결론: 4bit PTQ는 ‘적용’보다 ‘검증’이 핵심

PyTorch 2.0 환경에서 LLM 4bit PTQ는 이제 “한 줄 옵션”처럼 적용할 수 있는 시대가 됐지만, 실무에서 중요한 건 다음입니다.

• 내 트래픽 패턴(입력 길이, 배치, 응답 길이)에서 토큰/초와 지연시간이 실제로 개선되는지
• 내 도메인 데이터에서 정확도/안전성/일관성이 허용 범위인지
• 운영 환경에서 버전/커널/드라이버 차이로 재현성이 깨지지 않는지

추천 액션 플랜은 단순합니다.

1. W4 로딩으로 VRAM 절감 효과를 먼저 확인
2. 도메인 대표 데이터로 정량/정성 평가
3. 부족하면 nf4, double quant, compute dtype 조합을 바꿔가며 재평가
4. 그래도 부족하면 GPTQ/AWQ 같은 더 정교한 PTQ로 확장

이 과정을 “지표와 코드로 재현 가능하게” 만들어두면, LLM 모델 교체나 스케일 아웃 시에도 비용과 품질을 안정적으로 관리할 수 있습니다.