(LoRA) 在消費級硬體上微調 FLUX.1-dev

Derek Liu

derekl35

關注

Marc Sun

marcsun13

關注

Sayak Paul

sayakpaul

關注

merve

merve

關注

Linoy Tsaban

linoyts

關注

在上一篇文章《探索 Diffusers 中的量化後端》中，我們深入探討了各種量化技術如何縮小像 FLUX.1-dev 這樣的擴散模型，使其在不大幅影響效能的情況下，大大提高了進行**推理**的可訪問性。我們看到了 bitsandbytes、torchao 等如何減少生成影像的記憶體佔用。

執行推理很酷，但要真正讓這些模型成為我們自己的，我們還需要能夠微調它們。因此，在這篇文章中，我們將探討**高效**地**微調**這些模型，單 GPU 峰值記憶體使用量低於約 10 GB 視訊記憶體。本文將指導您使用 diffusers 庫透過 QLoRA 微調 FLUX.1-dev。我們將展示 NVIDIA RTX 4090 的結果。我們還將強調如何透過 torchao 進行 FP8 訓練，以在相容硬體上進一步最佳化速度。

目錄

• 資料集
• FLUX 架構
• 使用 diffusers 對 FLUX.1-dev 進行 QLoRA 微調
  • 關鍵最佳化技術
  • 設定與結果
• 使用 torchao 進行 FP8 微調
• 使用訓練好的 LoRA 介面卡進行推理
  • 選項 1：載入 LoRA 介面卡
  • 選項 2：將 LoRA 合併到基礎模型中
• 在 Google Colab 上執行
• 結論

資料集

我們旨在微調 black-forest-labs/FLUX.1-dev，使其採用阿爾豐斯·穆夏的藝術風格，使用一個小型資料集。

FLUX 架構

該模型由三個主要元件組成

• 文字編碼器（CLIP 和 T5）
• Transformer（主模型 - Flux Transformer）
• 變分自編碼器（VAE）

在我們的 QLoRA 方法中，我們**只**關注**微調 Transformer 元件**。文字編碼器和 VAE 在整個訓練過程中保持凍結狀態。

使用 Diffusers 對 FLUX.1-dev 進行 QLoRA 微調

我們使用了一個 diffusers 訓練指令碼（在此處稍作修改，旨在用於 FLUX 模型的 DreamBooth 風格 LoRA 微調。此外，此處提供了一個簡化版本，用於復現本博文中的結果（並在 Google Colab 中使用）。讓我們檢查 QLoRA 和記憶體效率的關鍵部分

關鍵最佳化技術

LoRA（低秩適應）深入探討： LoRA 透過使用低秩矩陣跟蹤權重更新，使模型訓練更高效。LoRA 不會更新完整的權重矩陣 $W$，而是學習兩個較小的矩陣 $A$ 和 $B$。模型權重的更新為 $\Delta W = B A$，其中 $A \in \mathbb{R}^{r \times k}$ 和 $B \in \mathbb{R}^{d \times r}$。數字 $r$（稱為*秩*）遠小於原始維度，這意味著需要更新的引數更少。最後，$\alpha$ 是 LoRA 啟用的縮放因子。它影響 LoRA 對更新的影響程度，通常設定為與 $r$ 相同的值或其倍數。它有助於平衡預訓練模型和 LoRA 介面卡的影響。有關該概念的總體介紹，請檢視我們之前的博文：《使用 LoRA 進行高效 Stable Diffusion 微調》。

QLoRA：效率利器： QLoRA 透過首先以量化格式（通常透過 bitsandbytes 以 4 位格式）載入預訓練的基礎模型來增強 LoRA，從而大幅削減基礎模型的記憶體佔用。然後，它在此量化基礎之上訓練 LoRA 介面卡（通常為 FP16/BF16）。這顯著降低了儲存基礎模型所需的視訊記憶體。

例如，在HiDream 的 DreamBooth 訓練指令碼中，使用 bitsandbytes 進行 4 位量化將 LoRA 微調的峰值記憶體使用量從約 60GB 降低到約 37GB，而質量退化可忽略不計。我們在此處應用相同的原理來在消費級硬體上微調 FLUX.1。

8 位最佳化器 (AdamW)： 標準 AdamW 最佳化器以 32 位（FP32）維護每個引數的一階和二階矩估計，這會消耗大量記憶體。8 位 AdamW 使用塊級量化以 8 位精度儲存最佳化器狀態，同時保持訓練穩定性。與標準 FP32 AdamW 相比，此技術可將最佳化器記憶體使用量減少約 75%。在指令碼中啟用它非常簡單

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

optimizer = optimizer_class(
    params_to_optimize,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

梯度檢查點： 在前向傳播過程中，通常會儲存中間啟用以用於後向傳播梯度計算。梯度檢查點透過僅儲存某些*檢查點啟用*並在反向傳播期間重新計算其他啟用來權衡計算與記憶體。

if args.gradient_checkpointing:
    transformer.enable_gradient_checkpointing()

快取潛變數： 這種最佳化技術在訓練開始前，透過 VAE 編碼器預處理所有訓練影像。它將生成的潛變量表示儲存在記憶體中。在訓練期間，直接使用快取的潛變數，而不是即時編碼影像。這種方法提供兩個主要好處

1. 消除了訓練過程中冗餘的 VAE 編碼計算，加快了每個訓練步驟的速度
2. 允許 VAE 在快取後完全從 GPU 記憶體中移除。缺點是儲存所有快取的潛變數會增加 RAM 使用量，但這對於小型資料集通常是可控的。

# Cache latents before training if the flag is set
    if args.cache_latents:
        latents_cache = []
        for batch in tqdm(train_dataloader, desc="Caching latents"):
            with torch.no_grad():
                batch["pixel_values"] = batch["pixel_values"].to(
                    accelerator.device, non_blocking=True, dtype=weight_dtype
                )
                latents_cache.append(vae.encode(batch["pixel_values"]).latent_dist)
        # VAE is no longer needed, free up its memory
        del vae
        free_memory()

設定 4 位量化 (BitsAndBytesConfig)

本節演示了基礎模型的 QLoRA 配置

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

nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=bnb_4bit_compute_dtype,
)

transformer = FluxTransformer2DModel.from_pretrained(
    args.pretrained_model_name_or_path,
    subfolder="transformer",
    quantization_config=nf4_config,
    torch_dtype=bnb_4bit_compute_dtype,
)
# Prepare model for k-bit training
transformer = prepare_model_for_kbit_training(transformer, use_gradient_checkpointing=False)
# Gradient checkpointing is enabled later via transformer.enable_gradient_checkpointing() if arg is set

定義 LoRA 配置 (LoraConfig)： 介面卡被新增到量化後的 Transformer 中

transformer_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank, 
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"], # FLUX attention blocks
)
transformer.add_adapter(transformer_lora_config)
print(f"trainable params: {transformer.num_parameters(only_trainable=True)} || all params: {transformer.num_parameters()}")
# trainable params: 4,669,440 || all params: 11,906,077,760

只有這些 LoRA 引數是可訓練的。

預計算文字嵌入 (CLIP/T5)

在啟動 QLoRA 微調之前，我們可以透過一次快取文字編碼器的輸出來節省大量的視訊記憶體和掛鐘時間。

在訓練時，資料載入器只需讀取快取的嵌入，而無需重新編碼字幕，因此 CLIP/T5 編碼器無需佔用 GPU 記憶體。

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

import pandas as pd
import torch
from datasets import load_dataset
from huggingface_hub.utils import insecure_hashlib
from tqdm.auto import tqdm
from transformers import T5EncoderModel

from diffusers import FluxPipeline


MAX_SEQ_LENGTH = 77
OUTPUT_PATH = "embeddings.parquet"


def generate_image_hash(image):
    return insecure_hashlib.sha256(image.tobytes()).hexdigest()


def load_flux_dev_pipeline():
    id = "black-forest-labs/FLUX.1-dev"
    text_encoder = T5EncoderModel.from_pretrained(id, subfolder="text_encoder_2", load_in_8bit=True, device_map="auto")
    pipeline = FluxPipeline.from_pretrained(
        id, text_encoder_2=text_encoder, transformer=None, vae=None, device_map="balanced"
    )
    return pipeline


@torch.no_grad()
def compute_embeddings(pipeline, prompts, max_sequence_length):
    all_prompt_embeds = []
    all_pooled_prompt_embeds = []
    all_text_ids = []
    for prompt in tqdm(prompts, desc="Encoding prompts."):
        (
            prompt_embeds,
            pooled_prompt_embeds,
            text_ids,
        ) = pipeline.encode_prompt(prompt=prompt, prompt_2=None, max_sequence_length=max_sequence_length)
        all_prompt_embeds.append(prompt_embeds)
        all_pooled_prompt_embeds.append(pooled_prompt_embeds)
        all_text_ids.append(text_ids)

    max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024
    print(f"Max memory allocated: {max_memory:.3f} GB")
    return all_prompt_embeds, all_pooled_prompt_embeds, all_text_ids


def run(args):
    dataset = load_dataset("Norod78/Yarn-art-style", split="train")
    image_prompts = {generate_image_hash(sample["image"]): sample["text"] for sample in dataset}
    all_prompts = list(image_prompts.values())
    print(f"{len(all_prompts)=}")

    pipeline = load_flux_dev_pipeline()
    all_prompt_embeds, all_pooled_prompt_embeds, all_text_ids = compute_embeddings(
        pipeline, all_prompts, args.max_sequence_length
    )

    data = []
    for i, (image_hash, _) in enumerate(image_prompts.items()):
        data.append((image_hash, all_prompt_embeds[i], all_pooled_prompt_embeds[i], all_text_ids[i]))
    print(f"{len(data)=}")

    # Create a DataFrame
    embedding_cols = ["prompt_embeds", "pooled_prompt_embeds", "text_ids"]
    df = pd.DataFrame(data, columns=["image_hash"] + embedding_cols)
    print(f"{len(df)=}")

    # Convert embedding lists to arrays (for proper storage in parquet)
    for col in embedding_cols:
        df[col] = df[col].apply(lambda x: x.cpu().numpy().flatten().tolist())

    # Save the dataframe to a parquet file
    df.to_parquet(args.output_path)
    print(f"Data successfully serialized to {args.output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--max_sequence_length",
        type=int,
        default=MAX_SEQ_LENGTH,
        help="Maximum sequence length to use for computing the embeddings. The more the higher computational costs.",
    )
    parser.add_argument("--output_path", type=str, default=OUTPUT_PATH, help="Path to serialize the parquet file.")
    args = parser.parse_args()

    run(args)

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

透過將其與快取的 VAE 潛變數 (--cache_latents) 結合使用，您可以將活動模型縮減為僅包含量化後的 Transformer + LoRA 介面卡，從而使整個微調過程在 10 GB 的 GPU 記憶體下舒適執行。

設定與結果

為了本次演示，我們利用了 NVIDIA RTX 4090 (24GB 視訊記憶體) 來探索其效能。使用 accelerate 的完整訓練命令如下所示。

# 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 \
  train_dreambooth_lora_flux_miniature.py \
  --pretrained_model_name_or_path="black-forest-labs/FLUX.1-dev" \
  --data_df_path="embeddings_alphonse_mucha.parquet" \
  --output_dir="alphonse_mucha_lora_flux_nf4" \
  --mixed_precision="bf16" \
  --use_8bit_adam \
  --weighting_scheme="none" \
  --width=512 \
  --height=768 \
  --train_batch_size=1 \
  --repeats=1 \
  --learning_rate=1e-4 \
  --guidance_scale=1 \
  --report_to="wandb" \
  --gradient_accumulation_steps=4 \
  --gradient_checkpointing \ # can drop checkpointing when HW has more than 16 GB.
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --cache_latents \
  --rank=4 \
  --max_train_steps=700 \
  --seed="0"

RTX 4090 配置： 在我們的 RTX 4090 上，我們使用了 train_batch_size 為 1，gradient_accumulation_steps 為 4，mixed_precision="bf16"，gradient_checkpointing=True，use_8bit_adam=True，LoRA rank 為 4，解析度為 512x768。潛變數透過 cache_latents=True 進行快取。

記憶體佔用 (RTX 4090)

• QLoRA： QLoRA 微調的峰值視訊記憶體使用量約為 9GB。
• BF16 LoRA： 在相同設定下執行標準 LoRA（基礎 FLUX.1-dev 為 FP16）消耗 26 GB 視訊記憶體。
• BF16 全量微調： 在不進行記憶體最佳化的前提下，估計大約需要 120 GB 視訊記憶體。

訓練時間 (RTX 4090)： 在 RTX 4090 上，使用 train_batch_size 為 1，解析度為 512x768，對阿爾豐斯·穆夏資料集進行 700 步的微調大約需要 41 分鐘。

輸出質量： 最終的衡量標準是生成的藝術作品。以下是我們在 derekl35/alphonse-mucha-style 資料集上使用 QLoRA 微調模型的樣本

此表比較了主要的 bf16 精度結果。微調的目標是讓模型學習阿爾豐斯·穆夏獨特的風格。

提示	基礎模型輸出	QLoRA 微調輸出（穆夏風格）
“寧靜的黑髮女人，月光下的百合，旋渦狀的植物圖案，阿爾豐斯·穆夏風格”
“池塘裡的小狗，阿爾豐斯·穆夏風格”
“華麗的狐狸，戴著秋葉和漿果的項圈，置身於森林樹葉的掛毯之中，阿爾豐斯·穆夏風格”

微調後的模型很好地捕捉了穆夏標誌性的新藝術風格，這從裝飾圖案和獨特的調色盤中顯而易見。QLoRA 過程在學習新風格的同時保持了出色的保真度。

點選檢視 fp16 對比

結果幾乎相同，表明 QLoRA 在 fp16 和 bf16 混合精度下都表現出色。

模型比較：基礎模型 vs. QLoRA 微調模型 (fp16)

提示	基礎模型輸出	QLoRA 微調輸出（穆夏風格）
“寧靜的黑髮女人，月光下的百合，旋渦狀的植物圖案，阿爾豐斯·穆夏風格”
“池塘裡的小狗，阿爾豐斯·穆夏風格”
“華麗的狐狸，戴著秋葉和漿果的項圈，置身於森林樹葉的掛毯之中，阿爾豐斯·穆夏風格”

使用 TorchAO 進行 FP8 微調

對於擁有計算能力 8.9 或更高（例如 H100、RTX 4090）的 NVIDIA GPU 使用者，可以透過 torchao 庫利用 FP8 訓練實現更高的速度效率。

我們使用略微修改的 diffusers-torchao 訓練指令碼，在 H100 SXM GPU 上對 FLUX.1-dev LoRA 進行了微調。使用的命令如下

accelerate launch train_dreambooth_lora_flux.py \
  --pretrained_model_name_or_path=black-forest-labs/FLUX.1-dev \
  --dataset_name=derekl35/alphonse-mucha-style --instance_prompt="a woman, alphonse mucha style" --caption_column="text" \
  --output_dir=alphonse_mucha_fp8_lora_flux \
  --mixed_precision=bf16 --use_8bit_adam \
  --weighting_scheme=none \
  --height=768 --width=512 --train_batch_size=1 --repeats=1 \
  --learning_rate=1e-4 --guidance_scale=1 --report_to=wandb \
  --gradient_accumulation_steps=1 --gradient_checkpointing \
  --lr_scheduler=constant --lr_warmup_steps=0 --rank=4 \
  --max_train_steps=700 --checkpointing_steps=600 --seed=0 \
  --do_fp8_training --push_to_hub

訓練執行時，**峰值記憶體使用量為 36.57 GB**，並在大約 **20 分鐘**內完成。

此 FP8 微調模型的定性結果也已提供：

啟用 FP8 訓練與 torchao 的關鍵步驟包括

1. 使用 torchao.float8 中的 convert_to_float8_training 將 FP8 層**注入**模型。
2. **定義 module_filter_fn** 來指定哪些模組應該轉換為 FP8，哪些不應該。

如需更詳細的指南和程式碼片段，請參閱此要點和diffusers-torchao 儲存庫。

使用訓練好的 LoRA 介面卡進行推理

訓練完 LoRA 介面卡後，您有兩種主要的推理方法。

選項 1：載入 LoRA 介面卡

一種方法是在基礎模型之上載入您訓練好的 LoRA 介面卡。

載入 LoRA 的好處

• 靈活性： 無需重新載入基礎模型即可輕鬆切換不同的 LoRA 介面卡
• 實驗： 透過交換介面卡來測試多種藝術風格或概念
• 模組化： 使用 set_adapters() 組合多個 LoRA 介面卡以實現創意融合
• 儲存效率： 維護一個基礎模型和多個小型介面卡檔案

from diffusers import FluxPipeline, FluxTransformer2DModel, BitsAndBytesConfig
import torch 

ckpt_id = "black-forest-labs/FLUX.1-dev"
pipeline = FluxPipeline.from_pretrained(
    ckpt_id, torch_dtype=torch.float16
)
pipeline.load_lora_weights("derekl35/alphonse_mucha_qlora_flux", weight_name="pytorch_lora_weights.safetensors")

pipeline.enable_model_cpu_offload()

image = pipeline(
    "a puppy in a pond, alphonse mucha style", num_inference_steps=28, guidance_scale=3.5, height=768, width=512, generator=torch.manual_seed(0)
).images[0]
image.save("alphonse_mucha.png")

選項 2：將 LoRA 合併到基礎模型中

當您想要以單一風格實現最大效率時，可以將LoRA 權重合並到基礎模型中。

合併 LoRA 的好處

• 視訊記憶體效率： 推理過程中沒有介面卡權重的額外記憶體開銷
• 速度： 推理速度略快，因為無需執行介面卡計算
• 量化相容性： 可以對合並後的模型重新量化，以實現最大記憶體效率

from diffusers import FluxPipeline, AutoPipelineForText2Image, FluxTransformer2DModel, BitsAndBytesConfig
import torch 

ckpt_id = "black-forest-labs/FLUX.1-dev"
pipeline = FluxPipeline.from_pretrained(
    ckpt_id, text_encoder=None, text_encoder_2=None, torch_dtype=torch.float16
)
pipeline.load_lora_weights("derekl35/alphonse_mucha_qlora_flux", weight_name="pytorch_lora_weights.safetensors")
pipeline.fuse_lora()
pipeline.unload_lora_weights()

pipeline.transformer.save_pretrained("fused_transformer")

bnb_4bit_compute_dtype = torch.bfloat16

nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=bnb_4bit_compute_dtype,
)
transformer = FluxTransformer2DModel.from_pretrained(
    "fused_transformer",
    quantization_config=nf4_config,
    torch_dtype=bnb_4bit_compute_dtype,
)

pipeline = AutoPipelineForText2Image.from_pretrained(
    ckpt_id, transformer=transformer, torch_dtype=bnb_4bit_compute_dtype
)
pipeline.enable_model_cpu_offload()

image = pipeline(
    "a puppy in a pond, alphonse mucha style", num_inference_steps=28, guidance_scale=3.5, height=768, width=512, generator=torch.manual_seed(0)
).images[0]
image.save("alphonse_mucha_merged.png")

在 Google Colab 上執行

雖然我們展示了 RTX 4090 上的結果，但相同的程式碼可以在更易於訪問的硬體上執行，例如 Google Colab 中免費提供的 T4 GPU。

在 T4 上，相同的步數下，微調過程預計會顯著延長，大約需要 4 小時。這是為了可訪問性而做出的權衡，但它使得無需高階硬體即可進行自定義微調成為可能。如果在 Colab 上執行，請注意使用限制，因為 4 小時的訓練執行可能會超出限制。

結論

QLoRA 與 diffusers 庫相結合，極大地普及了定製 FLUX.1-dev 等最先進模型的能力。正如在 RTX 4090 上所演示的，高效微調唾手可得，並能產生高質量的風格適應。此外，對於擁有最新 NVIDIA 硬體的使用者，torchao 透過 FP8 精度實現了更快的訓練。

在 Hub 上分享您的創作！

分享您微調的 LoRA 介面卡是為開源社群做出貢獻的絕佳方式。它讓其他人可以輕鬆嘗試您的風格，在您的工作基礎上繼續發展，並有助於建立充滿活力的創意 AI 工具生態系統。

如果您已經訓練了 FLUX.1-dev 的 LoRA，我們鼓勵您分享它。最簡單的方法是將 --push_to_hub 標誌新增到訓練指令碼中。另外，如果您已經訓練了一個模型並希望上傳它，您可以使用以下程式碼片段。

# Prereqs:
# - pip install huggingface_hub diffusers
# - Run `huggingface-cli login` (or set HF_TOKEN env-var) once.
# - save model

from huggingface_hub import create_repo, upload_folder

repo_id = "your-username/alphonse_mucha_qlora_flux"
create_repo(repo_id, exist_ok=True)

upload_folder(
    repo_id=repo_id,
    folder_path="alphonse_mucha_qlora_flux",
    commit_message="Add Alphonse Mucha LoRA adapter"
)

檢視我們的穆夏 LoRA 和 TorchAO FP8 LoRA。您可以在此集合中找到這兩者以及其他介面卡。

我們迫不及待地想看到您的創作！