beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM 해결 9가지

Authors

서빙이든 실험이든 로컬에서 LLM을 돌리다 보면 결국 한 번은 OOM(Out Of Memory)을 만납니다. 문제는 OOM이 단순히 "모델이 커서"만이 아니라, dtype, KV 캐시, 컨텍스트 길이, 배치 전략, 로더 옵션, 파편화(fragmentation) 같은 여러 요인이 합쳐져 터진다는 점입니다.

이 글은 Transformers 기반 로컬 추론에서 OOM을 9가지 체크포인트로 쪼개서, 어떤 레버를 당기면 VRAM과 RAM이 얼마나 줄어드는지 감을 잡게 해줍니다. 이미 4bit 양자화와 KV 캐시 튜닝을 깊게 다룬 글이 있다면 아래 글도 함께 보세요.

0. 먼저: OOM이 "어디"에서 났는지 분류하기

OOM은 크게 3종류로 나뉩니다.

  1. GPU VRAM OOM: CUDA out of memory 혹은 torch.cuda.OutOfMemoryError
  2. CPU RAM OOM: 프로세스가 죽거나 Killed, 혹은 OS OOM killer
  3. 메모리 파편화: 남은 VRAM이 있는데도 할당 실패, 혹은 실행마다 들쭉날쭉

최소한 아래 정도는 로그에 찍어두면 원인 파악이 빨라집니다.

import os, torch

def mem_report(tag=""):
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**2
        reserved = torch.cuda.memory_reserved() / 1024**2
        print(f"[{tag}] allocated={allocated:.1f}MiB reserved={reserved:.1f}MiB")

print("cuda:", torch.cuda.is_available())
mem_report("start")

또한, 같은 코드가 어떤 날은 되고 어떤 날은 터진다면 파편화 가능성이 높습니다. 이 경우 아래 8번을 먼저 적용하는 게 효과적입니다.

1) 모델을 더 작게가 아니라, 먼저 dtype를 낮춰라

가장 빠른 1차 처방은 fp16 또는 bf16으로 내리는 것입니다. 특히 Ampere 이후 GPU는 bf16이 안정적인 경우가 많습니다.

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

체크포인트

  • fp32로 로딩되면 VRAM이 즉시 2배 가까이 늘어납니다.
  • 일부 모델은 bf16에서 더 안정적이고, 일부는 fp16이 더 안전합니다. 둘 다 시도해 보세요.

2) device_map="auto"와 오프로딩으로 VRAM을 "쪼개서" 쓰기

단일 GPU VRAM이 부족하면 모델 일부를 CPU로 오프로딩해 OOM을 회피할 수 있습니다. 속도는 느려지지만, "일단 실행"이 목표라면 강력합니다.

from transformers import AutoModelForCausalLM

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

더 적극적으로는 accelerate 오프로딩 옵션을 씁니다.

from transformers import AutoModelForCausalLM

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

체크포인트

  • 오프로딩은 VRAM OOM을 RAM OOM으로 "이동"시킬 수 있습니다. RAM도 함께 모니터링하세요.
  • NVMe 오프로딩은 가능하면 SSD를 쓰고, 폴더는 충분한 여유 공간을 확보합니다.

3) 양자화: 8bit 또는 4bit로 가는 게 정석

추론 목적이라면 bitsandbytes 기반 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.bfloat16,
)

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

체크포인트

  • 4bit는 가중치 메모리를 크게 줄이지만, KV 캐시는 별개입니다. 컨텍스트가 길면 여전히 OOM이 납니다.
  • compute_dtype를 bf16 또는 fp16으로 바꾸는 것만으로도 피크 메모리가 달라질 수 있습니다.

4) 진짜 범인은 KV 캐시다: max_new_tokens와 컨텍스트를 줄여라

로컬 추론에서 OOM을 가장 자주 터뜨리는 건 긴 입력 컨텍스트 + 긴 생성 길이로 인해 커지는 KV 캐시입니다.

즉, 같은 7B 모델이라도

  • input_length=8192 + max_new_tokens=1024
  • input_length=2048 + max_new_tokens=256

은 메모리 사용량이 완전히 다릅니다.

import torch

inputs = tokenizer("긴 프롬프트...", return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,        # 줄이기
        do_sample=False,
        use_cache=True,            # 보통은 True가 빠름
    )

추가 팁: 프롬프트를 줄이는 실전 전략

  • RAG를 쓴다면 top-k를 무작정 키우지 말고, 문서 chunk를 줄이거나 rerank를 넣습니다.
  • 벡터 검색이 느려서 top-k를 못 줄이는 상황이라면 인덱스 튜닝이 먼저일 수 있습니다.

5) 배치와 패딩이 VRAM을 먹는다: 동적 패딩과 배치 1부터

추론 서버를 만들다 보면 배치를 묶고 싶어집니다. 하지만 LLM은 배치가 커질수록 KV 캐시와 활성화 메모리가 같이 커져 OOM이 쉽게 납니다.

먼저 배치 1로 안정화한 뒤, 아래를 적용하세요.

  • 배치 크기 줄이기
  • 입력 길이별로 버킷팅해서 패딩 낭비 줄이기
from transformers import DataCollatorWithPadding

collator = DataCollatorWithPadding(tokenizer=tokenizer, padding=True)

# DataLoader에서 collate_fn=collator 사용

체크포인트

  • 배치 2가 배치 1의 정확히 2배 메모리가 아닐 수 있습니다. 피크가 더 커질 수 있습니다.
  • 패딩이 많은 배치는 "실제 토큰"보다 "패딩 토큰" 때문에 OOM이 납니다.

6) torch.inference_mode()와 eval()은 필수, 그래프가 남지 않게

추론인데도 실수로 그래프를 쌓으면 GPU 메모리가 계속 증가합니다.

model.eval()

with torch.inference_mode():
    out = model(**inputs)

체크포인트

  • torch.no_grad()도 좋지만, 추론 전용이라면 torch.inference_mode()가 더 강합니다.
  • 스트리밍 구현에서 토큰마다 model(...)을 호출하면서 텐서를 리스트에 쌓아두면, 그 자체가 메모리 누수처럼 보일 수 있습니다. 디코딩용 텐서는 즉시 CPU로 옮기거나 버리세요.

7) attn_implementation로 메모리 효율을 챙기기

가능한 환경에서는 메모리 효율이 좋은 attention 구현을 선택합니다. 모델과 GPU, PyTorch 버전에 따라 지원이 달라질 수 있습니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2",  # 지원 시
)

지원이 안 되면 sdpa로도 이득을 볼 수 있습니다.

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",
)

체크포인트

  • flash attention 계열은 대체로 메모리와 속도 모두 이득이지만, 설치 조합이 까다롭습니다.
  • 적용 후에도 OOM이면 컨텍스트 길이와 KV 캐시를 다시 보세요.

8) 메모리 파편화 완화: PYTORCH_CUDA_ALLOC_CONF와 캐시 정리

"분명 여유 VRAM이 있는데" 할당이 실패한다면 파편화일 수 있습니다. 가장 흔한 처방은 max_split_size_mb 조정입니다.

셸에서 환경변수로 설정합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"

또는 상황에 따라 더 작게도 시도합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:64"

실험 루프에서 커널을 반복 실행한다면 중간에 캐시를 정리해 피크를 낮출 수 있습니다.

import gc, torch

def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

# 실험 반복 사이에 호출
cleanup()

체크포인트

  • empty_cache()는 "할당된 텐서"를 지우지 않습니다. 참조가 남아있으면 효과가 없습니다.
  • Jupyter에서 특히 파편화가 심해질 수 있으니, 재시작이 가장 빠른 해결책인 경우도 많습니다.

9) 로딩 단계 OOM: low_cpu_mem_usage, safetensors, 체크포인트 분할

모델이 실행 중이 아니라 로딩 중에 OOM이 난다면, CPU RAM 또는 GPU VRAM으로 순간 피크가 튈 가능성이 큽니다.

아래 옵션은 로딩 피크를 줄이는 데 도움이 됩니다.

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

또한 가능하면 safetensors 포맷을 우선 사용합니다(모델 리포지토리에 함께 제공되는 경우가 많습니다). safetensors는 로딩이 더 안전하고 예측 가능한 경우가 많습니다.

체크포인트

  • 로딩 OOM은 종종 "RAM이 부족한데 스왑이 없어서" 발생합니다. NVMe 스왑을 임시로라도 켜면 로딩은 통과하는 경우가 있습니다.
  • 여러 GPU가 있다면 device_map으로 레이어를 분산해 로딩 피크를 낮출 수 있습니다.

실전용: OOM 대응 템플릿 코드

아래는 로컬 추론에서 OOM을 덜 만나도록 기본값을 보수적으로 잡은 템플릿입니다.

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128")

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

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

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

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    quantization_config=bnb_config,
    low_cpu_mem_usage=True,
    attn_implementation="sdpa",
)
model.eval()

prompt = "요약: 로컬 LLM OOM을 줄이는 방법을 5줄로 정리해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

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

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

마무리: 9가지 중 무엇부터 적용할까

우선순위를 정하면 시행착오가 줄어듭니다.

  1. torch_dtype를 bf16 또는 fp16으로
  2. max_new_tokens와 입력 컨텍스트 길이부터 줄이기(KV 캐시)
  3. 4bit 양자화 적용
  4. device_map="auto"로 분산 및 오프로딩
  5. attn_implementation 최적화
  6. 배치와 패딩 최적화
  7. 파편화 완화 환경변수 적용
  8. 로딩 OOM이면 low_cpu_mem_usage와 safetensors 확인

이 순서로 하면 대부분의 로컬 OOM은 "원인"이 보이기 시작합니다. 그래도 계속 터진다면, 사용 중인 GPU VRAM, 목표 컨텍스트 길이, 모델 파라미터 크기(예: 7B, 13B, 34B)를 기준으로 아키텍처(모델 크기 축소, 멀티 GPU, vLLM 전환 등) 자체를 재설계해야 합니다.

Discuss on Twitter • View on GitHub