beautifulsoups.dev
Published on

Stable Diffusion LoRA 학습 NaN·검은 이미지 해결

Authors

서로 다른 증상이지만, LoRA 학습에서 NaN loss 와 검은 이미지(또는 거의 단색) 는 대부분 같은 계열의 원인에서 출발합니다. 즉, 수치 안정성(precision/optimizer/learning rate) 이 깨지거나, 데이터/전처리(해상도·알파 채널·정규화) 가 모델이 기대하는 범위를 벗어날 때 발생합니다.

이 글은 kohya-ss 계열 학습 스크립트(LoRA/LyCORIS)와 diffusers 기반 학습 모두에 통하는 형태로, 원인별로 “어떤 로그/증상으로 의심하고, 무엇을 바꿔서 확인할지”를 체크리스트로 정리합니다.

1) 증상으로 빠르게 분류하기

A. 학습 중 loss 가 갑자기 NaN 으로 변함

  • 보통 특정 step에서 갑자기 발생
  • 직전까지는 정상적으로 감소하다가 폭발
  • 이후는 계속 NaN 으로 유지되거나 학습이 멈춤

의심 우선순위

  1. 학습률 과다 또는 스케줄러/워밍업 부재
  2. fp16/bf16 혼용 + 옵티마이저 조합 문제
  3. gradient 폭발(clip 미설정)
  4. 데이터에 손상 이미지/알파 채널/비정상 픽셀 값

B. 학습은 끝났는데 샘플이 검게 나오거나 거의 단색

  • 결과가 전체적으로 어둡거나 완전 블랙
  • 특정 프롬프트에서만 아니라 전반적으로 발생
  • 학습 중 loss 는 정상처럼 보일 수도 있음

의심 우선순위

  1. VAE 문제(잘못된 VAE, VAE 미스매치, half precision 디코딩 문제)
  2. 데이터 전처리(알파 채널, 색공간, 16-bit PNG, EXIF 회전)
  3. 너무 강한 LoRA(네트워크 차원/알파/학습률 과다)로 출력이 붕괴
  4. 샘플링 설정(잘못된 CFG/steps/negative)이라기보다, 대개 학습/데이터 이슈

2) 가장 흔한 원인 1: 학습률 과다와 스케줄러

LoRA는 “파라미터 수가 적다”는 이유로 학습률을 과하게 주는 실수가 잦습니다. 특히 UNet 과 Text Encoder 를 같이 학습할 때, 텍스트 인코더가 먼저 불안정해지며 NaN 으로 튀는 경우가 많습니다.

권장 시작점(안전한 베이스라인)

  • UNet lr: 1e-4 또는 5e-5
  • Text Encoder lr: 5e-6 ~ 1e-5 (가능하면 더 낮게)
  • scheduler: cosine 또는 constant_with_warmup
  • warmup: 전체 step의 3% ~ 10%
  • gradient clipping: 1.0

kohya-ss 예시 커맨드

아래 예시는 “일단 NaN부터 막는” 보수적 설정입니다.

accelerate launch train_network.py \
  --pretrained_model_name_or_path "./sd15.safetensors" \
  --train_data_dir "./train" \
  --resolution 512,512 \
  --network_module networks.lora \
  --network_dim 16 \
  --network_alpha 16 \
  --output_dir "./output" \
  --output_name "my_lora" \
  --learning_rate 1e-4 \
  --text_encoder_lr 1e-5 \
  --unet_lr 1e-4 \
  --lr_scheduler "cosine" \
  --lr_warmup_steps 200 \
  --max_grad_norm 1.0 \
  --train_batch_size 2 \
  --mixed_precision "bf16" \
  --save_every_n_epochs 1 \
  --sample_every_n_steps 200

포인트

  • --max_grad_norm 은 NaN 방지에 체감이 큽니다.
  • 가능하면 fp16 보다 bf16 이 안정적인 경우가 많습니다(지원 GPU에서).

3) 가장 흔한 원인 2: precision, 옵티마이저, 그리고 “조합 문제”

fp16 에서 NaN이 잘 나는 패턴

  • AdamW 계열에서 모멘텀/분산 추정이 불안정
  • loss scale이 자동으로 튜닝되다가 특정 step에서 언더/오버플로

대응 순서(위에서부터 시도)

  1. bf16 으로 변경(가능한 경우)
  2. 옵티마이저를 AdamW8bit 로 변경(메모리 절감 + 종종 안정성 개선)
  3. 학습률을 2배~10배 낮추기
  4. gradient clipping 추가 또는 강화

bitsandbytes 8bit 옵티마이저 예시

accelerate launch train_network.py \
  --optimizer_type "AdamW8bit" \
  --learning_rate 8e-5 \
  --max_grad_norm 1.0 \
  --mixed_precision "bf16"

주의

  • 8bit 옵티마이저는 환경에 따라 설치/호환 이슈가 있을 수 있습니다.
  • fp16 을 꼭 써야 한다면, 학습률을 더 낮추고 워밍업을 늘리는 쪽이 안전합니다.

4) 가장 흔한 원인 3: 데이터셋 문제(알파 채널/손상 파일/비정상 포맷)

검은 이미지의 “진짜 범인”이 데이터인 경우가 생각보다 많습니다.

체크리스트

  • PNG에 알파 채널 이 포함되어 있고, 학습 파이프라인이 이를 잘못 처리
  • 16-bit PNG, CMYK, HDR 등 일반적인 8-bit RGB 가 아닌 포맷
  • 파일 일부가 손상(열리긴 열리지만 디코딩 결과가 비정상)
  • EXIF 회전 정보로 인해 크롭/리사이즈가 예상과 다르게 적용

빠른 검증: 학습 폴더를 전부 RGB 8-bit로 재인코딩

아래는 Pillow로 학습 이미지를 일괄 변환하는 스크립트입니다.

from pathlib import Path
from PIL import Image

src = Path("./train")
dst = Path("./train_rgb")
dst.mkdir(parents=True, exist_ok=True)

for p in src.rglob("*"):
    if p.suffix.lower() not in [".png", ".jpg", ".jpeg", ".webp"]:
        continue
    try:
        img = Image.open(p)
        img = img.convert("RGB")
        out = dst / (p.stem + ".jpg")
        img.save(out, quality=95)
    except Exception as e:
        print("FAIL", p, e)

이 변환만으로도

  • 알파 채널로 인한 검은 배경 학습
  • 일부 손상 파일
  • 특이한 색공간 문제가 한 번에 정리되는 경우가 많습니다.

5) VAE 미스매치가 만드는 “검은 이미지”

학습은 정상인데 결과가 검게 나오는 대표 케이스가 VAE입니다.

흔한 상황

  • SD 1.5 기반인데, SDXL VAE를 끼워서 디코딩
  • 특정 커스텀 VAE를 쓰다가 half precision 디코딩에서 깨짐
  • 학습은 A VAE로 했는데, 추론은 B VAE로 해서 색/명암이 붕괴

해결 가이드

  • 학습과 추론에서 같은 VAE 를 사용
  • 검은 출력이 나오면, 우선 VAE를 “기본 VAE” 로 되돌려 비교
  • 가능하면 VAE는 fp32 로 디코딩(툴에 옵션이 있으면)

툴마다 옵션명이 다르지만, 핵심은 “VAE만이라도 float로” 입니다.

6) LoRA 과적합/과강화로 인한 출력 붕괴(검은색/단색)

학습이 진행될수록 샘플이 점점 단색으로 무너지는 경우는 보통 다음 조합입니다.

  • network_dim 이 너무 큼
  • 학습률이 높음
  • 데이터가 적음
  • epoch가 많음

조정 순서

  1. network_dim 을 줄이기: 16 또는 8 부터
  2. network_alpha 를 dim 과 동일 또는 더 낮게
  3. epoch를 줄이고, 대신 좋은 데이터/캡션 품질에 투자
  4. UNet lr을 먼저 내리고, 그 다음 Text Encoder lr을 내리기

권장 감각

  • 데이터가 수십 장 수준이면 dim 4~16 사이가 안전합니다.

7) 캡션/토큰 문제로 NaN이 나는 경우도 있다

드물지만 캡션에 특수문자/제어문자/비정상 유니코드가 섞여 토크나이저에서 예외적 동작을 유발하는 경우가 있습니다.

체크리스트

  • 캡션 파일 인코딩을 UTF-8 로 통일
  • 줄바꿈/탭/제로폭 문자 제거
  • 너무 긴 캡션은 잘라내기

간단 정리 스크립트(제어문자 제거)

import re
from pathlib import Path

cap_dir = Path("./train")
for p in cap_dir.rglob("*.txt"):
    s = p.read_text(encoding="utf-8", errors="ignore")
    s2 = re.sub(r"[\x00-\x1f\x7f]", " ", s)
    s2 = re.sub(r"\s+", " ", s2).strip()
    if s2 != s:
        p.write_text(s2, encoding="utf-8")

8) “한 번에” 원인 좁히는 최소 실험 플로우

시간을 아끼려면 변경을 한 번에 많이 하지 말고, 아래 순서로 10~20분 단위 실험을 권합니다.

  1. 데이터 재인코딩: RGB 8-bit JPG로 변환 후 재학습
  2. 학습률 절반: UNet lr 0.5배, Text Encoder lr 0.5배
  3. warmup 추가 + grad clip: warmup 200~500, clip 1.0
  4. precision 변경: bf16 가능하면 전환, 아니면 fp16 유지하되 lr 더 낮춤
  5. VAE 통일: 학습/추론 VAE 동일화, 검은 출력이면 기본 VAE로 비교
  6. dim 축소: dim 32 이상이면 16 또는 8 로 낮춰보기

이 과정을 거치면 대부분의 NaN/검은 이미지 이슈는 “어느 범주인지”가 드러납니다.

9) 로그/모니터링 팁: NaN을 조기에 잡기

학습이 NaN으로 터질 때는 이미 늦은 경우가 많습니다. 아래 지표를 같이 보세요.

  • step별 loss가 갑자기 큰 값으로 점프하는지
  • grad norm이 비정상적으로 커지는지
  • 샘플 이미지가 특정 step부터 급격히 어두워지는지

운영 관점에서의 “재시작 루프” 대응은 인프라에서도 자주 겪는 문제입니다. 학습 프로세스가 크래시 후 자동 재시작되며 원인을 숨기는 경우가 있는데, 이런 패턴은 서비스 운영에서의 무한 재시작과 유사합니다. 필요하면 systemd 서비스 무한 재시작 - StartLimit 해결처럼 “재시작 정책과 로그 보존”을 먼저 정리하는 접근이 디버깅 시간을 크게 줄입니다.

또한 긴 학습을 돌릴수록 “지연/처리량”도 중요해집니다. 추론/서빙까지 염두에 둔다면 양자화나 런타임 최적화 관점도 참고할 만합니다. 예를 들어 ONNX Runtime로 LLM INT4 양자화와 지연 개선의 접근은 모델은 다르지만, 병목을 수치화하고 줄이는 방식 자체는 비슷합니다.

10) 자주 묻는 설정 조합 Q&A

Q1. Text Encoder를 같이 학습하면 NaN이 더 잘 나나요?

네. 일반적으로 TE가 더 민감합니다. 먼저 UNet만 학습해 안정성을 확인한 뒤, TE는 낮은 lr로 추가하는 방식이 안전합니다.

Q2. 검은 이미지가 나오면 무조건 VAE 문제인가요?

아닙니다. VAE는 “가장 흔한 원인 중 하나”지만, 알파 채널/비정상 포맷/데이터 손상도 매우 흔합니다. 그래서 데이터 재인코딩을 1순위로 권합니다.

Q3. dim 을 키우면 더 잘 학습되니 크게 잡는 게 좋은가요?

데이터가 충분하고 목적이 명확할 때만 그렇습니다. 데이터가 적을수록 큰 dim 은 과적합과 출력 붕괴를 부르기 쉽습니다.

11) 결론: NaN과 검은 이미지는 “수치 안정성 + 데이터”로 푼다

정리하면,

  • NaN loss 는 대부분 학습률, precision, optimizer, grad clip로 해결됩니다.
  • 검은 이미지 는 VAE 미스매치 또는 데이터 포맷(특히 알파 채널)에서 자주 발생합니다.

가장 효율적인 해결 루틴은 다음 3줄입니다.

  1. 학습 이미지를 전부 RGB 8-bit 로 재인코딩
  2. lr 낮추고 warmup 과 max_grad_norm 추가
  3. 학습/추론 VAE를 동일하게 맞추고 기본 VAE로 비교

이 3가지만 해도 재현되는 문제의 대부분이 사라지거나, 최소한 원인 범위가 급격히 좁혀집니다.

Discuss on Twitter • View on GitHub