Transformers 로컬 LLM OOM 해결 - KV캐시·PagedAttention

로컬 GPU에서 Hugging Face Transformers로 LLM을 추론하다가 CUDA out of memory를 맞는 순간, 대부분의 사람은 모델 크기(파라미터)만 의심합니다. 하지만 실제로는 KV 캐시(Key/Value cache) 가 메모리의 주범이 되는 경우가 훨씬 많습니다. 특히 컨텍스트가 길어지거나, 배치가 커지거나, 여러 요청을 동시에 처리하면 KV 캐시가 선형으로 커지면서 VRAM을 빠르게 잠식합니다.

이 글에서는

• OOM이 발생하는 메커니즘을 KV 캐시 관점에서 분해하고
• Transformers에서 당장 적용할 수 있는 메모리 절감 옵션을 정리하며
• 근본적으로는 PagedAttention(= paged KV 캐시) 로 문제를 풀어가는 방법(vLLM 중심)을 실전 코드와 함께 설명합니다.

추론 서버를 운영하다가 “처음엔 되다가 요청이 쌓이면 터진다” 같은 증상을 겪었다면, 이 글의 핵심은 메모리 파편화 + KV 캐시의 연속 할당 요구를 이해하는 데 있습니다. (비슷한 맥락으로 캐시/자원 고갈을 다룬 글로 Spring Boot 3에서 HikariCP 커넥션 고갈 원인 9가지도 참고할 만합니다.)

로컬 LLM OOM의 진짜 원인: 파라미터가 아니라 KV 캐시

KV 캐시가 왜 커지나

Autoregressive 디코딩에서 다음 토큰을 생성할 때, 매 스텝마다 과거 토큰들의 Attention을 다시 계산하면 비용이 너무 큽니다. 그래서 각 레이어의 Key/Value를 저장해 두고 재사용하는데, 이게 KV 캐시입니다.

KV 캐시 메모리 사용량은 대략 아래에 비례합니다.

• 배치 크기 B
• 레이어 수 L
• 현재 시퀀스 길이 S (프롬프트 길이 + 생성된 길이)
• 헤드 수 H 및 헤드 차원 D
• dtype 크기(예: fp16은 2바이트)

즉, 토큰이 늘어날수록 선형으로 증가하고, 동시 요청이 늘면 배치처럼 누적됩니다.

“처음엔 되는데 점점 터진다”의 정체

• 단일 요청은 통과하지만, 여러 요청을 동시에 처리하면 각 요청의 KV 캐시가 쌓입니다.
• 요청마다 프롬프트 길이가 다르면(가변 길이), KV 캐시가 연속 메모리로 깔끔하게 떨어지지 않고 파편화(fragmentation) 가 심해집니다.
• 결과적으로 아직 총 여유 VRAM이 남아 있어도, “연속된 큰 덩어리”를 못 구해서 OOM이 납니다.

이때 흔히 보이는 메시지가 “reserved memory가 큰데 allocated가 상대적으로 작다” 같은 형태입니다.

Transformers에서 당장 적용할 수 있는 OOM 완화 체크리스트

여기서는 “근본 해결(PagedAttention)”로 가기 전에, Transformers에서 현실적으로 자주 쓰는 완화책을 정리합니다.

1) max_new_tokens와 컨텍스트 길이 상한을 강제

가장 확실한 방법은 생성 길이를 제한하는 것입니다.

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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").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))

• use_cache=True는 속도를 위해 기본적으로 켜는 편이지만, KV 캐시가 커지므로 길이 상한이 없으면 OOM이 빨리 옵니다.
• 대화형 서비스라면 “대화 히스토리 무한 누적”이 OOM의 지름길입니다. 서버에서 히스토리를 요약하거나 슬라이딩 윈도우를 적용하세요.

2) dtype/양자화로 파라미터와 KV 캐시를 함께 줄이기

• fp16/bf16: 파라미터와 KV 캐시 모두 2바이트 기반
• int8/int4 양자화: 파라미터는 크게 줄지만, KV 캐시는 보통 fp16/bf16으로 남는 경우가 많아 기대만큼 안 줄 수도 있습니다.

BitsAndBytes 4bit 예시:

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import torch

model_id = "mistralai/Mistral-7B-Instruct-v0.2"

bnb_cfg = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_cfg,
    device_map="cuda",
)

포인트는 “모델 로딩은 가벼워졌는데도 긴 컨텍스트에서 OOM이 나는” 상황이 여전히 가능하다는 점입니다. 그때는 KV 캐시가 병목입니다.

3) FlashAttention/SDPA로 속도는 개선되지만, OOM은 별개

PyTorch의 SDPA나 FlashAttention은 attention 연산의 효율을 크게 올려주지만, KV 캐시의 총량 자체를 줄이는 해결책은 아닙니다. (일부 구현에서 메모리 피크를 낮추는 효과는 있을 수 있습니다.)

Transformers에서 SDPA를 쓰는 전형적인 형태:

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    attn_implementation="sdpa",
)

4) 동시성 제어: 배치/큐잉이 OOM을 만든다

로컬 추론 서버에서 진짜 위험한 건 “동시 요청을 무제한으로 받아서 내부적으로 배치를 키우는” 패턴입니다.

• 동시성 제한(세마포어)
• 요청 큐잉
• 긴 요청과 짧은 요청 분리

이건 OOM을 늦추는 데 매우 효과적입니다. 다만 “GPU를 최대한 활용하면서도 OOM 없이”라는 목표에는 근본적으로 한계가 있습니다.

KV 캐시 파편화 문제: 왜 PagedAttention이 필요한가

Transformers의 일반적인 추론은 요청마다 KV 캐시를 비교적 단순하게 할당합니다. 문제는 실제 서비스 트래픽이 아래처럼 가변적이라는 점입니다.

• 프롬프트 길이가 제각각
• 생성 길이도 제각각
• 중간에 취소되는 요청도 존재

이때 GPU 메모리에서 “큰 연속 블록”을 계속 요구하면 파편화가 누적되고, 결국 OOM이 발생합니다.

PagedAttention의 핵심 아이디어

PagedAttention은 KV 캐시를 OS의 페이징처럼 고정 크기 블록(page) 으로 쪼개 관리합니다.

• KV 캐시를 연속된 큰 덩어리로 잡지 않아도 됨
• 요청이 늘고 줄어도 page 단위로 재사용 가능
• 가변 길이 시퀀스에서도 메모리 효율이 좋아짐
• 결과적으로 동시 처리량(throughput)을 올리면서 OOM을 줄임

이 접근이 대중화된 구현이 vLLM입니다.

vLLM로 PagedAttention 적용하기 (로컬 추론 실전)

Transformers에서 “한 방에” KV 캐시 구조를 바꾸기는 어렵습니다. 실무에서는 추론 엔진을 vLLM로 교체하는 방식이 가장 빠르고 효과적입니다.

1) vLLM 설치

환경에 따라 CUDA 휠이 달라질 수 있으니, vLLM 공식 설치 가이드를 확인하는 것이 안전합니다. 보통은 아래처럼 시작합니다.

pip install vllm

2) OpenAI 호환 서버로 띄우기

python -m vllm.entrypoints.openai.api_server \
  --model mistralai/Mistral-7B-Instruct-v0.2 \
  --dtype half \
  --max-model-len 8192

• --max-model-len은 KV 캐시 상한과 직결됩니다. “가능한 크게”가 아니라 “서비스에 필요한 만큼”으로 잡는 게 OOM 방지에 유리합니다.
• 이 서버는 OpenAI 호환 API를 제공하므로, 기존 클라이언트를 크게 바꾸지 않고도 붙일 수 있습니다. (레이트 리밋/재시도 정책은 OpenAI API 429·Rate Limit 재시도 백오프 설계 같은 패턴을 로컬에도 적용하면 안정성이 좋아집니다.)

3) Python 클라이언트 호출 예시

OpenAI SDK 스타일로 호출하거나, 단순히 HTTP로 호출할 수 있습니다.

import requests

base_url = "http://127.0.0.1:8000/v1"

payload = {
    "model": "mistralai/Mistral-7B-Instruct-v0.2",
    "messages": [
        {"role": "user", "content": "KV 캐시가 왜 OOM을 유발하는지 5줄로 설명해줘"}
    ],
    "temperature": 0,
    "max_tokens": 200,
}

r = requests.post(f"{base_url}/chat/completions", json=payload, timeout=60)
print(r.json()["choices"][0]["message"]["content"])

이렇게 바꾸면, 동일 GPU에서 Transformers 대비 동시 처리량이 늘고 OOM 빈도가 줄어드는 경우가 많습니다. 특히 “여러 요청이 섞여 들어오는 환경”에서 체감이 큽니다.

OOM을 더 줄이는 운영 팁: 길이·동시성·캐시 정책

PagedAttention을 도입해도 “무한 컨텍스트, 무한 동시성”은 불가능합니다. 아래는 로컬 운영에서 효과가 큰 정책들입니다.

1) max_model_len과 입력 길이 컷

• 서버에서 입력 토큰 수를 측정해 상한을 넘으면 거절하거나 요약 후 재시도
• 사용자가 붙여 넣는 로그/문서가 길어질수록 KV 캐시가 급증

2) 프롬프트 캐싱(prompt caching)과 RAG 전처리

• 시스템 프롬프트가 길다면, 매 요청마다 같은 프롬프트를 반복하지 않도록 구조를 바꾸세요.
• RAG로 문서를 통째로 넣기보다, 검색으로 필요한 청크만 넣어 토큰을 줄이세요.

3) 관측: 메모리 파편화와 피크를 분리해서 보기

단순히 nvidia-smi만 보면 “사용량”만 보이고, 파편화나 reserved/allocated의 관계는 놓치기 쉽습니다.

• PyTorch 메모리 요약을 주기적으로 로그로 남기기
• OOM 직전의 배치 크기, 평균 입력 길이, 평균 생성 길이 기록

예시:

import torch

print(torch.cuda.memory_summary(device=0, abbreviated=True))

이 로그를 남겨두면 “총량 부족”인지 “파편화/피크”인지 구분이 빨라집니다.

Transformers를 계속 써야 한다면: 현실적인 타협점

조직/프로덕트 제약으로 vLLM을 바로 못 쓰는 경우도 있습니다. 그럴 때의 타협점을 정리하면 아래와 같습니다.

• 컨텍스트/생성 길이 상한을 강제하고, 대화 히스토리는 요약
• 동시성 제한과 요청 큐잉으로 피크를 제어
• 모델은 4bit로 내려 파라미터 여유를 확보하되, KV 캐시가 병목임을 인지
• 가능하면 attn_implementation을 sdpa로 두고 속도/피크를 약간이라도 개선

다만 “동시 요청이 많은 서비스”라면, 결국 KV 캐시 관리가 승부처라서 PagedAttention 계열로 가는 것이 장기적으로 유리합니다.

정리: OOM은 ‘모델 크기’보다 ‘토큰과 KV 캐시’ 문제다

• 로컬 LLM OOM의 핵심 원인은 시퀀스 길이와 동시성에 의해 커지는 KV 캐시인 경우가 많습니다.
• Transformers에서 할 수 있는 1차 대응은 max_new_tokens, 입력 길이 컷, 동시성 제한, dtype/양자화입니다.
• 하지만 가변 길이 요청이 섞이는 실제 트래픽에서는 메모리 파편화가 OOM을 앞당깁니다.
• PagedAttention(vLLM)은 KV 캐시를 page 단위로 관리해 파편화를 줄이고, 같은 GPU에서 더 많은 동시 요청을 안정적으로 처리하게 해줍니다.

로컬에서 “간단히 한두 번 돌리는” 단계는 Transformers로도 충분하지만, 조금이라도 서버 형태로 운영하거나 동시 처리가 필요하다면 PagedAttention 기반 엔진으로 전환하는 것이 가장 확실한 OOM 해법입니다.