beautifulsoups.dev
Published on

Transformers 로컬 LLM VRAM 부족 - 4bit·KV캐시 튜닝

Authors

로컬 GPU에서 Transformers로 LLM을 돌리다 보면 가장 흔한 실패가 CUDA out of memory 입니다. 특히 7B, 13B급 모델을 8GB~12GB VRAM에서 돌릴 때는 “가중치만 줄이면 되겠지”라고 접근했다가, 실제로는 KV 캐시가 길이와 배치에 비례해 폭증하면서 OOM이 나는 경우가 많습니다.

이 글에서는 VRAM 사용량을 가중치(Weights), KV 캐시, 활성화(Activations) 로 나눠 이해하고, Transformers에서 바로 적용 가능한 4bit 양자화(bitsandbytes) 와 KV 캐시 튜닝(use_cache, max_new_tokens, sliding window 등) 을 중심으로 로컬 추론을 안정화하는 방법을 정리합니다.

추가로 운영 환경에서 OOM이 반복될 때는 “한 번 터지고 끝”이 아니라 프로세스 재시작, 워커 크래시로 이어질 수 있으니, 장애 관점 진단 루틴도 함께 가져가면 좋습니다. 비슷한 맥락의 운영 팁은 Ray Serve 배포 시 모델 로딩 지연·OOM 해결법, 컨테이너가 반복 재시작되는 상황은 K8s CrashLoopBackOff 10분 원인별 진단법도 같이 참고하면 연결이 됩니다.

VRAM이 어디서 새는지: 3가지 덩어리로 쪼개기

1) 가중치 메모리(Weights)

  • 모델 파라미터가 GPU에 올라가며 차지하는 메모리
  • FP16/BF16 기준 대략 params * 2 bytes
  • 7B는 FP16이면 대략 14GB 수준(여기에 오버헤드 포함)

즉 12GB VRAM에서 7B를 FP16으로는 원천적으로 어렵고, 4bit/8bit 양자화가 사실상 필수입니다.

2) KV 캐시(Key/Value cache)

  • 디코딩(생성) 단계에서 과거 토큰의 attention을 재사용하기 위해 저장하는 캐시
  • 시퀀스 길이(프롬프트 길이 + 생성 길이) 에 비례해 증가
  • 배치 크기(동시 요청 수)에도 비례

체감상 “프롬프트가 길어질수록 갑자기 OOM”이 나는 주범이 KV 캐시입니다.

3) 활성화 메모리(Activations)

  • forward 중간 텐서(특히 학습/미세튜닝에서 큼)
  • 순수 추론에서도 배치가 커지면 증가
  • torch.no_grad() 혹은 torch.inference_mode()로 줄일 수 있음

먼저 확인할 것: 지금 OOM이 가중치인지 KV 캐시인지

OOM을 해결하려면 “줄여야 할 대상”을 먼저 맞춰야 합니다.

  • 모델 로딩 직후 OOM: 가중치가 원인일 확률이 큼(4bit/8bit 필요)
  • 짧은 프롬프트는 되는데 길어지면 OOM: KV 캐시가 원인일 확률이 큼
  • 동시 요청을 늘리면 OOM: KV 캐시와 활성화가 함께 원인일 확률이 큼

런타임에서 VRAM을 대략적으로 관찰하는 최소 코드는 아래처럼 잡을 수 있습니다.

import torch

def vram(msg=""):
    if not torch.cuda.is_available():
        print("CUDA not available")
        return
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"{msg} allocated={alloc:.1f}MiB reserved={reserved:.1f}MiB")

vram("start")
# 모델 로드
vram("after load")
# 생성 수행
vram("after generate")

allocated는 실제 사용량, reserved는 캐싱 allocator가 잡아둔 영역이라 둘 다 같이 보는 게 좋습니다.

4bit 양자화로 가중치부터 줄이기(Transformers + bitsandbytes)

로컬 LLM 추론에서 가장 실전적인 조합은 다음입니다.

  • transformers
  • bitsandbytes
  • (가능하면) accelerate

핵심은 BitsAndBytesConfig로 4bit 로딩을 지정하는 것입니다.

4bit 로딩 예시(NF4 + double quant)

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,  # GPU가 bf16 지원하면 권장
)

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

model.eval()

옵션 해설

  • load_in_4bit=True: 가중치 4bit 양자화 로딩
  • bnb_4bit_quant_type="nf4": 일반적으로 성능-품질 균형이 좋다고 많이 쓰는 방식
  • bnb_4bit_use_double_quant=True: 추가 압축으로 VRAM을 더 줄이는 경우가 많음
  • bnb_4bit_compute_dtype: 연산 dtype. bfloat16이 가능하면 안정적인 편이고, 안 되면 float16로

4bit로도 부족할 때 체크리스트

  • device_map="auto"가 일부 레이어를 CPU로 오프로딩할 수 있음(속도는 감소)
  • GPU가 8GB 이하라면 7B도 프롬프트 길이에 따라 빡빡할 수 있음(이때는 KV 캐시 튜닝이 더 중요)

KV 캐시가 진짜 폭탄이다: 길이와 동시성의 곱

KV 캐시는 대략적으로 다음에 비례합니다.

  • 배치 크기 batch_size
  • 레이어 수 num_layers
  • 헤드 수/헤드 차원
  • 시퀀스 길이 seq_len (프롬프트 + 생성 토큰)
  • dtype(보통 FP16/BF16)

즉, 동시 요청을 2배로 늘리면 KV 캐시도 거의 2배가 되고, max_new_tokens를 2배로 늘려도 거의 2배가 됩니다.

그래서 “4bit로 로딩했는데도 OOM”이면, 그 다음은 KV 캐시를 건드려야 합니다.

KV 캐시 튜닝 1: max_new_tokens를 먼저 제한

가장 즉각적인 효과는 생성 길이를 줄이는 것입니다.

inputs = tok("한국어로 짧게 요약해줘: ...", return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,   # 1024 같은 큰 값은 KV 캐시 폭증
        do_sample=False,
        use_cache=True,
    )

print(tok.decode(out[0], skip_special_tokens=True))
  • 제품/서비스에서는 “무제한 생성”을 허용하면 최악의 경우 한 요청이 GPU를 고갈시킵니다.
  • 실제 운영에서는 max_new_tokens 상한을 두고, 필요하면 “계속 생성” UX로 분할하는 편이 안전합니다.

KV 캐시 튜닝 2: 긴 프롬프트를 줄이는 전략(요약, RAG, chunk)

프롬프트 길이가 길면, 생성 길이가 짧아도 KV 캐시가 커집니다.

실전에서는 다음이 효과적입니다.

  • 대화 히스토리: 최근 N턴만 유지하고 이전은 요약본으로 교체
  • 문서 입력: chunking 후 필요한 chunk만 넣는 RAG
  • 시스템 프롬프트: 중복되는 지시문을 최소화

이건 “메모리 튜닝”이면서 동시에 latency도 줄여줍니다.

KV 캐시 튜닝 3: use_cache=False는 최후의 수단

use_cache=False는 KV 캐시를 저장하지 않으므로 메모리를 크게 줄일 수 있지만, 매 토큰마다 과거를 다시 계산해 속도가 급격히 느려집니다.

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        use_cache=False,  # 메모리는 줄지만 속도는 크게 손해
    )
  • 디버깅 또는 “어떻게든 돌아가게” 만들 때는 유용
  • 일반적인 채팅/생성 서비스에는 권장하지 않습니다

KV 캐시 튜닝 4: Sliding Window Attention(모델이 지원할 때)

일부 모델/아키텍처는 슬라이딩 윈도우(최근 토큰만 attention)로 KV 캐시 증가를 제한할 수 있습니다.

다만 이건 모델 설정과 구현 지원 여부에 따라 다릅니다. Transformers에서는 모델 config에 sliding_window 같은 필드가 있는지 확인하고, 지원하는 모델에서만 적용해야 합니다.

주의: 지원하지 않는 모델에 억지로 설정을 넣으면 성능 저하가 아니라 동작 자체가 꼬일 수 있습니다.

성능과 VRAM을 같이 잡는 실전 파라미터 조합

아래 조합은 “로컬 GPU 8GB~12GB에서 7B급을 현실적으로”를 목표로 할 때 자주 쓰는 세트입니다.

  • 가중치: 4bit(NF4 + double quant)
  • 연산 dtype: bfloat16 가능하면 사용, 아니면 float16
  • 생성 제한: max_new_tokens를 128~512 등으로 상한
  • 프롬프트 제한: 히스토리 요약 + RAG로 입력 길이 관리
  • 동시성 제한: 배치/동시 요청 수를 VRAM에 맞춰 제한

예시 코드:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

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

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

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

prompt = "다음 요구사항을 5줄로 정리해줘: ..."
inputs = tok(prompt, return_tensors="pt").to(model.device)

gen_kwargs = dict(
    max_new_tokens=256,
    temperature=0.2,
    top_p=0.9,
    do_sample=True,
    use_cache=True,
)

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

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

OOM이 반복될 때 운영 관점 체크(재시작, 단편화, 워커 격리)

로컬 개발이든 단일 서버 서빙이든, OOM은 종종 “한 번의 실패”로 끝나지 않고 프로세스 상태를 망가뜨립니다.

1) PyTorch 캐시 단편화와 reserved 폭증

  • OOM 직전 reserved가 과도하게 커져 있으면 캐시 allocator 영향일 수 있습니다.
  • 요청 패턴이 들쭉날쭉(짧은 요청과 긴 요청이 섞임)하면 단편화가 심해질 수 있습니다.

가능하면:

  • 입력 길이/생성 길이를 구간화해서 워커를 분리
  • “긴 요청 전용 워커”를 따로 두기

2) 동시 요청 제어가 사실상 필수

KV 캐시는 동시성에 선형 비례합니다. 따라서 서버에서는 다음 중 하나가 필요합니다.

  • 큐잉(동시 실행 제한)
  • 배치 전략(가능하면 길이가 비슷한 요청끼리)
  • 워커 수를 VRAM에 맞게 강제

서빙 환경에서 모델 로딩 지연과 OOM이 같이 난다면 위의 운영 팁을 좀 더 확장한 글인 Ray Serve 배포 시 모델 로딩 지연·OOM 해결법이 직접적으로 도움이 됩니다.

3) 컨테이너/프로세스가 계속 죽는다면 원인 분리

OOM은 결과이고, 원인은 “요청 길이 폭주”, “동시성 폭주”, “모델 교체 시 메모리 누수” 등 다양합니다. 쿠버네티스에서 반복 재시작으로 보이면 K8s CrashLoopBackOff 10분 원인별 진단법처럼 이벤트와 로그를 기준으로 원인부터 분리하는 게 빠릅니다.

문제 해결 순서(추천)

  1. 가중치부터: 4bit 로딩으로 “모델이 뜨는 상태” 만들기
  2. 생성 길이 제한: max_new_tokens 상한 설정
  3. 프롬프트 길이 관리: 히스토리 요약, RAG, chunking
  4. 동시성 제한: 큐잉 또는 워커 분리
  5. 그래도 안 되면: use_cache=False 같은 극단 옵션은 마지막에

마무리

Transformers 로컬 LLM에서 VRAM 부족을 해결하려면 “4bit면 끝”이 아니라, KV 캐시가 시퀀스 길이와 동시성에 의해 폭증한다는 구조를 이해하고 제어해야 합니다.

  • 4bit 양자화는 가중치를 줄이는 가장 강력한 1차 수단
  • KV 캐시는 max_new_tokens, 프롬프트 길이, 동시성 제한으로 다스리는 게 핵심
  • 운영에서는 OOM을 장애로 보고 재현 가능한 제한(상한, 큐잉, 워커 분리)으로 설계하는 게 재발을 막습니다

원하는 GPU VRAM(예: 8GB, 12GB, 24GB)과 목표 모델(예: 7B, 14B), 목표 컨텍스트 길이(예: 4k, 8k)를 알려주면, 그 조건에 맞춰 권장 max_new_tokens, 동시성, 프롬프트 정책까지 더 구체적인 튜닝 가이드를 제안할 수 있습니다.

Discuss on Twitter • View on GitHub