Published on

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

Authors

로컬에서 LLM을 돌릴 때 가장 먼저 부딪히는 벽은 OOM, 즉 GPU 메모리 부족입니다. 특히 transformers로 7B~13B 모델을 그대로 올리면 가중치만으로 VRAM이 꽉 차고, 여기에 KV 캐시와 활성화 텐서까지 더해져 추론 중에도 터집니다.

이 글은 OOM을 “모델을 더 작은 걸로 바꾸자”가 아니라, 같은 모델을 최대한 유지하면서 4bit 양자화 + 오프로딩(offloading)으로 버티게 만드는 방법을 단계별로 정리합니다. 목표는 다음 두 가지입니다.

  • VRAM이 부족해도 모델 로딩과 추론이 되게 만들기
  • 속도는 다소 느려져도, 재현 가능한 안정성을 확보하기

관련해서 KV 캐시 최적화까지 함께 보고 싶다면 아래 글도 참고하면 좋습니다.

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

OOM은 보통 “모델 파라미터가 커서”라고 끝나지만, 실제로는 다음 항목이 함께 VRAM을 잠식합니다.

  • 모델 가중치(Weights)
  • 옵티마이저 상태(학습 시) 및 그래디언트(학습 시)
  • 활성화 텐서(학습 시 특히 큼)
  • KV 캐시(추론 시 토큰이 길어질수록 증가)
  • 프레임워크 오버헤드 및 메모리 단편화

로컬 추론 기준으로는 가중치와 KV 캐시가 핵심입니다. 그래서 전략도 두 갈래로 나뉩니다.

  • 가중치를 줄이기: 8bit, 4bit 양자화
  • KV 캐시 부담을 줄이기: 짧은 컨텍스트, 캐시 정책, flash-attn 계열, GQA 모델 선택 등

이 글은 그중에서도 “가중치를 줄이고도 VRAM이 모자라면 일부를 GPU 밖으로 빼자”에 집중합니다.

4bit 양자화 개념: 왜 4bit가 체감이 큰가

bitsandbytes 기반 4bit 양자화는 대략적으로 가중치 메모리를 크게 줄입니다.

  • FP16: 파라미터당 2바이트
  • 8bit: 파라미터당 1바이트 수준
  • 4bit: 파라미터당 0.5바이트 수준(추가 스케일/제로포인트 오버헤드 존재)

실무 체감은 “7B 모델이 FP16로는 빡빡하던 GPU에서도 4bit면 로딩이 된다” 쪽입니다. 다만 4bit만으로도 부족한 GPU가 많습니다. 예를 들어 8GB VRAM에서는 7B도 상황에 따라 KV 캐시에서 터질 수 있습니다.

그때 필요한 것이 오프로딩입니다.

오프로딩이란: GPU 밖으로 일부를 밀어내는 선택

오프로딩은 한마디로 “모델 파라미터 또는 중간 상태 일부를 CPU RAM 또는 디스크로 옮기고, 필요할 때만 GPU로 가져오는 방식”입니다.

transformers에서는 크게 다음 두 계층이 엮입니다.

  • accelerate의 디바이스 맵(device map)과 오프로딩
  • bitsandbytes의 4bit/8bit 로딩

오프로딩 종류는 보통 아래로 나뉩니다.

  • CPU 오프로딩: RAM을 사용. 속도는 느려지지만 디스크보단 훨씬 낫습니다.
  • 디스크 오프로딩: RAM도 부족할 때 최후의 수단. SSD라도 체감 느립니다.

핵심은 “GPU VRAM을 넘는 부분을 어디에 둘지”를 명시하고, “레이어를 어떻게 배치할지”를 자동 또는 수동으로 결정하는 것입니다.

준비물: 패키지 버전과 CUDA 확인

환경이 엇갈리면 오프로딩은 설정은 맞는데도 로딩에서 터지거나 성능이 비정상적으로 나옵니다. 최소한 아래 조합을 맞추는 걸 권장합니다.

  • transformers
  • accelerate
  • bitsandbytes
  • torch는 CUDA 빌드

설치 예시는 아래처럼 시작하면 됩니다.

pip install -U "transformers>=4.39" "accelerate>=0.27" bitsandbytes

CUDA 가용 여부는 다음으로 확인합니다.

import torch
print(torch.cuda.is_available())
print(torch.version.cuda)
print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu")

기본 해법 1: 4bit 로딩 + 자동 디바이스 맵

가장 먼저 시도할 조합은 load_in_4bit=Truedevice_map="auto"입니다. 이 조합은 GPU에 최대한 올리고, 부족하면 CPU로 일부를 밀어내는 방향으로 동작합니다.

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.float16,
)

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 = "한국어로 로컬 LLM OOM을 줄이는 방법을 요약해줘."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

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

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

여기서 중요한 포인트는 다음입니다.

  • device_map="auto"는 편하지만, 항상 최적 배치를 보장하지 않습니다.
  • 일부 모델은 pad_token 설정이 필요할 수 있습니다.
  • 4bit에서도 KV 캐시가 길어지면 OOM이 날 수 있습니다.

그래도 “일단 로딩조차 안 되던 상태”를 벗어나는 데 가장 효과적인 첫 단계입니다.

기본 해법 2: 오프로딩 폴더 지정으로 디스크까지 사용

CPU RAM도 빡빡하면 디스크 오프로딩이 필요합니다. offload_folder를 지정하면 레이어 일부를 디스크에 내리고 필요할 때 읽어옵니다.

import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Llama-2-13b-chat-hf"  # 예시
os.makedirs("./offload", exist_ok=True)

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

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

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

디스크 오프로딩은 “돌긴 돈다”를 만들어주지만, 속도 저하가 큽니다. 운영 관점에서는 다음을 권장합니다.

  • NVMe SSD 사용
  • 오프로딩 폴더를 느린 네트워크 스토리지에 두지 않기
  • 컨텍스트 길이와 동시 요청 수를 강하게 제한하기

실전 해법: max_memory로 VRAM 상한을 강제

자동 디바이스 맵은 종종 “VRAM을 끝까지 채우려다” 런타임에서 터집니다. 이때 max_memory로 GPU와 CPU에 쓸 수 있는 메모리 상한을 명시하면 안정성이 크게 올라갑니다.

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_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

max_memory = {
    0: "7GiB",      # GPU 0은 7GiB까지만 사용
    "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,
)

포인트는 “GPU 메모리를 100퍼센트 쓰지 않게” 만드는 것입니다. 커널 워크스페이스, 캐시, 단편화 때문에 0.5~1GiB 정도의 여유가 실제로 중요합니다.

오프로딩을 더 통제하고 싶다면: 수동 device_map 개념

device_map을 수동으로 구성하면 “어떤 레이어는 GPU, 어떤 레이어는 CPU”를 직접 배치할 수 있습니다. 다만 모델 아키텍처마다 모듈 이름이 달라서, 우선 모델 구조를 출력해 이름을 확인해야 합니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.2",
    device_map="auto",
)

for name, _ in list(model.named_modules())[:50]:
    print(name)

실전에서는 완전 수동보다는 max_memory로 자동 배치를 유도하는 방식이 유지보수에 유리합니다.

추론 중 OOM을 줄이는 추가 체크리스트

4bit와 오프로딩으로 “로딩 OOM”은 줄였는데, 생성 중에 터지는 경우가 많습니다. 그때는 아래를 같이 점검해야 합니다.

1) 컨텍스트 길이와 max_new_tokens 제한

KV 캐시는 토큰 수에 비례해 늘어납니다. 긴 대화 로그를 그대로 넣으면 4bit라도 터질 수 있습니다.

  • 입력 프롬프트를 요약하거나 최근 N턴만 유지
  • max_new_tokens를 상한으로 강제
out = model.generate(
    **inputs,
    max_new_tokens=256,
    do_sample=True,
    temperature=0.7,
)

2) 배치 크기와 동시성 제한

로컬 서버로 띄우면 “동시 요청”이 곧 배치 크기입니다. 동시성이 늘면 KV 캐시도 그만큼 늘어납니다.

  • 단일 GPU 로컬은 동시 요청을 1~2로 제한
  • 큐잉을 두고 처리량을 관리

3) torch.inference_mode()로 불필요한 그래프 방지

추론인데도 그래프가 남으면 메모리 사용량이 불필요하게 커집니다.

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

4) PyTorch 메모리 단편화 완화

반복 실행 중에 단편화로 OOM이 나는 경우가 있습니다.

  • 요청 사이에 torch.cuda.empty_cache()를 남발하는 건 권장하지 않지만, 개발 단계에서 누수 확인에는 도움이 됩니다.
  • 프로세스를 재시작하는 방식으로 안정성을 확보하는 경우도 많습니다.
import torch

torch.cuda.empty_cache()

성능과 품질 트레이드오프 정리

4bit와 오프로딩은 만능이 아닙니다. 운영에서 기대치를 맞추려면 트레이드오프를 명확히 알아야 합니다.

  • 4bit 양자화: VRAM 절약 효과 큼, 품질 저하가 약간 있을 수 있음, 모델마다 민감도 다름
  • CPU 오프로딩: 속도 저하 중간, 안정성 높음
  • 디스크 오프로딩: 속도 저하 큼, “돌아가는 게 목표”일 때만 선택

로컬 환경에서 목표가 “빠른 응답”이라면 오프로딩보다 “더 작은 모델” 또는 “GQA 지원 모델” 또는 “컨텍스트 제한”이 더 낫습니다. 반대로 목표가 “특정 모델을 꼭 로컬에서”라면 오프로딩은 현실적인 해법입니다.

자주 겪는 실패 패턴과 해결 팁

패턴 1: 로딩은 되는데 첫 generate에서 OOM

  • 원인: KV 캐시가 생성되면서 폭증
  • 해결: 입력 길이 축소, max_new_tokens 축소, 동시 요청 제한

패턴 2: device_map="auto"인데도 GPU에 너무 많이 올림

  • 원인: 자동 배치가 여유분을 고려하지 않음
  • 해결: max_memory로 GPU 상한을 낮춰 강제

패턴 3: 디스크 오프로딩 후 지나치게 느림

  • 원인: I/O 병목
  • 해결: CPU RAM 오프로딩으로 바꾸거나, 오프로딩 자체를 줄이기 위해 GPU 상향 또는 더 작은 모델 선택

마무리: 추천 조합 3단계

정리하면, 로컬 LLM OOM을 줄이는 우선순위는 아래 3단계가 가장 실전적입니다.

  1. 4bit 양자화 + device_map="auto"
  2. max_memory로 GPU 상한을 낮춰 안정화
  3. 그래도 부족하면 offload_folder로 디스크 오프로딩까지 확장

그리고 생성 중 OOM은 대부분 KV 캐시와 동시성에서 터지므로, 토큰 길이 제한과 요청 큐잉이 같이 들어가야 합니다.

추가로 KV 캐시 자체를 더 깊게 최적화하는 방법은 아래 글에서 이어서 다룹니다.