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

로컬 GPU에서 Transformers로 LLM을 돌리다 보면 한 번쯤은 CUDA out of memory를 맞습니다. 문제는 “모델이 너무 커서”만이 아니라, 로딩 방식·dtype·KV 캐시·배치/시퀀스 길이·파편화 같은 운영 요소들이 겹쳐서 터진다는 점입니다.

이 글은 로컬 LLM 추론(또는 간단한 서빙)에서 OOM을 줄이는 7가지 방법을, 바로 적용 가능한 코드 중심으로 정리합니다. 특히 transformers + accelerate 조합을 기준으로 설명합니다.

참고: OOM은 한 번 발생하면 파이썬 프로세스가 살아 있어도 CUDA 메모리 상태가 꼬일 수 있습니다. 재현/튜닝 시엔 프로세스를 깔끔히 재시작하는 편이 더 빠를 때가 많습니다.

0) 먼저: 어디서 얼마나 쓰는지부터 측정

OOM 최적화는 감이 아니라 측정으로 시작합니다. 아래는 현재/최대 사용량을 확인하는 최소 코드입니다.

import torch

def cuda_mem(prefix=""):
    if not torch.cuda.is_available():
        print(prefix, "CUDA not available")
        return
    alloc = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"{prefix} alloc={alloc:.2f}GB reserved={reserved:.2f}GB peak={peak:.2f}GB")

cuda_mem("start")

추론 루프에서 cuda_mem("after generate") 같은 식으로 찍어보면, 모델 로딩이 문제인지(KV 캐시/출력 토큰이 문제인지) 빠르게 갈라집니다.

1) torch_dtype를 낮추고, 가능하면 bfloat16/float16로 고정

가장 즉효는 dtype입니다. FP32는 로컬 LLM에선 거의 금지 수준으로 비쌉니다.

• Ampere 이상(A100/RTX30/RTX40 등)에서 bfloat16은 안정성이 좋고, 메모리도 FP32 대비 절반
• float16은 더 넓게 지원되지만, 일부 모델/연산에서 수치 불안정이 있을 수 있음

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

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

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,  # 또는 torch.float16
    device_map="cuda"
)

추가로, torch.set_default_dtype 같은 전역 변경은 예기치 않은 메모리/정확도 문제를 만들 수 있어 권장하지 않습니다. 모델 로딩 시점에 명시하는 편이 안전합니다.

2) 4-bit/8-bit 양자화(bitsandbytes)로 가중치 메모리 압축

VRAM이 부족한 로컬 환경에서 가장 강력한 옵션이 양자화입니다. 특히 4-bit는 “돌아가게 만드는” 수준의 차이를 만듭니다.

설치(환경에 맞게 CUDA 버전 확인 필요):

pip install -U transformers accelerate bitsandbytes

4-bit 로딩 예시:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

주의할 점:

• 양자화는 VRAM을 크게 줄이지만, 속도/정확도/호환성이 모델과 GPU에 따라 달라집니다.
• 특정 연산이 FP16/BF16으로 올라오면서 예상보다 메모리가 늘 수 있으니, compute_dtype를 명시하고 측정하세요.

3) device_map="auto" + CPU 오프로딩으로 “완전 OOM”을 피하기

모델이 VRAM에 안 들어갈 때, 일부 레이어를 CPU로 내리는 방식이 있습니다. 속도는 느려지지만, 로컬에서 “일단 동작”이 목표라면 매우 유용합니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto",
    offload_folder="./offload",  # 오프로딩 파일 저장
)

팁:

• NVMe SSD가 있다면 오프로딩 체감이 훨씬 낫습니다.
• accelerate가 설치되어 있어야 device_map이 제대로 동작합니다.

4) max_new_tokens와 입력 길이(max_length)를 줄여 KV 캐시 폭발 막기

추론 OOM의 주범 중 하나는 KV 캐시입니다. 생성 토큰이 늘어날수록 레이어별 K/V 텐서가 누적되고, 배치가 커질수록 곱으로 증가합니다.

• 입력 프롬프트가 길다: 이미 KV 캐시가 큰 상태에서 시작
• max_new_tokens가 크다: 생성 중 계속 증가
• 배치가 크다: 요청을 묶을수록 선형 이상으로 커짐

import torch

prompt = "긴 문서를 요약해줘..."
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).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~512 범위에서 시작
• 입력은 truncation=True와 max_length를 반드시 명시

추론 품질을 유지하면서 길이를 줄이는 전략으로는 RAG가 대표적입니다. 로컬 LLM에서도 “문서 전체를 다 넣는” 대신 검색으로 컨텍스트를 줄이면 VRAM과 지연이 같이 줄어듭니다. 관련 패턴은 CoT 누출 없이 추론 강화하는 RAG+검증자 패턴도 참고할 만합니다.

5) 배치/동시성을 줄이고, 토크나이저 패딩을 최적화

여러 요청을 한 번에 처리하려고 배치를 키우면, KV 캐시와 어텐션 연산이 같이 커져 OOM이 빨리 옵니다. 로컬에서는 특히 “동시성 1~2”가 체감상 안정적인 경우가 많습니다.

또한 패딩이 비효율적이면 짧은 문장도 긴 문장 길이에 맞춰 메모리를 먹습니다. 가능한 한 배치 내 길이를 비슷하게 묶거나, 패딩 전략을 신경 써야 합니다.

texts = ["짧은 질문", "조금 더 긴 질문..." * 200]

batch = tokenizer(
    texts,
    return_tensors="pt",
    padding=True,           # 배치 패딩
    truncation=True,
    max_length=1024,
)

batch = {k: v.to(model.device) for k, v in batch.items()}

with torch.inference_mode():
    out = model.generate(**batch, max_new_tokens=128)

가능하면:

• 길이가 비슷한 요청끼리 묶기(서빙이면 큐에서 정렬)
• 배치 크기를 줄이고, 대신 캐시/스루풋은 다른 방식으로 확보

이때 “캐시”라는 관점은 GPU 메모리뿐 아니라 빌드/배포 파이프라인에도 통합니다. 모델 이미지/의존성 빌드가 느려서 자주 재시작이 어렵다면, Docker 빌드 70% 단축 - 멀티스테이지·캐시 전략처럼 운영 시간을 줄여 튜닝 반복을 빠르게 만드는 것도 도움이 됩니다.

6) 어텐션 최적화: SDPA/FlashAttention로 메모리 압박 완화

PyTorch 2.x에서는 SDPA(Scaled Dot-Product Attention)가 기본 최적화를 제공하며, 환경에 따라 FlashAttention 계열 커널이 사용됩니다. 이 경로는 어텐션의 중간 텐서 사용량을 줄여 OOM 가능성을 낮출 수 있습니다.

Transformers에서는 설정으로 켤 수 있습니다.

from transformers import AutoModelForCausalLM

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

주의:

• 모든 모델/버전에서 동일하게 동작하지 않습니다.
• 드라이버, PyTorch, CUDA 조합에 따라 커널 선택이 달라집니다.
• 문제가 생기면 attn_implementation을 제거하고 기본값으로 되돌려 비교하세요.

7) 메모리 파편화/캐시 관리: PYTORCH_CUDA_ALLOC_CONF와 안전한 정리

“분명 여유가 있어 보이는데 OOM”은 파편화일 수 있습니다. PyTorch의 CUDA allocator가 큰 블록을 못 잡아 실패하는 경우가 있습니다.

가장 흔한 처방은 max_split_size_mb 조정입니다. 실행 전에 환경변수로 지정합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"

워크로드에 따라 64/128/256 등으로 바꿔가며 재현 테스트를 해보세요.

또한, 추론 서버/노트북에서 모델을 여러 번 갈아끼우며 실험하면, 텐서 참조가 남아 메모리가 반환되지 않는 경우가 흔합니다. 실험용으로는 아래 패턴이 유용합니다.

import gc
import torch

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

# 모델 교체/실험 사이에 호출
cleanup_cuda()

단, empty_cache()는 “사용 중인 텐서를 줄이는” 게 아니라 캐싱된 블록을 반환하는 성격이라, 근본 해결은 여전히 “큰 텐서를 만들지 않는 설정”입니다.

자주 터지는 조합과 빠른 처방 요약

로딩 단계에서 바로 OOM

• FP32로 로딩됨
• 모델이 VRAM보다 큼

처방:

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

생성 도중 OOM(몇 토큰 뽑다가 터짐)

• max_new_tokens가 큼
• 입력이 너무 김
• 배치가 큼

처방:

• max_new_tokens 축소
• truncation=True, max_length=... 적용
• 동시성/배치 축소, 길이 비슷한 요청끼리 묶기

“여유 있어 보이는데” OOM

• 파편화/캐시 블록 문제

처방:

• PYTORCH_CUDA_ALLOC_CONF의 max_split_size_mb 튜닝
• 실험 중간에 참조 제거, gc.collect() 후 캐시 정리

결론: OOM은 ‘모델 크기’가 아니라 ‘운영 변수’의 곱

로컬 LLM에서 CUDA OOM을 줄이는 핵심은 다음 순서로 접근하는 것입니다.

1. dtype과 양자화로 가중치 메모리를 먼저 줄이고
2. max_new_tokens/입력 길이/배치로 KV 캐시 성장을 제한하며
3. 필요하면 오프로딩과 어텐션 최적화로 마진을 확보하고
4. 마지막으로 allocator 튜닝으로 파편화 리스크를 낮춥니다.

위 7가지를 적용하면, 같은 GPU에서도 “겨우 로딩”에서 “안정적으로 추론”까지 체감이 크게 바뀝니다. 다음 단계로는 vLLM/TensorRT-LLM 같은 서빙 스택으로 넘어가며 KV 캐시/페이지드 어텐션을 최적화하는 선택지도 검토해볼 만합니다.