beautifulsoups.dev
Published on

LoRA 학습 속도 2배 - kohya + bf16 + 8bit AdamW

Authors

서로 다른 병목을 동시에 건드리면 LoRA 학습 체감 속도가 확 올라갑니다. 특히 kohya_ss 기반 학습에서 bf16(연산/메모리 효율)과 8bit AdamW(옵티마이저 상태 메모리 절감)를 같이 쓰면, 같은 GPU에서 더 큰 배치나 더 높은 해상도를 밀어 넣을 수 있고 결과적으로 step 처리량이 늘어납니다.

이 글은 다음을 목표로 합니다.

  • bf16로 연산을 가볍게 만들고(가능한 GPU에서)
  • 8bit AdamW로 VRAM을 아껴 배치 크기를 키우고
  • 그 결과 wall-clock 기준 학습 시간을 줄이는 구성

아래 설정은 “무조건 2배”를 보장하는 마법은 아니지만, 많은 환경에서 가장 재현성 있게 체감 성능을 끌어올리는 조합입니다.

관련해서 VRAM이 자주 터지는 환경이라면 이 글도 같이 보면 좋습니다: Stable Diffusion VRAM OOM 해결 - xFormers·SDPA·VAE 타일링

왜 bf16 + 8bit AdamW가 빨라지나

1) bf16: 속도와 안정성의 균형점

혼합정밀은 크게 fp16과 bf16을 많이 씁니다.

  • fp16: 빠르지만 오버/언더플로우에 민감해서 스케일링 이슈가 생길 수 있음
  • bf16: 지수부가 넓어 안정성이 좋고, Ampere 이후 GPU에서 학습 안정성이 높은 편

특히 LoRA는 전체 UNet을 풀 파인튜닝하는 것보다 안정적이지만, 데이터가 거칠거나 LR을 공격적으로 잡으면 fp16에서 손실이 튀는 경우가 있습니다. bf16은 이런 상황에서 더 “그냥 돌아가는” 확률이 높습니다.

2) 8bit AdamW: 옵티마이저 상태가 줄면 배치가 커진다

Adam 계열은 파라미터 외에 1차/2차 모멘트 상태를 들고 있어 VRAM을 꽤 먹습니다. LoRA 자체는 학습 파라미터 수가 적지만, 그래도 해상도나 배치를 올리면 메모리 압박이 생깁니다.

bitsandbytes의 8bit 옵티마이저는 이 상태 메모리를 압축해 VRAM을 절약합니다. 절약된 VRAM을 배치 크기 증가로 전환하면,

  • step당 처리 이미지 수가 증가
  • 같은 목표 step을 더 빨리 소화

로 이어져 wall-clock 기준 학습 시간이 줄어듭니다.

3) “2배”가 나오는 전형적인 케이스

다음 조건에서 체감 2배에 가깝게 나오는 경우가 많습니다.

  • 원래 설정이 VRAM 한계에 걸려 train_batch_size=1 고정
  • 8bit AdamW로 VRAM을 확보해 train_batch_size=2 또는 4로 상승
  • bf16로 연산 효율까지 챙겨 step time도 소폭 감소

즉, 핵심은 “배치 증가”가 만드는 레버리지입니다.

준비물: GPU/드라이버/파이썬 패키지 체크

GPU가 bf16을 제대로 지원하는지

일반적으로 다음 세대에서 bf16이 유리합니다.

  • NVIDIA Ampere(RTX 30xx, A10, A100) 이상
  • Ada(RTX 40xx, L4)에서도 좋음

다만 프레임워크/드라이버 조합에 따라 bf16이 예상만큼 빠르지 않거나, 특정 커널이 느려지는 경우도 있습니다. 반드시 간단한 실험으로 step time을 비교하세요.

bitsandbytes 설치

8bit AdamW는 보통 bitsandbytes가 필요합니다.

pip install -U bitsandbytes

설치가 됐는데도 런타임에서 모듈을 못 찾거나, 가상환경이 꼬이면 아래 체크리스트가 빠릅니다.

kohya_ss에서 권장 조합: 핵심 옵션 맵

koyha GUI/CLI는 버전마다 플래그가 조금씩 다를 수 있으니, 아래는 “의도” 중심으로 보되 실제 플래그명은 본인 리포지토리의 train_network.py --help로 확인하는 것을 권장합니다.

목표 조합

  • 정밀도: bf16
  • 옵티마이저: AdamW8bit 또는 bitsandbytes 기반 8bit AdamW
  • 메모리/속도 보조: gradient_checkpointing, xformers 또는 sdpa

CLI 예시: LoRA 학습(UNet 위주) 템플릿

아래는 실전에서 많이 쓰는 형태의 예시입니다. 경로/이름은 환경에 맞게 수정하세요.

accelerate launch \
  --mixed_precision bf16 \
  --num_cpu_threads_per_process 8 \
  train_network.py \
  --pretrained_model_name_or_path "./models/sd15.safetensors" \
  --train_data_dir "./train" \
  --output_dir "./output" \
  --output_name "my_lora" \
  --network_module "networks.lora" \
  --network_dim 16 \
  --network_alpha 16 \
  --resolution "512,512" \
  --train_batch_size 2 \
  --gradient_accumulation_steps 1 \
  --max_train_steps 3000 \
  --learning_rate 1e-4 \
  --lr_scheduler "cosine" \
  --optimizer_type "AdamW8bit" \
  --max_grad_norm 1.0 \
  --save_every_n_steps 500 \
  --logging_dir "./logs" \
  --log_with "tensorboard" \
  --caption_extension ".txt" \
  --shuffle_caption \
  --keep_tokens 0 \
  --gradient_checkpointing \
  --xformers

포인트는 다음입니다.

  • --mixed_precision bf16: bf16을 켜서 연산/메모리 효율 확보
  • --optimizer_type "AdamW8bit": 옵티마이저 상태 VRAM 절감
  • --train_batch_size: 절약된 VRAM을 배치 증가로 전환(여기서 속도 체감이 큼)
  • --gradient_checkpointing: VRAM을 더 줄이지만 약간 느려질 수 있음. 배치 증가로 상쇄되는지 측정 필요

xformers/sdpa는 환경에 따라 유불리가 있으니, VRAM 압박이 심할 때 우선 켜고 step time을 비교하세요. VRAM 관련 튜닝은 위에서 언급한 OOM 글을 함께 참고하면 시행착오가 줄어듭니다.

“속도 2배”를 위해 가장 먼저 만져야 할 레버

1) 배치 크기 우선, 그 다음이 step time

학습 총 시간은 대략 아래로 결정됩니다.

  • 총 시간 = step 수 * step당 시간

그런데 같은 데이터량을 처리하는 관점에서는, 배치를 키워 step 수를 줄이거나(혹은 같은 step이면 더 많은 이미지를 처리) throughput을 올리는 게 효과적입니다.

실전 체크 순서:

  1. 8bit AdamW 적용
  2. VRAM이 남으면 train_batch_size를 1에서 2로
  3. 여유가 더 있으면 4까지(품질/일반화는 데이터에 따라 달라짐)
  4. 배치가 커지면 LR도 함께 재조정

2) 배치가 커지면 LR을 그대로 두면 안 되는 경우

배치를 올리면 gradient 추정이 더 안정적이어서 LR을 약간 올려도 되는 경우가 많습니다. 하지만 LoRA는 데이터가 적을 때 과적합이 빨리 오기도 해서, 다음 중 하나를 택합니다.

  • 보수적: 배치만 올리고 LR 유지
  • 공격적: 배치 증가 비율만큼 LR을 소폭 증가(예: 배치 2배면 LR 1.0x에서 1.3x 정도)

정답은 데이터셋 크기/캡션 품질/목표 스타일에 따라 달라서, 최소 2개 설정으로 짧게 A/B 테스트하는 게 가장 빠릅니다.

품질을 망치지 않는 범위에서 더 빨라지는 추가 팁

network_dim과 alpha는 속도보다 “용량”에 영향

network_dim을 크게 하면 파라미터가 늘어 학습이 약간 느려지고, VRAM도 늘 수 있습니다. 속도만 보고 dim을 과하게 줄이면 표현력이 부족해질 수 있으니, 다음 정도가 무난합니다.

  • 스타일/캐릭터: dim=8 또는 16에서 시작
  • 범용 개념/의상/포즈: dim=16 또는 32

resolution은 step time에 직격탄

512에서 768로 올리면 연산량이 크게 증가합니다. “2배 빠르게”가 목표라면, 우선 512에서 개념을 잡고 필요할 때만 고해상도로 재학습하거나, 후반에만 해상도를 올리는 전략이 효율적입니다.

데이터 로더 병목 확인

GPU는 놀고 있는데 step이 느리면 I/O 병목일 수 있습니다.

  • 이미지가 네트워크 드라이브에 있음
  • 캡션 파일이 너무 많고 작은 파일 I/O가 잦음
  • CPU 스레드가 부족함

이때는 --num_cpu_threads_per_process를 늘리거나, 데이터 위치를 NVMe로 옮기는 것만으로도 체감이 납니다.

자주 겪는 오류와 빠른 진단

1) bf16 켰더니 오히려 느리다

가능한 원인:

  • 특정 커널이 bf16 최적화가 덜 됨
  • xformers/sdpa 조합이 비효율적으로 동작

대응:

  • bf16과 fp16을 각각 200 step 정도만 돌려 step time 비교
  • xformers on/off 비교

2) 8bit AdamW에서 NaN이나 loss 폭주

가능한 원인:

  • LR이 과함
  • 데이터가 너무 적거나 캡션이 불안정

대응:

  • LR을 0.5x로 낮춰 재시도
  • max_grad_norm을 유지하거나 더 낮춰보기(예: 1.0에서 0.5)

3) 설치는 됐는데 bitsandbytes가 로드 실패

가상환경/인터프리터가 꼬였을 가능성이 큽니다. 특히 conda와 venv를 섞어 쓰면 흔합니다. 아래 글의 체크리스트대로 정리하면 대부분 해결됩니다.

벤치마크 방법: “진짜 2배”인지 확인하는 측정 루틴

속도 최적화 글에서 가장 흔한 함정은 “느낌상 빨라졌다”입니다. 최소한 아래는 기록하세요.

  • GPU: 모델명, VRAM
  • 설정: 정밀도, 옵티마이저, 배치, 해상도
  • 지표: step/sec 또는 sec/step, GPU utilization, VRAM peak

TensorBoard 로그 외에, 콘솔 로그에서 step time을 뽑아도 됩니다.

예를 들어 다음처럼 두 실험을 비교합니다.

  • A: fp16 + AdamW + batch 1
  • B: bf16 + AdamW8bit + batch 2

B에서 step time이 비슷하거나 약간 느려져도, batch가 2배면 throughput이 이깁니다. “총 학습 시간” 관점에서 B가 더 짧아지는지 확인하는 게 핵심입니다.

결론: 가장 재현성 있는 2배 가속 공식

정리하면, kohya 기반 LoRA 학습에서 속도를 크게 올리는 가장 실전적인 공식은 다음입니다.

  1. bf16으로 혼합정밀을 설정해 안정성과 효율을 확보하고
  2. 8bit AdamW로 VRAM을 절약해
  3. 절약한 VRAM을 train_batch_size 증가로 전환한다

여기에 xformers/sdpa, gradient_checkpointing을 상황에 따라 조합하면, “OOM 없이 더 큰 배치”를 달성하기 쉬워지고 결과적으로 학습 시간이 눈에 띄게 줄어듭니다.

다만 최종 목표가 속도만이 아니라 “원하는 LoRA 품질”이라면, 배치와 LR을 함께 튜닝하면서 짧은 A/B 실험을 반복하는 것이 가장 빠른 길입니다.

Discuss on Twitter • View on GitHub