beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM, bitsandbytes 4bit로 해결

Authors

로컬 GPU에서 transformers로 LLM을 올리면 가장 먼저 마주치는 벽이 OOM입니다. 특히 7B급 모델도 프롬프트가 길어지거나 배치가 조금만 커져도 VRAM이 급격히 치솟습니다. 이 글에서는 bitsandbytes 4bit 양자화 로딩으로 가중치 메모리를 줄여 OOM을 완화하는 방법을, 재현 가능한 코드와 함께 정리합니다.

가중치 메모리를 줄이면 “모델이 아예 안 올라가는” 문제를 먼저 해결할 수 있고, 그 다음 단계로는 KV 캐시(컨텍스트 길이, 생성 길이)에 의한 폭증을 다루는 흐름이 일반적입니다. KV 캐시 절감은 별도 글인 Transformers 로컬 LLM OOM? KV 캐시 절감 5가지에서 더 깊게 다뤘으니 함께 보면 좋습니다.

OOM의 구조: 가중치 vs KV 캐시

로컬 추론에서 VRAM을 크게 쓰는 축은 크게 2개입니다.

  1. 가중치(Weights): 모델 파라미터 자체. FP16/BF16로 올리면 파라미터 수에 비례해 VRAM이 고정적으로 들어갑니다.
  2. KV 캐시(Key/Value Cache): 디코더 계열 LLM은 생성 중 각 토큰의 어텐션 키/밸류를 저장합니다. 컨텍스트 길이와 생성 길이가 늘수록 선형적으로 커집니다.

따라서 “모델 로딩 단계에서 OOM”이면 가중치가 원인일 확률이 높고, “프롬프트 길이 또는 max_new_tokens를 늘리면 OOM”이면 KV 캐시가 원인일 확률이 큽니다.

이번 글의 목표는 가중치 메모리를 4bit로 줄여 첫 번째 병목을 해결하는 것입니다.

bitsandbytes 4bit가 해결하는 것

bitsandbytes는 GPU 상에서 8bit/4bit 양자화 가중치로 모델을 로드해 추론 메모리를 크게 줄입니다. 특히 4bit는 다음 조합이 실전에서 많이 쓰입니다.

  • load_in_4bit=True
  • bnb_4bit_quant_type="nf4" (NF4)
  • bnb_4bit_use_double_quant=True (더블 양자화)
  • bnb_4bit_compute_dtype=torch.bfloat16 또는 torch.float16

핵심은 저장(가중치) 포맷은 4bit로 두되, 연산은 BF16/FP16으로 수행해 품질 저하를 완화하는 방식입니다.

사전 체크: 환경과 호환성

4bit 로딩이 실패하는 케이스는 대부분 “CUDA/드라이버/라이브러리 조합” 문제입니다. 아래를 먼저 확인하세요.

  • NVIDIA 드라이버 및 CUDA 런타임이 정상인지
  • torch가 CUDA 빌드인지
  • transformers, accelerate, bitsandbytes 버전이 너무 오래되지 않았는지

간단 점검 코드입니다.

import torch
print("torch:", torch.__version__)
print("cuda available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("gpu:", torch.cuda.get_device_name(0))
    print("capability:", torch.cuda.get_device_capability(0))

설치 예시는 다음과 같습니다.

pip install -U "transformers" "accelerate" "bitsandbytes"

운영체제나 CUDA 조합에 따라 bitsandbytes가 추가 설정을 요구할 수 있습니다. 설치 단계에서부터 에러가 난다면, 먼저 pip show bitsandbytes로 버전을 확인하고, GPU가 아닌 CPU 전용 환경에서 설치하려는 것은 아닌지 점검하세요.

4bit 로딩: 가장 많이 쓰는 표준 레시피

아래 예시는 Hugging Face 모델을 4bit로 로드하고, 텍스트 생성을 수행합니다. device_map="auto"를 사용하면 GPU 메모리에 맞춰 자동으로 배치합니다.

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.bfloat16,
)

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

prompt = "한국어로 짧게, 로컬 LLM OOM 원인을 3가지로 정리해줘."
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,
        top_p=0.9,
    )

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

compute dtype는 BF16이 항상 정답일까

  • GPU가 BF16을 잘 지원하면 torch.bfloat16이 안정적인 편입니다.
  • BF16 지원이 애매하거나, 특정 커널에서 문제가 나면 torch.float16으로 바꿔보세요.
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

VRAM이 얼마나 줄어드나: 대략적인 감

모델 파라미터 가중치가 FP16일 때는 파라미터당 2바이트입니다. 4bit는 이상적으로는 0.5바이트 수준이지만, 실제로는 스케일/제로포인트/메타데이터 및 커널 구현 오버헤드가 있어 체감은 “정확히 4배”가 아니라 2배~3배대로 느껴지는 경우도 있습니다.

그럼에도 불구하고 “모델이 아예 GPU에 안 올라가던 상황”을 “로딩은 되고, 컨텍스트/생성 길이를 조절하며 쓸 수 있는 상황”으로 바꾸는 데 효과가 큽니다.

정량 확인을 위해 로드 전후 VRAM을 찍는 습관이 좋습니다.

import torch

def vram_mb():
    return torch.cuda.memory_allocated() / 1024 / 1024

torch.cuda.reset_peak_memory_stats()
print("allocated(MB):", vram_mb())

# 모델 로드 후
print("allocated(MB):", vram_mb())
print("peak(MB):", torch.cuda.max_memory_allocated() / 1024 / 1024)

OOM이 계속 난다면: 4bit 이후의 체크리스트

4bit로 가중치를 줄였는데도 OOM이 난다면, 다음 중 하나일 가능성이 큽니다.

1) KV 캐시가 병목이다

프롬프트가 길거나 max_new_tokens가 크면 KV 캐시가 폭증합니다. 해결책은 컨텍스트를 줄이거나, 생성 길이를 제한하거나, 캐시 정책을 조정하는 것입니다. 구체적인 테크닉은 Transformers 로컬 LLM OOM? KV 캐시 절감 5가지에서 이어서 보세요.

2) 배치 또는 동시 요청이 숨어 있다

로컬에서 API 서버로 감싸서 쓰는 경우, 동시 요청이 쌓이면서 배치처럼 동작해 메모리가 튈 수 있습니다. 이 경우는 “한 번에 한 요청만 처리”하도록 제한하거나, 큐잉을 넣는 방식이 먼저입니다.

3) device_map 오프로딩이 의도와 다르다

device_map="auto"는 편하지만, 특정 레이어가 CPU로 오프로딩되거나 예상 외로 분산되면 속도/메모리 패턴이 달라집니다. 메모리를 더 아끼려면 CPU 오프로딩을 명시적으로 쓰는 것도 방법입니다.

아래는 accelerate가 추론 시 레이어를 어떻게 배치했는지 확인하는 힌트입니다.

print(model.hf_device_map)

4bit에서 자주 겪는 에러와 대응

CUDA out of memory가 로딩 중에 난다

  • 같은 GPU에서 다른 프로세스가 VRAM을 잡고 있는지 확인합니다.
  • device_map="auto"를 쓰고도 로딩 OOM이면 모델 자체가 너무 크거나, GPU가 너무 작을 수 있습니다.
  • 그래도 안 되면 더 작은 모델로 내려가거나, CPU 오프로딩 전략을 섞어야 합니다.

bitsandbytes 관련 import 에러 또는 커널 로드 실패

  • bitsandbytes 버전과 CUDA 조합 문제일 가능성이 큽니다.
  • 먼저 pip install -U bitsandbytes로 올리고, 그래도 실패하면 torch CUDA 빌드와 드라이버를 재점검합니다.

결과 품질이 떨어진다

4bit는 품질 손실이 “완전히 0”은 아닙니다. 완화 팁은 다음과 같습니다.

  • bnb_4bit_quant_type를 "nf4"로 유지
  • bnb_4bit_use_double_quant=True 유지
  • 연산 dtype를 가능하면 BF16로
  • 샘플링 파라미터를 보수적으로(너무 높은 temperature는 환각을 키울 수 있음)

실전 팁: 로컬 RAG와 같이 쓸 때 메모리 예산 잡기

로컬 LLM을 RAG로 엮으면, 임베딩 모델, 벡터DB, 리랭커까지 함께 떠서 메모리 예산이 더 빡빡해집니다. 이때는 “LLM 가중치 4bit + KV 캐시 관리 + 검색 지연 튜닝”을 같이 최적화해야 전체 체감이 좋아집니다.

검색 레이어 쪽은 예를 들어 HNSW 파라미터 튜닝으로 지연을 줄여, LLM이 기다리는 시간을 줄이고 동시 처리 압박을 낮출 수 있습니다. 관련해서는 Qdrant HNSW 튜닝으로 RAG 검색지연 50% 줄이기도 함께 참고할 만합니다.

또한 양자화 자체를 더 넓은 관점에서 보면, 추론만이 아니라 배포/운영 비용을 줄이는 핵심 수단입니다. PyTorch 기반의 PTQ 흐름이 궁금하다면 PyTorch 2.x PTQ로 int8 양자화해 4배 경량화도 연결해서 읽으면 “왜 양자화가 효과적인가”를 더 체계적으로 잡을 수 있습니다.

정리

  • 로컬 LLM OOM은 크게 가중치 메모리와 KV 캐시가 원인입니다.
  • bitsandbytes 4bit 로딩은 가중치 메모리를 크게 줄여, “모델이 아예 안 올라감” 문제를 먼저 해결합니다.
  • 표준 설정은 NF4 + double quant + BF16 compute 조합이 무난합니다.
  • 4bit 이후에도 OOM이면 KV 캐시(컨텍스트/생성 길이), 동시성, 오프로딩 맵을 점검해야 합니다.

다음 단계로는, 4bit로 모델을 올린 상태에서 컨텍스트 길이와 생성 길이를 어떻게 설계하면 OOM 없이 품질을 유지할지를 다루는 것이 실전에서 가장 큰 효과를 냅니다. KV 캐시 절감 테크닉은 위에 링크한 글을 기준으로 설정을 맞춰보면 시행착오를 크게 줄일 수 있습니다.

Discuss on Twitter • View on GitHub