beautifulsoups.dev
Published on

Transformers 로컬 LLM 로드 VRAM OOM 7가지

Authors

로컬 GPU에서 transformers로 LLM을 로드할 때 가장 흔한 실패가 CUDA out of memory입니다. 문제는 “모델이 커서”만이 아니라, 로드 방식·dtype·KV 캐시·디바이스 매핑·추론 파라미터·파편화 같은 복합 요인으로 VRAM이 예상보다 빨리 소진된다는 점입니다.

이 글은 Transformers 기반 로컬 LLM 로드 시 VRAM OOM을 만드는 7가지 원인과, 각 원인에 대한 즉시 적용 가능한 해결 코드를 한 번에 정리합니다.

참고: 원인 진단을 할 때는 “모델 로드(가중치) VRAM”과 “추론 중(활성화/캐시) VRAM”을 분리해서 봐야 합니다. 로드는 되는데 첫 토큰에서 터지면 KV 캐시·배치·컨텍스트가 원인일 확률이 큽니다.

OOM 빠른 진단 체크리스트

아래 3가지만 먼저 확인해도 절반은 해결됩니다.

  1. dtype이 float32로 로드되고 있지 않은가
  2. device_map이 의도대로 GPU/CPU에 분산되고 있는가
  3. max_new_tokens, max_length, batch_size, num_beams가 과도하지 않은가

디버깅용으로는 아래처럼 현재 VRAM 사용량을 자주 찍어보는 습관이 좋습니다.

import torch

def vram(tag=""):
    if torch.cuda.is_available():
        alloc = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"[{tag}] allocated={alloc:.2f}GB reserved={reserved:.2f}GB")

vram("start")

1) dtype 기본값이 float32로 로드됨

가장 흔한 실수입니다. 특히 torch_dtype를 지정하지 않거나, 모델 설정이 FP32로 고정된 경우 VRAM이 2배 이상 튈 수 있습니다.

해결

  • Ampere 이상이면 보통 bfloat16이 안정적입니다.
  • 그 외에는 float16을 우선 시도합니다.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

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

tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,  # 또는 torch.float16
    device_map="auto",
)

추가로, 추론 시에도 autocast를 쓰면 일부 연산이 더 안전하게 낮은 precision으로 내려갑니다.

import torch

with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    out = model.generate(**tok("Hello", return_tensors="pt").to("cuda"), max_new_tokens=64)

2) device_map 미설정 또는 잘못된 디바이스 매핑

device_map을 지정하지 않으면 기본적으로 “통째로 GPU 0”에 올라가려 하면서 OOM이 납니다. 또는 device_map="auto"를 썼는데도 특정 레이어가 GPU에 과하게 몰릴 수 있습니다.

해결 A: 기본은 device_map="auto"

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

해결 B: 일부를 CPU로 오프로딩

GPU VRAM이 애매할 때는 CPU 오프로딩이 현실적인 타협입니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    offload_folder="./offload",
)

해결 C: max_memory로 상한을 명시

max_memory = {
    0: "20GiB",  # GPU 0
    "cpu": "64GiB"
}

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    max_memory=max_memory,
    offload_folder="./offload",
)

3) 4bit/8bit 양자화 로드 미사용(또는 잘못 사용)

로컬에서 가장 강력한 OOM 해법은 양자화입니다. 다만 “양자화했다고 생각했는데 실제로는 FP16/FP32로 올라간” 케이스가 많습니다. 예를 들어 bitsandbytes 미설치, 잘못된 파라미터, 혹은 특정 모델이 양자화 로드를 지원하지 않는 경우입니다.

해결: 4bit 로드(QLoRA 스타일)

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

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

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

양자화는 가중치 VRAM을 크게 줄이지만, KV 캐시/활성화 메모리는 별개로 남습니다. 즉, “로드는 되는데 생성에서 터짐”은 다음 원인(4, 5)을 같이 봐야 합니다.

4) 컨텍스트 길이 증가로 KV 캐시가 폭증

요즘 모델은 context_length가 8k, 32k, 128k까지도 갑니다. 하지만 컨텍스트가 길어질수록 KV 캐시가 선형적으로 증가하고, 레이어 수가 많을수록 그 비용이 커집니다.

특히 다음 조합은 OOM을 매우 쉽게 만듭니다.

  • 긴 입력(input_ids 길이 큼)
  • 큰 max_new_tokens
  • 배치가 2 이상
  • num_beams 사용

해결 A: 입력을 줄이고 max_new_tokens를 보수적으로

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

out = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
)

해결 B: 빔서치 피하기(또는 빔 수 축소)

num_beams는 KV 캐시를 사실상 빔 수만큼 늘리는 방향으로 작동해 OOM을 유발하기 쉽습니다.

out = model.generate(
    **inputs,
    max_new_tokens=128,
    num_beams=1,  # OOM 시 1로
)

해결 C: 배치 추론 시 마이크로배치

# pseudo: prompts를 작은 chunk로 나눠 순차 처리
for chunk in chunks(prompts, size=1):
    inputs = tok(chunk, return_tensors="pt", padding=True, truncation=True).to(model.device)
    out = model.generate(**inputs, max_new_tokens=128)

5) use_cache=True로 인한 메모리 상주(특히 긴 생성)

기본적으로 causal LM 생성은 속도를 위해 KV 캐시를 사용합니다(use_cache=True). 긴 문장을 생성할수록 캐시가 계속 쌓이니 VRAM이 부족한 환경에서는 OOM이 날 수 있습니다.

해결: 캐시 비활성화(속도는 느려짐)

model.config.use_cache = False

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

현실적으로는 캐시를 끄기보다, 4번의 컨텍스트/토큰/배치/빔을 줄이는 방식이 더 흔한 정답입니다. 다만 “무조건 OOM을 피해야 한다”면 캐시 비활성화가 마지막 카드가 될 수 있습니다.

6) PyTorch CUDA 메모리 파편화 및 reserved 폭증

에러 메시지에서 allocated는 낮은데 reserved가 비정상적으로 높거나, 같은 코드를 반복 실행할 때 점점 OOM이 빨라지면 파편화를 의심해야 합니다.

해결 A: 프로세스 재시작이 가장 확실

노트북/서버에서 반복 실험 시, 한 번 꼬이면 깔끔하게 프로세스를 재시작하는 게 가장 효과적입니다.

해결 B: PYTORCH_CUDA_ALLOC_CONF로 분할 정책 조정

환경 변수 설정 예시입니다. max_split_size_mb를 조정하면 큰 블록 분할로 인한 파편화를 완화하는 데 도움이 될 때가 있습니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"

해결 C: 불필요 텐서 참조 제거 후 캐시 정리

import gc, torch

del inputs
# del model  # 모델까지 내릴 때만

gc.collect()
torch.cuda.empty_cache()

캐시 정리는 “이미 파편화된 reserved를 완전히 되돌리는” 만능은 아니지만, 반복 실험 중 일시적으로 숨통을 틔우는 데는 도움이 됩니다.

7) 학습/미세조정 코드가 섞여 그래디언트·옵티마이저 상태가 VRAM을 잡아먹음

로컬에서 추론만 한다고 생각했는데, 코드 어딘가에 model.train()이 있거나, torch.set_grad_enabled(True) 상태로 돌아가거나, 심지어 옵티마이저를 만들면서 상태 텐서가 GPU에 올라가면 VRAM이 급격히 부족해집니다.

해결: 추론 모드 고정

import torch

model.eval()

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

해결: 미세조정이라면 메모리 절약 옵션 사용

학습이 목적이라면 아래가 정석입니다.

  • gradient checkpointing
  • LoRA/QLoRA
  • optimizer를 8bit로
model.gradient_checkpointing_enable()
model.config.use_cache = False  # 체크포인팅과 같이 쓰는 경우가 많음

실전: OOM을 피하는 “안전한” 로드 템플릿

아래 템플릿은 로컬에서 실패 확률을 크게 낮추는 기본 조합입니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

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

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

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

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

prompt = "요약: Transformers에서 VRAM OOM을 줄이는 방법은?"
inputs = tok(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)

model.eval()
with torch.inference_mode(), torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        num_beams=1,
    )

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

운영 관점 팁: 재현 가능한 실험 환경 만들기

OOM은 “어떤 커밋에서, 어떤 드라이버/쿠다/파이토치 조합에서” 갑자기 발생하는 경우가 많습니다. 로컬 실험을 Docker로 고정하면 재현성이 좋아집니다. 특히 멀티 GPU/서버를 오가며 테스트한다면 이미지 빌드 체계가 중요합니다.

마무리: OOM을 “모델 크기 문제”로만 보지 말 것

정리하면 VRAM OOM은 보통 아래 순서로 해결됩니다.

  1. dtype을 bfloat16 또는 float16로 고정
  2. device_map="auto"와 max_memory로 분산/상한 설정
  3. 4bit/8bit 양자화 적용
  4. 컨텍스트/생성 토큰/배치/빔을 줄여 KV 캐시 폭증 억제
  5. 파편화가 의심되면 프로세스 재시작 + 할당 정책 조정
  6. 추론 모드(inference_mode)로 그래디언트 차단

원하는 GPU VRAM 용량(예: 8GB, 12GB, 24GB)과 모델 크기(예: 7B, 13B, 34B), 목표 컨텍스트 길이를 알려주면, 해당 조건에 맞춘 현실적인 로드 구성(양자화/오프로딩/토큰 제한)을 더 구체적으로 제안할 수 있습니다.

Discuss on Twitter • View on GitHub