beautifulsoups.dev
Published on

Stable Diffusion LoRA 학습 OOM 해결 가이드

Authors

LoRA 학습을 돌리다 보면 가장 먼저 부딪히는 벽이 CUDA OOM(out of memory)입니다. 특히 8GB~12GB급 GPU에서 batch size 를 조금만 올리거나, 해상도를 768 로 올리거나, 텍스트 인코더까지 같이 학습하는 순간 VRAM이 터지는 경우가 흔합니다.

이 글은 “왜 OOM이 나는지”를 이론으로만 설명하지 않고, Stable Diffusion LoRA 학습 파이프라인에서 VRAM을 실제로 줄이는 옵션 조합을 우선순위대로 정리합니다. 대상은 kohya-ss/sd-scripts 계열(LoRA 학습에서 가장 흔함)이지만, diffusers 기반 학습에도 동일한 원리가 적용됩니다.

또한 OOM이 GPU 내부 메모리 부족인지, 시스템 메모리 부족으로 커널 OOM Killer가 개입하는 케이스인지도 구분해야 합니다. 리눅스에서 프로세스가 통째로 죽는다면 아래 글의 방식으로 원인을 추적해보는 게 좋습니다.

OOM의 3가지 대표 원인(LoRA 학습 관점)

1) 활성화(activation) 메모리 폭증

UNet의 중간 feature map이 가장 큰 비중을 차지합니다. 해상도(512 → 768)가 올라가면 픽셀 수가 2.25x 로 늘고, activation도 비슷한 비율로 커집니다. 여기에 batch size 까지 올리면 곱으로 증가합니다.

2) 옵티마이저 상태(optimizer states)

Adam 계열은 파라미터 외에 1차·2차 모멘트 상태를 들고 있어 메모리를 많이 씁니다. LoRA는 학습 파라미터 수가 상대적으로 적지만, 그래도 fp32 상태로 들고 있으면 VRAM을 체감할 정도로 차지할 수 있습니다. 이때 8bitAdam 이 효과적입니다.

3) 어텐션 구현(Attention backend)

기본 PyTorch attention은 메모리 사용량이 큽니다. xformers 또는 PyTorch의 SDPA(scaled dot product attention)를 쓰면 attention 메모리 사용량을 크게 줄일 수 있습니다.

1순위: xformers로 어텐션 메모리 절감

xformers 는 LoRA 학습에서 OOM 해결 효과가 가장 즉각적인 옵션 중 하나입니다. 특히 batch size 를 1에서 2로 올릴 때, 혹은 resolution 을 올릴 때 생기는 OOM을 버텨주는 경우가 많습니다.

설치(환경에 따라 다름)

CUDA/torch 버전 호환이 민감합니다. 가장 안전한 흐름은 다음과 같습니다.

# (권장) 가상환경에서 torch 먼저 설치
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# xformers 설치(버전에 맞는 wheel이 있으면 빠르게 설치됨)
pip install xformers

설치가 꼬이면 pip 가 소스 빌드를 시도하면서 시간이 오래 걸리거나 실패할 수 있습니다. 이때는 torch 와 CUDA 조합에 맞는 xformers wheel을 찾아 설치하거나, 차라리 PyTorch SDPA로 우회하는 것도 방법입니다(아래 참고).

kohya-ss에서 사용

대부분 학습 스크립트에서 --xformers 플래그로 켤 수 있습니다.

accelerate launch train_network.py \
  --pretrained_model_name_or_path "runwayml/stable-diffusion-v1-5" \
  --train_data_dir "./data" \
  --output_dir "./out" \
  --network_module "networks.lora" \
  --xformers

PyTorch SDPA로 대체(대안)

환경상 xformers 가 불안정하면, PyTorch 2.x에서 SDPA를 켜는 방식이 있습니다. 코드 레벨에서 다음과 같이 설정하는 패턴이 흔합니다.

import torch

# PyTorch 2.x: SDPA 커널 선택(환경에 따라 동작이 달라질 수 있음)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

프로젝트에 따라 옵션이 이미 잡혀 있거나, diffusers 내부에서 처리되는 경우도 있으니 “내가 켰는데도 적용이 안 된다”면 실제 attention backend가 무엇인지 로그로 확인하는 편이 안전합니다.

2순위: 8bitAdam(또는 Adafactor)로 옵티마이저 메모리 줄이기

LoRA는 학습 파라미터가 작지만, 옵티마이저 상태는 여전히 VRAM을 잡아먹습니다. 특히 text encoder까지 같이 학습하면(LoRA를 UNet+TE에 걸면) 효과가 더 큽니다.

bitsandbytes 설치

pip install bitsandbytes

kohya-ss에서 8bit Adam 사용

accelerate launch train_network.py \
  --optimizer_type "AdamW8bit" \
  --learning_rate 1e-4 \
  --network_dim 16 \
  --network_alpha 16

환경에 따라 optimizer 이름이 AdamW8bit 또는 AdamW8bit 계열로 다를 수 있으니, 스크립트의 옵션 목록을 확인하세요.

언제 효과가 큰가

  • network_dim 을 크게 잡았을 때(예: 64, 128)
  • UNet뿐 아니라 text encoder도 같이 학습할 때
  • batch size 를 2 이상으로 올리려 할 때

3순위: Gradient Checkpointing으로 activation 메모리 절감

gradient checkpointing 은 activation을 저장하지 않고 역전파 시 재계산하는 방식이라, VRAM을 줄이는 대신 속도가 느려집니다. OOM이 자주 난다면 속도 손해를 감수할 가치가 큽니다.

kohya-ss 계열에서는 보통 다음 플래그로 켭니다.

accelerate launch train_network.py \
  --gradient_checkpointing

체감상 batch size 를 올리기 위한 마지막 한 끗을 만들어주는 옵션인 경우가 많습니다.

4순위: Mixed Precision과 저장 dtype 정리(fp16, bf16)

mixed precision

학습을 fp16 또는 bf16 으로 돌리면 activation과 일부 연산 메모리가 줄어 OOM 가능성이 내려갑니다.

accelerate launch train_network.py \
  --mixed_precision "fp16" \
  --save_precision "fp16"
  • Ampere 이상(NVIDIA RTX 30xx/40xx, A100 등)은 bf16 도 좋은 선택입니다.
  • 다만 일부 조합에서 fp16 은 overflow/underflow로 학습이 불안정할 수 있어, loss가 튀면 bf16 또는 learning rate 조정이 필요합니다.

주의: save_precision 만 바꿔서는 VRAM이 크게 줄지 않음

저장 dtype은 주로 체크포인트 크기에 영향이 있고, 학습 중 VRAM은 mixed precision과 attention/activation 쪽이 좌우합니다.

5순위: 배치 전략(gradient accumulation으로 우회)

OOM이 나는 가장 흔한 지점은 batch size 입니다. 하지만 LoRA는 배치 1에서도 학습은 됩니다. 대신 gradient accumulation 으로 “유효 배치”를 키우는 방식이 안전합니다.

예를 들어 VRAM 때문에 train_batch_size=1 만 가능하다면:

accelerate launch train_network.py \
  --train_batch_size 1 \
  --gradient_accumulation_steps 4

이러면 유효 배치는 4 가 됩니다(속도는 느려짐). OOM을 피하면서 품질을 어느 정도 유지하는 현실적인 타협입니다.

6순위: 해상도·버킷팅(bucketing)·크롭 설정 튜닝

해상도는 VRAM에 직격탄

  • 512 는 대부분의 GPU에서 가장 안정적
  • 768 은 activation이 커져 OOM 빈도가 급증

가능하면 512 로 시작하고, 정말 필요할 때만 768 로 올리세요.

버킷팅으로 낭비 픽셀 줄이기

데이터셋 이미지 비율이 제각각이면, 무작정 정사각 리사이즈는 픽셀 낭비가 생깁니다. bucket 을 켜면 비율별로 묶어서 패딩/낭비를 줄일 수 있습니다.

accelerate launch train_network.py \
  --enable_bucket \
  --min_bucket_reso 256 \
  --max_bucket_reso 1024 \
  --bucket_reso_steps 64

버킷은 “품질”에도 영향이 있지만, 여기서는 VRAM 낭비를 줄여 OOM을 완화하는 관점에서 유용합니다.

7순위: UNet만 학습하고 Text Encoder는 끄기

Text Encoder까지 학습하면 품질이 좋아지는 경우도 있지만, 그만큼 VRAM과 시간이 늘어납니다. OOM이 심하면 우선 UNet LoRA만으로 시작하세요.

kohya-ss에서는 보통 text encoder 학습을 끄는 옵션이 따로 있거나, 관련 플래그를 주지 않으면 UNet만 학습하는 구성인 경우가 많습니다. 본인이 사용하는 스크립트에서 --train_text_encoder 같은 옵션을 확인하고, OOM 상황에서는 끄는 쪽으로 시도하세요.

OOM 디버깅: “진짜로 어디서 터지는지” 확인하는 최소 코드

학습 스크립트를 통째로 보기 어렵다면, 최소한 아래 패턴으로 피크 VRAM 을 확인할 수 있습니다.

import torch

def report(tag: str):
    torch.cuda.synchronize()
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] alloc={alloc:.0f}MiB reserved={reserved:.0f}MiB peak={peak:.0f}MiB")

torch.cuda.reset_peak_memory_stats()
report("start")

# 여기에 forward/backward/optimizer.step 호출

report("after step")
  • allocated 는 실제 텐서가 점유 중인 메모리
  • reserved 는 캐싱 allocator가 확보해 둔 메모리
  • OOM은 보통 reserved 가 높아서 생기는 게 아니라, “추가 할당이 필요한데 연속 블록이 안 나오는” 단편화 상황에서도 발생합니다

단편화가 의심되면, 실행 전 환경 변수로 allocator 설정을 조정하는 것도 도움이 됩니다.

export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128,garbage_collection_threshold:0.8"

자주 쓰는 OOM 회피 조합(현실적인 프리셋)

8GB GPU(예: RTX 3060 8GB) 기준

  • resolution=512
  • train_batch_size=1
  • gradient_accumulation_steps=4 또는 8
  • --xformers
  • --gradient_checkpointing
  • --mixed_precision fp16
  • --optimizer_type AdamW8bit
accelerate launch train_network.py \
  --pretrained_model_name_or_path "runwayml/stable-diffusion-v1-5" \
  --train_data_dir "./data" \
  --output_dir "./out" \
  --resolution 512 \
  --train_batch_size 1 \
  --gradient_accumulation_steps 8 \
  --mixed_precision "fp16" \
  --save_precision "fp16" \
  --gradient_checkpointing \
  --optimizer_type "AdamW8bit" \
  --xformers \
  --network_module "networks.lora" \
  --network_dim 16 \
  --network_alpha 16

12GB GPU(예: RTX 3060 12GB) 기준

  • 위 설정에서 gradient_accumulation_steps 를 줄여 속도 개선
  • 데이터가 크지 않다면 batch size=2 도 도전 가능(단, OOM 나면 다시 1로)

OOM이 아니라 “학습이 멈추거나 죽는” 케이스도 점검

OOM은 CUDA 에러로 명확히 뜨지만, 다음 같은 경우는 겉으로 OOM처럼 보일 수 있습니다.

  • 데이터 로더 워커가 죽어서 학습이 멈춤(num_workers 이슈)
  • 시스템 RAM 부족으로 커널 OOM Killer가 프로세스를 종료
  • 드라이버/전원/온도 문제로 GPU 리셋

특히 컨테이너 환경에서 학습한다면, 이미지 빌드가 느려서 반복 실험이 고통스러울 수 있습니다. 학습 환경을 자주 갈아엎는다면 BuildKit 캐시 최적화도 같이 해두면 효율이 좋아집니다.

결론: OOM은 “옵션 하나”가 아니라 “조합”으로 잡는다

Stable Diffusion LoRA 학습 OOM은 대개 다음 우선순위로 해결됩니다.

  1. --xformers 로 attention 메모리 절감
  2. AdamW8bit 로 옵티마이저 상태 축소
  3. --gradient_checkpointing 으로 activation 절감
  4. mixed_precision 을 fp16 또는 bf16 로
  5. batch size 를 내리고 gradient_accumulation_steps 로 유효 배치 확보
  6. 해상도는 512 부터, 버킷팅으로 픽셀 낭비 최소화

이 조합으로도 안 잡히면, 그때는 “내 GPU에서 가능한 목표치(해상도/배치/TE 학습 여부)”를 재설정해야 합니다. 무리해서 768 과 큰 배치를 고집하기보다, 안정적으로 오래 학습시키는 쪽이 결과가 더 좋게 나오는 경우가 많습니다.

Discuss on Twitter • View on GitHub