beautifulsoups.dev
Published on

Transformers 로컬 LLM OOM 7분 진단·해결

Authors

로컬에서 transformers로 LLM을 띄우다 보면, 잘 되던 코드가 어느 순간 갑자기 OOM(Out Of Memory)으로 터집니다. 특히 GPU VRAM은 여유가 있어 보이는데도 CUDA out of memory가 나거나, 매번 다른 지점에서 죽는 경우가 많습니다. 이 글은 7분 안에 OOM 원인을 가중치(모델 파라미터), KV 캐시(시퀀스 길이), 배치/동시성, 메모리 파편화 중 어디에 두어야 하는지 판단하고, 바로 적용 가능한 처방을 제공합니다.

아래 내용은 단일 GPU 기준으로 설명하지만, 멀티 GPU에서도 동일한 원리로 접근합니다.

0분: OOM 유형부터 구분하기

먼저 에러가 GPU VRAM OOM인지, 시스템 RAM OOM인지 구분해야 합니다.

  • GPU OOM 전형 메시지: RuntimeError: CUDA out of memory. Tried to allocate ...
  • CPU/RAM OOM 전형 증상: 프로세스가 갑자기 종료되거나, 리눅스 dmesg에 OOM killer 로그

컨테이너/쿠버네티스 환경이라면 OOMKilled 여부까지 같이 봐야 합니다. OOMKilled 진단 흐름은 이 글과 유사한 방식으로 접근하면 빠릅니다: EKS Pod CrashLoopBackOff? OOMKilled 진단법

1분: 현재 VRAM이 “어디에” 쓰이는지 계측하기

nvidia-smi로 큰 그림 확인

  • 모델 로드 직후 VRAM
  • 첫 토큰 생성 직후 VRAM
  • 긴 문장 생성 중 VRAM 증가 추세

이 3개만 봐도 가중치가 큰지, KV 캐시가 커지는지 감이 옵니다.

PyTorch에서 정밀 계측

아래는 단계별로 VRAM 피크를 찍는 최소 코드입니다.

import torch

def vram(tag: str):
    torch.cuda.synchronize()
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] alloc={alloc:.1f}MiB reserved={reserved:.1f}MiB peak={peak:.1f}MiB")

torch.cuda.reset_peak_memory_stats()

이걸 모델 로드/추론 전후에 끼워 넣으면, OOM 직전에 무엇이 급증했는지가 보입니다.

2분: OOM 원인 1위, KV 캐시 폭증을 먼저 의심

로컬 LLM에서 가장 흔한 함정은 **max_new_tokens**나 입력 길이가 커지면서 KV 캐시가 VRAM을 잡아먹는 경우입니다.

  • “모델 로드는 되는데 생성만 하면 터짐”이면 KV 캐시 가능성이 큼
  • “짧은 프롬프트는 되는데 긴 문서 넣으면 터짐”도 동일

즉시 처방: 토큰 길이 상한을 강제

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

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

tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    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))

핵심은 truncation=True와 max_length로 입력 길이를 통제하고, max_new_tokens도 서비스 정책으로 상한을 두는 것입니다.

고급 처방: 배치/동시성에서 KV 캐시가 곱으로 커진다

동시에 여러 요청을 처리하면 KV 캐시가 거의 동시 요청 수만큼 늘어납니다.

  • 단일 요청은 되는데, 멀티스레드/멀티프로세스로 돌리면 터진다
  • 혹은 FastAPI로 띄우고 동시 요청에서만 OOM

이 경우는 요청 큐잉 또는 동시성 제한이 정답입니다.

3분: OOM 원인 2위, “가중치 자체가 안 들어감”

모델 로드 단계에서 바로 터지면, 대부분 가중치가 VRAM에 들어가지 않는 상황입니다.

즉시 처방 A: torch_dtype를 낮춰서 로드

  • FP16 또는 BF16: 일반적인 첫 선택
  • FP32는 로컬 LLM에 거의 불필요하게 무겁습니다.
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,  # 또는 torch.float16
    device_map="cuda",
)

즉시 처방 B: 4bit/8bit 로드로 체급을 바꾸기

bitsandbytes를 쓰면 가중치 메모리를 크게 줄일 수 있습니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

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

8bit/4bit 양자화는 “돌아가게 만들기”에 가장 강력한 카드입니다. 다만 환경별로 오류가 다양해서, 양자화 에러/성능 튜닝은 아래 글을 함께 참고하면 시행착오를 줄일 수 있습니다: PyTorch 2.1 int8 양자화 오류별 해결과 속도

즉시 처방 C: device_map="auto"로 CPU 오프로딩

GPU가 작을 때는 일부 레이어를 CPU로 넘기는 방식도 가능합니다.

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

단, CPU 오프로딩은 속도 하락이 크므로 “개발/검증용”으로 쓰고, 운영은 양자화나 더 큰 GPU가 보통 정답입니다.

4분: OOM 원인 3위, 파편화와 reserved 함정

PyTorch는 메모리 풀을 잡아두기 때문에 nvidia-smi에서 여유가 있어 보여도, PyTorch 관점에서 “연속된 큰 덩어리”가 없어 OOM이 날 수 있습니다.

증상

  • 같은 입력인데도 어떤 때는 되고 어떤 때는 터짐
  • 에러 메시지에 reserved memory가 크게 나오고 allocated는 상대적으로 작음

즉시 처방 A: 프로세스 단위로 워커를 재시작

장시간 서비스에서 파편화가 누적되면, 가장 확실한 해결은 워커 재시작입니다.

  • gunicorn/uvicorn worker 수를 줄이고 주기적 recycle
  • 요청 수 기준으로 재시작

즉시 처방 B: PYTORCH_CUDA_ALLOC_CONF 조정

환경변수로 allocator 동작을 바꿔 파편화를 완화할 수 있습니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"

워크로드에 따라 128 또는 64 등으로 조정합니다. 값이 너무 작으면 오히려 오버헤드가 늘 수 있어, “OOM 재현 입력”으로 A/B 테스트하는 게 안전합니다.

즉시 처방 C: torch.inference_mode()와 불필요 텐서 제거

추론에서 no_grad보다 더 강한 inference_mode를 쓰고, 중간 텐서를 오래 들고 있지 않게 합니다.

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

# 큰 텐서를 오래 잡지 않게 참조를 끊고 캐시 정리
 del out
 torch.cuda.empty_cache()

empty_cache()는 “진짜 사용 중인 메모리”를 줄이지는 못하지만, 캐시 반환으로 다음 큰 할당이 성공하는 경우가 있어 응급 처치로는 의미가 있습니다.

5분: OOM 원인 4위, 배치 크기와 텐서 형태 실수

LLM 추론에서 배치 크기는 곧 메모리입니다. 그런데 배치가 커지는 실수는 종종 “의도치 않게” 발생합니다.

흔한 실수

  • 토크나이저 결과를 리스트로 쌓아 한 번에 generate 호출
  • 데이터로더에서 batch_size를 키운 채 그대로 추론
  • 대화 히스토리를 계속 누적해서 입력이 기하급수적으로 증가

빠른 점검 코드: 실제 입력 토큰 수 출력

input_ids = inputs["input_ids"]
print("batch=", input_ids.shape[0], "seq=", input_ids.shape[1])

여기서 seq가 생각보다 크면, OOM은 거의 확정입니다.

6분: “7분 안에 끝내는” 체크리스트

아래 순서대로 보면 보통 1~2번에서 결론이 납니다.

  1. OOM이 GPU인지 CPU인지 확인 (nvidia-smi, dmesg)
  2. 모델 로드에서 터지는지, 생성에서 터지는지 분리
  3. 생성에서 터지면 입력 길이와 max_new_tokens를 줄여 재현 확인
  4. 동시 요청/배치가 있는지 확인 (동시성 N이면 KV 캐시도 대략 N배)
  5. 로드에서 터지면 FP16/BF16, 8bit/4bit, device_map을 순서대로 적용
  6. 같은 조건에서 들쭉날쭉하면 파편화 의심 (reserved가 큰지 확인) 후 allocator/워커 recycle

7분: 실전용 “안 터지는” 기본 템플릿

아래 템플릿은 로컬에서 재현 가능한 범용 설정입니다.

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 파편화 완화(필요 시)
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"

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

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

def generate(text: str, max_input_tokens: int = 2048, max_new_tokens: int = 256):
    inputs = tok(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=max_input_tokens,
    ).to("cuda")

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

print(generate("요약해줘: ..."))

이 상태에서 OOM이 난다면, 남는 선택지는 대개 둘입니다.

  • 더 공격적으로 줄이기: max_input_tokens 및 max_new_tokens 축소, 동시성 제한
  • 체급 바꾸기: 4bit/8bit 양자화 또는 더 큰 GPU

자주 묻는 질문(현장에서 가장 많이 만나는 케이스)

Q1. VRAM 24GB인데 7B도 OOM이 납니다

대부분은 긴 입력 + 긴 생성 + 동시 요청 조합입니다. 7B 가중치 자체는 들어가도, KV 캐시가 폭증하면 24GB도 금방 찹니다. 먼저 max_length와 max_new_tokens를 강제하고, 동시성을 1로 내려 재현 여부를 보세요.

Q2. empty_cache()를 계속 호출하면 해결되나요

근본 해결은 아닙니다. 파편화/캐시 반환으로 “운 좋게” 다음 할당이 성공할 수는 있으나, 입력 길이나 동시성이 원인이라면 결국 다시 터집니다.

Q3. 4bit로 로드했는데도 느립니다

양자화는 메모리를 줄이지만, 설정/커널/드라이버 조합에 따라 속도 이득이 제한될 수 있습니다. 특히 compute dtype, attention 구현, PyTorch 버전에 따라 차이가 큽니다. 양자화 관련 트러블슈팅은 PyTorch 2.1 int8 양자화 오류별 해결과 속도 흐름대로 점검하는 게 빠릅니다.

마무리

Transformers 로컬 LLM OOM은 “모델이 커서”만이 아니라, **KV 캐시(토큰 길이)**와 동시성, 그리고 파편화가 결합되며 발생하는 경우가 훨씬 많습니다. 7분 진단의 핵심은

  • 로드 OOM인지 생성 OOM인지 분리하고
  • 입력/출력 토큰 상한으로 KV 캐시를 통제하며
  • 필요하면 4bit/8bit로 체급을 낮추고
  • 장시간 서비스는 파편화를 전제로 워커 recycle을 설계

하는 것입니다.

다음 단계로는, “내 GPU에서 안전한 토큰 예산(입력 길이, 생성 길이, 동시성)”을 표로 만들어 운영 정책으로 고정해두면 OOM 재발을 크게 줄일 수 있습니다.

Discuss on Twitter • View on GitHub