beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM과 KV 캐시 최적화

Authors

로컬 GPU 한 장에서 Transformers 기반 LLM을 돌리다 보면, 모델 로딩은 되는데 생성(generate) 순간에 터지거나, 배치나 컨텍스트 길이를 조금만 올려도 바로 CUDA OOM이 나는 경우가 많습니다. 특히 최근 모델들은 파라미터 메모리보다 KV 캐시가 지배적이 되는 구간이 쉽게 오기 때문에, OOM을 “가중치가 커서”라고만 보면 해결이 안 됩니다.

이 글은 (1) 어떤 메모리가 어디서 터지는지 구조적으로 쪼개고, (2) KV 캐시를 줄이는 옵션과 트레이드오프, (3) 운영 관점에서 재현 가능한 체크리스트와 코드 패턴을 제공합니다.

또한 로컬에서 잘 돌린 뒤 서버로 올릴 때는 서빙 프레임워크/롤백 전략이 필요합니다. 배포 파이프라인은 BentoML로 GPU 모델 서빙·롤백 1시간 구축 글도 함께 참고하면 좋습니다.

OOM을 만드는 메모리 4종 세트

Transformers 추론에서 GPU 메모리는 대략 아래로 나뉩니다.

  1. 가중치(Weights): 모델 파라미터 자체. 로딩 순간에 대부분 결정됩니다.
  2. 활성화(Activations): 추론에서는 학습보다 훨씬 작지만, 일부 연산(특히 긴 시퀀스)에서 늘어날 수 있습니다.
  3. KV 캐시(Key/Value cache): 오토리그레시브 생성에서 과거 토큰의 K/V를 레이어별로 저장. 컨텍스트 길이와 생성 길이에 비례해 증가합니다.
  4. 프래그먼테이션/버퍼: CUDA allocator 조각화, 커널 워크스페이스, attention 구현체의 임시 버퍼 등.

로컬 LLM에서 “프롬프트가 길어질수록” 혹은 “max_new_tokens를 올릴수록” OOM이 나는 패턴이면, 거의 항상 3번 KV 캐시가 주범입니다.

KV 캐시가 왜 그렇게 커지나

KV 캐시는 레이어마다, 헤드마다, 토큰마다 쌓입니다. 대략적인 감을 잡는 데 중요한 건 다음입니다.

  • 토큰 수가 늘면 선형으로 증가합니다. 프롬프트 길이 + 생성된 토큰 수가 모두 포함됩니다.
  • 레이어 수가 많을수록 증가합니다.
  • num_key_value_heads(GQA/MQA) 여부에 따라 증가폭이 줄기도 합니다.
  • dtype가 핵심입니다. fp16/bf16이면 토큰당 저장 비용이 크고, fp8/int8 계열 KV 캐시를 쓰면 크게 줄일 수 있습니다(지원 여부는 모델/커널/옵션에 따라 다름).

결론적으로 “7B는 가벼우니까 8GB에서 넉넉하겠지”가 성립하지 않습니다. 7B라도 컨텍스트 8k, 생성 2k, 배치 4 같은 조합이면 KV 캐시가 먼저 터질 수 있습니다.

먼저 확인할 것: 지금 OOM은 어디서 나나

1) generate 직전/직후 메모리 스냅샷

아래처럼 단계별로 메모리를 찍어두면, 가중치/캐시/프래그먼테이션 중 어디가 문제인지 빨리 좁힐 수 있습니다.

import torch

def mem(tag: str):
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserv = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] allocated={alloc:.1f}MiB reserved={reserv:.1f}MiB peak={peak:.1f}MiB")

# 사용 예
# mem("after_load")
# mem("before_generate")
# mem("after_generate")
  • 로딩 직후부터 reserved가 과도하게 크면 allocator 설정/조각화 가능성을 봅니다.
  • before_generate에서 괜찮다가 generate 시작하자마자 급증하면 KV 캐시 가능성이 큽니다.

2) 컨텍스트 길이(프롬프트 토큰 수)부터 측정

문장 길이와 토큰 수는 다릅니다. 반드시 토크나이즈 결과 길이를 기준으로 보세요.

from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("your-model")
text = "..."
inputs = tok(text, return_tensors="pt")
print("prompt_tokens=", inputs["input_ids"].shape[-1])

KV 캐시 최적화 1순위: 생성 길이와 배치 전략

max_new_tokens를 상한으로 두고, 스트리밍으로 UX를 맞추기

OOM을 막는 가장 확실한 방법은 “생성 토큰 상한”을 강제하는 것입니다. 긴 답변이 필요하면 한 번에 길게 생성하기보다, 적절한 상한과 스트리밍으로 체감 지연을 낮추는 편이 안정적입니다.

서버에서 스트리밍을 붙이면 프록시 버퍼링/타임아웃 때문에 499/502가 늘어나는 경우가 있습니다. 운영 체크리스트는 LLM SSE 스트리밍 499 502 급증과 응답 끊김을 잡는 프록시 튜닝 체크리스트도 같이 보면 좋습니다.

배치 크기와 동시성은 KV 캐시에 직격탄

동시 요청을 배치로 묶으면 처리량은 좋아지지만, KV 캐시는 요청 수만큼 별도로 쌓입니다.

  • 단일 GPU에서 OOM이 잦다면: 배치 크기부터 줄이고, 동시성은 큐로 제어합니다.
  • “프롬프트가 길고 생성도 길다”면: 배치 1로도 터질 수 있으니 컨텍스트/생성 상한을 먼저 조절합니다.

KV 캐시 최적화 핵심: cache_implementation와 양자화

Transformers는 버전에 따라 KV 캐시 구현체를 선택할 수 있습니다. 환경에 따라 옵션 이름/지원이 달라질 수 있으니, 아래 코드는 “패턴”으로 이해하고 본인 버전에 맞게 확인하세요.

1) Static KV cache로 allocator 압박 줄이기

동적 확장 방식은 토큰이 늘 때마다 할당이 반복되어 조각화나 피크 메모리를 키울 수 있습니다. 가능하면 static 캐시로 미리 잡아두는 게 안정적일 때가 많습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "your-model"

tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    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,
        cache_implementation="static",  # 지원되는 버전에서만 사용
    )

print(tok.decode(out[0], skip_special_tokens=True))
  • 장점: 메모리 피크와 조각화 리스크 감소
  • 단점: 최대 길이를 넉넉히 잡으면 그만큼 예약 메모리가 커질 수 있음

2) Quantized KV cache로 메모리 절감

지원되는 조합에서는 KV 캐시를 int8 등으로 저장해 메모리를 크게 줄일 수 있습니다. 이 경우 품질 영향, 커널 지원, 속도 변화가 있을 수 있어 벤치마크가 필수입니다.

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        use_cache=True,
        cache_implementation="quantized",  # 지원되는 환경에서만
    )

실무 팁:

  • “가중치 4bit 양자화”만 하고 KV 캐시는 fp16인 경우가 많습니다. 이러면 프롬프트가 길 때 여전히 OOM이 납니다.
  • KV 캐시 양자화는 환경 의존성이 크니, 옵션을 켰을 때 실제로 메모리가 줄었는지 memory_allocated로 검증하세요.

가중치 메모리 줄이기: 4bit 로딩과 dtype 정리

KV 캐시 이전에, 기본 체급(가중치)을 낮추는 것도 중요합니다.

bitsandbytes 4bit 로딩 예시

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "your-model"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

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

inputs = tok("...", return_tensors="pt").to("cuda")

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

주의할 점:

  • compute dtype를 fp16/bf16 중 무엇으로 둘지에 따라 속도/안정성이 바뀝니다.
  • 4bit로 가중치를 줄여도, 긴 컨텍스트에서는 KV 캐시가 다시 병목이 됩니다.

컨텍스트를 줄이는 실전 패턴: 잘라내기, 요약, RAG

KV 캐시는 “지나간 모든 토큰”을 들고 있는 구조라, 컨텍스트 자체를 줄이는 게 가장 직접적인 최적화입니다.

  • 슬라이딩 윈도우: 최근 대화 N토큰만 유지
  • 요약 메모리: 오래된 대화는 요약으로 압축
  • RAG: 필요한 문서 조각만 가져와 프롬프트를 짧게 유지

특히 RAG를 쓰면 “항상 긴 프롬프트” 문제를 줄일 수 있고, 결과적으로 KV 캐시 OOM도 줄어듭니다.

PyTorch/Transformers에서 자주 하는 OOM 유발 실수

1) torch.no_grad() 대신 torch.inference_mode()를 안 쓰는 경우

추론에서는 torch.inference_mode()가 더 강하게 그래프/버퍼 생성을 억제해 메모리와 속도에 유리합니다.

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

2) 출력 텐서를 GPU에 계속 쌓아두는 경우

여러 요청을 루프에서 처리할 때, 디코딩을 늦게 하거나 GPU 텐서를 리스트에 담아두면 메모리가 회수되지 않습니다.

texts = []
for prompt in prompts:
    inputs = tok(prompt, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=128)
    # 즉시 CPU 문자열로 변환해 GPU 텐서 수명을 끊기
    texts.append(tok.decode(out[0].to("cpu"), skip_special_tokens=True))
    del inputs, out
    torch.cuda.empty_cache()  # 빈번 호출은 비추지만, 디버깅 단계에선 유용

empty_cache()는 근본 해결책이 아니라 “조각화/캐시”를 관찰하기 위한 도구로 쓰는 편이 좋습니다.

3) device_map="auto"가 의도치 않게 CPU offload를 만드는 경우

GPU가 모자라면 일부 레이어가 CPU로 넘어가면서 속도가 급락하거나, 전송 버퍼로 메모리 사용 패턴이 꼬일 수 있습니다. 로컬 추론에서 목표가 “OOM 회피”인지 “속도 유지”인지에 따라 전략을 분리하세요.

  • OOM 회피가 최우선이면: CPU offload도 고려
  • 속도가 최우선이면: 모델 크기/컨텍스트/배치를 줄여 GPU 단독으로 맞추기

OOM 디버깅 체크리스트(재현 가능 버전)

  1. 프롬프트 토큰 수, max_new_tokens, 배치 크기 기록
  2. 모델 로딩 직후/생성 직전/직후 memory_allocated, memory_reserved 기록
  3. dtype 확인: 가중치 dtype, KV 캐시 dtype(가능하면)
  4. use_cache 켰을 때와 껐을 때 비교(속도는 느려지지만 원인 분리에 도움)
  5. cache_implementation을 static으로 바꿔 피크 변화 확인
  6. 가능하면 KV 캐시 양자화 옵션 적용 후 실제 메모리 감소 확인
  7. 여전히 터지면: 컨텍스트 상한을 먼저 낮추고, 그 다음 배치/동시성 조절

마무리: “가중치 최적화”만으로는 부족하다

Transformers 로컬 LLM OOM은 보통 2단계로 옵니다.

  • 1단계: 모델이 커서 로딩부터 힘듦 → 4bit/8bit, device_map, offload로 해결
  • 2단계: 로딩은 되는데 생성에서 터짐 → KV 캐시, 컨텍스트, 배치가 핵심

특히 컨텍스트가 길어지는 제품(대화, 문서 요약, 코드 분석)은 KV 캐시가 곧 비용이자 안정성입니다. max_new_tokens 상한, 컨텍스트 관리(RAG/요약), 그리고 가능한 환경에서는 cache_implementation 최적화까지 적용하면, 같은 GPU에서도 “돌아가게”가 아니라 “안정적으로 운영 가능하게” 만들 수 있습니다.

Discuss on Twitter • View on GitHub