beautifulsoups.dev
Published on

ONNX Runtime로 LLM 4bit 양자화해 2배 가속

Authors

서빙 환경에서 LLM 추론 병목은 대개 두 가지로 수렴합니다. 첫째, KV 캐시와 가중치가 차지하는 메모리로 인해 배치나 동시성이 막히는 문제. 둘째, 메모리 대역폭 한계 때문에 GPU나 CPU 코어가 놀면서도 토큰 생성 속도가 기대만큼 나오지 않는 문제입니다. 4bit 양자화는 이 두 축을 동시에 건드리는 가장 강력한 카드 중 하나입니다. 특히 ONNX Runtime은 실행 그래프 최적화와 양자화 도구 체인이 잘 갖춰져 있어, PyTorch 기반 모델을 비교적 짧은 경로로 실서비스 가능한 형태로 만들 수 있습니다.

이 글은 ONNX Runtime에서 LLM을 4bit로 양자화해 체감 성능을 2배 수준으로 끌어올리는 것을 목표로 합니다. 다만 모델, 하드웨어, 시퀀스 길이, 배치, 프롬프트 길이에 따라 이득은 달라질 수 있으니, 재현 가능한 벤치마크 방법까지 같이 다룹니다.

왜 4bit가 LLM에서 특히 잘 먹히나

LLM 추론은 오토리그레시브 특성 때문에 토큰을 한 개씩 생성합니다. 이때 매 스텝마다 대부분의 연산은 "가중치 읽기"에 의해 지배됩니다. 즉, 연산량보다 메모리 대역폭이 먼저 한계에 닿는 경우가 많습니다.

4bit로 가중치를 줄이면 다음 효과가 생깁니다.

  • 가중치 메모리 풋프린트 감소: 같은 VRAM에서 더 큰 모델을 올리거나, 동일 모델로 더 큰 배치와 더 많은 동시 요청을 처리 가능
  • 메모리 대역폭 절감: 가중치 로딩이 줄어 토큰 생성 속도 개선 가능
  • 캐시 압박 완화: 시스템 메모리나 L3 캐시, GPU 캐시 관점에서 유리

하지만 4bit는 정확도 저하 리스크가 큽니다. 그래서 "어떤 방식의 4bit"인지가 중요합니다. 단순한 per-tensor 양자화는 손실이 커질 수 있고, 일반적으로는 group-wise, per-channel, 그리고 스케일/제로포인트 설계가 성패를 가릅니다.

ONNX Runtime에서의 선택지: 정적 양자화 vs 4bit weight-only

ONNX Runtime의 양자화는 크게 다음 범주로 이해하면 편합니다.

  • INT8 정적 양자화: activation까지 포함하는 경우가 많고, 캘리브레이션 데이터가 필요합니다. CPU에서 특히 강력합니다.
  • 4bit weight-only 양자화: activation은 FP16 또는 FP32로 두고 가중치만 4bit로 내립니다. LLM에서 실전성이 높습니다.

LLM에서는 보통 "가중치가 크고 activation은 상대적으로 작다"는 특성이 있으므로, weight-only 4bit가 현실적인 타협점입니다. activation까지 극단적으로 내리면 품질이 무너질 가능성이 커지고, 커널 지원도 제한됩니다.

참고로 INT8에서 정확도 하락을 잡는 일반 원칙은 4bit에서도 그대로 도움이 됩니다. 특히 캘리브레이션 프롬프트 구성, 레이어별 민감도 확인, 일부 연산을 FP16로 남기는 전략은 공통입니다. 관련해서는 PyTorch→ONNX→INT8 양자화 정확도 하락 잡기 글의 체크리스트도 같이 보면 시행착오를 크게 줄일 수 있습니다.

전체 파이프라인 개요

실전 파이프라인은 보통 아래 순서로 갑니다.

  1. PyTorch 모델을 ONNX로 export
  2. ONNX 그래프 최적화(상수 폴딩, 불필요 노드 제거, 패턴 퓨전)
  3. 4bit weight-only 양자화 적용
  4. ONNX Runtime에서 실행 provider 설정(CUDA, TensorRT, CPU 등)
  5. 벤치마크로 토큰 처리량과 지연 측정
  6. 품질 검증(정확도, 퍼플렉서티, 태스크 기반 평가)

이 글에서는 1, 3, 4, 5에 집중하고, 2는 필요한 지점만 짚겠습니다.

준비물: 버전과 환경 주의점

ONNX Runtime과 관련 패키지는 버전 조합에 따라 커널 지원이나 성능이 크게 달라질 수 있습니다.

  • Python 3.10 이상 권장
  • onnxruntime 또는 onnxruntime-gpu 최신 계열 권장
  • CUDA 사용 시 CUDA Toolkit, 드라이버, cuDNN 호환 확인

환경 설치 예시는 아래처럼 시작하면 됩니다.

python -m venv .venv
source .venv/bin/activate

pip install -U pip
pip install onnx onnxruntime onnxruntime-gpu numpy transformers torch

서버 CI에서 캐시가 제대로 안 먹으면 빌드와 배포 시간이 늘어나고, 결국 모델 아티팩트 재생성이 잦아져 운영이 불안정해집니다. 파이프라인을 CI로 굳힐 계획이라면 GitHub Actions 캐시가 안 먹을 때 - key·restore-keys 완전 정복도 같이 참고해두면 좋습니다.

1) PyTorch에서 ONNX로 내보내기: LLM export의 핵심 포인트

LLM을 통째로 export하는 방식은 모델 구조와 목적에 따라 다릅니다.

  • "프롬프트 한번, 로그잇 한번" 형태로 내보내면 단순하지만, KV 캐시 최적화가 어렵습니다.
  • 디코더 스텝을 KV 캐시 입력과 출력으로 분리하면, 오토리그레시브 루프를 외부에서 돌리며 성능을 끌어올릴 수 있습니다.

여기서는 개념 전달을 위해 단순화된 export 코드 예시를 보여줍니다. 실제로는 모델별로 입력 시그니처가 다르므로, export 전에 입력 텐서 이름과 shape를 확정해야 합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "gpt2"  # 예시

tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
model.eval()

text = "Hello"
inputs = tok(text, return_tensors="pt")

# 동적 축 지정: 배치와 시퀀스 길이는 동적으로
dynamic_axes = {
    "input_ids": {0: "batch", 1: "seq"},
    "logits": {0: "batch", 1: "seq"},
}

with torch.no_grad():
    torch.onnx.export(
        model,
        (inputs["input_ids"],),
        "llm_fp32.onnx",
        input_names=["input_ids"],
        output_names=["logits"],
        dynamic_axes=dynamic_axes,
        opset_version=17,
    )

실전에서는 다음을 꼭 점검하세요.

  • opset 버전: 너무 낮으면 최적화나 커널 매칭이 깨질 수 있습니다.
  • 동적 축: 시퀀스 길이가 고정되면 실제 서비스 트래픽에서 비효율이 커집니다.
  • KV 캐시 분리: 가능한 모델이라면 export 구조를 캐시 친화적으로 설계해야 2배 가속에 가까워집니다.

2) 4bit weight-only 양자화 적용

ONNX에서 4bit 양자화는 보통 "가중치만"을 대상으로 합니다. 핵심은 다음입니다.

  • Linear 계열 연산(MatMul, Gemm)에 적용했을 때 효과가 큼
  • attention 자체보다 MLP 블록의 비중이 커서, MLP 쪽 양자화가 성능에 크게 기여
  • 레이어별로 민감도가 달라 일부 레이어는 FP16 또는 INT8로 남기는 하이브리드가 유리할 수 있음

아래 코드는 ONNX Runtime의 양자화 도구를 호출하는 형태의 예시입니다. 실제 함수 이름과 옵션은 설치된 패키지 버전에 따라 다를 수 있으니, 실행 전 python -c "import onnxruntime; print(onnxruntime.__version__)"로 버전을 고정하고 문서에 맞춰 조정하세요.

from pathlib import Path

fp_model = Path("llm_fp32.onnx")
q4_model = Path("llm_q4.onnx")

# 버전별로 API가 달라질 수 있어, 개념적 예시로 이해하세요.
# 일반적으로는 weight-only 4bit, group size, per-channel 여부를 지정합니다.

from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    model_input=str(fp_model),
    model_output=str(q4_model),
    weight_type=QuantType.QInt8,  # 일부 구현은 내부적으로 4bit packing을 사용
)

print("saved:", q4_model)

위 예시는 "동적 양자화" 형태로 보이지만, 실무에서 원하는 4bit weight-only는 별도 패키지나 확장 옵션이 필요할 수 있습니다. 중요한 건 코드 한 줄보다 다음 체크리스트입니다.

4bit에서 정확도 저하를 줄이는 체크리스트

  • group-wise quantization 사용: 너무 큰 그룹은 손실이 커지고, 너무 작은 그룹은 오버헤드가 늘 수 있습니다.
  • per-channel 스케일: 특히 Linear weight는 채널별 스케일이 유리한 경우가 많습니다.
  • 민감 레이어 제외: 임베딩, LM head, 첫 블록 또는 마지막 블록이 민감한 경우가 잦습니다.
  • 클리핑/아웃라이어 처리: LLM weight에는 아웃라이어가 존재할 수 있어, 일부 채널만 별도 처리하는 방식이 필요할 수 있습니다.

정확도 이슈가 발생하면, 먼저 "어디서" 무너지는지부터 확인해야 합니다. 레이어별 cosine similarity나 logits KL divergence를 찍어보면, 양자화 설정을 바꿔야 할 지점이 빨리 보입니다.

3) ONNX Runtime 실행: provider와 세션 옵션으로 성능을 끌어올리기

양자화된 모델을 올렸는데도 속도가 안 나오면, 원인은 대부분 실행 provider 설정이나 세션 옵션입니다.

  • GPU라면 CUDAExecutionProvider 또는 TensorRT 연동을 고려
  • CPU라면 스레드 수, affinity, arena 설정 등 런타임 옵션이 중요

아래는 기본적인 세션 생성 예시입니다.

import onnxruntime as ort

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

providers = [
    "CUDAExecutionProvider",
    "CPUExecutionProvider",
]

sess = ort.InferenceSession("llm_q4.onnx", sess_options=so, providers=providers)

print(sess.get_providers())

자주 놓치는 성능 함정

  • 그래프 최적화 레벨이 기본값: 최적화가 꺼져 있거나 낮으면 퓨전이 안 일어나 성능이 급락할 수 있습니다.
  • 입출력 복사 비용: CPU에서 만든 입력을 GPU로 매번 복사하면 작은 배치에서 지연이 커집니다.
  • 동적 shape 과도 사용: 필요 이상으로 shape가 동적이면 커널 선택이 비효율적일 수 있습니다.

4) 벤치마크: "2배"를 주장하려면 무엇을 재야 하나

LLM 성능은 평균 지연 하나로 말하기 어렵습니다. 최소 아래 3가지는 같이 봐야 합니다.

  • TTFT: Time To First Token
  • TPOT: Time Per Output Token
  • TPS: Tokens Per Second

그리고 반드시 조건을 고정해야 합니다.

  • 프롬프트 길이, 출력 길이
  • 배치 크기
  • KV 캐시 사용 여부
  • warm-up 횟수

아래는 매우 단순화한 벤치마크 스켈레톤입니다. 실제 LLM은 디코딩 루프, 샘플링, KV 캐시가 포함되므로, 이 코드는 "런타임 호출 오버헤드" 감을 잡는 용도로 보세요.

import time
import numpy as np
import onnxruntime as ort

sess = ort.InferenceSession("llm_q4.onnx", providers=["CPUExecutionProvider"])

# 예시 입력: batch=1, seq=16
input_ids = np.random.randint(0, 1000, size=(1, 16), dtype=np.int64)

# warm-up
for _ in range(10):
    _ = sess.run(None, {"input_ids": input_ids})

# measure
t0 = time.perf_counter()
iters = 100
for _ in range(iters):
    _ = sess.run(None, {"input_ids": input_ids})

t1 = time.perf_counter()

print("avg ms:", (t1 - t0) * 1000 / iters)

2배 가속이 잘 나오는 조건

  • 메모리 대역폭이 병목인 경우: 특히 큰 모델, 긴 시퀀스, 낮은 배치에서 효과가 큽니다.
  • CPU 추론: INT8, 4bit weight-only가 CPU에서 더 체감되는 경우가 많습니다.
  • 동시성 증가: 단일 요청 TPS가 1.5배여도, VRAM 여유로 동시 요청을 늘려 전체 처리량이 2배가 되는 경우가 흔합니다.

반대로, 이미 TensorRT로 잘 튜닝된 FP16 경로가 있고 compute-bound라면 4bit가 항상 2배를 보장하진 않습니다.

품질 검증: "빠른데 이상해요"를 막는 최소 절차

4bit에서 흔한 실패는 "문장 생성은 되는데, 논리/수리/코딩 태스크가 무너진다"입니다. 따라서 단순 샘플 몇 개로 끝내지 말고, 최소 아래를 권장합니다.

  • 고정 시드로 50개 내외 프롬프트 세트 구성
  • 태스크별로 분리: 요약, 번역, 코드, 수리, 지식 QA
  • FP16 또는 FP32 기준 출력과의 차이를 정량화

정량화 지표는 상황에 따라 다르지만, 운영 관점에서는 "사용자 불만이 터지는 유형"을 미리 잡는 게 더 중요합니다.

운영 팁: 배포와 롤백을 전제로 설계하기

양자화 모델은 환경에 따라 예기치 못한 커널 미지원, 성능 급락, 품질 저하가 발생할 수 있습니다. 그래서 배포는 반드시 롤백 가능한 형태로 설계해야 합니다.

  • 모델 아티팩트를 버전으로 분리하고, 라우팅 레이어에서 스위치 가능하게 구성
  • canary 트래픽으로 품질과 지연 모니터링
  • 장애 시 FP16 경로로 즉시 우회

모델을 서빙하는 플랫폼이 Triton이라면, 모델 교체 과정에서 핫스왑이 실패하는 케이스를 자주 만나게 됩니다. 운영 체크리스트는 Triton 모델 핫스왑 실패 원인과 해결 체크리스트도 도움이 됩니다.

마무리: 4bit는 "한 방"이 아니라 "조합"이다

ONNX Runtime로 LLM을 4bit로 양자화하면, 메모리와 대역폭을 동시에 줄여 추론 성능을 크게 끌어올릴 수 있습니다. 다만 2배 가속을 안정적으로 얻으려면 다음 조합이 필요합니다.

  • export 단계에서 입력 시그니처와 KV 캐시 구조를 성능 친화적으로 설계
  • 4bit weight-only를 적용하되 group-wise, per-channel 등 손실을 줄이는 설정 선택
  • provider, 세션 옵션, I/O 복사 비용까지 포함해 엔드투엔드로 측정
  • 품질 검증을 태스크 기반으로 수행하고, 운영에서 롤백 가능한 배포 전략 마련

이 과정을 한 번 자동화해두면, 이후 모델 교체나 하드웨어 변경에도 "성능과 품질"을 빠르게 회귀 테스트할 수 있고, 그 자체가 가장 큰 비용 절감으로 이어집니다.

Discuss on Twitter • View on GitHub