Published on

로컬 LLM OOM·속도 해결 - bitsandbytes+KV 캐시

Authors

로컬에서 LLM을 실행하면 가장 먼저 부딪히는 벽이 두 가지입니다. 첫째는 CUDA out of memory 같은 OOM(메모리 부족)이고, 둘째는 토큰 생성 속도(TPS)가 기대보다 훨씬 느린 문제입니다. 특히 7B~13B급 모델을 노트북 GPU나 보급형 데스크톱 GPU에서 돌리면, "로딩은 됐는데 생성이 느리다" 혹은 "아예 로딩이 안 된다"가 흔합니다.

이 글에서는 **bitsandbytes(4bit/8bit 양자화)**로 VRAM을 크게 줄이고, **KV 캐시(Key/Value cache)**를 올바르게 활용해 생성 속도를 안정적으로 끌어올리는 방법을 정리합니다. 단순히 옵션 몇 개 켜는 수준이 아니라, 어떤 메모리가 어디서 터지고 왜 느려지는지 감을 잡을 수 있게 구성했습니다.

OOM을 원인별로 추적하는 접근 자체는 JVM에서도 비슷합니다. 힙덤프 분석처럼 “어디가 많이 먹는지”부터 잡는 게 정답입니다. 메모리 문제를 체계적으로 접근하는 관점은 Spring Boot OOM 원인추적과 힙덤프 분석 실전 글도 참고할 만합니다.

로컬 LLM에서 메모리가 터지는 지점 3가지

로컬 LLM의 GPU 메모리 사용은 보통 아래 3덩어리로 나뉩니다.

1) 모델 가중치(Weights)

가장 큰 비중을 차지합니다. FP16으로 로드하면 7B도 수 GB 단위로 VRAM을 먹습니다. 여기서 bitsandbytes 양자화가 가장 큰 효과를 냅니다.

2) KV 캐시(Attention cache)

생성 단계에서 매 토큰마다 과거 토큰들의 Key/Value를 저장해 재사용하는 캐시입니다. 컨텍스트 길이(max_new_tokens가 아니라 입력+생성 합의 시퀀스 길이)가 길어질수록 KV 캐시가 커집니다.

  • 긴 대화(챗) 시나리오에서 시간이 갈수록 VRAM이 늘어나는 주범
  • OOM이 “처음엔 되다가 몇 턴 뒤에 터지는” 형태로 나타나면 KV 캐시 가능성이 큼

3) 활성화/버퍼(Activations, temp buffers)

추론 중 일시적으로 잡히는 버퍼입니다. FlashAttention, SDPA, fused kernel, torch.compile 여부 등에 따라 변동폭이 큽니다.

bitsandbytes로 VRAM 줄이기: 8bit vs 4bit

bitsandbytes는 Hugging Face Transformers에서 가장 널리 쓰이는 양자화 백엔드입니다.

  • 8bit: 품질 손실이 비교적 적고 안정적. VRAM 절감은 중간.
  • 4bit(NF4): VRAM 절감이 매우 큼. 대신 환경/커널/모델에 따라 속도나 호환성 이슈가 있을 수 있음.

실무적으로는 다음 순서를 추천합니다.

  1. 8bit로 먼저 안정적으로 구동
  2. VRAM이 부족하면 4bit로 전환
  3. 4bit에서 속도가 오히려 느려지면(또는 CPU offload가 늘면) 다시 8bit+KV 튜닝으로 타협

실전 코드: Transformers + bitsandbytes 4bit 로딩

아래 예시는 transformers 기반으로 4bit 양자화를 적용해 모델을 로드하는 최소 구성입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, 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",
    torch_dtype=torch.bfloat16,
)

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=False,
        use_cache=True,
    )

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

체크포인트

  • use_cache=True는 KV 캐시 사용 여부입니다. 기본적으로 켜져 있는 경우가 많지만, 명시해 두는 편이 안전합니다.
  • device_map="auto"는 레이어를 GPU/CPU로 자동 배치합니다. VRAM이 부족하면 CPU로 떨어지면서 속도가 급락할 수 있습니다.

속도가 느린 진짜 이유: KV 캐시를 “쓰고” 있는가, “감당”하고 있는가

KV 캐시는 보통 속도를 올려줍니다. 그런데 체감상 느려지는 경우가 있습니다.

1) 컨텍스트가 길어져 KV 캐시가 너무 커짐

대화가 길어질수록 KV 캐시가 커지고, VRAM이 한계에 가까워지면 다음이 발생합니다.

  • 메모리 단편화로 할당이 불안정해짐
  • device_map에 의해 일부가 CPU로 오프로딩됨
  • 결과적으로 TPS가 급락하거나 OOM

해결 방향은 두 갈래입니다.

  • 컨텍스트 길이를 제한하거나, 대화 히스토리를 요약해 넣기
  • KV 캐시의 dtype/구현을 최적화(아래 참고)

2) KV 캐시가 비효율적인 dtype로 저장됨

모델 가중치는 4bit로 줄였는데 KV 캐시는 여전히 FP16/FP32로 쌓이면, “로딩은 되는데 생성 중에 터지는” 상황이 나옵니다.

모델/라이브러리 조합에 따라 KV 캐시 dtype을 bf16로 맞추거나, 더 나아가 quantized KV cache를 지원하는 스택을 고려할 수 있습니다(환경 의존성이 있어 여기서는 원리 중심으로 설명합니다).

KV 캐시 최적화 실전: 길이, 배치, 프리필 전략

KV 캐시는 크게 **프리필(prefill)**과 디코드(decode) 두 구간에서 성능 특성이 갈립니다.

  • 프리필: 입력 프롬프트 전체를 한 번에 처리(연산량 큼)
  • 디코드: 토큰을 1개씩 생성(반복 호출, 캐시 효율이 중요)

1) max_new_tokens보다 중요한 건 “총 시퀀스 길이”

OOM을 줄이려면 다음을 함께 관리해야 합니다.

  • 입력 길이: len(input_ids)
  • 생성 길이: max_new_tokens
  • 합: 입력+생성 = KV 캐시가 커지는 상한

운영에서 흔한 실수는 입력이 이미 길어진 상태에서 max_new_tokens만 줄이는 것입니다. 이미 쌓인 히스토리 때문에 KV 캐시가 커진 상태라 효과가 제한적입니다.

2) 배치 크기와 동시 요청 수를 낮춰라

로컬 서버로 LLM을 띄우고 여러 요청을 동시에 받으면, KV 캐시가 요청 수만큼 생깁니다. 즉,

  • 동시 요청 수가 2배면 KV 캐시도 거의 2배

단일 사용자용 로컬 챗이라면 동시성(큐/스레드)을 낮추는 게 가장 확실한 OOM 완화책입니다.

3) 프롬프트 캐싱(Static prefix caching)

시스템 프롬프트나 고정 지침이 매번 붙는다면, 매 요청마다 프리필을 다시 하는 낭비가 생깁니다.

서빙 프레임워크(vLLM, TGI 등)에서는 prefix caching을 제공하는 경우가 많고, 직접 구현하려면 “고정 prefix의 past key values를 재사용”하는 형태가 됩니다. 다만 Transformers 순정 API로는 구현 난이도가 있어, 로컬 서빙 목적이면 프레임워크 채택이 더 현실적입니다.

메모리/속도 진단: 숫자로 확인하기

감으로 튜닝하면 금방 한계가 옵니다. 아래 코드는 현재 GPU 메모리 사용량을 대략적으로 확인하는 용도입니다.

import torch

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

report_cuda_mem("before")
# 모델 로드/추론 수행
report_cuda_mem("after")

추가로, OOM이 났을 때는 “모델 로드 중 OOM”인지 “생성 중 OOM”인지부터 분리하세요.

  • 로드 중 OOM: 가중치가 원인인 경우가 많음 -> 양자화/모델 축소/오프로딩
  • 생성 중 OOM: KV 캐시/컨텍스트 길이가 원인인 경우가 많음 -> 히스토리 관리/동시성 제한

위 화살표 표기는 MDX에서 문제가 될 수 있으니, 문서/코드에서는 반드시 ->처럼 백틱 처리하는 습관이 안전합니다.

권장 튜닝 레시피(보급형 GPU 기준)

1) 8GB VRAM

  • 7B: 4bit + 짧은 컨텍스트(예: 2k~4k) + 동시성 1
  • 13B: 매우 빡빡하거나 불가능. CPU offload가 늘면 속도가 급락

2) 12GB VRAM

  • 7B: 8bit도 가능, 컨텍스트를 늘릴 여유
  • 13B: 4bit로 타협하면 “돌아가는” 수준은 가능, 속도는 환경 의존

3) 24GB VRAM

  • 13B: 8bit 또는 4bit로 여유, 컨텍스트/동시성 튜닝 폭이 큼

로컬 서빙이라면: bitsandbytes만으로 부족한 순간

노트북에서 개인 실험은 Transformers + bitsandbytes로 충분한 경우가 많습니다. 하지만 API 서버로 띄워 동시 요청을 받거나, 긴 컨텍스트를 다루거나, TPS를 끝까지 뽑아야 한다면 다음 계열을 검토하게 됩니다.

  • vLLM: PagedAttention으로 KV 캐시 메모리 효율 개선, 서빙 친화
  • TGI(Text Generation Inference): 배치/서빙 최적화

이 글의 범위는 bitsandbytes+KV 캐시에 집중하지만, 결론적으로 “속도”는 커널/서빙 구조의 영향을 크게 받습니다.

자주 터지는 함정 체크리스트

1) device_map="auto"가 CPU 오프로딩을 만들고 있지 않은가

VRAM이 부족하면 일부 레이어가 CPU로 내려가고, 그 순간 TPS가 크게 떨어집니다. 로드 로그를 확인하거나, 생성 속도가 비정상적으로 느리면 의심하세요.

2) 컨텍스트가 누적되는 챗 UI 구조

프론트/백엔드에서 히스토리를 계속 누적하면, KV 캐시가 커지는 건 필연입니다. 일정 턴마다 요약하거나, 최근 N턴만 유지하는 정책이 필요합니다.

3) OOM 이후 재시도 루프

OOM이 난 뒤 같은 조건으로 즉시 재시도하면, 메모리 단편화/캐시 상태 때문에 더 빨리 죽는 경우가 있습니다. 프로세스 재시작이 가장 깔끔한 경우도 많습니다.

K8s 환경에서라면 이런 문제는 곧바로 재시작 루프로 이어질 수 있어, 증상 관점에서 K8s CrashLoopBackOff 8가지 원인, 로그로 끝내기 같은 체크리스트가 도움이 됩니다.

정리

  • 로컬 LLM OOM의 큰 덩어리는 가중치, KV 캐시, 임시 버퍼입니다.
  • bitsandbytes 4bit/8bit 양자화는 가중치 VRAM을 크게 줄이는 가장 쉬운 방법입니다.
  • 생성 중에 터지거나 시간이 갈수록 느려지면, 대개 **KV 캐시(컨텍스트 길이/동시성)**가 원인입니다.
  • use_cache=True는 기본이지만, 컨텍스트 정책(요약/최근 턴 유지)과 동시성 제한이 함께 있어야 안정적으로 운영됩니다.

실험을 할 때는 “한 번에 모든 옵션을 바꾸지 말고”, 로드 단계생성 단계를 분리해 메모리 사용량을 숫자로 확인하면서 튜닝하면 훨씬 빠르게 답에 도달합니다.