beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM? 8bit·4bit 로딩 가이드

Authors

로컬 GPU에서 LLM을 transformers로 올리다 보면, 모델 로딩 단계에서 바로 OOM이 나거나(가중치가 안 올라감), 로딩은 되는데 첫 생성에서 터지거나(주로 KV 캐시), 배치/컨텍스트를 조금만 올려도 급격히 메모리가 증가하는 문제를 자주 만납니다.

이 글은 그 중에서도 가장 즉효가 큰 8bit·4bit 로딩을 중심으로, 어디서 메모리가 새는지(가중치 vs KV 캐시 vs 임시 버퍼), 그리고 어떤 옵션 조합이 안전한지 정리합니다. KV 캐시까지 포함한 OOM 최적화는 아래 글에서 더 깊게 다룹니다.

OOM이 나는 지점부터 구분하자

OOM은 크게 3가지 축으로 나뉩니다.

1) 가중치(Weights) 메모리

모델 파라미터가 GPU VRAM에 올라가는 비용입니다.

  • FP16/BF16: 파라미터당 2바이트
  • FP32: 파라미터당 4바이트
  • 8bit: 파라미터당 1바이트 수준(실제로는 스케일/제로포인트 등 오버헤드 존재)
  • 4bit: 파라미터당 0.5바이트 수준(오버헤드 존재)

대략적인 감만 잡아도 도움이 됩니다.

  • 7B 모델 FP16 가중치만 대략 7e9 * 2 바이트 ≈ 14GB 근처
  • 13B 모델 FP16은 ≈ 26GB 근처

여기에 레이어별 버퍼, 옵티마이저(학습 시), KV 캐시가 더해집니다.

2) KV 캐시(Key/Value cache)

추론 시 토큰이 늘어날수록 증가하는 메모리입니다. 컨텍스트 길이(max_new_tokens가 아니라 입력 길이 포함)와 배치 크기에 비례해서 커집니다.

  • “로딩은 되는데 첫 생성에서 OOM”은 KV 캐시가 원인인 경우가 많습니다.

3) 임시 버퍼/연산 그래프/프레임워크 오버헤드

  • torch.compile, flash-attn, fused kernel, attention 구현체에 따라 피크 메모리가 달라질 수 있습니다.
  • generate() 과정에서 logits 저장, beam search, return dict 옵션 등도 영향을 줍니다.

8bit 로딩: 가장 무난한 첫 번째 처방

8bit는 품질 하락이 상대적으로 적고, 4bit보다 디버깅 난이도가 낮습니다. “일단 로컬에서 돌아가게 만들기”에는 8bit가 가장 무난합니다.

설치 체크

bitsandbytes가 필요합니다.

pip install -U transformers accelerate bitsandbytes

CUDA/드라이버 조합에 따라 bitsandbytes가 실패할 수 있으니, 설치 후 간단히 import 테스트를 권장합니다.

import bitsandbytes as bnb
print(bnb.__version__)

8bit 로딩 예제

아래 예시는 AutoModelForCausalLM에 BitsAndBytesConfig를 적용합니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

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

bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
)

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

prompt = "요약: 로컬에서 LLM OOM을 줄이는 방법을 알려줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
    )

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

핵심 포인트

  • device_map="auto"는 VRAM이 부족하면 일부 레이어를 CPU로 오프로딩할 수 있습니다. 다만 속도는 크게 느려질 수 있습니다.
  • torch_dtype는 양자화 자체와 별개로, 일부 연산/버퍼 dtype에 영향을 줍니다.

4bit 로딩: VRAM을 더 줄이되 옵션을 제대로

4bit는 VRAM 절감 폭이 크지만, 설정을 잘못하면 속도/품질/안정성에서 손해를 볼 수 있습니다. 현재 실무에서 가장 흔한 조합은 NF4 + double quant + FP16 compute 입니다.

4bit 로딩 예제(NF4 권장)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

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

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

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

prompt = "질문: 4bit 양자화가 왜 VRAM을 줄이나요?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
    )

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

옵션 설명(실전 감각)

  • bnb_4bit_quant_type="nf4": 일반적인 fp4보다 LLM에서 품질 유지가 유리한 경우가 많습니다.
  • bnb_4bit_use_double_quant=True: 양자화 스케일을 한 번 더 양자화해 메모리를 추가로 줄입니다(대체로 이득).
  • bnb_4bit_compute_dtype=torch.float16: 연산 dtype입니다. GPU가 BF16에 강하면 torch.bfloat16도 후보입니다.

“4bit로도 OOM”이면: 실제 범인은 KV 캐시일 확률이 높다

가중치를 4bit로 줄였는데도 generate()에서 터진다면, 높은 확률로 KV 캐시가 원인입니다.

대표적인 트리거는 다음입니다.

  • 컨텍스트가 길다: 긴 프롬프트 + 큰 max_new_tokens
  • 배치가 크다: 여러 요청을 한 번에 처리
  • use_cache=True 기본값으로 KV 캐시가 쌓임

KV 캐시까지 포함한 튜닝(예: 캐시 dtype, max_position_embeddings 고려, paged attention 계열 접근)은 아래 글에서 이어서 보는 게 좋습니다.

OOM 디버깅: “얼마나 쓰는지”를 코드로 찍어라

OOM을 감으로 해결하면 같은 문제가 반복됩니다. 최소한 아래 정도는 측정해두면 원인 분리가 빨라집니다.

import torch

def report_cuda_mem(tag: str):
    if not torch.cuda.is_available():
        print(f"{tag}: CUDA not available")
        return
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    max_alloc = torch.cuda.max_memory_allocated() / 1024**2
    print(f"{tag}: allocated={alloc:.1f}MB reserved={reserved:.1f}MB max_alloc={max_alloc:.1f}MB")

# 사용 예시
report_cuda_mem("before load")
# model load...
report_cuda_mem("after load")
# generate...
report_cuda_mem("after generate")
  • after load에서 이미 한계면 가중치/오프로딩 전략 문제
  • after generate에서 급증하면 KV 캐시/컨텍스트/배치 문제

자주 하는 실수와 안전한 기본값

1) device_map을 안 쓰고 통째로 GPU에 올리기

VRAM이 애매한 환경에서는 device_map="auto"로 시작하는 게 안전합니다. 다만 CPU 오프로딩이 발생하면 속도는 느려집니다.

2) torch_dtype를 무의식적으로 FP32로 두기

특별한 이유가 없다면 FP16 또는 BF16을 기본으로 두는 편이 낫습니다.

  • BF16이 더 안정적인 GPU도 있습니다(특히 최신 아키텍처).

3) 컨텍스트 길이를 “모델이 지원하니까” 무작정 키우기

지원 여부와 “내 VRAM에서 감당 가능”은 다른 문제입니다. 로컬 추론에서는 컨텍스트를 키우는 순간 KV 캐시가 급증합니다.

4) 추론인데 torch.no_grad()나 torch.inference_mode()를 안 쓰기

추론에서는 autograd가 필요 없습니다.

with torch.inference_mode():
    out = model.generate(**inputs, max_new_tokens=128)

8bit vs 4bit 선택 기준(현업 기준)

  • VRAM이 16GB 이상이고 7B급: 8bit부터 시도, 여유 있으면 FP16도 고려
  • VRAM이 8GB~12GB이고 7B급: 4bit가 현실적인 선택
  • 13B 이상을 단일 GPU로: 4bit + 오프로딩/컨텍스트 제한이 사실상 필수

품질 측면에서는 일반적으로 FP16이 가장 좋고, 그 다음이 8bit, 그 다음이 4bit입니다. 다만 실제 서비스/개발에서는 프롬프트 튜닝, 샘플링 파라미터, RAG 품질이 더 큰 변수를 만들기도 합니다.

체크리스트: OOM을 빠르게 줄이는 순서

  1. 8bit 로딩 적용(load_in_8bit=True)
  2. 안 되면 4bit 로딩(NF4 + double quant)
  3. 그래도 OOM이면 컨텍스트/배치부터 줄이고, KV 캐시 최적화로 넘어가기
  4. device_map="auto"로 오프로딩을 허용(속도 트레이드오프)
  5. 메모리 측정 코드를 넣고 “로딩 OOM vs 생성 OOM”을 분리

로컬 LLM OOM은 대부분 “가중치”와 “KV 캐시” 중 하나가 범인입니다. 8bit·4bit 로딩은 가중치 문제를 빠르게 해결해주고, 그 다음 단계로 KV 캐시까지 다루면 같은 GPU로 훨씬 긴 컨텍스트와 더 안정적인 추론을 만들 수 있습니다.

Discuss on Twitter • View on GitHub