在消费级硬件上微调 FLUX.1-dev：QLoRA 与 FP8 实践

在上一篇文章探索Diffusers中的量化后端中，我们深入探讨了各种量化技术如何能压缩像FLUX.1-dev这样的扩散模型，使其在不大幅牺牲性能的情况下，显著更易于推理。我们看到了bitsandbytes、torchao等如何减少生成图像时的内存占用。

进行推理很酷，但要让这些模型真正属于我们自己，我们还需要能够微调它们。因此，在这篇文章中，我们将解决这些模型的高效微调问题，目标是在单个GPU上峰值内存使用低于约10 GB VRAM。本文将指导你使用diffusers库通过QLoRA微调FLUX.1-dev。我们将展示在NVIDIA RTX 4090上的结果。我们还将强调使用torchao进行FP8训练如何在兼容硬件上进一步优化速度。

目录

• 数据集
• FLUX架构
• QLoRA Fine-tuning FLUX.1-dev with diffusers
  • 关键优化技术
  • 设置与结果
• 使用torchao进行FP8微调
• Inference with Trained LoRA Adapters
  • 选项1：加载LoRA适配器
  • 选项2：将LoRA合并到基础模型
• 在Google Colab上运行
• 结论

数据集

我们的目标是微调black-forest-labs/FLUX.1-dev，以采用阿尔丰斯·穆夏的艺术风格，使用一个小型数据集。

FLUX架构

该模型由三个主要组件组成：

• 文本编码器（CLIP和T5）
• 变换器（主模型 - Flux Transformer）
• 变分自编码器（VAE）

在我们的QLoRA方法中，我们专注于微调变换器组件。文本编码器和VAE在整个训练过程中保持冻结。

使用Diffusers对FLUX.1-dev进行QLoRA微调

我们使用了一个diffusers训练脚本（稍作修改自此处，专为FLUX模型的DreamBooth风格LoRA微调设计。此外，一个简化版本可用于重现本博客中的结果（并用于Google Colab），可在此处获取。让我们检查QLoRA和内存效率的关键部分：

关键优化技术

LoRA（低秩适应）深入探讨：LoRA通过用低秩矩阵跟踪权重更新，使模型训练更高效。LoRA不是更新完整的权重矩阵W W ，而是学习两个较小的矩阵A A 和B B 。模型权重的更新为ΔW=BA \Delta W = B A ，其中A∈Rr×k A \in \mathbb{R}^{r \times k} 和B∈Rd×r B \in \mathbb{R}^{d \times r} 。数字r r （称为秩）远小于原始维度，这意味着要更新的参数更少。最后，α \alpha 是LoRA激活的缩放因子。这影响LoRA对更新的影响程度，通常设置为与r r 相同的值或其倍数。它有助于平衡预训练模型和LoRA适配器的影响。关于该概念的通用介绍，请查看我们之前的博客文章：使用LoRA进行高效稳定扩散微调。

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）： 在 Alphonse Mucha 数据集上微调 700 步，在 RTX 4090 上配合 train_batch_size 为 1 且分辨率为 512x768 时，耗时约 41 分钟。

输出质量： 最终的衡量标准是生成的画作。以下是我们基于 derekl35/alphonse-mucha-style 数据集进行 QLoRA 微调模型的样本：

该表比较了主要的 bf16 精度结果。微调的目标是让模型学会 Alphonse Mucha 的独特风格。

提示词	基础模型输出	QLoRA 微调输出（Mucha 风格）
"Serene raven-haired woman, moonlit lilies, swirling botanicals, alphonse mucha style"
"a puppy in a pond, alphonse mucha style"
"Ornate fox with a collar of autumn leaves and berries, amidst a tapestry of forest foliage, alphonse mucha style"

微调后的模型很好地捕捉到了 Mucha 标志性的新艺术运动风格，这从装饰图案和独特的色调中可见一斑。QLoRA 流程在学习新风格的同时保持了出色的保真度。

点击查看 fp16 对比

结果几乎完全一致，表明 QLoRA 在 fp16 和 bf16 混合精度下都能有效运行。

模型对比：基础模型 vs. QLoRA 微调（fp16）

提示词	基础模型输出	QLoRA 微调输出（Mucha 风格）
"Serene raven-haired woman, moonlit lilies, swirling botanicals, alphonse mucha style"
"a puppy in a pond, alphonse mucha style"
"Ornate fox with a collar of autumn leaves and berries, amidst a tapestry of forest foliage, alphonse mucha style"

使用 TorchAO 进行 FP8 微调

对于计算能力为 8.9 或更高的 NVIDIA GPU（例如 H100、RTX 4090）用户，通过 torchao 库利用 FP8 训练可以实现更高的速度效率。

我们在 H100 SXM GPU 上微调了 FLUX.1-dev LoRA，对diffusers-torchao 训练脚本稍作修改。使用了以下命令：

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 微调模型的定性结果也可供查看：

使用 torchao 启用 FP8 训练的关键步骤包括：

1. 注入 FP8 层：使用来自 torchao.float8 的 convert_to_float8_training 将 FP8 层注入模型。
2. 定义 module_filter_fn以指定哪些模块应转换为 FP8，哪些不应转换。

关于更详细的指南和代码片段，请参阅此 gist 以及 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"
)

看看我们的 Mucha LoRA 和 TorchAO FP8 LoRA。你可以在这个合集中找到这两个以及其他适配器。

我们迫不及待想看到你创作的作品！