Transformers 로컬 LLM 로딩 OOM 해결 8가지

서버나 로컬 워크스테이션에서 transformers로 LLM을 올리다 보면, 모델 파일 다운로드보다 더 자주 막히는 게 CUDA out of memory 입니다. 특히 7B~13B급을 기본 설정 그대로 from_pretrained로 올리거나, 로딩 직후 바로 generate를 돌리면 VRAM이 순식간에 고갈됩니다.

이 글은 “왜 OOM이 나는지”를 메모리 구성요소(가중치, KV 캐시, 활성화, 프래그먼테이션)로 쪼개고, 로컬 로딩 단계에서 바로 적용 가능한 8가지 처방을 정리합니다. 목표는 단순히 실행되게 만드는 것이 아니라, 재현 가능한 방식으로 “메모리 예산을 설계”하는 것입니다.

관련해서 추론 품질을 올리기 위한 프롬프트 설계는 별개 축이므로, 필요하면 Chain-of-Thought 없이 추론력 올리는 5가지 프롬프트도 함께 참고하면 좋습니다.

OOM을 만드는 4가지 덩어리

로컬 LLM 추론에서 GPU 메모리는 대략 아래로 소비됩니다.

1. 모델 가중치(Weights): 파라미터 수에 비례. dtype에 따라 2배, 4배 차이.
2. KV 캐시: batch_size * seq_len * num_layers * hidden_size에 비례. 긴 컨텍스트에서 폭증.
3. 활성화(Activations): 주로 학습에서 크지만, 추론에서도 일부 연산 버퍼가 생김.
4. 메모리 프래그먼테이션: “총량은 남아있는데 연속 공간이 없어” 실패하는 케이스.

따라서 해결책도 “가중치 줄이기”, “KV 캐시 줄이기”, “버퍼/프래그먼테이션 줄이기”, “GPU 밖으로 빼기”로 분류됩니다.

1) dtype 낮추기: bfloat16 또는 float16

가장 먼저 해야 할 일입니다. FP32로 로딩하면 7B도 VRAM이 빠듯해집니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

• Ampere 이상이면 bfloat16이 안정적인 경우가 많습니다.
• float16은 더 보편적이지만, 일부 연산에서 오버/언더플로로 품질 저하나 NaN이 날 수 있습니다.

체감: FP32 대비 대략 절반 VRAM 절감.

2) 8bit 또는 4bit 로딩: bitsandbytes 양자화

가중치 메모리를 직접적으로 크게 줄입니다. 로컬 환경에서 “일단 돌아가게” 만드는 최우선 카드입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

팁:

• 4bit는 VRAM을 크게 아끼지만, 일부 모델/환경에서 속도가 느려지거나 품질이 소폭 흔들릴 수 있습니다.
• device_map="auto"는 레이어를 분산 배치하지만, “GPU에 다 못 올릴 때” CPU 오프로딩까지 섞이면서 속도가 확 떨어질 수 있습니다. 그래도 OOM 방지가 우선이라면 유효합니다.

3) device_map과 오프로딩으로 VRAM 예산 맞추기

OOM을 피하려고 무작정 양자화만 하는 것보다, “어디에 올릴지”를 설계하면 더 안정적입니다.

from transformers import AutoModelForCausalLM

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

• offload_folder는 CPU/RAM 또는 디스크로 일부 텐서를 내릴 때 사용됩니다.
• 디스크 오프로딩은 매우 느릴 수 있으니, 가능하면 RAM에 여유가 있는 머신에서 사용하세요.

현장에서 흔한 실패 패턴은 “로딩은 되는데 첫 generate에서 OOM”입니다. 이 경우는 가중치가 아니라 KV 캐시가 원인일 확률이 큽니다. 다음 항목을 같이 적용해야 합니다.

4) 컨텍스트 길이와 max_new_tokens를 줄여 KV 캐시 폭발 막기

KV 캐시는 길이가 늘수록 선형으로 커집니다. 특히 “긴 시스템 프롬프트 + 긴 대화 히스토리 + 큰 max_new_tokens” 조합은 1회 요청만으로도 OOM을 만들 수 있습니다.

import torch

prompt = "요약: ..."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

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

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

실전 체크리스트:

• 로컬 테스트에서 max_new_tokens는 128~256부터 시작
• 대화형이면 과거 메시지를 무한히 누적하지 말고, 요약/슬라이딩 윈도우 적용
• 필요 이상으로 큰 max_length를 고정값으로 두지 않기

추론 품질이 필요해서 프롬프트를 길게 가져가야 한다면, 프롬프트 자체를 “짧고 강하게” 만드는 방향도 도움이 됩니다. 위에서 링크한 프롬프트 글이 그 축입니다.

5) Flash Attention 또는 SDPA로 어텐션 메모리 줄이기

PyTorch의 SDPA 또는 Flash Attention이 가능한 환경이면 어텐션의 메모리 사용량과 속도를 동시에 개선할 수 있습니다.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="cuda",
    attn_implementation="sdpa",  # 환경에 따라 "flash_attention_2"
)

주의:

• 모델/버전 조합에 따라 attn_implementation 지원 여부가 다릅니다.
• Flash Attention 2는 별도 설치 및 CUDA 호환성이 필요할 수 있습니다.

이 최적화는 “가중치 OOM”보다는 “긴 컨텍스트에서 터지는 OOM”에 특히 효과적입니다.

6) torch.inference_mode()와 불필요한 그래프/버퍼 제거

추론인데도 no_grad를 안 쓰거나, 실수로 model.train() 상태로 돌리면 메모리 사용량이 증가합니다. 또한 일부 환경에서 그래프가 남아 프래그먼테이션을 악화시키기도 합니다.

import torch

model.eval()

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

추가로, 반복 호출 루프에서 텐서를 리스트에 계속 쌓는 코드가 있으면 GPU 메모리가 “천천히” 증가하다가 터집니다. 디코딩 결과만 CPU 문자열로 보관하고, 중간 텐서는 유지하지 않도록 구조를 점검하세요.

7) 메모리 프래그먼테이션 완화: allocator 설정과 캐시 정리

nvidia-smi로 보면 여유가 있는데도 OOM이 나는 케이스가 있습니다. 이는 연속 메모리 블록 할당 실패(프래그먼테이션)일 수 있습니다.

실전에서는 아래 두 가지를 자주 씁니다.

7-1) PyTorch CUDA allocator 설정

프로세스 시작 전에 환경변수를 설정해야 합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,expandable_segments:True"
python run_infer.py

• max_split_size_mb는 큰 블록을 쪼개는 정책에 영향을 줍니다.
• expandable_segments는 특정 상황에서 단편화를 줄이는 데 도움이 됩니다.

7-2) 캐시 정리(응급처치)

import torch

torch.cuda.empty_cache()

이건 “진짜 사용 중인 텐서”를 지우는 게 아니라 캐시를 비우는 수준이라 근본 해결은 아닙니다. 하지만 실험을 반복하며 VRAM이 들쭉날쭉해진 상황에서는 일시적으로 도움이 됩니다.

8) 로딩 자체를 줄이기: safetensors, 체크포인트/샤드, 불필요한 모듈 배제

OOM이 “로딩 중”에 터지는 경우는, 모델이 GPU로 올라가는 순간뿐 아니라 CPU에서 GPU로 옮기는 과정의 피크 메모리(특히 여러 복사본) 때문에 발생하기도 합니다.

8-1) safetensors 사용

가능하면 safetensors 포맷을 쓰는 모델을 선택하세요. 로딩 안정성과 속도에서 이점이 있습니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    use_safetensors=True,
    device_map="auto",
    torch_dtype="auto",
)

8-2) 불필요한 기능 비활성화 및 저장소 정리

• 학습용 헤드나 어댑터가 섞인 체크포인트는 구성에 따라 메모리 피크가 커질 수 있습니다.
• PEFT 어댑터를 붙이는 경우에도, base 모델 로딩을 4bit로 줄이고 어댑터만 추가하는 식으로 예산을 잡는 게 좋습니다.

# 예: base는 4bit로, 어댑터만 추가 로딩하는 구조(개념 예시)
# 실제 코드는 사용하는 PEFT 방식에 맞춰 조정 필요

재현 가능한 “OOM 디버깅 루틴”

아래 순서로 보면 원인을 빨리 좁힐 수 있습니다.

1. 로딩 단계에서 OOM인지, generate 단계에서 OOM인지 구분
2. 로딩 OOM이면: dtype 낮추기, 4bit/8bit, device_map 오프로딩, safetensors
3. generate OOM이면: max_new_tokens와 입력 길이 축소, SDPA/Flash Attention, 배치 크기 축소
4. “여유가 있는데 OOM”이면: allocator 설정으로 프래그먼테이션 완화

간단한 메모리 로깅도 효과적입니다.

import torch

def report(tag: str):
    if not torch.cuda.is_available():
        return
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    print(f"[{tag}] allocated={alloc:.1f}MB reserved={reserved:.1f}MB")

report("start")
# model load
report("after_load")
# generate
report("after_generate")

마무리: 가장 효과가 큰 조합

로컬 LLM OOM 해결에서 “효과 대비 노력”이 큰 조합은 보통 아래입니다.

• torch_dtype=bfloat16 또는 4bit 양자화
• device_map="auto"로 분산 및 필요 시 오프로딩
• max_new_tokens와 입력 컨텍스트 길이 제한으로 KV 캐시 제어
• attn_implementation="sdpa" 적용 가능 여부 확인

여기까지 적용하면, 같은 GPU에서도 “모델을 올릴 수 있느냐”뿐 아니라 “어느 정도 컨텍스트로 얼마나 안정적으로 돌릴 수 있느냐”까지 예측 가능해집니다. 결국 OOM은 운이 아니라 예산 문제이고, 예산은 위 8가지 레버로 충분히 통제할 수 있습니다.