在消費者硬體上使用 QLoRA 微調 FLUX.1-dev
Hugging Face 詳細說明了一種使用 QLoRA 高效微調 FLUX.1-dev 的方法，使該過程能在消費級硬體上執行，峰值記憶體使用量低於 10 GB VRAM。透過結合 4 位元量化與低秩適應（LoRA），使用者可以在單一 GPU（例如 NVIDIA RTX 4090）上客製化高效能擴散模型，而無需企業級硬體。
FLUX.1-dev 架構與微調重點
FLUX.1-dev 由三個主要組件構成：文字編碼器（CLIP 與 T5）、Flux Transformer（主模型）以及變分自編碼器（VAE）。為了最大化效率，QLoRA 方法僅聚焦於微調 transformer component，而文字編碼器與 VAE 在整個訓練過程中保持凍結。
QLoRA 的記憶體最佳化技術
為了將記憶體佔用降低至 10 GB 以下，若干最佳化技術被整合至 diffusers 訓練流程中：
量化與適應
- QLoRA (Quantized LoRA)： 基礎模型透過
bitsandbytes以 4 位元量化格式 (NF4) 載入，極大降低保存基礎模型所需的記憶體。接著在此量化基礎上以 FP16 或 BF16 精度訓練 LoRA 適配器。 - LoRA (Low-Rank Adaptation)： LoRA 不是更新完整的權重矩陣，而是學習兩個較小的低秩矩陣 ($A$ 與 $B$)，顯著減少可訓練參數的數量。在示範的設定中，僅有 4,669,440 個參數可訓練，總參數量為 11,906,077,760。
計算效能提升
- 8-bit AdamW Optimizer： 透過區塊式量化將優化器狀態以 8 位元精度儲存，較標準 FP32 AdamW 可減少約 75% 的記憶體使用。
- Gradient Checkpointing： 此技術以計算換取記憶體，僅儲存特定的 checkpoint 激活，並在反向傳播時重新計算其他部分。
- Cache Latents： 訓練影像在訓練開始前先經過 VAE 編碼器預處理。此舉消除冗餘的 VAE 計算，並使 VAE 在訓練階段可完全從 GPU 記憶體中移除。
- Pre-computing Text Embeddings： 在訓練前一次性快取 CLIP 與 T5 文字編碼器的輸出。此舉防止文字編碼器在微調過程中佔用 GPU 記憶體。
NVIDIA RTX 4090 的效能基準
使用 NVIDIA RTX 4090（24GB VRAM）對 Alphonse Mucha 風格資料集（512x768 解析度、batch size 1、rank 4）進行 FLUX.1-dev 微調，觀測到以下記憶體與時間指標：
| 方法 | 峰值 VRAM 使用量 | 訓練時間（700 步） |
|---|---|---|
| QLoRA | ~9 GB | ~41 分鐘 |
| BF16 LoRA | 26 GB | 未指定 |
| BF16 Full Fine-tuning | ~120 GB（估計） | 未指定 |
對於具備計算能力 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 上亦可執行微調，只是所需時間顯著更長——相同的 700 步在 RTX 4090 上為 41 分鐘，而在 T4 上約需 4 小時。