Transformers 로컬 LLM OOM 해결 - 4bit+Offload

로컬 GPU로 LLM을 띄우다 보면 가장 먼저 부딪히는 벽이 OOM(Out Of Memory)입니다. 특히 7B~13B급 모델을 기본 설정(예: fp16)으로 로드하면, 추론 자체는 시작도 못 하고 로딩 단계에서 터지는 경우가 많습니다.

이 글에서는 Hugging Face Transformers 기준으로 4bit 양자화(bitsandbytes) + offload(accelerate) 조합으로 OOM을 해결하는 방법을 정리합니다. 단순히 “양자화하면 된다”가 아니라, 실제로 자주 겪는 실패 지점(로드는 되는데 생성에서 터짐, offload가 느려서 타임아웃, 디스크가 꽉 차서 캐시 실패 등)까지 함께 다룹니다.

관련해서 추론 속도까지 끌어올리는 최적화는 별도로 정리해 둔 글이 있으니, OOM을 해결한 뒤 성능을 더 올리고 싶다면 함께 읽는 것을 권합니다: 로컬 LLM 속도 2배 - KV 캐시·FlashAttention2

OOM이 나는 지점 3가지

Transformers 로컬 추론에서 메모리는 보통 아래 3군데에서 터집니다.

1) 모델 가중치 로딩 단계

가장 흔한 케이스입니다. 예를 들어 7B 모델을 fp16으로 올리면 가중치만 대략 14GB 수준(대략치)이고, 여기에 레이어 버퍼/오버헤드가 붙어 16GB VRAM에서도 아슬아슬하거나 실패합니다.

2) KV 캐시(생성 중) 단계

로드는 되었는데 generate()에서 OOM이 나는 경우입니다. 프롬프트 길이와 max_new_tokens가 늘어날수록 KV 캐시가 커집니다. 즉, “모델은 올라가는데 생성이 안 된다”면 KV 캐시가 원인일 가능성이 큽니다.

3) 파편화/캐시/동시성

같은 VRAM 용량이라도, 이미 다른 프로세스가 GPU를 점유 중이거나 파편화가 심하면 OOM이 빨리 납니다. 또한 torch.compile 같은 최적화가 오히려 초기 메모리 피크를 키워 실패하는 경우도 있습니다.

해결 전략 개요: 4bit + Offload

OOM을 해결하는 가장 실전적인 조합은 다음입니다.

• 4bit 양자화: 모델 가중치를 16bit 대신 4bit로 로드해 VRAM 사용량을 크게 줄입니다.
• Offload: GPU에 다 올리기 어려운 레이어를 CPU 메모리(또는 디스크)로 보내고, 필요할 때만 GPU로 가져옵니다.

이 조합의 장점은 “작은 GPU에서도 큰 모델을 띄울 수 있다”는 점이고, 단점은 “속도가 느려질 수 있다”는 점입니다. 따라서 목표는 보통 다음처럼 잡습니다.

• 목표 1: 일단 OOM 없이 로딩/생성 성공
• 목표 2: offload 비율을 줄이거나, 더 빠른 커널/캐시 최적화로 속도 회복

환경 준비(필수 패키지)

아래 패키지 조합이 가장 흔합니다.

• transformers
• accelerate
• bitsandbytes
• torch

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

GPU/드라이버/CUDA 조합에 따라 bitsandbytes가 설치는 되지만 로딩 시 실패하는 경우가 있으니, 문제가 생기면 먼저 python -c "import bitsandbytes as bnb; print(bnb.__version__)" 같은 식으로 import 테스트를 권합니다.

4bit 양자화로 로드하기(가장 먼저 시도)

bitsandbytes 기반 4bit 양자화는 BitsAndBytesConfig로 설정합니다. 핵심은 load_in_4bit=True와 bnb_4bit_quant_type(보통 nf4)입니다.

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,  # Ampere 이상이면 bfloat16 선호
)

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 = "한국어로 4bit 양자화의 장단점을 5줄로 요약해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

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

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

여기서도 OOM이 난다면?

• device_map="auto"가 “가능한 한 GPU에 많이 올리려는” 경향이 있어서, 카드에 따라 로딩 피크가 여전히 높을 수 있습니다.
• 이 경우 명시적 offload를 섞어야 합니다.

Offload로 VRAM 한계를 넘기기

Transformers는 accelerate와 함께 device_map을 통해 레이어를 CPU로 보낼 수 있습니다. 또한 CPU로만 보내는 게 아니라, CPU 메모리도 부족하면 디스크(offload 폴더)를 활용할 수 있습니다.

아래는 “GPU에 다 못 올리는 일부를 CPU로 offload”하는 전형적인 패턴입니다.

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

offload_dir = "./offload"
os.makedirs(offload_dir, exist_ok=True)

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)

# max_memory를 명시하면 "auto"가 더 보수적으로 배치합니다.
# 예: GPU 12GB 카드라면 10~11GiB 정도로 제한을 걸어 피크를 낮춥니다.
max_memory = {
    0: "11GiB",
    "cpu": "48GiB",
}

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    max_memory=max_memory,
    offload_folder=offload_dir,
    offload_state_dict=True,
)

prompt = "13B 모델을 12GB GPU에서 굴릴 때 offload가 필요한 이유를 설명해줘."
inputs = tokenizer(prompt, return_tensors="pt")

# 입력은 모델의 첫 디바이스로 보내는 것이 안전합니다.
# device_map이 섞여 있으면 model.device가 단일 값이 아닐 수 있어,
# 보통은 아래처럼 첫 파라미터의 device를 사용합니다.
first_device = next(model.parameters()).device
inputs = {k: v.to(first_device) for k, v in inputs.items()}

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

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

offload_folder를 쓸 때 주의점

• offload는 디스크 I/O에 민감합니다. 가능하면 NVMe SSD를 권장합니다.
• 디스크 공간이 부족하면 모델 로딩이 실패하거나, 캐시가 깨져서 재다운로드가 반복됩니다.
• 로컬 개발 환경에서 디스크가 꽉 찬 상태로 모델 캐시까지 겹치면 문제가 커지므로, 디스크 사용량/inode 상태도 함께 점검하는 습관이 좋습니다. (관련 문제 해결 흐름은 이 글이 도움이 됩니다: 리눅스 디스크 100%? inode 고갈 진단·복구 실전)

“로드는 되는데 generate에서 OOM” 해결 체크리스트

4bit+offload로 로딩은 성공했는데 생성에서 OOM이 나면, 대부분 KV 캐시가 원인입니다. 아래 순서대로 줄여보면 빠르게 안정화됩니다.

1. max_new_tokens 줄이기

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

2. 입력 프롬프트 길이 줄이기(특히 RAG로 긴 컨텍스트를 넣는 경우)

3. 배치 크기 1로 고정하기(동시 요청 처리 중이면 큐잉)

4. 샘플링 옵션이 메모리를 크게 쓰진 않지만, 실험 중이라면 단순화

out = model.generate(
    **inputs,
    do_sample=False,
    num_beams=1,
)

5. 가능하면 flash-attn과 같은 커널 최적화는 “OOM 해결 후” 적용

속도 최적화는 메모리 패턴을 바꿔 초기 피크를 키울 수 있어, 먼저 안정화한 다음 단계적으로 켜는 편이 안전합니다. 속도 튜닝은 앞서 언급한 글(로컬 LLM 속도 2배 - KV 캐시·FlashAttention2)에 더 자세히 정리했습니다.

흔한 에러/증상별 처방

CUDA out of memory가 로딩 중 발생

• 4bit 적용 여부 확인: BitsAndBytesConfig(load_in_4bit=True)
• device_map="auto" + max_memory로 GPU 상한을 낮춰 피크를 줄이기
• 그래도 안 되면 더 공격적으로 CPU offload(또는 더 작은 모델)로 전환

로딩은 되는데 응답이 지나치게 느림

• offload 비율이 높을수록 느립니다. 가능한 만큼 GPU에 더 올리도록 max_memory를 재조정합니다.
• offload 디렉터리를 HDD가 아니라 SSD/NVMe로 옮깁니다.
• 토크나이저가 느리면 use_fast=True 확인

디스크가 급격히 차고 캐시 문제가 생김

• 모델 캐시(~/.cache/huggingface)와 offload 폴더가 동시에 커질 수 있습니다.
• 특히 컨테이너 환경에서 ephemeral storage가 작으면 “디스크 압박” 형태로 장애가 납니다. 쿠버네티스에서 비슷한 류의 문제를 다룬 글도 참고할 만합니다: EKS Pod Evicted(디스크 압박) 원인·해결 가이드

실전 권장 조합(요약)

• VRAM이 넉넉하지 않다: 4bit는 기본값으로 둔다
• 그래도 OOM: device_map="auto"에 max_memory를 주고, offload_folder로 디스크 offload까지 허용
• 로드는 되는데 generate OOM: max_new_tokens와 입력 길이를 줄여 KV 캐시를 먼저 안정화
• 속도가 너무 느리다: offload 비율을 줄이고, 이후 FlashAttention2/KV 캐시 최적화로 회복

마무리: OOM은 “용량”이 아니라 “피크” 문제다

로컬 LLM OOM은 단순히 “VRAM이 부족해서”라기보다, 로딩/생성 순간의 메모리 피크와 KV 캐시 증가가 겹치면서 발생하는 경우가 많습니다. 그래서 해결도 “무조건 더 큰 GPU”가 아니라,

• 4bit로 가중치 자체를 줄이고
• offload로 피크를 분산시키고
• 생성 파라미터로 KV 캐시 폭증을 통제

하는 식으로 접근하면, 작은 GPU에서도 꽤 큰 모델을 실용적으로 운용할 수 있습니다.

다음 단계로는 “OOM 없이 띄운 모델을 더 빠르게” 만드는 것이 남습니다. 위에서 연결한 속도 최적화 글을 기반으로, offload 비율을 줄이면서 FlashAttention2 같은 커널 최적화를 단계적으로 적용해 보세요.