beautifulsoups.dev
Published on

Transformers 로컬 LLM CUDA OOM 해결 10가지

Authors

로컬에서 transformers로 LLM을 띄우다 보면 가장 먼저 만나는 벽이 CUDA out of memory입니다. 문제는 OOM이 단순히 “VRAM이 부족”해서만 발생하지 않는다는 점입니다. 모델 파라미터 외에도 KV 캐시, 활성화(activation), 옵티마이저 상태, 파편화(fragmentation), 그리고 잘못된 로딩/디바이스 매핑까지 다양한 요인이 겹칩니다.

이 글은 “어떻게든 돌아가게”가 아니라, 왜 OOM이 나는지를 빠르게 진단하고 VRAM을 실제로 줄이는 10가지 방법을 정리합니다. 예시는 transformers 기준이지만, accelerate/bitsandbytes/torch 전반에 그대로 적용됩니다.

참고로, 리소스 누수 관점에서의 접근은 다른 문제(예: 파일 핸들, 세션, 제너레이터)에서도 통합니다. 비슷한 사고방식이 필요하다면 데코레이터+컨텍스트 매니저로 리소스 누수 0, Python Generator close·throw로 누수 잡기도 함께 보면 도움이 됩니다.

0. 먼저: OOM을 “원인별로” 분류하기

OOM이 났을 때 로그에 흔히 보이는 패턴은 다음과 같습니다.

  • allocated는 낮은데 reserved가 큰 경우: 캐싱/파편화/allocator 이슈 가능성
  • 첫 토큰 생성 전부터 OOM: 로딩/가중치/디바이스 매핑 문제
  • 프롬프트가 길어질수록 OOM: KV 캐시가 원인인 경우가 많음
  • 배치가 커질수록 OOM: 활성화/attention 메모리 증가

아래 코드를 습관처럼 붙여서 “지금 VRAM이 어디로 갔는지”를 확인하세요.

import torch

def cuda_mem(prefix: str = ""):
    if not torch.cuda.is_available():
        print(prefix, "CUDA not available")
        return
    free, total = torch.cuda.mem_get_info()
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    print(
        f"{prefix} free={free/1e9:.2f}GB total={total/1e9:.2f}GB "
        f"allocated={allocated/1e9:.2f}GB reserved={reserved/1e9:.2f}GB"
    )

cuda_mem("before")
# load / generate ...
cuda_mem("after")

allocated는 실제 사용 중인 텐서, reserved는 PyTorch가 잡아둔 캐시 영역입니다. OOM을 줄이려면 둘 다 관찰해야 합니다.

1) 가장 효과적: 4bit/8bit 양자화로 가중치 VRAM을 줄이기

모델 파라미터가 VRAM의 대부분을 차지한다면, 양자화가 1순위입니다. bitsandbytes 기반으로 4bit 로딩 시, 7B급 모델을 8GB~12GB에서도 현실적으로 다루는 경우가 많습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

주의점:

  • GPU가 bfloat16을 잘 지원하지 않으면 torch.float16으로 바꾸는 게 안전합니다.
  • 4bit는 속도가 느려질 수 있습니다. 하지만 “돌아가냐/못 돌아가냐” 단계에서는 최고의 카드입니다.

2) device_map="auto" + 오프로딩으로 VRAM 상한을 강제하기

단일 GPU에 다 못 올리는 모델이라면, accelerate 방식의 디바이스 매핑과 오프로딩이 강력합니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

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

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

팁:

  • NVMe가 빠르면 CPU/NVMe 오프로딩의 체감이 좋아집니다.
  • 오프로딩은 “VRAM OOM”은 줄이지만, “속도”는 크게 희생될 수 있습니다.

3) max_new_tokens와 입력 길이 제한으로 KV 캐시 폭발 막기

생성형 모델에서 OOM의 숨은 주범은 KV 캐시입니다. 대략적으로 KV 캐시는 레이어 수, 헤드 수, hidden size, 그리고 무엇보다 시퀀스 길이에 비례해 커집니다.

즉, 프롬프트가 길거나 max_new_tokens가 크면 VRAM이 계속 늘어납니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "microsoft/Phi-3-mini-4k-instruct"

tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda")

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

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        use_cache=True,
    )
print(tok.decode(out[0], skip_special_tokens=True))

실무적으로는 다음 조합이 자주 통합니다.

  • max_length를 명시적으로 제한
  • max_new_tokens를 작게 시작해서 점진적으로 확대
  • 긴 문맥이 필요하면 RAG로 외부 컨텍스트를 “요약해서” 넣기

4) 배치 크기와 num_beams를 줄여 attention 메모리 급증 방지

OOM은 종종 “배치 크기”가 아니라 “사실상 배치처럼 동작하는 옵션” 때문에 납니다.

  • num_beams는 빔 개수만큼 후보 시퀀스를 유지하므로 메모리 사용이 커집니다.
  • num_return_sequences도 비슷한 효과를 냅니다.
with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        num_beams=1,              # 4, 8 같은 값이면 OOM 위험 증가
        num_return_sequences=1,
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
    )

빔서치가 꼭 필요하면:

  • num_beams를 2부터 시작
  • max_new_tokens를 줄이고
  • 양자화 또는 오프로딩과 같이 사용

5) torch_dtype를 강제해서 FP32 로딩 사고를 막기

의외로 자주 있는 실수는 FP16/BF16 모델을 기대했는데 FP32로 올라간 경우입니다. FP32는 파라미터만 2배 VRAM을 먹습니다.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype=torch.float16,   # 또는 torch.bfloat16
)

추가로, 일부 체크포인트는 torch_dtype="auto"가 안전하지만, “확실히 줄이려면” 명시가 낫습니다.

6) inference_mode와 no_grad로 그래프/활성화 메모리 제거

추론인데도 학습 그래프가 잡히면 활성화 텐서가 남아 OOM이 빨리 옵니다. 추론은 아래 패턴을 기본값으로 두세요.

import torch

model.eval()

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

torch.inference_mode()는 torch.no_grad()보다 더 공격적으로 autograd 관련 오버헤드를 줄입니다.

7) Flash Attention 또는 SDPA로 attention 메모리 최적화

PyTorch 2.x 환경에서는 scaled dot product attention 경로가 최적화되어 VRAM과 속도 모두 개선될 수 있습니다. 모델에 따라 attn_implementation 옵션이 지원됩니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="cuda",
    torch_dtype="auto",
    attn_implementation="sdpa",  # 지원 시
)

주의:

  • 모델/버전에 따라 sdpa가 미지원이거나, 특정 GPU에서 이득이 제한적일 수 있습니다.
  • Flash Attention 계열은 설치/호환성 이슈가 있을 수 있으니, 먼저 sdpa부터 시도하는 편이 안정적입니다.

8) PyTorch allocator 파편화 해결: PYTORCH_CUDA_ALLOC_CONF

reserved는 큰데 allocated는 작은데도 OOM이 나면, 메모리 블록이 잘게 쪼개진 파편화를 의심해야 합니다. 이때는 allocator 설정이 효과적입니다.

실행 전에 환경변수를 설정합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.8"

설명:

  • max_split_size_mb를 낮추면 큰 블록이 과도하게 쪼개지는 것을 완화할 수 있습니다.
  • garbage_collection_threshold는 캐시 회수 타이밍에 영향을 줍니다.

이건 “모델 자체가 너무 큰 문제”를 해결하진 못하지만, 애매하게 OOM이 나는 케이스(특히 여러 번 로드/언로드 반복)에서 체감이 큽니다.

9) 생성 루프에서 텐서/출력을 누적하지 않기 (파이썬 레벨 누수)

대화형 루프를 만들 때, 매 턴의 inputs나 past_key_values를 잘못 들고 있거나, GPU 텐서를 리스트에 쌓아두면 VRAM이 계속 증가합니다.

안티 패턴 예시(하지 말 것):

# 나쁜 예: GPU 텐서를 history에 계속 보관
history = []
for step in range(100):
    out = model.generate(**inputs, max_new_tokens=64)
    history.append(out)  # out은 GPU 텐서일 수 있음

개선 패턴:

history_text = []
for step in range(100):
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=64)

    text = tok.decode(out[0].detach().cpu(), skip_special_tokens=True)
    history_text.append(text)

    # 다음 입력을 만들 때도 GPU 텐서를 오래 붙잡지 않게 주의
    inputs = tok(text[-1000:], return_tensors="pt", truncation=True, max_length=2048).to("cuda")

핵심은 “GPU 텐서는 필요한 순간에만” 들고, 기록은 CPU 문자열로 남기는 것입니다. 리소스가 누적되는 패턴을 구조적으로 끊는 접근은 데코레이터+컨텍스트 매니저로 리소스 누수 0에서 다룬 방식과도 통합니다.

10) 정말 급할 때: 캐시 비우기와 프로세스 격리 전략

10-1. torch.cuda.empty_cache()의 정확한 의미

torch.cuda.empty_cache()는 “할당된 텐서를 지우는” 함수가 아니라, PyTorch가 잡아둔 캐시(reserved)를 드라이버에 반환하는 성격입니다. 즉, 코드가 GPU 텐서를 계속 참조하고 있으면 효과가 없습니다.

import gc
import torch

# 참조를 끊은 뒤
del out
 del inputs
gc.collect()

torch.cuda.empty_cache()

(위 코드에서 del inputs처럼 들여쓰기/공백은 실제로는 한 줄로 작성하세요.)

10-2. 서비스/툴에서는 “프로세스 1회성”이 가장 확실

주피터/REPL/서버처럼 같은 프로세스에서 모델을 여러 번 갈아끼우면 파편화와 잔존 참조로 OOM이 더 잘 납니다. 가장 확실한 방법은 다음입니다.

  • 모델 로딩과 추론을 별도 워커 프로세스로 격리
  • 요청 단위로 워커를 재시작하거나, 일정 횟수 후 롤링 재시작

Kubernetes 환경에서 이런 증상은 결국 재시작 루프로 이어지기도 합니다. 운영에서 “왜 계속 죽는지” 진단하는 관점은 K8s CrashLoopBackOff 원인 10분 진단법과도 맞닿아 있습니다.

체크리스트: OOM이 나면 이 순서로 줄여라

  1. torch_dtype가 FP16/BF16인지 확인하고 강제
  2. max_new_tokens와 입력 max_length를 줄여 KV 캐시부터 잡기
  3. num_beams를 1로, 배치를 1로
  4. 8bit, 그래도 안 되면 4bit 양자화
  5. device_map="auto" + 오프로딩
  6. attn_implementation="sdpa" 등 최적화 경로 적용
  7. PYTORCH_CUDA_ALLOC_CONF로 파편화 완화
  8. 루프에서 GPU 텐서를 누적하지 않는지 점검
  9. 마지막으로 캐시 비우기/프로세스 격리

마무리

CUDA OOM은 단일 처방으로 끝나는 문제가 아니라, 가중치(정적), KV 캐시(시퀀스 의존), 활성화(배치/서치 의존), **파편화/누수(런타임 습관)**를 분리해서 다뤄야 빨리 해결됩니다.

만약 현재 환경의 GPU VRAM, 모델 이름, 목표 컨텍스트 길이(예: 4k, 8k), 그리고 OOM이 나는 시점(로딩 중인지, 첫 generate인지, 몇 토큰 후인지)을 알려주면 위 10가지 중 “가장 비용 대비 효과가 큰 조합”으로 구체적인 설정 값을 추천해 드릴 수 있습니다.

Discuss on Twitter • View on GitHub