Transformers 로컬 LLM OOM·속도 최적화 가이드

로컬에서 LLM을 돌리다 보면 거의 항상 두 가지 벽에 부딪힙니다. 첫째는 CUDA out of memory 같은 OOM이고, 둘째는 “돌긴 도는데 너무 느리다”입니다. 이 글은 Hugging Face Transformers 기준으로 OOM을 줄이면서도 처리량을 끌어올리는 방법을, 원인 분해와 실전 코드 중심으로 정리합니다.

특히 OOM은 단순히 “모델이 커서”가 아니라, 가중치(Weights) + KV 캐시 + 활성화(Activations) + 프레임워크 오버헤드의 합으로 터집니다. 속도도 마찬가지로 커널 선택, dtype, attention 구현, 배치 전략, 토크나이저 병목이 겹쳐서 느려집니다.

추가로, 메모리 문제를 원인 추적하는 관점은 JVM이든 Python이든 비슷합니다. 힙 덤프 기반으로 원인을 좁혀가는 접근은 아래 글도 함께 참고할 만합니다.

Spring Boot OOM 원인추적과 힙덤프 분석 실전

OOM을 이해하는 4가지 메모리 덩어리

1) 가중치 메모리(Weights)

모델 파라미터 자체가 차지하는 공간입니다.

FP16/BF16: 파라미터 1개당 2바이트
FP32: 4바이트
8bit/4bit quantization: 대폭 감소(대신 약간의 오버헤드와 성능/정확도 트레이드오프)

대략적인 감으로는 **파라미터 수 x dtype 바이트**가 바닥 비용입니다. 예를 들어 7B를 FP16으로 올리면 가중치만 약 7e9 x 2 바이트 수준이므로 수십 GB급이 됩니다(실제는 레이어 구조/버퍼로 더 늘 수 있음).

2) KV 캐시(Key/Value cache)

생성형 추론에서 가장 자주 OOM을 유발합니다. 입력 프롬프트가 길거나 max_new_tokens가 크면 KV 캐시가 계속 쌓입니다.

컨텍스트 길이가 커질수록 선형으로 증가
배치 크기(동시 요청 수)가 커질수록 선형으로 증가
모델의 num_layers, num_heads, head_dim에 비례

즉, “모델은 8GB에 올라가는데 2천 토큰만 생성해도 터진다”는 대부분 KV 캐시 문제입니다.

3) 활성화 메모리(Activations)

학습/파인튜닝에서 크게 문제지만, 추론에서도 설정에 따라 늘 수 있습니다.

use_cache=True는 보통 추론에서 필수지만 KV 캐시를 만듭니다.
torch.no_grad() 또는 torch.inference_mode()를 안 쓰면 그래프가 쌓여 메모리가 급증합니다.

4) 프레임워크/커널 오버헤드

PyTorch allocator 캐시, fragmentation, attention 구현(FlashAttention 여부), 커널 선택 등에 따라 “여유 VRAM이 있어 보이는데도” OOM이 날 수 있습니다.

OOM 최적화 1순위: dtype와 quantization

BF16/FP16로 먼저 내리기

GPU가 BF16을 잘 지원하면 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="cuda",
)
model.eval()

8bit/4bit quantization(bitsandbytes)

가장 체감이 큰 OOM 해결책입니다. 특히 4bit는 “안 올라가던 모델이 올라가는” 레벨로 바뀝니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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()

주의점:

4bit는 compute dtype를 BF16로 두는 편이 일반적으로 안정적입니다.
일부 GPU/드라이버 조합에서 커널 호환 이슈가 날 수 있어, 문제가 생기면 8bit 또는 FP16로 후퇴하는 전략이 필요합니다.

OOM 최적화 2순위: offload와 device_map을 제대로 쓰기

GPU VRAM이 부족하면 일부 레이어를 CPU로 내리는 offload가 가능합니다. 속도는 느려지지만 “일단 돌아가게” 만드는 데 효과적입니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

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

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    offload_folder="./offload",
    low_cpu_mem_usage=True,
)
model.eval()

체크포인트:

device_map="auto"는 편하지만, 특정 레이어가 GPU에 몰리면 OOM이 날 수 있습니다.
offload_folder는 디스크 IO가 관여하므로 NVMe가 유리합니다.

KV 캐시로 터질 때: 컨텍스트/생성 길이/배치를 줄여라

max_new_tokens, max_length를 보수적으로

max_new_tokens를 크게 잡아 놓고 “알아서 멈추겠지”라고 두면 메모리가 계속 예약됩니다.

import torch

prompt = "요약: Transformers에서 OOM을 줄이는 핵심은"
inputs = tokenizer(prompt, return_tensors="pt").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))

배치 처리 시 padding 전략

여러 요청을 묶으면 throughput은 좋아지지만, padding이 길어지면 KV 캐시가 불필요하게 커집니다.

가능하면 길이가 비슷한 요청끼리 묶기(bucketing)
padding_side와 truncation 정책 명확히

tokenizer.padding_side = "left"  # decoder-only 모델에서 흔히 사용
batch = tokenizer(
    ["짧은 질문", "훨씬 더 긴 질문 ..."],
    return_tensors="pt",
    padding=True,
    truncation=True,
    max_length=1024,
).to(model.device)

속도 최적화: 느린 원인을 분해해서 고치기

속도는 “모델 크기”만의 문제가 아닙니다. 아래 순서로 점검하면 삽질이 줄어듭니다.

토크나이저가 병목인지
attention 커널이 최적인지(FlashAttention 등)
dtype이 맞는지(BF16/FP16)
generate 파라미터가 과도한지(beam search, repetition penalty 등)
배치 전략이 올바른지

토크나이저 병목 줄이기

Fast tokenizer를 쓰고, 가능한 한 토크나이징을 요청 경로 밖으로 빼거나 캐시합니다.

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

# 반복 요청이라면 prompt 템플릿의 고정 부분은 미리 토큰화해 두는 전략도 가능

torch.compile로 커널 최적화(조건부)

환경에 따라 효과가 크거나, 오히려 이득이 없을 수도 있습니다. 그래도 한 번은 시도할 가치가 있습니다.

import torch

model = model.to("cuda")
model.eval()

# PyTorch 2.x
compiled_model = torch.compile(model)

inputs = tokenizer("Hello", return_tensors="pt").to("cuda")
with torch.inference_mode():
    _ = compiled_model.generate(**inputs, max_new_tokens=64)

주의:

컴파일 초기 1회 비용이 큽니다. 서버라면 warm-up을 별도로 수행하세요.
dynamic shape이 많으면 기대만큼 안 나올 수 있습니다.

FlashAttention/SDPA 사용 여부 확인

Transformers는 PyTorch의 SDPA(Scaled Dot-Product Attention) 경로를 타면 빨라지는 경우가 많습니다. 모델/버전에 따라 설정이 다르지만, 핵심은 “빠른 attention 구현을 타고 있는가”입니다.

실무 팁:

PyTorch, CUDA, GPU 아키텍처 조합에 따라 사용 가능한 커널이 달라집니다.
같은 모델이라도 attention 구현에 따라 체감 속도가 크게 변합니다.

generate 설정이 속도를 망치는 패턴

num_beams를 키우면 품질은 좋아질 수 있지만 속도는 급락합니다.
do_sample=True와 높은 top_p, temperature는 경우에 따라 속도보다 품질에 초점을 둡니다.

권장 시작점:

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
        num_beams=1,
        use_cache=True,
    )

메모리 파편화와 “남아 있는데도 OOM” 대응

PyTorch는 메모리 allocator가 캐시를 잡고 있고, fragmentation 때문에 큰 연속 블록을 못 구해 OOM이 날 수 있습니다.

실전 대응:

추론 서버 프로세스를 장시간 띄우면 fragmentation이 누적될 수 있으니, 배포/롤링 전략을 고려합니다.
요청 단위로 텐서를 오래 잡고 있지 않게 하고, 전역 리스트에 누적되는 참조를 제거합니다.

디버깅에 유용한 코드:

import torch

print(torch.cuda.memory_summary())
# 또는
print(torch.cuda.mem_get_info())

OOM이 간헐적이라면, “특정 길이의 요청”이 트리거일 때가 많습니다. 이 경우는 애플리케이션 레벨에서 입력 길이 제한과 동시성 제한이 가장 확실한 해법입니다.

배포 관점 최적화: 캐시와 워밍업, 그리고 빌드 속도

로컬 LLM은 모델 다운로드, 컨테이너 빌드, 런타임 워밍업 시간이 길어 운영 비용을 키웁니다.

모델 파일은 볼륨 캐시로 재사용
컨테이너는 레이어 캐시 최적화
서버 시작 시 warm-up 요청 1회로 커널/컴파일 비용 상쇄

CI에서 이미지 빌드를 자주 한다면 아래 글의 Docker 레이어 캐시 전략도 그대로 응용할 수 있습니다.

GitHub Actions Docker 레이어 캐시로 배포 3배 빠르게

최소 체크리스트(실무용)

OOM이 나면

torch.inference_mode() 적용 여부 확인
dtype를 FP16 또는 BF16으로
8bit 또는 4bit quantization 적용
max_new_tokens, 입력 max_length, 동시성(배치) 줄이기
device_map="auto" + offload로 일단 생존
메모리 summary로 fragmentation/누수 의심 지점 찾기

느리면

fast tokenizer 사용, 토크나이징 캐시
beam search 끄고(num_beams=1) baseline 측정
SDPA/FlashAttention 경로 확인
torch.compile 시도(서버는 warm-up 필수)
길이 비슷한 요청끼리 batching

마무리

Transformers 로컬 LLM의 OOM과 속도 문제는 “옵션 몇 개”로 끝나지 않고, 가중치 크기, KV 캐시 성장, 커널 구현, 배치 전략이 함께 얽혀 있습니다. 하지만 원인을 덩어리로 나누고(Weights, KV, Activations, Overhead), 우선순위대로 적용하면 대부분의 케이스는 안정적으로 해결됩니다.

다음 단계로는 “내 GPU/모델 조합에서 최적점이 어디인지”를 찾는 벤치마크가 중요합니다. 동일 프롬프트로 tokens/sec, VRAM peak, p95 latency를 기록해 두면, 설정 변경이 실제로 개선인지 빠르게 판단할 수 있습니다.