beautifulsoups.dev
Published on

Transformers 로컬 LLM CUDA OOM 줄이는 9가지

Authors

로컬에서 LLM을 transformers로 실행하다 보면 가장 먼저 부딪히는 문제가 CUDA out of memory입니다. 특히 7B~13B급 모델을 단일 GPU에서 돌리거나, 배치/컨텍스트를 조금만 키워도 OOM이 쉽게 발생합니다.

OOM은 단순히 “VRAM이 부족하다”로 끝나지 않습니다. 같은 모델·같은 GPU에서도 dtype, KV cache, max_new_tokens, device_map, allocator 설정에 따라 성공/실패가 갈립니다. 이 글에서는 로컬 LLM 추론 기준으로, 바로 적용 가능한 CUDA OOM 줄이는 9가지를 정리합니다.

운영 환경에서 컨테이너가 메모리/리소스 문제로 재시작되는 패턴과도 유사합니다. 쿠버네티스에서 원인 추적이 필요하다면 Kubernetes CrashLoopBackOff 10가지 원인과 15분 진단도 함께 참고하면 좋습니다.

1) FP16/BF16로 로드하고 FP32를 피하기

가장 기본이지만 효과가 큽니다. 모델을 FP32로 로드하면 파라미터만으로도 VRAM이 급격히 늘어납니다. GPU가 BF16을 잘 지원하면 BF16이 안정적인 경우가 많고, 그렇지 않으면 FP16을 사용합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

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

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    device_map="cuda",
)

체크 포인트

  • BF16 사용 시: GPU 아키텍처에 따라 성능/안정성이 달라집니다.
  • FP16 사용 시: 일부 모델/연산에서 nan이 생기면 BF16 또는 flash-attn 계열 최적화를 고려하세요.

2) 8bit/4bit 양자화로 파라미터 메모리 줄이기

VRAM이 빠듯하다면 양자화가 가장 강력합니다. bitsandbytes를 이용한 8bit/4bit 로딩은 로컬 추론에서 사실상 표준이 됐습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

tokenizer = 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,
    quantization_config=bnb_config,
    device_map="auto",
)

트레이드오프

  • 4bit는 VRAM 절감이 크지만, 모델/프롬프트에 따라 품질 저하가 있을 수 있습니다.
  • compute_dtype를 BF16으로 두면 속도와 안정성에 도움이 되는 경우가 많습니다.

3) max_new_tokens와 max_length를 보수적으로 잡기

OOM의 많은 경우는 “모델 로드”가 아니라 생성 중에 터집니다. 이유는 KV cache가 토큰 수에 비례해 커지기 때문입니다.

  • 입력 길이 input_tokens
  • 생성 길이 max_new_tokens
  • 합계가 길어질수록 KV cache가 선형으로 증가
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

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

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
)

prompt = "요약해줘: ..."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,  # 1024 같은 큰 값은 OOM을 유발하기 쉬움
        do_sample=False,
        use_cache=True,
    )

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

실전 팁

  • 먼저 max_new_tokens=128로 안정 구간을 찾고 점진적으로 늘리세요.
  • 긴 문맥이 필요하면, 토큰 예산을 정해 입력 토큰과 출력 토큰을 함께 관리해야 합니다.

4) 배치를 줄이거나, 마이크로배치로 쪼개기

로컬 추론에서 배치를 크게 잡으면 activation 및 KV cache가 같이 늘어 OOM이 빠르게 발생합니다.

  • 동시 요청을 한 번에 처리하려면 배치가 필요하지만, VRAM이 한계라면 요청을 큐잉하거나 마이크로배치로 나누는 편이 안정적입니다.

간단히는 “한 번에 하나씩” 처리하는 것이 가장 안전합니다.

# 예시: 여러 프롬프트를 배치로 넣는 대신 순차 처리
prompts = ["질문1...", "질문2...", "질문3..."]

for p in prompts:
    inputs = tokenizer(p, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        out = model.generate(**inputs, max_new_tokens=128)
    print(tokenizer.decode(out[0], skip_special_tokens=True))
    torch.cuda.empty_cache()

torch.cuda.empty_cache()는 “메모리 사용량을 0으로 만든다”가 아니라 캐시를 반환하는 동작이라 만능은 아니지만, 배치/요청 사이에 피크를 낮추는 데 도움이 되는 경우가 있습니다.

5) device_map="auto"와 오프로딩으로 VRAM 피크 낮추기

단일 GPU VRAM이 부족하면 일부 레이어를 CPU로 오프로딩해 “일단 돌아가게” 만들 수 있습니다. 속도는 느려지지만 OOM을 피하는 현실적인 방법입니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

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

tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype="auto",
    offload_folder="./offload",  # 디스크 오프로딩 경로
)

주의

  • 오프로딩은 스토리지/CPU/RAM 상태에 크게 영향을 받습니다.
  • NVMe가 있으면 훨씬 낫고, 네트워크 디스크는 병목이 심할 수 있습니다.

6) attn_implementation 최적화로 메모리 절약

PyTorch scaled_dot_product_attention 또는 Flash Attention 계열은 attention 메모리를 줄여주는 경우가 많습니다. 모델/환경에 따라 옵션이 다르지만, 지원되는 경우 시도 가치가 큽니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
    attn_implementation="sdpa",  # 가능하면 시도
)

환경에 따라 flash_attention_2 같은 옵션이 가능할 수 있습니다. 다만 이는 설치/빌드 조건이 까다로울 수 있어, 먼저 sdpa부터 확인하는 편이 안전합니다.

7) 메모리 단편화 줄이기: PyTorch allocator 설정

OOM 로그를 자세히 보면 “free는 남아 있는데 할당 실패”처럼 보일 때가 있습니다. 이는 VRAM **단편화(fragmentation)**일 수 있습니다.

가장 흔히 쓰는 완화책 중 하나가 PYTORCH_CUDA_ALLOC_CONF 설정입니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"
python run.py

상황에 따라 64, 128, 256 등으로 조정하면서 재현 테스트를 해보세요.

추가로, 디버깅할 때는 아래 코드로 현재/최대 메모리 사용량을 확인하면 원인 파악이 빨라집니다.

import torch

print("allocated", torch.cuda.memory_allocated() / 1024**2, "MB")
print("reserved", torch.cuda.memory_reserved() / 1024**2, "MB")
print("max allocated", torch.cuda.max_memory_allocated() / 1024**2, "MB")

8) use_cache를 상황에 따라 끄기

생성에서 use_cache=True는 대체로 속도를 크게 올리지만, KV cache를 유지하는 만큼 메모리를 더 씁니다. 아주 긴 생성이나 VRAM이 극단적으로 부족한 상황에서는 use_cache=False가 OOM 회피에 도움이 될 수 있습니다.

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

현실적인 판단

  • VRAM이 조금 부족해서 OOM이라면 max_new_tokens를 줄이는 쪽이 보통 더 낫습니다.
  • 그래도 안 되면 use_cache=False로 “동작 우선” 모드로 전환합니다.

9) 불필요한 그래디언트/학습 모드 제거: inference_mode와 eval

추론인데도 학습 관련 그래프/버퍼가 남아 있으면 메모리 낭비가 발생합니다. 아래 3가지는 기본으로 고정하세요.

  • model.eval()
  • torch.inference_mode() 또는 torch.no_grad()
  • 필요 없다면 output_attentions, output_hidden_states 끄기
model.eval()

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

추가로, 파이프라인/서빙 코드에서 요청마다 모델을 다시 로드하는 실수가 의외로 많습니다. 모델 로드는 프로세스 생명주기에서 1회만 수행하고, 요청은 토크나이징과 generate만 반복하도록 구조를 잡는 것이 안정적입니다.

OOM을 재현 가능하게 줄이는 체크리스트

아래 순서대로 적용하면 “감으로 튜닝”하는 시간을 크게 줄일 수 있습니다.

  1. torch_dtype를 BF16 또는 FP16으로 고정
  2. max_new_tokens를 128부터 시작해 점진적으로 증가
  3. 배치 크기를 1로 낮춰 기준선 확보
  4. 그래도 OOM이면 4bit 양자화 적용
  5. 여전히 부족하면 device_map="auto" + 오프로딩
  6. attn_implementation="sdpa" 시도
  7. 단편화 의심 시 PYTORCH_CUDA_ALLOC_CONF 조정

마무리

CUDA OOM은 단일 원인보다 여러 요인의 합으로 발생합니다. 특히 로컬 LLM은 “모델 파라미터”뿐 아니라 “생성 길이로 커지는 KV cache”가 핵심 변수라서, max_new_tokens와 context length 관리가 가장 먼저입니다. 그 다음이 dtype, 양자화, 디바이스 매핑, 어텐션 최적화, allocator 설정 순으로 효과가 누적됩니다.

성능 최적화가 결국 병목을 줄이는 과정이라는 점에서, 프론트엔드의 체감 성능을 단계적으로 줄이는 접근과도 닮았습니다. 웹 성능 최적화 관점이 궁금하다면 Next.js LCP 4초→1초 - RSC·이미지·폰트 최적화도 함께 보면 “병목을 측정하고 줄이는” 사고방식에 도움이 됩니다.

Discuss on Twitter • View on GitHub