beautifulsoups.dev
Published on

로컬 LLM OOM 해결 - vLLM·KV 캐시 최적화

Authors

로컬에서 transformers로 LLM을 돌리다 보면, 모델 가중치는 잘 올라갔는데도 추론 중간에 CUDA out of memory가 터지는 경우가 많습니다. 특히 동시 요청이 늘거나, 컨텍스트 길이를 조금만 올려도 갑자기 메모리가 폭증합니다. 이 글은 OOM의 핵심 원인이 KV 캐시인 경우가 많다는 점을 짚고, vLLM의 PagedAttention 기반 KV 캐시 관리로 어떻게 해결하는지, 그리고 transformers 환경에서도 당장 적용 가능한 완화책을 함께 정리합니다.

또한 운영 관점에서 “왜 갑자기 OOM이 났는지”를 추적하는 방식은 LLM 서빙에서도 중요합니다. 장애 원인 추적 루틴은 systemd 서비스 무한 재시작 원인과 journalctl 추적 글의 관점(로그 기반 원인 좁히기)을 그대로 가져오면 도움이 됩니다.

OOM을 부르는 진짜 범인: KV 캐시

LLM 추론 메모리는 크게 3개 덩어리로 나뉩니다.

  1. 모델 가중치(Weights): 로딩 시 거의 고정. fp16/bf16/int8/int4로 줄일 수 있습니다.
  2. 활성화(Activations): 주로 프리필(prefill, prompt 처리) 구간에서 증가. 배치가 커지면 커집니다.
  3. KV 캐시(Key/Value Cache): 디코딩(decoding, 토큰 생성) 동안 누적. 컨텍스트가 길고 동시 요청이 많을수록 폭증합니다.

많은 팀이 “가중치가 GPU에 들어가니 끝”이라고 생각하지만, 실제 OOM은 KV 캐시가 GPU 메모리를 잠식하면서 발생하는 일이 많습니다.

KV 캐시가 왜 그렇게 큰가

Self-Attention은 과거 토큰들의 Key/Value를 재사용합니다. 그래서 토큰을 한 개 생성할 때마다 각 레이어에 대해 K/V 텐서를 저장합니다. 대략적인 스케일은 아래처럼 잡을 수 있습니다.

  • KV 캐시 메모리 ≈ batch_size * num_layers * seq_len * hidden_size * 2(K,V) * dtype_bytes (여기에 head 분할/패딩/오버헤드가 붙습니다)

즉, 다음이 동시에 커지면 위험합니다.

  • seq_len(컨텍스트 길이)
  • 동시 요청 수(서빙에서의 effective batch)
  • 레이어 수가 많은 모델
  • fp16/bf16라도 dtype 자체는 2바이트라 결코 작지 않음

Transformers에서 OOM이 잦은 구조적 이유

transformers의 기본 서빙 패턴은 보통 다음 중 하나입니다.

  • 요청마다 generate() 호출
  • 여러 요청을 한 프로세스에서 처리하되, 프레임워크 레벨의 효율적인 스케줄링이 약함

이때 자주 생기는 문제가 있습니다.

1) 동적 배칭 부재로 인한 비효율

요청이 동시에 들어오면 GPU를 잘 쓰기 위해 배치를 키우고 싶은데, 단순히 묶으면 서로 다른 길이의 prompt 때문에 padding이 늘고, 프리필 메모리가 커집니다.

2) KV 캐시의 “연속 할당”과 단편화

일반적인 구현은 시퀀스 길이가 늘어날수록 KV 캐시를 큰 덩어리로 잡는 경향이 있습니다. 요청이 끝나고 메모리가 반환되어도, CUDA allocator 관점에서 단편화가 생기면 “총 여유 메모리는 있는데도” OOM이 납니다.

3) max_new_tokens와 max_length를 넉넉히 잡는 습관

서빙에서 안전하게 하려고 상한을 크게 잡으면, 프레임워크가 내부적으로 더 큰 버퍼를 준비하거나, 최악 케이스에 맞춘 스케줄링으로 메모리 압박이 커질 수 있습니다.

vLLM이 OOM을 줄이는 핵심: PagedAttention과 KV 캐시 관리

vLLM의 강점은 단순히 “빠르다”가 아니라, KV 캐시를 페이지 단위로 관리해서 단편화와 과할당을 줄이는 데 있습니다.

PagedAttention이 하는 일

  • KV 캐시를 OS의 페이징처럼 고정 크기 블록(페이지) 으로 쪼갭니다.
  • 시퀀스가 늘어나면 필요한 페이지만 추가로 할당합니다.
  • 요청이 끝나면 페이지를 풀에 반환하여 재사용합니다.

결과적으로 다음이 개선됩니다.

  • 긴 컨텍스트에서의 메모리 효율
  • 동시 요청 처리(continuous batching)
  • 단편화로 인한 “애매한 OOM” 감소

실전: vLLM로 로컬 서빙 전환하기

가장 단순한 시작은 OpenAI 호환 서버 모드입니다.

pip install vllm

# 예: Llama 계열 모델
python -m vllm.entrypoints.openai.api_server \
  --model meta-llama/Llama-2-7b-chat-hf \
  --dtype float16 \
  --gpu-memory-utilization 0.90 \
  --max-model-len 4096

옵션 해설

  • --gpu-memory-utilization
    • vLLM이 KV 캐시 등으로 사용할 GPU 메모리 비율을 조절합니다.
    • 너무 높이면 다른 프로세스/드라이버 오버헤드 때문에 불안정해질 수 있어 0.85부터 올리는 편이 안전합니다.
  • --max-model-len
    • 운영 상한입니다. 크게 잡을수록 KV 캐시 풀이 커질 수 있습니다.
    • “필요한 만큼만” 잡는 게 OOM 예방의 기본입니다.

클라이언트 예시(OpenAI 호환)

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")

resp = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-chat-hf",
    messages=[{"role": "user", "content": "KV 캐시가 뭔지 3줄로 설명해줘"}],
    temperature=0.2,
    max_tokens=200,
)

print(resp.choices[0].message.content)

그래도 터진다: OOM 원인별 체크리스트

아래는 “vLLM로 바꿨는데도 OOM” 혹은 “Transformers를 유지해야 함” 상황에서 바로 점검할 항목입니다.

1) 컨텍스트 상한을 실제 트래픽에 맞추기

가장 흔한 실수는 max_model_len 또는 애플리케이션 레벨 컨텍스트를 과하게 잡는 것입니다.

  • RAG를 한다면, 검색 결과를 무조건 다 넣지 말고 상위 k와 chunk 길이를 제한하세요.
  • 검색 인덱스 튜닝으로 “필요한 문서만” 넣는 것도 효과가 큽니다. RAG가 느리거나 chunk가 과도해지는 문제는 pgvector RAG 느림? IVFFlat 튜닝 체크리스트처럼 검색단에서 해결하면 LLM 메모리 압박도 같이 줄어듭니다.

2) 동시성 제한: 배치가 커지면 KV 캐시가 직격탄

서빙에서 OOM은 보통 “특정 순간 동시 요청이 몰림”과 함께 옵니다.

  • vLLM을 써도 동시성이 무한이면 언젠가는 터집니다.
  • API 게이트웨이/서버에서 동시 요청 수 제한과 큐잉을 두세요.

예: asyncio.Semaphore로 간단 제한

import asyncio

sem = asyncio.Semaphore(4)  # GPU 1장 기준, 모델/컨텍스트에 맞게 조정

async def handle(req):
    async with sem:
        return await run_inference(req)

3) 프리필 구간 메모리: 긴 프롬프트 한 방이 터뜨린다

긴 문서를 한 번에 넣는 워크로드는 디코딩보다 프리필이 더 위험할 때가 있습니다.

  • 입력 토큰 길이 제한
  • 문서 요약 후 투입
  • 시스템 프롬프트를 과도하게 길게 쓰지 않기

4) dtype와 양자화 전략 재점검

가중치가 줄면 KV 캐시가 지배적인 상황에서도 “총량”이 줄어 안정성이 올라갑니다.

  • float16/bfloat16 기본
  • 가능하면 int8/int4 양자화(모델/품질/커널 지원 확인)

Transformers에서 bitsandbytes 양자화 예시

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_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
)

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

주의할 점은, 양자화는 가중치 메모리를 크게 줄이지만 KV 캐시는 여전히 fp16/bf16로 남는 경우가 많다는 것입니다. 그래서 “양자화했는데도 OOM”이 충분히 가능합니다.

5) 메모리 단편화 의심 시: 프로세스 전략

  • 한 프로세스에서 너무 오래 돌리면 단편화가 누적될 수 있습니다.
  • 워크로드가 들쭉날쭉하면, 일정 조건에서 워커를 재시작하는 전략도 실무에서 씁니다.

다만 무작정 재시작하면 장애가 커지므로, 재시작 루프/원인 로그를 반드시 잡아야 합니다. 이런 운영 패턴은 systemd 서비스 무한 재시작 원인과 journalctl 추적처럼 “왜 재시작됐는지”를 로그로 남기고, 재현 가능한 단서(요청 크기, 토큰 수, 동시성)를 수집하는 게 핵심입니다.

6) 요청 단위 토큰 예산제를 도입하기

OOM을 막는 가장 강력한 방법 중 하나는 “사용자가 GPU 메모리를 마음대로 쓰지 못하게” 하는 것입니다.

  • 입력 토큰 상한
  • 출력 토큰 상한
  • 입력과 출력의 합 상한
  • RAG 컨텍스트 합 상한

예: 토큰 길이 초과 시 컷

from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf", use_fast=True)

def clip_prompt(prompt: str, max_input_tokens: int = 2048) -> str:
    ids = tok.encode(prompt)
    ids = ids[-max_input_tokens:]
    return tok.decode(ids)

Transformers를 유지해야 한다면: OOM 완화 팁

vLLM로 옮기기 어렵거나, 특수한 커스텀 로직 때문에 transformers를 유지해야 하는 경우도 있습니다. 그때의 우선순위는 아래입니다.

  1. 컨텍스트 상한을 줄이고(입력 토큰 컷)
  2. 동시 요청을 제한하고
  3. 양자화로 가중치를 줄인 뒤
  4. 가능하면 flash-attn 등 메모리 효율 커널을 검토합니다(모델/환경 의존)

또한 generate() 호출 시 불필요한 옵션(예: 과도한 num_beams)은 메모리를 크게 올릴 수 있습니다. 빔 서치는 KV 캐시를 빔 수만큼 사실상 복제하는 효과가 나기 때문에, 로컬 서빙에서는 특별한 이유가 없으면 피하는 편이 안전합니다.

결론: OOM은 “모델이 커서”가 아니라 “캐시를 못 다뤄서” 난다

로컬 LLM OOM을 줄이는 데 가장 효과적인 레버는 다음 순서로 정리할 수 있습니다.

  • KV 캐시 관점으로 문제를 재정의한다
  • 컨텍스트 상한과 동시성을 먼저 통제한다
  • vLLM의 PagedAttention으로 KV 캐시 할당/재사용을 최적화한다
  • 그래도 부족하면 양자화와 워크로드(프롬프트/RAG) 자체를 줄인다

특히 “동시 요청이 늘면 갑자기 터진다”, “총 여유 메모리는 있어 보이는데도 OOM이다”, “긴 컨텍스트에서 불안정하다”라면 vLLM 전환이 체감 효과가 큰 편입니다.

원한다면 사용하는 GPU 용량, 목표 동시성, 평균 입력 토큰/출력 토큰, 모델(예: 7B/13B/34B) 정보를 기준으로 대략적인 KV 캐시 메모리 예산과 --gpu-memory-utilization, --max-model-len 권장값까지 같이 잡아드릴 수 있습니다.

Discuss on Twitter • View on GitHub