beautifulsoups.dev
Published on

Transformers 로컬 LLM KV 캐시 OOM 해결 가이드

Authors

로컬 GPU 한 장으로 Transformers 기반 LLM을 돌리다 보면, 모델 로딩은 성공하는데 생성 단계에서 갑자기 CUDA out of memory가 터지는 경우가 많습니다. 이때 범인은 대개 KV 캐시(Key/Value cache) 입니다. 특히 긴 컨텍스트, 큰 배치, 높은 max_new_tokens를 주면 KV 캐시는 선형으로 커지고, 어느 순간 VRAM을 밀어내며 OOM을 냅니다.

이 글에서는 KV 캐시가 왜 이렇게 빨리 커지는지(메모리 관점), 그리고 Transformers에서 바로 적용 가능한 OOM 완화/해결 옵션을 코드와 함께 정리합니다. 마지막에는 운영 환경에서 재발 방지 체크리스트까지 다룹니다.

KV 캐시 OOM이 생기는 구조적 이유

Transformers의 디코더 계열 LLM(예: Llama, Mistral, Qwen 계열)은 오토리그레시브 생성 시 매 토큰마다 self-attention을 수행합니다. 효율을 위해 과거 토큰의 attention 계산 결과인 Key/Value를 레이어별로 저장해두는데, 이것이 KV 캐시입니다.

핵심은 KV 캐시 메모리가 다음 요인에 대략 선형 비례한다는 점입니다.

  • 레이어 수(예: 32, 40, 80)
  • 배치 크기 batch_size
  • 시퀀스 길이 seq_len (프롬프트 길이 + 생성된 토큰 수)
  • 헤드 수 및 헤드 차원(num_heads, head_dim)
  • dtype 크기(예: fp16은 2바이트, bf16은 2바이트, fp32는 4바이트)

대략적인 감을 잡기 위한 근사식은 아래처럼 생각하면 됩니다.

  • KV 캐시는 레이어마다 Key와 Value 두 개가 있고
  • 각 토큰마다 hidden_size에 비례하는 텐서가 저장됩니다.

즉, batch_size가 1이라도 seq_len이 8k, 16k로 길어지면 KV 캐시만으로도 VRAM을 크게 잡아먹습니다. 여기에 다음 요소가 겹치면 OOM이 가속됩니다.

  • use_cache=True가 기본으로 켜져 있음(생성 성능을 위해 기본값)
  • max_new_tokens를 크게 줘서 생성 길이가 길어짐
  • 여러 요청을 한 번에 태우는 동적 배칭(서버 환경)
  • torch.compile 또는 일부 커널 선택으로 임시 버퍼가 추가로 생성
  • device_map="auto"로 레이어 오프로딩이 되더라도 KV 캐시는 GPU에 남는 구성(설정에 따라 다름)

증상으로 구분하는 KV 캐시 OOM 시나리오

1) 프롬프트는 짧은데 max_new_tokens가 클 때

처음 몇 토큰 생성하다가 일정 시점에 OOM이 납니다. 생성이 진행될수록 KV 캐시가 누적되기 때문입니다.

2) 프롬프트가 길 때(특히 RAG)

생성 시작 직후 또는 첫 스텝에서 OOM이 날 수 있습니다. 이미 프롬프트 토큰만으로 KV 캐시가 크게 잡히기 때문입니다.

3) 동시 요청이 늘수록 갑자기 터질 때

서버에서 배치가 커지면서 batch_size가 증가하고 KV 캐시가 배치에 비례해 커집니다. 이 경우는 단일 요청 테스트에서는 멀쩡하다가 운영에서 터집니다.

운영형 배포에서 GPU 자원 문제를 다룰 때는 오토스케일과 부하 패턴까지 함께 봐야 합니다. GPU 서빙을 K8s로 운영한다면 KServe+Knative로 GPU 모델 오토스케일 배포처럼 스케일 전략도 같이 점검하는 편이 안전합니다.

가장 먼저 적용할 수 있는 7가지 해결책

아래는 효과가 큰 순서대로 정리했습니다. 여러 개를 조합하는 것이 일반적입니다.

1) max_new_tokens와 입력 길이 상한을 강제하기

가장 확실한 방법은 총 시퀀스 길이(입력 + 출력) 를 제한하는 것입니다.

  • 입력 토큰이 길면: 프롬프트를 자르거나 요약
  • 출력 토큰이 길면: max_new_tokens 제한
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)

prompt = "..."  # 길어질 수 있는 입력
inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=2048).to("cuda")

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        use_cache=True,
    )

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

포인트는 truncation=True와 max_length를 명시해 입력 상한을 두는 것입니다. RAG 파이프라인이면, 검색 결과를 그대로 붙이지 말고 문서 수/문단 수를 제한하거나, 중복 제거와 압축을 넣으세요.

2) 배치를 키우지 말고 요청을 쪼개기(서버라면 동적 배칭 제한)

batch_size는 KV 캐시를 정직하게 곱합니다. 로컬에서 단일 프롬프트는 되는데, 서버에서 동시 요청이 늘면 터지는 이유입니다.

  • 동적 배칭을 쓰면 max_batch_size를 낮추기
  • 긴 요청과 짧은 요청을 같은 배치로 묶지 않기
  • 스트리밍 응답으로 사용자 체감 지연을 줄이되, 내부적으로는 동시성을 제한

3) use_cache=False로 KV 캐시 자체를 끄기(최후의 보루)

KV 캐시를 끄면 메모리는 줄지만, 생성 속도가 크게 느려집니다. 그래도 “어떻게든 돌아가게” 만드는 응급 처치로는 유효합니다.

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

운영에서는 보통 권장되지 않지만, 디버깅이나 VRAM이 극도로 부족한 환경에서는 선택지입니다.

4) dtype을 낮추고(또는 양자화) KV 캐시 크기를 줄이기

KV 캐시도 텐서이므로 dtype 영향을 받습니다.

  • fp16 또는 bf16 사용(이미 대부분 이렇게 함)
  • 4bit, 8bit 양자화로 가중치 메모리를 줄여 KV 캐시를 위한 여유 VRAM을 확보

Transformers에서 bitsandbytes 4bit 양자화 예시는 아래처럼 구성합니다.

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_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

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

주의할 점은 “양자화는 KV 캐시를 직접 줄이는 게 아니라” 가중치가 차지하던 VRAM을 줄여서 KV 캐시가 들어갈 공간을 만든다는 것입니다. 긴 컨텍스트에서는 여전히 KV 캐시가 병목입니다.

5) FlashAttention 또는 SDPA로 attention 메모리/속도 최적화

KV 캐시 OOM은 아니더라도, attention 자체가 큰 임시 버퍼를 만들면서 OOM을 유발할 수 있습니다. PyTorch의 SDPA 또는 FlashAttention 경로를 타면 개선되는 경우가 많습니다.

Transformers에서는 버전에 따라 다음처럼 설정합니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
    attn_implementation="sdpa",  # 또는 "flash_attention_2"
)

환경에 따라 커널 선택이 달라지고, 드라이버 및 CUDA 버전 영향도 큽니다.

6) 프롬프트 캐싱 전략을 바꾸기: 멀티턴 대화는 “전체 재전송”을 피하기

채팅 UI에서 흔히 하는 실수는, 매 턴마다 이전 대화를 전부 프롬프트로 다시 보내는 것입니다. 그러면 컨텍스트가 매번 길어지고 KV 캐시도 커집니다.

  • 시스템 프롬프트와 정책 텍스트는 짧게 유지
  • 오래된 대화는 요약해서 교체
  • 필요하면 “대화 메모리”를 RAG로 전환(요약 + 검색)

7) 메모리 파편화 완화와 관측

OOM이 “남은 VRAM이 있어 보이는데도” 터지면 파편화일 수 있습니다.

  • PYTORCH_CUDA_ALLOC_CONF 설정으로 split 크기 조정
  • 주기적으로 메모리 사용량을 로그로 남겨 원인을 특정
# 예: 파편화 완화에 도움이 되는 경우가 있음
export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"

그리고 코드에서 관측을 넣어두면, KV 캐시가 커지는 패턴을 빨리 찾을 수 있습니다.

import torch

def log_cuda_mem(tag: str):
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"[{tag}] allocated={alloc:.1f}MB reserved={reserved:.1f}MB")

log_cuda_mem("before")
# generate 호출
log_cuda_mem("after")

실전: “긴 컨텍스트 RAG”에서 OOM을 줄이는 레시피

로컬 LLM에서 KV 캐시 OOM이 가장 자주 터지는 조합은 RAG입니다. 검색 결과를 많이 붙이면 입력 토큰이 급격히 늘고, 생성 길이까지 더해져 총 토큰이 폭증합니다.

권장 레시피는 다음과 같습니다.

  1. 검색 결과 개수 제한(예: top-k 3~5)
  2. 문서 chunk 길이 제한(예: 256~512 토큰)
  3. 중복 문장 제거, 표/코드 블록 과다 포함 방지
  4. 최종 프롬프트는 max_length로 강제 truncate
  5. max_new_tokens는 제품 요구사항에 맞게 상한 설정

벡터 검색 품질을 올려 “적은 문서로도 답을 만들게” 하면 입력 토큰을 줄일 수 있습니다. HNSW 기반 벡터DB를 튜닝 중이라면 RAG 리콜 급락? HNSW 파라미터 튜닝 가이드도 같이 보면, 불필요하게 많은 문서를 붙이는 상황을 줄이는 데 도움이 됩니다.

자주 하는 오해: device_map="auto"면 OOM이 해결된다?

device_map="auto"는 레이어를 GPU와 CPU로 나눠 올려 가중치 메모리를 줄이는 데 효과적입니다. 하지만 KV 캐시는 생성 중에 계속 커지며, 구현/설정에 따라 GPU에 남는 경우가 많습니다.

즉,

  • 모델 로딩 OOM은 해결됐는데
  • 생성 OOM은 여전히 발생

이라는 패턴이 흔합니다. 이때는 위에서 설명한 것처럼 총 토큰 수, 배치, max_new_tokens 를 먼저 줄여야 합니다.

디버깅 체크리스트

아래 순서로 보면 원인 파악이 빠릅니다.

  1. 입력 토큰 수를 로그로 찍기
    • len(tok(prompt).input_ids) 같은 방식
  2. max_new_tokens와 실제 생성 길이 확인
  3. 동시 요청 수, 배치 정책 확인
  4. dtype 확인(fp16 또는 bf16인지)
  5. attention 구현 경로 확인(sdpa 또는 flash_attention_2)
  6. 파편화 의심 시 reserved가 과도하게 큰지 확인

운영 환경에서 장애가 났다면 “로그로 원인을 좁히는 절차”가 중요합니다. 인프라 레벨의 장애 분석 흐름은 다르지만, 접근법 자체는 비슷합니다. 예를 들어 VM이 부팅 실패할 때도 콘솔 로그로 원인을 좁히듯이, GPU OOM도 메모리 지표와 입력/출력 토큰 로그가 핵심 단서가 됩니다. 관련 사고 대응 방식은 Azure VM 부팅 실패, Serial Console 로그로 복구 같은 글의 문제 분해 방식이 참고가 됩니다.

결론: KV 캐시는 “토큰 수의 세금”이다

Transformers 로컬 LLM에서 KV 캐시 OOM은 버그라기보다 구조적인 비용입니다. 해결의 우선순위는 명확합니다.

  • 1순위: 입력 길이와 출력 길이 상한 설정(총 토큰 제한)
  • 2순위: 배치/동시성 제한
  • 3순위: 양자화로 여유 VRAM 확보
  • 4순위: attention 커널 최적화(SDPA, FlashAttention)
  • 응급: use_cache=False

이 조합으로 대부분의 로컬 추론 OOM은 실용 수준으로 안정화할 수 있습니다. 특히 “RAG에서 검색 결과를 많이 붙이는 설계”는 KV 캐시 비용을 폭발시키는 지름길이므로, 검색 품질 튜닝과 프롬프트 압축을 함께 가져가는 것이 가장 효과적입니다.

Discuss on Twitter • View on GitHub