(LoRA) 소비자 하드웨어에서 FLUX.1-dev 미세 조정

이전 글인 Diffusers 의 양자화 백엔드 탐구 (Exploring Quantization Backends in Diffusers) 에서 우리는 다양한 양자화 기법이 FLUX.1-dev 와 같은 확산 모델의 크기를 줄여 추론 (inference) 성능을 크게 희생하지 않으면서 훨씬 더 접근 가능하게 만들 수 있음을 살펴보았습니다. bitsandbytes, torchao 및 기타 도구들이 이미지 생성을 위한 메모리 발자국을 줄이는 방법을 확인했습니다.

추론 수행은 흥미롭지만, 이러한 모델을 우리 자신의 것으로 만들기 위해서는 미세 조정 (fine-tuning) 도 가능합니다. 따라서 이 글에서는 단일 GPU 에서 ~10 GB VRAM 의 최대 메모리 사용량을 유지하면서 이러한 모델의 미세 조정 (efficient fine-tuning) 을 효율적으로 해결합니다. 이 글은 diffusers 라이브러리를 사용하여 FLUX.1-dev 를 미세 조정하는 방법을 안내할 것입니다. NVIDIA RTX 4090 에서의 결과를 소개할 것이며, torchao를 사용한 FP8 학습이 호환 가능한 하드웨어에서 속도를 더욱 최적화할 수 있음을 강조할 것입니다.

• 데이터셋 (Dataset)
• FLUX 아키텍처 (FLUX Architecture)
• diffusers 를 사용한 QLoRA 미세 조정 FLUX.1-dev
• torchao를 사용한 FP8 미세 조정
• 학습된 LoRA 어댑터로 추론 수행 (Inference with Trained LoRA Adapters)
• Google Colab 실행 (Running on Google Colab)
• 결론 (Conclusion)

우리는 Alphonse Mucha 의 예술적 스타일을 채택하도록 black-forest-labs/FLUX.1-dev를 미세 조정할 계획이며, 이를 위해 작은 데이터셋을 사용합니다.

모델은 세 가지 주요 구성 요소로 이루어져 있습니다:

• 텍스트 인코더 (CLIP 및 T5)
• 트랜스포머 (주 모델 - Flux Transformer)
• 변분 자동 인코더 (VAE)

우리의 QLoRA 접근법에서는 트랜스포머 구성 요소에만 미세 조정을 집중합니다. 텍스트 인코더와 VAE 는 훈련 전후로 고정 (frozen)됩니다.

우리는 diffusers 훈련 스크립트를 사용했습니다 (DreamBooth 스타일의 FLUX 모델 LoRA 미세 조정용으로 설계된 것으로 약간 수정됨). 또한, 이 블로그 글의 결과를 재현하기 위한 축소 버전 (Google Colab 에서 사용함) 은 여기에서 이용 가능합니다. QLoRA 및 메모리 효율성을 위한 핵심 부분을 살펴보겠습니다:

LoRA (Low-Rank Adaptation) 심층 탐구: LoRA 는 저랭크 행렬로 가중치 업데이트를 추적함으로써 모델 훈련을 더 효율적으로 만듭니다. 전체 가중치 행렬을 업데이트하는 대신, LoRA 는 두 개의 작은 행렬인 와 를 학습합니다. 모델의 가중치 업데이트는 , 여기서 와 입니다. 랭크 (rank) 는 원래 차원보다 훨씬 작아, 업데이트해야 할 파라미터가 줄어듭니다. 마지막으로, 는 LoRA 활성화에 대한 스케일링 인자입니다. 이는 LoRA 가 업데이트에 미치는 영향을 조절하며, 종종 사전 학습된 모델의 또는 그 배수와 같은 값을 설정합니다. 이는 사전 학습된 모델과 LoRA 어댑터의 영향력을 균형을 맞추는 데 도움이 됩니다. 개념에 대한 일반적인 소개를 원하시면 이전 블로그 글인 Using LoRA for Efficient Stable Diffusion Fine-Tuning 을 확인하세요.

QLoRA: 효율성 파워하우스: QLoRA 는 양자화 형식 (일반적으로 bitsandbytes 를 통한 4 비트) 으로 사전 학습된 기본 모델을 먼저 로드하여 기본 모델의 메모리 발자국을 크게 줄이는 방식으로 LoRA 를 향상시킵니다. 그런 다음 이 양자화된 기본 위에 LoRA 어댑터 (일반적으로 FP16/BF16) 를 훈련합니다. 이는 기본 모델을 유지하는 데 필요한 VRAM 을 크게 줄입니다.

예를 들어, HiDream 4 비트 양자화 및 bitsandbytes를 사용한 DreamBooth 훈련 스크립트에서는 LoRA 미세 조정의 최대 메모리 사용량을 ~60GB 에서 ~37GB 로 줄이며 품질 저하가 거의 없거나 없습니다. 이 원리는 여기에서 소비자급 하드웨어에서 FLUX.1 을 미세 조정하는 데 적용됩니다.

8-bit Optimizer (AdamW):
Standard AdamW optimizer maintains first and second moment estimates for each parameter in 32-bit (FP32), which consumes a lot of memory. The 8-bit AdamW uses block-wise quantization to store optimizer states in 8-bit precision, while maintaining training stability. This technique can reduce optimizer memory usage by ~75% compared to standard FP32 AdamW. Enabling it in the script is straightforward:

# Check for the --use_8bit_adam flag
if args.use_8bit_adam:
    optimizer_class = bnb.optim.AdamW8bit
...

Gradient Checkpointing:
During forward pass, intermediate activations are typically stored for backward pass gradient computation. Gradient checkpointing trades computation for memory by only storing certain checkpoint activations and recomputing others during backpropagation.

if args.gradient_checkpointing:
    transformer.enable_gradient_checkpointing()

Cache Latents:
This optimization technique pre-processes all training images through the VAE encoder before the beginning of the training. It stores the resulting latent representations in memory. During the training, instead of encoding images on-the-fly, the cached latents are directly used. This approach offers two main benefits:

• eliminates redundant VAE encoding computations during training, speeding up each training step
• allows the VAE to be completely removed from GPU memory after caching. The trade-off is increased RAM usage to store all cached latents, but this is typically manageable for small datasets.

# Cache latents before training if the flag is set
if args.cache_latents:
    latents_cache = []
...

Setting up 4-bit Quantization (BitsAndBytesConfig):

This section demonstrates the QLoRA configuration for the base model:

# Determine compute dtype based on mixed precision
bnb_4bit_compute_dtype = torch.float32
if args.mixed_precision == "fp16":
...

Defining LoRA Configuration (LoraConfig):
Adapters are added to the quantized transformer:

transformer_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
...

Only these LoRA parameters become trainable.

Before we launch the QLoRA fine-tune we can save a huge chunk of VRAM and wall-clock time by caching outputs of text encoders once.

At training time the dataloader simply reads the cached embeddings instead of re-encoding the caption, so the CLIP/T5 encoder never has to sit in GPU memory.

Code

# https://github.com/huggingface/diffusers/blob/main/examples/research_projects/flux_lora_quantization/compute_embeddings.py
import argparse
import pandas as pd
...

python compute_embeddings.py \
--max_sequence_length 77 \
--output_path embeddings_alphonse_mucha.parquet

By combining this with cached VAE latents (--cache_latents)
you whittle the active model down to just the quantized transformer + LoRA adapters, keeping the whole fine-tune comfortably under 10 GB of GPU memory.

For this demonstration, we leveraged an NVIDIA RTX 4090 (24GB VRAM) to explore its performance. The full training command using accelerate
is shown below.

# You need to pre-compute the text embeddings first. See the diffusers repo.
# https://github.com/huggingface/diffusers/tree/main/examples/research_projects/flux_lora_quantization
accelerate launch --config_file=accelerate.yaml \
...

Configuration for RTX 4090:
On our RTX 4090, we used a train_batch_size of 1, `gradient_accumulation_steps"

4 중 1 개, `mixed_precision=

VRAM 효율성: 추론 (inference) 시 어댑터 가중치로 인한 추가 메모리 오버헤드 없음
속도: 어댑터 계산을 적용할 필요가 없으므로 추론이 약간 더 빠름
양자화 호환성: 최대 메모리 효율성을 위해 병합된 모델을 재양자화 (re-quantize) 가능

코드

from diffusers import FluxPipeline, AutoPipelineForText2Image, FluxTransformer2DModel, BitsAndBytesConfig
import torch
ckpt_id = "black-forest-labs/FLUX.1-dev"
...

RTX 4090 에서 결과를 소개했지만, 이 코드는 Google Colab 의 무료 T4 GPU 와 같은 더 접근 가능한 하드웨어에서도 실행할 수 있습니다.

T4 에서는 동일한 단계 수를 위해 세밀 조정 (fine-tuning) 프로세스가 약 4 시간 동안 훨씬 더 오래 걸릴 것으로 예상됩니다. 이는 접근성을 위한 트레이드 오프이지만, 고급 하드웨어 없이 커스텀 세밀 조정이 가능하게 합니다. Colab 에서 실행할 경우 사용 제한을 주의해야 하며, 4 시간 훈련 실행은 이를 초과할 수 있습니다.

QLoRA 와 diffusers 라이브러리의 결합은 최첨단 모델 (예: FLUX.1-dev) 을 커스터마이징하는 능력을 크게 민주화했습니다. RTX 4090 에서 입증된 바와 같이, 효율적인 세밀 조정은 충분히 도달 가능하며, 고품질의 스타일적 적응을 제공합니다. 또한 최신 NVIDIA 하드웨어를 가진 사용자는 torchao 를 통해 FP8 정밀도를 통해 더 빠른 훈련을 가능하게 합니다.

세밀 조정된 LoRA 어댑터를 공유하는 것은 오픈 소스 커뮤니티에 기여하는 훌륭한 방법입니다. 이는 다른 사람들이 당신의 스타일을 쉽게 시도하고, 당신의 작업 위에 구축하며, 창의적인 AI 도구의 활기찬 생태계를 조성하는 데 도움을 줍니다.

FLUX.1-dev 로 LoRA 를 훈련했다면 공유를 권장합니다. 가장 쉬운 방법은 훈련 스크립트에 --push_to_hub 플래그를 추가하는 것입니다. 대안적으로, 이미 모델을 훈련하고 업로드하고 싶다면 다음 스니펫을 사용할 수 있습니다.

# Prereqs:
# - pip install huggingface_hub diffusers
# - Run `huggingface-cli login` (또는 HF_TOKEN 환경 변수 설정) 한 번.
...

Mucha LoRA 와 TorchAO FP8 LoRA 를 확인해 보세요. 이 컬렉션에는 두 가지와 다른 어댑터가 포함되어 있습니다.

당신이 만들어낸 것을 보고 기다릴 수 없습니다!