Transformers 로컬 LLM OOM? 4bit+오프로딩 실전

서빙이든 개인 개발이든 로컬에서 LLM을 돌리다 보면 가장 먼저 부딪히는 벽이 CUDA out of memory 입니다. 특히 Transformers로 7B~13B급 모델을 불러오면, “가중치만 올렸는데도” 터지거나, 로딩은 되는데 첫 generate()에서 갑자기 죽는 일이 흔합니다.

이 글은 OOM을 단순히 “VRAM이 부족해서”로 끝내지 않고, 메모리가 어디서 터지는지(가중치, KV 캐시, 활성화, 버퍼), 그리고 4bit 양자화 + 오프로딩(offload) 조합으로 실제로 돌아가는 세팅을 만드는 방법을 정리합니다.

추가로, 양자화는 성능과 정확도 트레이드오프가 있으니, 비슷한 맥락의 글로 PyTorch PTQ INT8 양자화 정확도 급락 해결법도 함께 보면 “왜 품질이 흔들리는지” 감을 잡는 데 도움이 됩니다.

OOM의 진짜 원인: 가중치보다 KV 캐시가 먼저 터진다

LLM 메모리 사용량은 크게 4가지로 나뉩니다.

모델 가중치(Weights): 로딩 순간에 VRAM을 크게 점유합니다.
KV 캐시(Key/Value cache): generate()에서 토큰이 늘어날수록 선형으로 커집니다. 길게 생성하면 여기서 먼저 터집니다.
활성화(Activations): 학습/미세조정에서 더 크지만, 추론에서도 일부 연산 버퍼가 생깁니다.
프레임워크/커널 버퍼: attention 구현, torch.compile, cudnn/cublas 워크스페이스 등.

많은 사람이 “7B면 8GB에서 4bit로 되겠지”라고 생각하지만, 실제로는 컨텍스트 길이와 배치, 동시 요청 수, 샘플링 설정에 따라 KV 캐시가 급격히 커져서 OOM이 납니다.

대략적인 감각 잡기

정확한 계산은 모델 구조에 따라 다르지만, 실전에서는 아래 항목을 먼저 의심하면 됩니다.

max_new_tokens를 크게 잡았다
max_length 또는 입력 프롬프트가 길다
batch_size 또는 동시 요청이 많다
use_cache=True(기본값)로 KV 캐시가 쌓인다

KV 캐시는 품질과 속도에 중요해서 무작정 끄면(use_cache=False) 속도가 크게 느려질 수 있습니다. 그래서 보통은 가중치부터 줄이고(4bit), 그래도 부족하면 오프로딩으로 버틴다가 현실적인 접근입니다.

4bit 양자화 기본: bitsandbytes + Transformers

가장 흔한 조합은 Transformers의 BitsAndBytesConfig로 4bit 로딩(nf4)을 하는 방식입니다.

설치 체크

아래는 최소 구성입니다.

pip install -U "transformers>=4.39" "accelerate>=0.26" bitsandbytes

환경에 따라 bitsandbytes가 CUDA 버전과 맞지 않으면 로딩 단계에서 에러가 날 수 있습니다. 이 경우는 OOM이 아니라 바이너리 호환성 문제라서, CUDA/드라이버/파이토치 버전을 먼저 맞춰야 합니다.

4bit 로딩 코드

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 지원하면 권장
)

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

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

prompt = "요약: Transformers에서 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,
        top_p=0.9,
    )

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

여기서 핵심은 device_map="auto" 입니다. 이 설정은 accelerate가 레이어를 GPU에 “가능한 만큼” 올리고, 남는 건 CPU로 분산시키는 기반이 됩니다. 하지만 이것만으로는 충분하지 않은 경우가 많습니다.

그래도 OOM이면: 오프로딩을 설계한다 (CPU / NVMe)

4bit로도 OOM이 나면 “VRAM이 정말 작다”거나 “컨텍스트가 길다”거나 “동시성이 있다”는 의미입니다. 이때는 오프로딩을 적극적으로 써야 합니다.

오프로딩은 크게 두 가지 축이 있습니다.

CPU offload: 일부 레이어를 CPU RAM에 올림
Disk(NVMe) offload: CPU RAM조차 부족하면 디스크로 스왑

Transformers/accelerate에서는 offload_folder와 max_memory로 이를 제어합니다.

max_memory로 상한을 걸어 OOM을 회피

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, 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.float16,
)

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

max_memory = {
    0: "10GiB",      # GPU 0에 10GiB까지만 사용
    "cpu": "48GiB"  # CPU RAM 상한
}

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    max_memory=max_memory,
    offload_folder="./offload",  # 필요 시 일부가 디스크로 내려감
)

max_memory를 명시하면 “일단 다 GPU에 넣다가 터지는” 패턴을 줄이고, 로딩 단계에서부터 분산을 유도할 수 있습니다.

offload_folder를 쓸 때의 현실적인 주의점

NVMe offload는 속도가 매우 느려질 수 있습니다. “돌긴 도는” 수준을 목표로 할 때 사용하세요.
./offload 폴더는 충분한 디스크 용량이 필요합니다.
오프로딩이 걸리면 지연시간이 튀므로, API 서빙에서는 타임아웃/큐잉 전략을 같이 설계해야 합니다.

OOM을 줄이는 생성 파라미터 튜닝

모델을 어떻게 올렸는지와 별개로, 생성 설정이 KV 캐시를 폭발시키는 경우가 많습니다.

1) 입력 길이와 생성 길이를 강제 제한

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

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

truncation=True와 max_length는 “프롬프트가 길어져서 KV 캐시 + attention 비용이 폭발”하는 상황을 예방합니다.

2) 배치/동시성: 로컬은 1부터 시작

로컬에서 OOM이 나면 먼저 아래를 지키는 게 좋습니다.

배치는 1
동시에 여러 요청을 돌리지 않기
스트리밍 출력으로 max_new_tokens를 과도하게 키우지 않기

동시성은 VRAM뿐 아니라 CPU RAM(오프로딩 시)도 같이 압박합니다.

메모리 상태를 눈으로 확인하는 체크리스트

OOM 해결은 “추측”보다 “측정”이 빠릅니다.

nvidia-smi로 피크 확인

watch -n 0.5 nvidia-smi

로딩 시점 피크
첫 generate() 시점 피크
토큰이 길어질수록 우상향하는지(KV 캐시)

이 3가지를 분리해서 보면 원인이 명확해집니다.

PyTorch에서 간단히 메모리 찍기

import torch

def report(tag: str):
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    print(f"[{tag}] allocated={allocated:.2f}GiB reserved={reserved:.2f}GiB")

report("after_load")
# generate 전후로 report("before_generate"), report("after_generate") 찍기

reserved가 큰데 allocated는 상대적으로 작으면, 캐싱/프래그먼테이션 성격이 강할 수 있습니다. 이때는 프로세스 재시작이 가장 확실한 해결책인 경우도 많습니다.

자주 겪는 함정 5가지

1) `torch_dtype`와 4bit compute dtype 혼동

4bit로 로딩해도 내부 연산 dtype(bnb_4bit_compute_dtype)이 float32면 느리고 메모리도 더 씁니다. 가능하면 float16 또는 bfloat16을 선택하세요.

2) `device_map="auto"`가 만능이 아니다

auto는 “대충 분산”이지 “최적”은 아닙니다. 특정 GPU에서만 OOM이 나면 max_memory로 상한을 강제하거나, 필요하면 수동 device_map을 고려해야 합니다.

3) 컨텍스트를 길게 쓰면 4bit여도 터진다

가중치를 줄여도 KV 캐시는 줄지 않습니다. 긴 문맥이 필요하면 다음 중 하나를 선택해야 합니다.

더 큰 VRAM
더 공격적인 오프로딩
컨텍스트를 줄이고 RAG로 외부 지식을 붙이기

4) CPU RAM이 부족하면 “GPU OOM”처럼 보일 수 있다

오프로딩 중에 CPU 메모리가 부족하면 시스템이 스왑으로 넘어가며 전체가 느려지고, 결국 타임아웃이나 커널 OOM 킬이 날 수 있습니다. GPU만 보지 말고 htop도 같이 보세요.

5) Windows에서 특히 오프로딩이 불안정한 케이스

드라이버/파이토치/bitsandbytes 조합에 따라 예민할 수 있습니다. 가능하면 WSL2 또는 리눅스 환경이 운영 난이도가 낮습니다.

마무리: OOM은 설정 문제가 아니라 설계 문제다

Transformers 로컬 LLM OOM은 보통 “모델이 커서”가 아니라, **가중치 로딩 전략(4bit), 레이어 배치(device_map), 메모리 상한(max_memory), 오프로딩 경로(offload_folder), 그리고 생성 길이(KV 캐시)**가 함께 맞물려 터집니다.

정리하면 아래 순서로 접근하면 시행착오가 줄어듭니다.

4bit nf4로 가중치부터 줄이기
device_map="auto"에 max_memory로 상한 걸기
부족하면 offload_folder로 CPU 또는 NVMe 오프로딩
max_length, max_new_tokens로 KV 캐시 폭발 방지
nvidia-smi와 torch.cuda.memory_*로 피크 지점 확인

이 조합만 제대로 잡아도, “로딩은 되는데 생성에서 죽는” 문제의 대부분은 재현 가능하게 통제할 수 있고, 로컬 환경에서도 충분히 실험 가능한 수준으로 안정화할 수 있습니다.