Transformers 로컬 LLM OOM - 4bit+오프로딩

로컬 GPU에서 transformers로 LLM을 띄우다 보면 가장 먼저 부딪히는 벽이 OOM(CUDA out of memory)입니다. 특히 7B~13B급 모델을 기본 설정(대개 float16)으로 로드하면, 추론만 하더라도 가중치 + KV 캐시 + 활성화 텐서가 VRAM을 빠르게 잠식합니다.

이 글에서는 OOM을 줄이는 데 효과가 큰 두 가지 축을 실전 관점에서 정리합니다.

4bit 양자화: 가중치 메모리를 크게 줄여 “일단 로드”를 가능하게 함
오프로딩(offloading): GPU에 전부 올리지 않고 CPU RAM(필요 시 디스크)로 일부를 내림

추가로, OOM을 재현/진단하는 체크리스트와 흔한 함정(예: device_map 실수, KV 캐시 폭증, max_new_tokens 과다)을 함께 다룹니다.

OOM이 나는 진짜 이유: 가중치만 문제가 아니다

OOM이라고 하면 모델 파라미터(가중치) 크기만 떠올리기 쉽지만, 로컬 추론에서 메모리를 크게 쓰는 항목은 다음과 같습니다.

모델 가중치
- float16이면 파라미터당 2바이트
- 7B면 대략 14GB 수준(단순 계산 기준)이라 16GB VRAM에서도 빡빡해질 수 있음
KV 캐시(Key/Value cache)
- 디코딩(생성) 과정에서 토큰이 늘어날수록 선형으로 증가
- 배치 크기(batch_size)와 max_new_tokens, 프롬프트 길이가 직접적인 영향
활성화(activation) 및 임시 버퍼
- attention 연산, matmul, 커널별 workspace 등
- torch.compile/FlashAttention/SDPA 설정에 따라 피크 메모리가 달라짐

따라서 “4bit로 로드하면 끝”이 아니라, KV 캐시를 포함한 전체 메모리 예산을 잡아야 안정적으로 운영할 수 있습니다.

4bit 양자화: bitsandbytes로 VRAM 절약하기

가장 보편적인 접근은 bitsandbytes 기반 4bit 로딩입니다. 이 방식은 가중치를 4bit로 저장하고, 연산 시 일부를 적절히 변환하여 추론합니다.

설치(환경 준비)

CUDA 환경이 맞는다는 전제 하에 다음을 준비합니다.

pip install -U transformers accelerate bitsandbytes

transformers: 모델/토크나이저
accelerate: device_map/오프로딩/분산 로딩
bitsandbytes: 8bit/4bit 양자화 백엔드

4bit 로딩 코드 예제

아래 예시는 AutoModelForCausalLM에 4bit 설정을 적용하고, 자동 장치 배치를 사용합니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

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

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # 일반적으로 nf4 권장
    bnb_4bit_use_double_quant=True,       # 추가 압축(상황에 따라 유리)
    bnb_4bit_compute_dtype=torch.float16  # GPU 연산 dtype
)

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

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

prompt = "로컬에서 OOM을 줄이는 방법을 3가지로 요약해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
    )

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

포인트

device_map="auto"는 GPU VRAM이 부족하면 일부 레이어를 CPU로 내릴 수 있습니다(오프로딩의 시작점).
torch_dtype는 4bit 자체와 별개로, 연산 dtype(예: float16)에 영향을 줍니다.

오프로딩: GPU가 부족하면 CPU RAM을 쓰자

4bit로도 VRAM이 모자라면, 다음 단계는 오프로딩입니다. 오프로딩에는 보통 두 종류가 있습니다.

CPU 오프로딩: 일부 레이어를 CPU 메모리에 상주시킴
디스크 오프로딩: CPU RAM도 부족하면 디스크에 내려서 필요 시 로드(가장 느림)

transformers + accelerate 조합에서는 device_map과 max_memory를 통해 “GPU에 얼마나 올릴지”를 통제할 수 있습니다.

max_memory로 장치별 메모리 예산 잡기

아래는 GPU VRAM을 10GB 정도만 사용하도록 제한하고, 나머지는 CPU로 넘기는 패턴입니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, 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.float16,
)

max_memory = {
    0: "10GiB",        # GPU 0에 10GiB까지만
    "cpu": "48GiB"    # CPU RAM 예산
}

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

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

기대 효과

“조금 느려져도 OOM 없이 돌아가게” 만드는 데 매우 효과적입니다.
특히 8GB~~12GB VRAM GPU에서 7B~~13B급을 억지로 돌릴 때 현실적인 타협점이 됩니다.

디스크 오프로딩(최후의 수단)

CPU RAM도 빡빡하다면 디스크 오프로딩을 고려할 수 있습니다. 다만 지연(latency)이 크게 늘어납니다.

from transformers import AutoModelForCausalLM

offload_dir = "./offload"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    offload_folder=offload_dir,
    offload_state_dict=True,
)

offload_folder는 오프로딩 파일이 저장될 경로입니다.
NVMe SSD가 아니면 체감 성능이 크게 떨어질 수 있습니다.

OOM을 악화시키는 대표 원인 6가지(체크리스트)

같은 모델/같은 GPU인데 “어제는 됐는데 오늘은 OOM”이 나는 경우가 많습니다. 아래 항목을 우선 점검하세요.

1) 프롬프트가 길고 max_new_tokens가 크다

KV 캐시는 토큰 수에 비례해서 커집니다.

프롬프트 길이: 시스템 프롬프트/대화 히스토리 누적
생성 길이: max_new_tokens 과다

대응:

히스토리 요약/윈도잉 적용
max_new_tokens를 현실적으로 제한

2) batch_size를 무심코 키웠다

여러 요청을 배치로 묶으면 처리량은 오르지만 KV 캐시도 배치만큼 늘어납니다.

대응:

로컬 단일 사용자 추론이면 배치를 1로 고정

3) device_map이 의도와 다르게 잡혔다

device_map="auto"는 편하지만, 특정 환경에서는 레이어 배치가 비효율적으로 잡힐 수 있습니다.

대응:

max_memory로 예산을 명시
필요하면 device_map을 수동으로 조정(고급)

4) dtype 조합이 꼬였다

예를 들어 일부 텐서가 float32로 올라가면 메모리 사용량이 크게 증가할 수 있습니다.

대응:

torch_dtype=torch.float16 또는 환경에 따라 bfloat16
커스텀 코드에서 불필요한 float() 호출 제거

5) PyTorch CUDA 메모리 단편화

OOM 로그에 “reserved memory”가 크고 “allocated”가 상대적으로 작으면 단편화 가능성이 있습니다.

대응:

프로세스 재시작이 가장 확실
실험적으로 PYTORCH_CUDA_ALLOC_CONF 조정

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

6) 다른 프로세스가 VRAM을 잡고 있다

브라우저, 다른 노트북 커널, 다른 추론 서버 등이 VRAM을 점유할 수 있습니다.

대응:

nvidia-smi로 점유 프로세스 확인 후 정리

메모리 절약을 더 당기는 실전 팁

inference_mode 사용

추론에서는 autograd가 필요 없습니다.

import torch

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

no_grad보다 더 강하게 추론 최적화가 적용되는 경우가 있습니다.

attention 구현 선택(SDPA 등)

환경에 따라 SDPA가 더 메모리 효율적일 수 있습니다. 다만 모델/버전에 따라 차이가 있습니다.

import torch
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

이 영역은 조합에 따라 성능/메모리 특성이 달라서, “정답 설정”이라기보다 벤치마크 기반 튜닝에 가깝습니다.

생성 파라미터를 보수적으로

max_new_tokens를 줄이면 KV 캐시가 즉시 줄어듭니다.
num_beams(빔서치)는 메모리를 크게 늘립니다.

out = model.generate(
    **inputs,
    max_new_tokens=128,
    num_beams=1,
    do_sample=True,
)

에러 로그를 읽는 법: 어디서 터졌는지 확인

OOM 로그에는 보통 “추가로 필요한 메모리”, “현재 할당/예약”이 나옵니다. 여기서 중요한 건 다음입니다.

allocated가 큰데 OOM이면: 진짜로 메모리 부족
reserved가 큰데 allocated는 상대적으로 작으면: 단편화 또는 캐시 영향

또한 OOM이 generate에서만 난다면 KV 캐시/생성 길이 이슈일 가능성이 큽니다. 반대로 from_pretrained에서 바로 터지면 가중치 로딩/장치 배치가 원인인 경우가 많습니다.

운영 관점: “로컬 추론도 결국 리소스 관리”다

로컬 LLM 추론은 단순히 모델만 잘 고르면 끝이 아니라, GPU/CPU/RAM/디스크를 어떻게 나눠 쓰는지가 안정성을 좌우합니다. 이 점은 클라우드 운영에서 TLS나 런타임 이슈를 추적하는 것과 비슷하게, 증상(OOM) 뒤의 병목 자원을 찾아내는 과정입니다.

클러스터에서 특정 조건에서만 실패하는 네트워크 이슈를 파고드는 방식은 EKS에서 TLS 1.3만 실패할 때 - OpenSSL·ALPN 같은 사례와 닮아 있습니다.
“자원이 모자라서 생기는 실패”를 로그와 지표로 구조화해 해결하는 관점은 GitLab CI Docker 빌드 no space left 해결 가이드에서 디스크 자원을 다루는 방식과도 유사합니다.

결론: 4bit로 로드, 오프로딩으로 완주

정리하면 로컬 transformers 환경에서 OOM을 줄이는 가장 현실적인 순서는 다음입니다.

4bit 양자화(bitsandbytes)로 가중치 메모리 절감
device_map="auto" + max_memory로 CPU 오프로딩 설계
그래도 부족하면 디스크 오프로딩(성능 저하 감수)
마지막으로 KV 캐시를 줄이는 생성 파라미터 튜닝(max_new_tokens, num_beams, 프롬프트 길이)

이 조합만 제대로 잡아도, “VRAM이 애매해서 로컬 LLM은 불가능”하다고 느꼈던 환경에서 의외로 안정적인 추론이 가능합니다.