beautifulsoups.dev
Published on

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

Authors

로컬 환경에서 Hugging Face Transformers로 LLM을 올리다 보면 가장 먼저 부딪히는 벽이 OOM(Out Of Memory)입니다. 특히 from_pretrained() 한 줄이 GPU VRAM을 순식간에 소진하거나, CPU RAM이 터지면서 프로세스가 죽는 일이 흔합니다.

이 글은 “왜 OOM이 나는지”를 빠르게 분류하고, 바로 적용 가능한 9가지 해결책을 코드와 함께 정리합니다. 대상은 주로 PyTorch + Transformers 기반 로컬 로딩이며, 단일 GPU, 노트북 GPU, 서버 GPU 모두에 적용 가능합니다.

아래 해결책은 대체로 “메모리 사용량 자체를 줄이기”와 “피크 메모리(순간 최대치)를 낮추기” 두 축으로 나뉩니다. OOM은 보통 피크에서 터지므로, 단순히 모델을 더 작은 걸로 바꾸기 전에 피크를 먼저 낮추는 게 효과적입니다.

관련해서 운영체제 레벨에서 OOM이 발생했는지(리눅스 OOM Killer)까지 확인하고 싶다면 아래 글도 함께 참고하세요.

0. 먼저: OOM이 VRAM인지 RAM인지부터 구분

해결책을 적용하기 전에, 터지는 메모리가 GPU인지 CPU인지부터 확인해야 합니다.

  • GPU VRAM OOM: 보통 에러에 CUDA out of memory 문구가 뜹니다.
  • CPU RAM OOM: 파이썬이 Killed로 끝나거나, RuntimeError: DefaultCPUAllocator: can't allocate memory 류가 보입니다.

빠른 체크 코드

import torch

def mem_report(tag=""):
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        print(tag, "GPU free/total (GB):", free/1024**3, total/1024**3)

mem_report("before")
# model load here
mem_report("after")

CPU RAM은 htop, free -h, 컨테이너라면 cgroup limit까지 같이 봐야 합니다.

1) dtype 낮추기: float16 / bfloat16로 로딩

가장 먼저 적용할 수 있고, 효과가 큰 방법입니다. FP32로 로딩하면 메모리를 2배 이상 더 먹습니다.

  • NVIDIA Ampere 이상이면 bfloat16도 안정적입니다.
  • 구형 GPU나 일부 연산에서 float16이 더 호환이 좋을 때도 있습니다.
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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.float16,  # 또는 torch.bfloat16
    device_map="cuda"
)

체감 포인트: “로딩은 되는데 추론에서 OOM”이라면, dtype만 낮춰도 해결되는 경우가 많습니다.

2) device_map="auto" + accelerate로 레이어 분산

단일 GPU에 다 못 올리면 CPU로 일부를 오프로딩해서라도 로딩을 성공시킬 수 있습니다. device_map="auto"는 내부적으로 가능한 배치를 찾습니다.

from transformers import AutoModelForCausalLM

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

CPU 오프로딩을 명시적으로

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16,
    offload_folder="./offload"  # CPU 오프로딩 캐시
)

주의: CPU 오프로딩은 “로딩 성공”에는 도움이 되지만, 추론 latency가 크게 증가할 수 있습니다. 그래도 개발/검증 단계에서 유용합니다.

3) 8bit/4bit 양자화 로딩: bitsandbytes

VRAM이 부족한 로컬 환경에서 가장 강력한 카드입니다. 7B~13B급은 4bit로 현실적으로 돌아가는 경우가 많습니다.

설치(환경에 따라 다름):

pip install -U bitsandbytes accelerate transformers

4bit 로딩 예시

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

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

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

실전 팁:

  • 4bit는 VRAM을 크게 줄이지만, 일부 모델/환경에서 커널 호환성 문제가 날 수 있습니다.
  • 문제가 나면 8bit로 먼저 안정성을 확보한 뒤 4bit로 내려가세요.

4) low_cpu_mem_usage=True로 CPU 피크 줄이기

모델 로딩 과정에서 CPU 메모리가 순간적으로 크게 튀면서 OOM이 나는 경우가 많습니다. 특히 큰 모델을 다운로드한 뒤 “로딩 중”에 죽는 케이스가 여기에 해당합니다.

from transformers import AutoModelForCausalLM

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

왜 효과가 있나: 로딩 과정에서 가중치를 한 번에 크게 복사/변환하지 않도록 경로를 최적화해 CPU RAM 피크를 낮춥니다.

5) Flash Attention / SDPA로 어텐션 메모리 절감

“로딩은 되는데, 첫 토큰 생성에서 OOM”이 나는 경우는 어텐션 KV 캐시나 어텐션 구현이 원인인 경우가 많습니다. PyTorch SDPA 또는 Flash Attention 계열을 쓰면 메모리 효율이 좋아집니다.

PyTorch SDPA 사용(가능한 경우)

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
    attn_implementation="sdpa",  # 지원 모델에서
)

Flash Attention 2(환경 의존)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cuda",
    attn_implementation="flash_attention_2",
)

주의: 모델/Transformers 버전에 따라 attn_implementation 지원 여부가 다릅니다. 지원되지 않으면 에러가 나니 버전 확인이 필요합니다.

6) 생성 파라미터로 KV 캐시 폭발 막기

추론에서 메모리를 가장 빨리 갉아먹는 건 대개 KV 캐시입니다. 다음 값들이 커지면 VRAM이 급격히 늘어납니다.

  • max_new_tokens (또는 max_length)
  • num_beams (빔 서치)
  • 배치 크기

안전한 기본값 예시

import torch

inputs = tok("Explain OOM in Transformers.", return_tensors="pt").to("cuda")

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        num_beams=1,  # 빔 서치는 메모리 크게 증가
        use_cache=True,
    )

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

정말 급하면 use_cache=False

out = model.generate(**inputs, max_new_tokens=128, use_cache=False)

use_cache=False는 속도가 느려질 수 있지만, OOM 회피에 즉효가 있습니다.

7) 메모리 조각화 대응: PYTORCH_CUDA_ALLOC_CONF

VRAM이 “총량은 남아 있는데도” OOM이 나는 경우가 있습니다. 이때 흔한 원인이 메모리 조각화(fragmentation)입니다. PyTorch CUDA allocator 설정으로 완화할 수 있습니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,expandable_segments:True"

그리고 코드에서 불필요한 텐서를 빨리 해제하고 캐시를 비우는 것도 도움이 됩니다.

import gc
import torch

del inputs
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()

현실적인 조언: 조각화는 “긴 실행 세션”에서 누적되기 쉬우므로, 반복 실험 중이면 커널/프로세스를 재시작하는 게 가장 확실한 해결이기도 합니다.

8) safetensors 사용 + 캐시/스왑/컨테이너 제한 점검

safetensors 우선

같은 모델이라도 safetensors가 있으면 로딩이 더 안전하고 빠른 경우가 많습니다(메모리 매핑, 안전한 역직렬화 등).

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

컨테이너/WSL/서버에서 “RAM이 있는데도” 죽는 경우

  • Docker라면 메모리 limit이 낮을 수 있습니다.
  • WSL은 기본 메모리 제한이 걸려 있을 수 있습니다.
  • swap이 꺼져 있으면 CPU OOM에 더 취약합니다.

이런 시스템 레벨 이슈를 추적하는 방법은 아래 글이 도움이 됩니다.

9) 로딩/추론 파이프라인을 “작게 쪼개기”: 스트리밍·배치·프롬프트 다이어트

모델 자체를 줄이지 않고도, 입력과 실행 방식을 바꾸면 OOM을 피할 수 있습니다.

배치 크기 줄이기

동시에 여러 문장을 넣는 배치 추론은 VRAM을 선형으로 올립니다. 우선 배치 1로 안정화하세요.

프롬프트 길이 줄이기

컨텍스트 길이가 길수록 KV 캐시가 커집니다. “긴 시스템 프롬프트 + 긴 히스토리”가 누적되면 작은 모델도 OOM이 납니다.

스트리밍 출력으로 UX 개선(메모리 자체는 동일해도 체감 개선)

from transformers import TextStreamer

streamer = TextStreamer(tok, skip_prompt=True)

_ = model.generate(
    **inputs,
    max_new_tokens=256,
    streamer=streamer,
)

핵심: OOM은 메모리 문제지만, 사용자는 “멈춘 것처럼 보이는” 상황을 더 싫어합니다. 스트리밍은 운영 관점에서 실패율을 낮추는 데 도움이 됩니다.

자주 겪는 OOM 패턴별 빠른 처방표

  • 로딩 즉시 CUDA out of memory

    • torch_dtype=float16 또는 bfloat16
    • 8bit/4bit 양자화
    • device_map="auto"로 CPU 오프로딩

  • 로딩 중 CPU가 터짐(프로세스 Killed)

    • low_cpu_mem_usage=True
    • use_safetensors=True
    • 컨테이너 메모리 limit, swap 확인

  • 첫 생성에서 OOM

    • max_new_tokens 줄이기
    • num_beams=1
    • attn_implementation="sdpa" 또는 flash_attention_2
    • 필요 시 use_cache=False

  • “남은 VRAM이 있는데도” OOM

    • PYTORCH_CUDA_ALLOC_CONF로 조각화 완화
    • 커널/프로세스 재시작

마무리: 최적의 순서(추천)

OOM을 가장 빠르게 줄이는 적용 순서는 보통 다음이 효율적입니다.

  1. torch_dtype=float16 또는 bfloat16
  2. low_cpu_mem_usage=True
  3. device_map="auto" (+ 필요 시 offload_folder)
  4. 8bit/4bit 양자화
  5. max_new_tokens, num_beams 등 생성 파라미터 조정
  6. SDPA/Flash Attention 적용
  7. 조각화 대응 환경변수

LLM 로딩 OOM은 “한 가지 해결책”으로 끝나는 문제가 아니라, 환경(GPU 세대, 드라이버, PyTorch/Transformers 버전, 컨테이너 제한)과 워크로드(컨텍스트 길이, 배치, 빔서치)에 의해 결정됩니다. 위 9가지를 체크리스트처럼 적용하면 대부분의 로컬 OOM은 재현 가능하게 줄일 수 있습니다.

Discuss on Twitter • View on GitHub