Transformers 로컬 LLM OOM 해결 - QLoRA+4bit

로컬 GPU 한 장으로 LLM을 돌리다 보면, 가장 먼저 부딪히는 벽이 CUDA out of memory 입니다. 특히 transformers + Trainer 조합으로 미세튜닝을 시도하거나, 컨텍스트 길이를 늘려 추론을 돌리는 순간 VRAM이 순식간에 바닥납니다.

이 글에서는 OOM이 왜 발생하는지(파라미터, 옵티마이저 상태, 활성화 메모리, KV cache), 그리고 이를 현실적으로 해결하는 대표 조합인 QLoRA + 4bit(bitsandbytes 기반)를 재현 가능한 코드로 정리합니다. 목표는 “성능을 크게 잃지 않으면서도 12GB~24GB급 GPU에서 미세튜닝이 가능하게 만들기”입니다.

추가로, 메모리 이슈는 LLM뿐 아니라 프레임워크 전반에서 반복됩니다. Next.js 서버에서도 비슷한 패턴의 누수가 자주 나오니, 원인 분석 관점이 필요하면 Next.js 14 App Router 메모리 누수 9가지도 같이 참고해두면 좋습니다.

OOM의 진짜 원인: 파라미터보다 무서운 것들

OOM을 “모델이 커서”라고만 생각하면 해결이 어렵습니다. 실제 VRAM 사용량은 대략 아래 항목의 합입니다.

• 모델 파라미터: FP16이면 파라미터 수 N에 대해 대략 2 * N 바이트
• 그라디언트: 학습 시 FP16/FP32 구성에 따라 추가
• 옵티마이저 상태: Adam 계열은 m, v 상태가 있어 파라미터의 2배 이상을 더 먹기 쉬움
• 활성화(activation) 메모리: 배치, 시퀀스 길이, 레이어 수에 비례하며 체크포인팅 여부에 따라 급변
• KV cache: 추론에서 길게 생성할수록 증가(특히 긴 컨텍스트)

즉, “7B 모델이니까 14GB면 되겠지” 같은 계산은 자주 틀립니다. 학습에서는 옵티마이저 상태와 활성화가 폭발하고, 추론에서는 KV cache가 발목을 잡습니다.

해결 전략 개요: 무엇을 줄일 것인가

OOM을 해결하는 방법은 크게 네 가지 축입니다.

1. 가중치 메모리 줄이기: 8bit/4bit 양자화, CPU offload
2. 옵티마이저 상태 줄이기: QLoRA(LoRA만 학습), 8bit optimizer
3. 활성화 줄이기: gradient checkpointing, sequence length 조절, micro-batch + grad accumulation
4. KV cache 줄이기(추론): max_new_tokens 제한, 컨텍스트 길이 관리, flash_attention_2 같은 효율 커널 사용

이 중 로컬 환경에서 “효과 대비 구현 난이도”가 가장 좋은 조합이 4bit 로드 + LoRA 학습(QLoRA)입니다.

QLoRA + 4bit가 OOM에 강한 이유

QLoRA는 핵심적으로 아래를 합니다.

• 베이스 모델 가중치는 4bit로 로드해서 VRAM을 크게 절약
• 학습 가능한 파라미터는 LoRA 어댑터만 두어, 그라디언트/옵티마이저 상태가 작은 범위에서만 생김
• 결과적으로 “베이스 모델 전체를 FP16으로 업데이트”하는 비용을 피함

주의할 점은, 4bit로 로드하더라도 연산은 내부적으로 특정 dtype(bf16/fp16)로 수행될 수 있고, 활성화/옵티마이저/배치 구성에 따라 OOM이 여전히 날 수 있습니다. 하지만 출발점 자체가 훨씬 낮아집니다.

환경 준비: 필수 패키지 버전 체크

아래 조합이 가장 무난합니다.

• transformers
• accelerate
• peft
• bitsandbytes
• (권장) torch는 CUDA 버전에 맞게 설치

pip install -U transformers accelerate peft bitsandbytes datasets trl

버전 충돌이 있으면 bitsandbytes 로딩에서 막히는 경우가 많습니다. 특히 리눅스가 아닌 환경(예: Windows)에서는 제약이 있을 수 있어, 가능하면 WSL2 또는 리눅스 환경을 권장합니다.

4bit로 모델 로드하기: BitsAndBytesConfig

가장 중요한 1단계는 “모델을 4bit로 안전하게 로드”하는 것입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

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

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

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

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

model.config.use_cache = False  # 학습 시 메모리 절약(중요)

설정 포인트

• bnb_4bit_quant_type는 보통 nf4가 성능/안정성 균형이 좋습니다.
• bnb_4bit_compute_dtype는 GPU가 지원하면 bfloat16을 권장합니다(Ampere 이상에서 안정적인 편).
• 학습 시 model.config.use_cache = False는 거의 필수입니다. KV cache는 추론엔 유용하지만 학습 메모리를 크게 잡아먹습니다.

QLoRA 적용: PEFT LoRA 어댑터 붙이기

이제 베이스 모델은 4bit로 고정하고, LoRA만 학습하도록 구성합니다.

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

target_modules 팁

모델 아키텍처마다 모듈 이름이 다릅니다. LLaMA 계열은 보통 q_proj, k_proj, v_proj, o_proj가 잘 맞고, MPT/Falcon 등은 다를 수 있습니다. 모듈명을 잘못 지정하면 “학습되는 파라미터가 0”이 되거나 효과가 급감합니다.

학습 루프: TRL SFTTrainer로 OOM 줄이기

Trainer로도 가능하지만, 텍스트 SFT는 trl의 SFTTrainer가 편합니다. OOM을 줄이는 핵심은 아래입니다.

• per_device_train_batch_size를 작게(예: 1)
• gradient_accumulation_steps로 유효 배치를 키우기
• gradient_checkpointing=True
• max_seq_length를 현실적으로 제한

from datasets import load_dataset
from transformers import TrainingArguments
from trl import SFTTrainer

dataset = load_dataset("json", data_files="./train.jsonl", split="train")

args = TrainingArguments(
    output_dir="./qlora-out",
    num_train_epochs=1,
    learning_rate=2e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    logging_steps=10,
    save_steps=200,
    bf16=True,
    optim="paged_adamw_8bit",
    gradient_checkpointing=True,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=args,
    max_seq_length=1024,
    packing=False,
)

trainer.train()
trainer.save_model("./qlora-adapter")

왜 paged_adamw_8bit인가

bitsandbytes의 8bit 옵티마이저는 옵티마이저 상태 메모리를 크게 줄여줍니다. 특히 LoRA 학습이라도 시퀀스가 길거나 accumulation이 크면 이 차이가 체감됩니다.

흔한 OOM 패턴과 즉시 적용 가능한 처방

1) 시퀀스 길이만 늘리면 터진다

• 원인: 활성화 메모리와 attention 연산 비용이 O(n^2)로 커짐
• 처방:
  • max_seq_length부터 줄이기(예: 2048에서 1024로)
  • packing=True로 짧은 샘플을 묶어 낭비를 줄이기(데이터 특성에 따라)
  • gradient_checkpointing=True 유지

2) 배치 1인데도 터진다

• 원인: 모델이 너무 크거나, dtype/커널 문제, 또는 use_cache=True
• 처방:
  • model.config.use_cache = False
  • bf16이 불안정하면 fp16=True로 변경해 비교
  • device_map="auto"로 일부 레이어가 CPU로 오프로드되는지 확인(속도는 느려질 수 있음)

3) 학습은 되는데 저장/평가에서 터진다

• 원인: 평가 시 generate가 길게 돌면서 KV cache가 증가
• 처방:
  • 평가/샘플링에 max_new_tokens를 작게
  • do_sample=False로 짧고 결정적인 출력만 확인

추론에서 OOM 줄이기: KV cache 관리

학습을 끝냈다면, 어댑터를 붙여 추론할 때도 메모리 관리가 필요합니다.

from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
)
model = PeftModel.from_pretrained(base, "./qlora-adapter")
model.eval()

prompt = "요약: QLoRA가 왜 메모리를 줄이나?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        use_cache=True,
    )

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

추론에서는 use_cache=True가 속도에 유리하지만, 긴 생성에서는 OOM을 유발할 수 있습니다. “길게 생성해야 하는데 OOM”이라면 max_new_tokens를 제한하거나, 컨텍스트를 줄이거나, 더 공격적인 양자화/오프로드 전략을 검토해야 합니다.

OOM 디버깅 체크리스트(실전)

아래 순서대로 보면 대부분 원인을 잡습니다.

1. 현재 VRAM 사용량 확인: nvidia-smi로 학습 시작 전/후 비교
2. 캐시 비활성화 확인: 학습 시 use_cache가 False인지
3. 시퀀스 길이 확인: 데이터 전처리에서 예상보다 길어지지 않았는지
4. 배치와 accumulation 확인: 실효 배치가 너무 커져 있지 않은지
5. gradient checkpointing 적용 여부
6. 옵티마이저가 8bit인지: optim="paged_adamw_8bit"
7. 불필요한 텐서 참조 제거: 커스텀 루프라면 loss/outputs를 리스트에 쌓지 않기

메모리 문제는 “한 번에 한 가지씩” 바꿔야 원인-결과가 명확합니다. 서버 메모리 누수도 같은 원리로 접근하는데, 진단 관점은 Spring Boot 3 가상스레드 장애 - 블로킹 I/O 진단 같은 글의 트러블슈팅 방식과도 통합니다.

결론: 로컬에서 현실적인 최적점은 QLoRA+4bit

로컬 GPU 환경에서 transformers로 LLM을 다룰 때 OOM은 피할 수 없는 관문입니다. 하지만 접근을 “모델을 줄인다”가 아니라 “메모리 구성 요소를 분해해서 줄인다”로 바꾸면 해결책이 명확해집니다.

• 4bit 로드로 베이스 가중치 메모리를 크게 절감
• LoRA만 학습하는 QLoRA로 그라디언트/옵티마이저 상태를 제한
• checkpointing, seq length, accumulation으로 활성화 메모리를 제어

이 조합은 단순히 OOM을 피하는 수준이 아니라, 로컬에서도 반복 실험(데이터/하이퍼파라미터 튜닝)을 가능하게 만들어 개발 속도를 올려줍니다. 다음 단계로는 flash_attention_2, rope scaling, 더 정교한 데이터 packing, 멀티 GPU(accelerate)까지 확장해보면 좋습니다.