beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM - 4bit+KV 캐시 튜닝

Authors

로컬 GPU에서 transformers로 LLM을 추론하다 보면, 모델을 4bit로 내렸는데도 CUDA out of memory가 터지는 경우가 많습니다. 이때 대부분의 병목은 모델 가중치(weight) 가 아니라 KV 캐시(Key/Value cache) 입니다. 특히 컨텍스트 길이를 늘리거나 배치/동시 요청을 올리면 KV 캐시가 선형으로 커지면서 VRAM을 잠식합니다.

이 글에서는

  • OOM을 재현/측정하는 방법
  • 4bit 로딩으로 가중치 메모리를 줄이는 방법
  • KV 캐시가 왜 커지는지(대략적인 메모리 모델)
  • transformers에서 적용 가능한 KV 캐시/추론 튜닝 포인트
  • 운영 관점에서 “동시성”과 “최대 토큰”을 어떻게 잡을지

를 한 번에 정리합니다.

관련해서 “원인 재현 후 해결” 흐름이 중요하다는 점은 네트워크 타임아웃을 다룬 글과도 결이 비슷합니다. 예: OpenAI Responses API 408 타임아웃 재현과 해결 실전 가이드

OOM의 2가지 큰 축: 가중치 vs KV 캐시

LLM 추론 시 GPU 메모리는 크게 다음이 차지합니다.

  1. 모델 가중치: 로딩 시 거의 고정. 4bit/8bit로 크게 줄일 수 있음.
  2. KV 캐시: 생성 길이, 입력 컨텍스트 길이, 배치 크기, 동시 요청 수에 비례해 증가.

4bit로 내렸는데도 OOM이 나는 전형적인 패턴은 이겁니다.

  • 가중치는 4bit로 줄어서 “처음 로드”는 성공
  • 프롬프트가 길거나 max_new_tokens를 크게 주거나
  • 동시에 여러 요청을 처리(배치 또는 멀티 세션)하면서
  • KV 캐시가 누적되어 OOM

즉, 4bit는 시작점이고, 실제 운영 안정성은 KV 캐시를 어떻게 제한/관리하느냐에 달려 있습니다.

먼저: 현재 VRAM 사용량을 “수치로” 확인하기

OOM 튜닝은 감으로 하면 끝이 없습니다. 최소한 아래 2가지는 확인하세요.

  • 로딩 직후 VRAM
  • 추론(특히 긴 출력) 이후 VRAM

nvidia-smi로 스냅샷 보기

watch -n 0.5 nvidia-smi

파이썬에서 피크 메모리 확인

import torch

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

torch.cuda.reset_peak_memory_stats()
report("start")

여기서 allocated는 실제 할당, reserved는 캐싱 allocator가 잡아둔 풀입니다. OOM은 보통 reserved 관점에서 먼저 벽을 만납니다.

4bit 양자화로 가중치 메모리 줄이기 (bitsandbytes)

가중치 메모리를 줄이는 가장 흔한 방법은 bitsandbytes 기반 4bit 로딩입니다.

설치

pip install -U transformers accelerate bitsandbytes

CUDA/드라이버 조합에 따라 bitsandbytes가 까다로울 수 있습니다. 파이썬 환경 문제로 설치가 꼬이면, 패키지 의존성을 정리하는 접근이 도움이 됩니다: Python 3.12에서 pkg_resources 에러 근본 해결

4bit 로딩 예제

아래 예제는 “가중치 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,
)
model.eval()

자주 하는 실수

  • torch_dtype를 무심코 float32로 두면 연산/일부 버퍼가 커질 수 있습니다.
  • GPU가 bfloat16을 잘 못 받으면 float16으로 바꾸세요.

가중치가 줄어도 KV 캐시는 보통 float16 또는 bfloat16로 잡히기 때문에, 긴 컨텍스트/긴 생성에서 OOM이 계속 날 수 있습니다.

KV 캐시가 왜 이렇게 빨리 커질까 (대략적인 메모리 감)

KV 캐시는 토큰이 늘어날수록 레이어마다 Key/Value 텐서를 저장합니다.

대충 이런 요소에 비례합니다.

  • 레이어 수
  • 히든 차원(헤드 수, 헤드 차원)
  • 시퀀스 길이(입력 토큰 + 생성 토큰)
  • 배치 크기
  • dtype 크기(float16이면 2바이트)

정확한 공식은 모델 구조에 따라 다르지만, 실전에서는 아래 결론만 기억해도 충분합니다.

  • 컨텍스트 2배 = KV 캐시 2배
  • 동시 요청 2배 = KV 캐시 2배
  • 배치 2배 = KV 캐시 2배

그래서 “4bit로 줄였는데도” OOM이 나는 겁니다. 가중치는 줄었지만, KV 캐시는 그대로 커지기 때문입니다.

transformers에서 OOM을 줄이는 KV 캐시/추론 튜닝 포인트

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

가장 즉효입니다. 운영에서는 “무제한 생성”을 허용하면 언젠가 터집니다.

prompt = "긴 프롬프트..."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        use_cache=True,
    )

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

추가로 입력 토큰도 잘라야 합니다.

inputs = tokenizer(
    prompt,
    return_tensors="pt",
    truncation=True,
    max_length=2048,
).to(model.device)
  • “입력 8k + 출력 2k” 같은 조합은 24GB급 GPU에서도 위험해질 수 있습니다.

2) 배치/동시성 제한: OOM의 가장 흔한 트리거

로컬 서버로 띄워두고 여러 요청을 동시에 받으면 KV 캐시가 세션별로 생깁니다. 단일 요청은 되는데, 동시에 2개만 들어와도 터지는 형태가 이 케이스입니다.

해결은 단순합니다.

  • 동시 요청 수를 제한(큐잉)
  • 요청당 max_new_tokens 상한
  • 길이 기반 비용 모델로 admission control

이건 API 레벨에서도 “백오프/큐잉”이 핵심인데, 로컬 LLM도 똑같습니다. 트래픽 제어 관점은 다음 글이 참고됩니다: OpenAI 429·Rate Limit - 백오프·큐잉 실전

3) use_cache=False는 최후의 수단 (속도 급락)

KV 캐시를 끄면 메모리는 줄 수 있지만, 디코딩이 매 토큰마다 전체 시퀀스를 다시 보게 되어 속도가 크게 느려집니다.

out = model.generate(
    **inputs,
    max_new_tokens=128,
    use_cache=False,
)
  • 디버깅/검증에는 유용
  • 운영 기본값으로는 비추천

4) Flash Attention / SDPA로 메모리·속도 최적화

transformers는 PyTorch의 SDPA(Scaled Dot Product Attention) 경로를 활용해 더 효율적인 attention을 태울 수 있습니다. 환경에 따라 효과가 크고, 특히 긴 컨텍스트에서 체감이 납니다.

import torch

torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

주의할 점

  • GPU 아키텍처/드라이버/torch 버전에 따라 사용 가능한 경로가 달라집니다.
  • 일부 조합에서는 오히려 느려지거나, 특정 dtype에서 이슈가 날 수 있어 A/B 테스트가 필요합니다.

5) device_map과 오프로딩으로 “OOM 대신 느려짐”을 선택하기

VRAM이 절대 부족하면 일부 레이어를 CPU로 오프로딩하는 방법이 있습니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    offload_folder="./offload",
)
  • 장점: OOM 회피
  • 단점: PCIe 전송 때문에 지연이 크게 늘 수 있음

“무조건 빠르게”가 아니라 “일단 안 죽게”가 목표인 환경(개발용 PC, 데모 머신)에서는 현실적인 타협입니다.

6) 긴 컨텍스트가 필요하면: 슬라이딩 윈도우/요약 전략

KV 캐시는 “전체 컨텍스트를 끝까지 들고 있는” 구조라서, 입력이 길어질수록 불리합니다. 애플리케이션에서 다음을 고려하세요.

  • 최근 N토큰만 유지(슬라이딩 윈도우)
  • 이전 대화는 요약해서 압축
  • RAG를 쓰더라도 “검색 결과를 무작정 많이 붙이지 않기” (top-k 제한)

이건 모델 튜닝이라기보다 제품 설계에 가깝고, OOM 방지에 매우 효과적입니다.

실전: OOM이 나는 설정을 안전하게 바꾸는 체크리스트

아래 순서로 바꾸면 원인 분리가 쉽습니다.

  1. 동시성 1로 고정 (멀티 요청/배치 제거)
  2. max_new_tokens를 128 같은 작은 값으로 낮춤
  3. 입력 max_length를 1024~2048로 제한
  4. 그래도 OOM이면 모델 자체가 VRAM에 비해 큼
    • 더 작은 모델로 변경 또는 오프로딩
  5. 위가 안정화되면
    • Flash/SDPA 최적화 적용
    • 동시성을 2, 3… 점진적으로 증가

그리고 각 단계에서 torch.cuda.max_memory_allocated() 피크를 기록해 “어디서 급증하는지”를 확인하세요.

예제: 4bit + 길이 제한 + 메모리 리포트까지 한 번에

아래는 최소한의 “안 터지게” 구성한 예시입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

def report(tag: str):
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated() / 1024**2
    alloc = torch.cuda.memory_allocated() / 1024**2
    print(f"[{tag}] alloc={alloc:.1f}MB peak={peak:.1f}MB")

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="cuda:0",
)
model.eval()

prompt = """You are a helpful assistant. Summarize the following text..."""

inputs = tokenizer(
    prompt,
    return_tensors="pt",
    truncation=True,
    max_length=2048,
).to(model.device)

torch.cuda.reset_peak_memory_stats()
report("after_load")

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

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

여기서도 OOM이 나면 다음을 의심하세요.

  • GPU VRAM 자체가 너무 작음(예: 8GB에서 7B도 빡빡)
  • 백그라운드에서 다른 프로세스가 VRAM을 점유
  • 드라이버/torch 조합 문제로 메모리 단편화가 심함

운영 팁: “OOM이 안 나는 서비스”를 위한 정책

로컬 LLM을 서비스 형태로 붙이면, 모델 튜닝보다 정책이 더 중요해집니다.

  • 요청당 상한: 입력 토큰, 출력 토큰
  • 동시성 상한: 워커 수 또는 세마포어
  • 길이 기반 큐잉: 긴 요청은 별도 큐
  • 타임아웃: 생성이 일정 시간 넘으면 중단

이건 DB 복제 지연이나 gRPC 타임아웃처럼 “병목을 숨기지 말고 제어”하는 문제와 유사합니다. 예: Go gRPC 데드라인 초과 해결 - context·LB·Keepalive

마무리

transformers 로컬 LLM에서 OOM을 잡는 핵심은 다음 한 줄로 요약됩니다.

  • 4bit는 가중치를 줄이고, KV 캐시는 길이/동시성을 줄여라.

실전에서는

  • max_new_tokens와 입력 max_length를 강제하고
  • 동시성을 제한하며
  • Flash/SDPA 같은 효율 경로를 켜고
  • 필요하면 오프로딩으로 “죽지 않는” 구성을 만든 뒤

점진적으로 성능을 올리는 방식이 가장 안정적입니다.

다음 단계로는 “내 GPU에서 이 모델이 컨텍스트 N, 동시성 M일 때 안전한가”를 수치화하는 것이 좋습니다. 피크 VRAM을 기록해 표로 만들면, OOM은 더 이상 운이 아니라 관리 가능한 지표가 됩니다.

Discuss on Twitter • View on GitHub