Transformers 로컬 LLM CUDA OOM·속도 튜닝

서버에 올리기 전, 개발 PC나 단일 GPU에서 Transformers로 로컬 LLM을 돌리다 보면 대부분 두 가지 벽을 만납니다. 첫째는 CUDA out of memory(OOM), 둘째는 “돌긴 도는데 너무 느림”입니다. 문제는 이 둘이 서로 얽혀 있다는 점입니다. 메모리를 줄이려다 속도가 떨어지고, 속도를 올리려다 메모리가 터집니다.

이 글은 OOM의 원인을 가중치(Weights), KV 캐시, 활성화(Activations), 메모리 프래그멘테이션으로 분해하고, Transformers 환경에서 바로 적용 가능한 튜닝을 코드로 정리합니다. 추가로 속도 튜닝의 핵심인 Flash Attention / SDPA, 양자화, 배치·프롬프트 전략까지 한 번에 연결합니다.

관련해서 KV 캐시와 양자화로 체감 속도를 올리는 방법은 아래 글도 함께 보면 좋습니다.

• 로컬 LLM 속도 2배 - KV 캐시·양자화 튜닝
• 더 공격적인 배포 최적화(ONNX, TensorRT)에서 막힐 때는: PyTorch→ONNX→TensorRT INT8 양자화 오류 해결

1) OOM부터 구조적으로 이해하기: 무엇이 VRAM을 먹나

로컬 LLM 추론에서 VRAM 사용량은 대략 아래 항목의 합입니다.

1. 모델 가중치(Weights): 파라미터 수와 dtype에 비례
2. KV 캐시: prefill(프롬프트 인코딩)과 decode(토큰 생성) 동안 쌓이는 키/밸류 텐서
3. 활성화(Activations): 학습이 아니라 추론이면 상대적으로 작지만, prefill 구간에서 배치가 크면 증가
4. 프래그멘테이션/캐싱: PyTorch CUDA allocator가 잡아둔 메모리, 조각난 블록 등

대부분의 “어제는 됐는데 오늘은 OOM”은 2번(컨텍스트 길이 증가) 또는 4번(프래그멘테이션)입니다.

가중치 메모리의 대략적인 감

• FP16/BF16: 파라미터 1개당 2바이트
• INT8: 대략 1바이트 수준이지만, 스케일/제로포인트 등 부가 메타데이터와 커널 구현에 따라 상이
• 4bit: 대략 0.5바이트 수준 + 오버헤드

예를 들어 7B 모델을 FP16으로 올리면 가중치만 대략 7e9 * 2 bytes ≒ 14GB 수준입니다. 여기서 KV 캐시가 추가되면 16GB GPU도 쉽게 터집니다.

2) 가장 흔한 OOM 원인 4가지와 처방

(1) 프롬프트가 길어져 KV 캐시가 폭증

KV 캐시는 토큰 수에 거의 선형으로 늘어납니다. “모델은 로드되는데 생성 시작하자마자 OOM”이면 대개 KV 캐시가 마지막 한 방을 날린 겁니다.

처방

• max_new_tokens를 줄이고, 필요하면 스트리밍으로 UX를 보완
• max_length 대신 max_new_tokens 중심으로 제어
• 긴 문서를 넣는다면 chunking + 요약/검색(RAG)로 컨텍스트를 줄이기

(2) dtype이 FP32로 올라가거나 섞임

Transformers에서 로딩 옵션을 잘못 주면 일부가 FP32로 올라가 VRAM이 급증합니다.

처방

• 로딩 시 torch_dtype=torch.bfloat16 또는 torch.float16를 명시
• 가능하면 Ampere 이후 GPU면 BF16 권장(수치 안정성)

(3) device_map 미사용으로 한 GPU에 몰림

멀티 GPU가 있어도 device_map을 안 쓰면 한 장에 다 올라가 OOM이 납니다.

처방

• device_map="auto" 또는 accelerate 기반 device placement 사용

(4) PyTorch 메모리 프래그멘테이션

“남은 VRAM이 있어 보이는데도 OOM”이면 프래그멘테이션을 의심합니다.

처방

• 프로세스를 재시작하는 게 가장 확실
• 환경변수 PYTORCH_CUDA_ALLOC_CONF로 allocator 동작을 조정

아래는 실전에서 자주 쓰는 설정입니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.8"

• max_split_size_mb: 큰 블록을 쪼개는 기준을 조절해 조각화를 완화
• garbage_collection_threshold: 캐시된 블록을 얼마나 공격적으로 회수할지

3) Transformers 기준: OOM 방지 로딩 템플릿

아래 코드는 “일단 안 터지게”를 목표로 한 안전한 로딩 예시입니다.

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128,garbage_collection_threshold:0.8"

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

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

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,          # Ampere+ 권장
    device_map="auto",                 # 단일 GPU면 자동으로 cuda:0
    low_cpu_mem_usage=True,             # 로딩 중 RAM 사용 감소
)

model.eval()

여기서도 OOM이면 다음 우선순위로 조정합니다.

1. BF16이 안 되면 FP16
2. 그래도 안 되면 8bit 또는 4bit 양자화
3. 컨텍스트 길이 및 max_new_tokens 축소

4) 8bit/4bit 양자화로 “로드 OOM”부터 해결

가중치가 VRAM 대부분을 차지하는 상황(특히 7B 이상 FP16)에서는 양자화가 가장 즉효입니다. Transformers에서는 bitsandbytes를 많이 씁니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

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

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

model.eval()

• nf4: LLM에서 성능/품질 밸런스가 좋은 편
• compute_dtype: BF16을 지원하면 BF16로 두는 것이 일반적으로 유리

주의할 점은, 양자화는 “무조건 빠름”이 아니라는 겁니다. GPU/드라이버/커널에 따라 4bit가 FP16보다 느릴 수도 있습니다. 하지만 OOM 해결 관점에서는 가장 강력합니다.

5) 속도 튜닝의 핵심 1: torch.inference_mode()와 생성 파라미터

기본 중 기본인데, 적용 여부에 따라 체감이 꽤 납니다.

import torch

prompt = "Explain KV cache in Transformers in Korean."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        do_sample=False,
        max_new_tokens=256,
        use_cache=True,   # KV 캐시 사용
    )

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

• torch.no_grad()보다 torch.inference_mode()가 더 공격적으로 autograd 상태를 끄고 최적화합니다.
• use_cache=True는 디코드 속도에 핵심입니다(대부분 기본값이지만 명시 추천).

6) 속도 튜닝의 핵심 2: SDPA/Flash Attention 사용

PyTorch 2 계열에서는 Scaled Dot-Product Attention(SDPA)이 기본 경로가 되며, 조건이 맞으면 Flash Attention 계열 커널을 탑니다. Transformers에서도 attn_implementation 옵션으로 경로를 고를 수 있습니다.

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",  # 가능하면 "flash_attention_2" 옵션도 모델/환경에 따라 고려
)

체크 포인트:

• PyTorch 버전이 낮으면 SDPA 최적화가 제한됩니다.
• 일부 모델/아키텍처는 특정 attention 구현과 궁합이 다릅니다.
• Flash Attention 2를 쓰려면 별도 패키지/호환성 이슈가 있을 수 있어, 먼저 sdpa로 안정적인 이득을 확인하는 편이 좋습니다.

7) OOM과 속도를 동시에 잡는 실전 파라미터 조합

(1) 컨텍스트 길이 관리: “프롬프트 토큰 수”를 수치로 보자

사람이 보기엔 짧아 보여도 토크나이즈하면 길 수 있습니다. 먼저 토큰 수를 찍어보세요.

text = """여기에 긴 문서..."""
ids = tokenizer(text).input_ids
print("prompt_tokens=", len(ids))

• prompt_tokens가 4k를 넘기기 시작하면, 7B급도 KV 캐시로 VRAM이 눈에 띄게 증가합니다.
• OOM이 “생성 중간”에 난다면, max_new_tokens가 너무 큰 경우가 많습니다.

(2) 배치 전략: 로컬은 배치보다 “동시성 제어”가 중요

로컬에서 여러 요청을 동시에 처리하면 KV 캐시가 요청 수만큼 쌓여 OOM이 빨라집니다. 서버로 운영한다면 동시성을 제한하거나 큐잉이 필요합니다.

• 단일 GPU: 동시 1~2부터 시작
• 프롬프트가 길면 동시 1로 제한하는 게 안전

(운영 환경에서 GPU 추론이 불안정하거나 콜드스타트가 문제라면, 인퍼런스 계층 이슈도 같이 봐야 합니다. 예를 들어 Knative 기반이면 503과 콜드스타트가 엮일 수 있는데, 이 경우는 이 글이 도움이 됩니다: KServe+Knative GPU 추론 503·콜드스타트 해결)

(3) temperature와 top_p는 속도 자체보다 “토큰 수”에 영향

샘플링을 켜면 결과가 길어지는 경향이 있어, 결국 KV 캐시가 커지고 시간이 늘어납니다. 속도/안정성 우선이면 아래처럼 시작하세요.

• do_sample=False
• max_new_tokens 보수적으로

8) 디버깅: 지금 무엇이 VRAM을 먹는지 측정하기

OOM을 “감”으로 잡으면 끝이 없습니다. 최소한 아래 지표는 찍어두면 원인 파악이 빨라집니다.

import torch

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

torch.cuda.reset_peak_memory_stats()
vram_report("after_load")

inputs = tokenizer("Hello", return_tensors="pt").to(model.device)
with torch.inference_mode():
    vram_report("before_generate")
    out = model.generate(**inputs, max_new_tokens=128, use_cache=True)
    vram_report("after_generate")

• allocated는 실제 텐서가 점유한 메모리
• reserved는 allocator가 잡아둔 풀(프래그멘테이션 이슈가 있으면 이 값이 과하게 큼)
• peak는 피크 사용량(특히 prefill에서 튀는지 확인)

9) 자주 터지는 케이스별 체크리스트

케이스 A: 모델 로딩 단계에서 바로 OOM

• FP16/BF16로도 안 되면 8bit/4bit로 로드
• device_map="auto" 적용 여부 확인
• 동일 GPU에 다른 프로세스가 VRAM 잡고 있는지 확인

케이스 B: 생성 시작 직후 OOM

• 프롬프트 토큰 수 확인
• max_new_tokens 과도한지 확인
• 동시 요청(멀티스레드/멀티프로세스) 여부 확인

케이스 C: “남은 VRAM 있는데” OOM

• 프래그멘테이션 의심
• 프로세스 재시작
• PYTORCH_CUDA_ALLOC_CONF 조정

10) 결론: 추천 튜닝 순서(가장 효율적인 루트)

1. 측정부터: prompt_tokens, allocated/reserved/peak 찍기
2. 안전 로딩: torch_dtype=bf16, device_map="auto", low_cpu_mem_usage=True
3. KV 캐시 관리: max_new_tokens 보수적으로, 긴 컨텍스트는 chunking
4. 속도 개선: torch.inference_mode(), use_cache=True, attn_implementation="sdpa"
5. 그래도 OOM이면 양자화: 8bit 또는 4bit로 가중치부터 줄이기
6. 운영 관점: 동시성 제한과 큐잉으로 KV 캐시 폭증 방지

로컬 LLM은 “모델 크기”보다 “컨텍스트와 동시성”이 더 자주 발목을 잡습니다. OOM을 단순히 VRAM 부족으로만 보지 말고, KV 캐시와 allocator 상태까지 포함해 구조적으로 보면 해결 속도가 확 올라갑니다.