소비자 하드웨어에서 QLoRA를 사용한 FLUX.1-dev 파인튜닝
Hugging Face는 QLoRA를 사용한 FLUX.1-dev의 효율적인 파인튜닝 방법을 상세히 제시했으며, 이를 통해 소비자 등급 하드웨어에서도 VRAM 피크 사용량이 10 GB 이하로 실행할 수 있습니다. 4비트 양자화와 Low-Rank Adaptation (LoRA)을 결합함으로써, 사용자는 NVIDIA RTX 4090과 같은 단일 GPU에서 엔터프라이즈급 하드웨어 없이도 고성능 디퓨전 모델을 맞춤화할 수 있습니다.
FLUX.1-dev 아키텍처 및 파인튜닝 초점
FLUX.1-dev는 텍스트 인코더(CLP와 T5), Flux Transformer(주 모델), 변분 오토인코더(VAE) 등 세 가지 주요 구성 요소로 이루어져 있습니다. 효율성을 극대화하기 위해 QLoRA 접근 방식은 transformer component만을 파인튜닝 대상으로 삼으며, 텍스트 인코더와 VAE는 학습 과정 내내 동결된 상태로 유지됩니다.
QLoRA를 위한 메모리 최적화 기법
메모리 사용량을 10 GB 이하로 만들기 위해, 여러 최적화 기법이 diffusers 학습 파이프라인에 통합됩니다:
양자화 및 적응
- QLoRA (Quantized LoRA): 기본 모델은
bitsandbytes를 통해 4비트 양자화 형식(NF4)으로 로드되며, 이는 기본 모델을 차지하는 메모리를 크게 줄여줍니다. 이후 LoRA 어댑터는 이 양자화된 기반 위에서 FP16 또는 BF16 정밀도로 학습됩니다. - LoRA (Low-Rank Adaptation): 전체 가중치 행렬을 업데이트하는 대신, LoRA는 두 개의 작은 저랭크 행렬($A$와 $B$)을 학습하여 학습 가능한 파라미터 수를 크게 감소시킵니다. 예시 설정에서는 전체 11,906,077,760개 중 4,669,440개만이 학습 가능한 파라미터입니다.
계산 효율성
- 8-bit AdamW Optimizer: 블록 단위 양자화를 사용해 옵티마이저 상태를 8비트 정밀도로 저장함으로써, 표준 FP32 AdamW에 비해 옵티마이저 메모리 사용량을 약 75% 절감합니다.
- Gradient Checkpointing: 이 기법은 특정 체크포인트 활성화값만 저장하고 역전파 시 나머지를 재계산함으로써, 계산량을 늘려 메모리를 절감합니다.
- Cache Latents: 학습 이미지가 훈련 시작 전에 VAE 인코더를 통해 사전 처리됩니다. 이를 통해 중복된 VAE 연산을 없애고, 학습 단계에서 VAE를 GPU 메모리에서 완전히 제거할 수 있습니다.
- Pre-computing Text Embeddings: CLIP 및 T5 텍스트 인코더의 출력이 훈련 전에 한 번 캐시됩니다. 이를 통해 파인튜닝 과정에서 텍스트 인코더가 GPU 메모리를 차지하는 것을 방지합니다.
NVIDIA RTX 4090에서의 성능 벤치마크
NVIDIA RTX 4090(24GB VRAM)을 사용해 Alphonse Mucha 스타일 데이터셋(512x768 해상도, 배치 크기 1, rank 4)에서 FLUX.1-dev를 파인튜닝한 결과, 다음과 같은 메모리 및 시간 지표가 관찰되었습니다:
| 방법 | 피크 VRAM 사용량 | 학습 시간 (700 스텝) |
|---|---|---|
| QLoRA | ~9 GB | ~41분 |
| BF16 LoRA | 26 GB | 명시되지 않음 |
| BF16 Full Fine-tuning | ~120 GB (est.) | 명시되지 않음 |
컴퓨팅 능력 8.9 이상인 NVIDIA GPU(예: H100, RTX 4090)를 사용하는 경우, **FP8 training via torchao**가 속도를 추가로 최적화합니다. H100 SXM GPU에서 실행한 결과, 피크 메모리 사용량이 36.57 GB에 도달했으며 700 스텝 학습을 약 20분에 완료했습니다.
학습된 어댑터를 위한 추론 전략
LoRA 어댑터가 학습되면 두 가지 방식으로 활용할 수 있습니다:
1. LoRA 어댑터 로드
어댑터는 기본 모델 위에 로드됩니다. 이는 최대한의 유연성을 제공하여, 사용자가 set_adapters()를 사용해 기본 모델을 다시 로드하지 않고도 다양한 스타일을 전환하거나 여러 어댑터를 결합할 수 있게 합니다.
2. LoRA를 기본 모델에 병합
LoRA 가중치는 기본 모델에 융합됩니다. 이 방식은 어댑터 가중치의 메모리 오버헤드를 없애고, 병합된 모델을 재양자화하여 추가 메모리를 절감할 수 있기 때문에 가장 높은 추론 효율성을 제공합니다.
하드웨어 접근성
RTX 4090에 최적화되어 있지만, 이 워크플로는 보다 접근하기 쉬운 하드웨어에서도 호환됩니다. Google Colab T4 GPU에서는 파인튜닝이 가능하지만, RTX 4090에서 41분이 걸리던 동일한 700 스텝을 수행하는 데 약 4시간이 소요됩니다.