Transformers 로컬 LLM OOM - 4bit·offload 최적화

서빙이 아니라 로컬에서 LLM을 직접 돌릴 때 가장 자주 마주치는 문제는 단연 OOM(out of memory)입니다. 특히 transformers + accelerate 조합에서 모델을 로드하는 순간, 혹은 첫 generate 호출에서 VRAM이 터지는 경우가 많습니다. 이 글은 “왜 OOM이 나는지”를 메커니즘 관점에서 정리하고, 4bit 양자화와 offload(CPU·디스크) 를 이용해 실제로 로드·추론을 안정화하는 설정을 단계별로 다룹니다.

(이미지 생성 쪽 OOM 튜닝이 필요하다면 Stable Diffusion VRAM OOM 해결 - xFormers·VAE 타일링도 함께 참고하면 메모리 감각을 잡는 데 도움이 됩니다.)

OOM이 터지는 지점 3가지

로컬 LLM OOM은 대개 아래 3구간 중 하나에서 발생합니다.

1) 모델 가중치 로드 단계

• FP16으로 7B만 되어도 가중치만 대략 ~14GB 수준(대략적)이라 12GB VRAM에서 바로 터질 수 있습니다.
• device_map="auto"를 쓰더라도, 레이어 배치가 한쪽으로 쏠리거나(특히 작은 VRAM) 로딩 중 피크 메모리가 순간적으로 커져 OOM이 날 수 있습니다.

2) KV cache(어텐션 캐시) 증가

• 추론 시 메모리는 가중치만이 아니라 KV cache가 크게 잡아먹습니다.
• max_new_tokens, context length, batch size(동시 요청 수)가 늘수록 KV cache가 선형에 가깝게 증가합니다.

3) 일시적 피크(temporary tensors)

• 샘플링(Top-k, Top-p), logits 처리, beam search, use_cache 설정 등에 따라 순간 피크가 생깁니다.
• PyTorch CUDA allocator 특성상 “남는 것처럼 보이는데도” 조각화로 OOM이 날 수 있습니다.

먼저 확인할 체크리스트

아래는 튜닝 전에 반드시 확인할 기본입니다.

• GPU 확인: nvidia-smi로 VRAM 총량, 사용량, 다른 프로세스 점유 여부 확인
• PyTorch 버전과 CUDA: 드라이버-런타임 불일치가 있으면 성능과 메모리 효율이 악화될 수 있음
• transformers, accelerate, bitsandbytes 버전 호환
• 모델의 기본 torch_dtype와 권장 로딩 옵션(모델 카드)

간단한 메모리 관측 코드:

import torch

def vram(msg=""):
    if not torch.cuda.is_available():
        print("CUDA not available")
        return
    a = torch.cuda.memory_allocated() / 1024**3
    r = torch.cuda.memory_reserved() / 1024**3
    print(f"{msg} allocated={a:.2f}GB reserved={r:.2f}GB")

vram("start")

4bit 양자화: 가장 큰 레버리지

가중치 메모리부터 줄이는 게 가장 확실합니다. bitsandbytes의 4bit(NF4) 양자화는 로컬 환경에서 체감 효과가 큽니다.

권장 로딩 패턴(Transformers + BitsAndBytes)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "your-model-id"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  # GPU가 bf16 지원이면 권장
)

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.bfloat16,
    low_cpu_mem_usage=True,
)

model.eval()

옵션 해설

• load_in_4bit=True: 가중치를 4bit로 로드
• nf4: 일반적으로 품질 대비 효율이 좋아 많이 사용
• double_quant: 추가 압축으로 VRAM 절감에 도움
• compute_dtype: 연산 dtype. bfloat16 지원 GPU면 안정적이고 오버플로우에 강한 편
• low_cpu_mem_usage=True: 로딩 중 CPU RAM 피크를 줄이는 데 유리

4bit인데도 OOM이 나는 이유

• KV cache는 여전히 FP16·BF16로 잡히는 경우가 많아, 긴 컨텍스트에서 OOM이 납니다.
• device_map이 자동 배치되면서 특정 레이어가 GPU에 과도하게 몰릴 수 있습니다.

즉, 4bit는 “가중치”를 줄여주지만, “추론 중 증가분(KV cache)”까지 해결하진 않습니다. 그래서 다음 단계가 offload입니다.

Offload 전략: VRAM을 RAM으로, RAM을 디스크로

accelerate의 device mapping과 offload는 “일부 레이어는 GPU, 일부는 CPU”로 보내 VRAM을 맞추는 방식입니다.

CPU offload(가장 현실적인 타협)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "your-model-id"

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=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,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    offload_folder="./offload",  # CPU offload 시에도 버퍼/체크포인트로 사용될 수 있음
    offload_state_dict=True,
    low_cpu_mem_usage=True,
)

• offload_folder: offload에 필요한 파일을 저장할 경로
• offload_state_dict=True: 로딩 과정에서 CPU RAM 피크를 줄이는 데 도움

CPU offload의 대가:

• 속도는 느려집니다. 특히 GPU-CPU 왕복이 잦으면 체감이 큽니다.
• 하지만 “아예 못 돌리는 것”을 “느리지만 돌아가게” 만드는 가장 강력한 안전장치입니다.

디스크 offload(최후의 수단)

RAM도 부족하면 디스크로 내려야 하는데, 이 경우 NVMe SSD가 사실상 필수입니다. SATA SSD나 HDD는 지연이 커서 추론이 매우 느려질 수 있습니다.

accelerate 설정을 통해 디스크 offload를 더 강하게 유도할 수 있지만, 먼저 아래의 현실적인 순서를 권합니다.

1. 4bit 적용
2. max_new_tokens와 컨텍스트를 줄여 KV cache를 관리
3. 그래도 안 되면 CPU offload
4. 마지막으로 디스크 offload

max_memory로 강제 상한을 걸어 OOM 방지

device_map="auto"는 편하지만, VRAM이 작은 환경에서는 종종 “배치 실패”가 납니다. 이때 max_memory로 상한을 명시하면 OOM을 예방하는 데 도움이 됩니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "your-model-id"

tokenizer = AutoTokenizer.from_pretrained(model_id)

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 VRAM 상한
    "cpu": "48GiB"  # 시스템 RAM 상한
}

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

포인트:

• GPU VRAM 상한을 실제보다 약간 낮게 잡으면 로딩 피크를 흡수할 여지가 생깁니다.
• CPU RAM도 상한을 두면 시스템 전체가 스왑으로 무너지는 상황을 줄일 수 있습니다.

추론 OOM 줄이기: KV cache와 생성 파라미터

모델이 “로드는 되는데 생성에서 터지는” 경우는 대부분 KV cache 문제입니다.

1) 생성 길이 제한

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

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

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

• max_new_tokens를 줄이면 KV cache가 즉시 줄어듭니다.
• num_beams를 늘리는 beam search는 메모리 사용량을 크게 늘릴 수 있어, 로컬에서는 가급적 피하는 편이 안전합니다.

2) 배치 크기(동시 처리) 줄이기

• 한 번에 여러 프롬프트를 넣으면 KV cache가 배치만큼 늘어납니다.
• 로컬 테스트에서는 배치를 1로 시작하고, 여유가 있을 때만 늘리세요.

3) 컨텍스트 길이 관리

• 긴 대화 히스토리를 그대로 넣지 말고 요약하거나, 최근 N턴만 유지하는 전략이 필요합니다.
• RAG를 붙이는 경우에도 retrieval 결과를 무제한으로 붙이면 바로 OOM이 납니다.

메모리 조각화와 PyTorch allocator 튜닝

“분명 VRAM이 남아 보이는데 OOM”이면 조각화 가능성이 있습니다.

자주 쓰는 환경 변수

아래는 프로세스 시작 전에 설정해야 합니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,expandable_segments:True"

• max_split_size_mb: 큰 블록 분할을 제한해 조각화를 줄이는 데 도움
• expandable_segments: 상황에 따라 메모리 세그먼트 확장을 더 유연하게

또한 실험 중에는 커널 재시작이 가장 확실한 “메모리 청소”입니다. torch.cuda.empty_cache()는 “캐시 반환”일 뿐, 이미 잡힌 텐서를 없애주지는 않습니다.

흔한 실수 패턴과 해결

실수 1) 로딩 dtype를 섞어서 오히려 피크 증가

• torch_dtype와 bnb_4bit_compute_dtype를 무작정 섞으면 예상과 다른 캐스팅이 발생할 수 있습니다.
• 우선은 compute_dtype=float16 또는 bfloat16로 단순하게 시작하고, 문제가 있을 때만 조정하세요.

실수 2) device_map을 맹신

• 자동 배치는 편하지만, “내 환경의 상한”을 모릅니다.
• 반드시 max_memory를 같이 써서 강제 상한을 걸어두는 편이 안전합니다.

실수 3) 컨텍스트를 무한히 늘림

• 로컬 OOM의 2대 원인은 “가중치”와 “KV cache”입니다.
• 4bit로 가중치를 줄였는데도 OOM이면, 다음은 거의 항상 KV cache입니다.

데이터 파이프라인에서 메모리 폭발이 함께 보인다면(예: 대용량 데이터셋을 Arrow로 처리), Arrow Invalid - offset overflow 에러 원인·해결처럼 입력 데이터 구조 자체를 점검하는 것도 중요합니다.

실전 레시피: 12GB VRAM에서 7B급을 “안정적으로”

아래 조합이 가장 실패 확률이 낮았습니다.

1. 4bit(NF4) + double_quant
2. device_map="auto" + max_memory로 VRAM 상한을 실제보다 조금 낮게
3. offload_folder 지정(필요 시 CPU offload 유도)
4. 생성은 max_new_tokens를 보수적으로 시작(예: 128~256)
5. 동시 요청은 배치 1부터

예시 코드(통합):

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = "your-model-id"

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",
    max_memory={0: "10GiB", "cpu": "48GiB"},
    offload_folder="./offload",
    offload_state_dict=True,
    low_cpu_mem_usage=True,
)

prompt = "Write a concise explanation of 4-bit quantization and offloading for local LLMs."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        use_cache=True,
    )

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

마무리: OOM은 “가중치”와 “캐시”를 분리해 보자

로컬 LLM OOM을 빠르게 해결하려면, 메모리를 두 덩어리로 나눠 생각하는 게 핵심입니다.

• 가중치 메모리: 4bit 양자화가 가장 큰 효과
• 추론 증가분(KV cache): max_new_tokens, 컨텍스트 길이, 배치 크기 관리가 핵심
• 그래도 부족하면 CPU offload, 최후에 디스크 offload
• max_memory로 상한을 걸어 “자동 배치의 실패”를 줄이기

이 순서대로 접근하면 “일단 돌아가게” 만드는 데 성공 확률이 크게 올라가고, 이후에 속도 최적화(FlashAttention, 더 나은 GPU, 더 짧은 컨텍스트 전략)로 확장하기도 쉬워집니다.