Transformers 로컬 LLM OOM? 8bit·Q4·KV캐시로 해결

로컬 GPU에서 transformers로 LLM을 띄우면 가장 흔한 실패가 CUDA out of memory입니다. 많은 분이 “모델이 너무 커서”라고만 생각하지만, 실제 VRAM은 크게 세 덩어리로 나뉩니다.

• 가중치(Weights): 모델 파라미터 자체
• KV 캐시(KV cache): 생성(디코딩) 중 토큰마다 누적되는 Key/Value 텐서
• 활성화(Activations): 프리필/디코딩 중 레이어를 통과하며 생기는 중간 텐서(추론에서는 상대적으로 작지만 설정에 따라 커짐)

이 글은 특히 **가중치(8bit, Q4)**와 KV 캐시를 중심으로, 로컬 LLM OOM을 “재현 가능한 방식”으로 줄이는 실전 레시피를 제공합니다. OOM 전반을 더 넓게 커버한 체크리스트가 필요하면 Transformers 로컬 LLM OOM? 7가지 즉시 해결도 함께 보세요.

OOM을 숫자로 쪼개서 이해하기

1) 가중치 메모리: dtype과 양자화가 곧 VRAM

대략적인 감은 다음과 같습니다.

• fp16/bf16: 파라미터 1개당 2바이트
• int8: 1바이트(스케일/제로포인트 등 오버헤드 존재)
• int4(Q4): 0.5바이트(블록 스케일 등 오버헤드 존재)

예를 들어 7B 모델은 파라미터가 약 70억 개이므로, 순수 계산만 하면 fp16 가중치만으로도 약 14GB가 필요합니다(여기에 오버헤드와 버퍼가 추가). 그래서 12GB GPU에서 7B가 “간당간당”하거나 바로 OOM이 나는 경우가 많습니다.

2) KV 캐시: “컨텍스트 길이 × 레이어 × 헤드”로 선형 증가

생성 단계에서 KV 캐시는 토큰이 늘어날수록 계속 쌓입니다. 즉,

• max_new_tokens를 늘리면 KV 캐시가 커지고
• 입력 프롬프트가 길어지면 프리필 단계에서 이미 KV 캐시가 크게 잡힙니다.

실무에서 “처음엔 되다가 대화가 길어지면 죽는” 패턴은 대부분 KV 캐시가 원인입니다.

3) 활성화: 추론에서는 상대적으로 작지만 설정에 따라 폭증

추론에서 활성화는 학습보다 작지만, 다음 조건에서 커질 수 있습니다.

• 배치가 커짐(batch_size 또는 동시 요청)
• use_cache=False로 디코딩을 매번 풀 패스 계산(비추천)
• 빔서치/샘플링 설정으로 동시에 유지되는 시퀀스 수 증가

전략 1: 8bit 로딩으로 “일단 띄우기”

bitsandbytes의 8bit 로딩은 가장 빠르게 VRAM을 줄이는 방법입니다. 품질 손실이 크지 않은 편이고, Q4보다 안정적으로 동작하는 환경이 많습니다.

설치

pip install -U transformers accelerate bitsandbytes

8bit 로딩 예제

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "meta-llama/Llama-2-7b-chat-hf"  # 예시

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    load_in_8bit=True,
    device_map="auto",
    torch_dtype=torch.float16,
)

prompt = "요약: KV 캐시가 왜 OOM을 만들지 설명해줘."
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,
        use_cache=True,
    )

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

포인트

• device_map="auto"는 GPU VRAM이 부족하면 일부 레이어를 CPU로 오프로딩할 수 있습니다. 다만 속도는 떨어집니다.
• torch_dtype는 8bit 로딩이라도 일부 연산 dtype에 영향을 줍니다.

전략 2: Q4(4bit)로 VRAM을 크게 줄이기

12GB 이하 GPU에서 7B~13B를 노릴 때는 Q4가 사실상 필수인 경우가 많습니다. transformers에서는 BitsAndBytesConfig로 4bit 양자화를 설정합니다.

4bit 로딩 예제(NF4 권장)

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = "meta-llama/Llama-2-13b-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",
)

prompt = "Q4 양자화가 VRAM에 어떤 영향을 주는지 5줄로."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

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

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

Q4에서 자주 겪는 이슈

• 드라이버/CUDA/bitsandbytes 조합에 따라 로딩이 실패할 수 있습니다.
• bnb_4bit_compute_dtype를 bf16으로 두면 성능/안정성이 좋아지는 경우가 많지만, GPU가 bf16을 잘 지원하지 않으면 fp16으로 바꿔야 합니다.

전략 3: KV 캐시를 줄여 “대화가 길어져도 안 죽게” 만들기

가중치를 Q4로 줄여도, 긴 컨텍스트와 긴 생성 길이에서는 KV 캐시가 다시 OOM을 유발합니다. 특히 챗봇처럼 누적 대화가 길어지는 워크로드에서 중요합니다.

1) 가장 효과적인 1차 처방: max_new_tokens와 입력 길이 제한

• max_new_tokens를 줄이면 KV 캐시 상한이 즉시 내려갑니다.
• 프롬프트 토큰 수가 길면 프리필 단계에서 이미 KV 캐시가 크게 잡힙니다.

아래처럼 토크나이저에서 입력을 잘라내는 것만으로도 OOM이 사라지는 경우가 많습니다.

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

out = model.generate(
    **inputs,
    max_new_tokens=256,
    use_cache=True,
)

2) 캐시 타입/동작을 명시적으로 관리하기

transformers는 버전이 올라가면서 캐시 구현이 다양해졌고, 모델에 따라 기본값도 다릅니다. 다음을 점검하세요.

• use_cache=True는 일반적으로 속도에 유리하지만, KV 캐시 메모리를 사용합니다.
• 반대로 use_cache=False는 KV 캐시를 쌓지 않지만, 디코딩이 매우 느려질 수 있어 실전에서는 잘 쓰지 않습니다.

“속도는 조금 손해 보더라도 OOM을 피하겠다”면, 특정 상황에서만 use_cache=False를 선택하는 전략도 가능합니다(예: 매우 짧은 생성).

out = model.generate(
    **inputs,
    max_new_tokens=64,
    use_cache=False,
)

3) 동시성(배치/서빙)이 KV 캐시를 폭발시키는 패턴

로컬에서 단일 요청은 되는데, 서버로 묶는 순간 OOM이 나는 대표 원인이 동시 요청 수만큼 KV 캐시가 복제되기 때문입니다.

• 동시 요청 N개
• 각 요청이 컨텍스트 C 토큰
• 각 요청이 생성 G 토큰

이면 KV 캐시 메모리는 사실상 N배로 증가합니다. 이때는 양자화보다 먼저 동시성 제한(큐잉)이나 마이크로배칭 튜닝이 필요합니다.

실전 조합 레시피: 12GB / 24GB에서 실패를 줄이는 설정

12GB GPU(예: RTX 3060) 권장 조합

• 모델: 7B급
• 로딩: Q4(nf4) 우선, 안 되면 8bit
• 입력 길이: 2k 내로 시작
• 생성 길이: max_new_tokens 128~256
• 동시성: 1부터 시작

24GB GPU(예: RTX 4090) 권장 조합

• 모델: 13B도 Q4면 여유, 7B는 fp16도 가능
• 긴 컨텍스트를 쓰려면 KV 캐시가 병목이 되므로 입력/생성 길이 정책이 여전히 중요

OOM을 “진짜로” 잡는 디버깅: 메모리 스냅샷과 재현

OOM은 감으로 고치면 재발합니다. 아래처럼 최소한의 계측을 넣어, 어떤 단계에서 VRAM이 치솟는지 확인하세요.

import torch

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

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

# prefill
with torch.inference_mode():
    _ = model(**inputs)
report("after_prefill")

# decode
with torch.inference_mode():
    out = model.generate(**inputs, max_new_tokens=256, use_cache=True)
report("after_generate")

• after_load가 크면 가중치/오프로딩 문제
• after_prefill에서 급증하면 긴 입력으로 KV 캐시가 크게 잡히는 패턴
• after_generate에서 계속 증가하면 생성 길이와 KV 캐시가 원인

OS 레벨에서 프로세스가 죽는 경우(리눅스 OOM Killer 개입)까지 의심된다면 리눅스 OOM Killer로 프로세스 죽을 때 원인 추적처럼 시스템 로그 기반으로 원인을 분리하는 것도 도움이 됩니다.

자주 하는 실수 5가지

1) max_length와 max_new_tokens를 혼용

• max_length는 “입력+출력 전체 길이” 상한입니다.
• max_new_tokens는 “출력 토큰 수” 상한입니다.

긴 프롬프트를 넣을 때는 max_new_tokens만 믿으면 전체 길이가 예상보다 커질 수 있으니, 입력은 tokenizer(..., truncation=True, max_length=...)로 별도 제한하는 편이 안전합니다.

2) device_map="auto"의 CPU 오프로딩을 성능 문제로만 봄

오프로딩은 속도를 희생해 OOM을 피하는 최후의 안전장치입니다. 로컬 개발/디버깅 단계에서는 유효하지만, 서빙에서는 지연이 급증할 수 있으니 목표 TPS에 맞춰 GPU 용량/모델 크기를 다시 산정해야 합니다.

3) 동시 요청을 늘리면서도 토큰 상한 정책이 없음

KV 캐시는 동시성에 매우 민감합니다. 서빙이라면 반드시 다음을 정책으로 못 박아야 합니다.

• 입력 토큰 상한
• 출력 토큰 상한
• 동시 요청 상한
• 초과 시 요약/리트라이/거절 전략

4) “Q4면 무조건 해결”이라고 생각

Q4는 가중치 VRAM을 크게 줄이지만, 긴 컨텍스트에서는 KV 캐시가 다시 지배적이 됩니다. 결국 토큰 정책이 없으면 재발합니다.

5) 파편화(fragmentation)를 무시

같은 VRAM 사용량이라도 할당 패턴이 나쁘면 OOM이 더 빨리 납니다. 반복 실행 시 갑자기 실패한다면 파편화를 의심하고, 프로세스 재시작 또는 캐시/할당 전략을 점검하세요.

결론: “가중치”와 “KV 캐시”를 따로 최적화해야 한다

로컬 LLM OOM은 보통 두 단계로 해결합니다.

1. 가중치 줄이기: 8bit로 일단 띄우고, 필요하면 Q4로 더 줄이기
2. KV 캐시 통제: 입력/출력 토큰 상한, 동시성 제한, 캐시 사용 정책으로 장기 안정성 확보

이 두 축을 분리해서 접근하면, 같은 GPU에서도 “처음만 되는 데모”가 아니라 “대화가 길어져도 안 죽는” 로컬 LLM 환경을 만들 수 있습니다.

추가로 더 많은 즉시 처방(배치, torch.compile, offload, attention 최적화 등)을 한 번에 보고 싶다면 Transformers 로컬 LLM OOM? 7가지 즉시 해결도 같이 참고해 보세요.