使用 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 组件 进行微调，而文本编码器和 VAE 在整个训练过程中保持冻结。

QLoRA 的内存优化技术

为了实现低于 10 GB 的内存占用，在 diffusers 训练流水线中集成了多种优化技术：

量化与适配

• QLoRA（量化 LoRA）: 基础模型通过 bitsandbytes 以 4 位量化格式（NF4）加载，显著降低了保存基础模型所需的内存。随后在此量化基础上以 FP16 或 BF16 精度训练 LoRA 适配器。
• LoRA（低秩适配）: LoRA 不更新完整的权重矩阵，而是学习两个更小的低秩矩阵（$A$ 和 $B$），显著减少可训练参数的数量。在演示的设置中，仅有 4,669,440 个参数可训练，而总参数量为 11,906,077,760。

计算效率

• 8 位 AdamW 优化器: 通过块级量化将优化器状态存储为 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 风格数据集（512×768 分辨率，批量大小 1，秩 4）上微调 FLUX.1-dev，观察到以下显存和时间指标：

方法	峰值显存使用	训练时间（700 步）
QLoRA	~9 GB	~41 分钟
BF16 LoRA	26 GB	未指定
BF16 Full Fine-tuning	~120 GB (est.)	未指定

对于计算能力 8.9 或更高的 NVIDIA GPU（例如 H100、RTX 4090），通过 torchao 的 FP8 训练进一步提升速度。在 H100 SXM GPU 上的运行实现了 36.57 GB 的峰值显存使用，并在约 20 分钟内完成了 700 步训练。

已训练适配器的推理策略

LoRA 适配器训练完成后，可通过两种方式使用：

1. 加载 LoRA 适配器

适配器加载在基础模型之上。这提供了最大的灵活性，允许用户在不同风格之间切换或使用 set_adapters() 组合多个适配器，而无需重新加载基础模型。

2. 将 LoRA 合并到基础模型

LoRA 权重被融合到基础模型中。此方法提供了最高的推理效率，因为它消除了适配器权重的内存开销，并且可以对合并后的模型进行重新量化以进一步节省内存。

硬件可及性

虽然针对 RTX 4090 进行了优化，但此工作流兼容更易获取的硬件。在 Google Colab T4 GPU 上也可以进行微调，只是耗时显著更长——相同的 700 步大约需要 4 小时，而在 RTX 4090 上仅需 41 分钟。