tdrussell/diffusion-pipe

diffusion-pipe – 用于扩散模型的管道并行训练框架

概述

• 一个开源训练脚本，利用DeepSpeed的管道并行处理，在多个GPU上训练非常大的扩散模型（图像到图像、文本到图像、文本到视频等）。
• 支持一系列最新的扩散模型家族（SDXL、Flux、Stable Diffusion 3、Qwen-Image、Cosmos、Hunyuan、MiniMax H3等），适用于图像和视频数据。

主要功能

功能	重要性
管道并行	将单个模型拆分到多个GPU上，使无法放在一张卡上的模型也能进行训练。
混合数据+管道并行	利用剩余的GPU进行数据并行，自动扩展批量大小。
TensorBoard日志记录与评估指标	训练期间可监控损失、VRAM使用量和验证分数。
检查点与恢复	保存完整的训练状态（权重、优化器、数据加载器），无需重新运行配置即可从任意检查点恢复。
潜在变量和文本嵌入的预缓存	将繁重的VAE和文本编码器工作移到预处理步骤，为主模型释放VRAM。
统一图像/视频管道	同一代码路径处理静态图像和视频帧，视频模型可选支持音频处理。
可扩展的模型接口	添加新的扩散模型只需实现一个小的子类。
低VRAM技巧	示例配置展示了如何使用8位AdamW、块交换、激活检查点和CUDA内存分配器技巧，在单个24GB GPU上训练140亿参数模型。

快速开始

1. 克隆仓库（包括子模块）

   git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe
   

2. 创建Conda环境（需要Python 3.12）

   conda create -n diffusion-pipe python=3.12
   conda activate diffusion-pipe
   

3. 安装与GPU/CUDA版本匹配的PyTorch（例如，RTX 4090使用torch==2.9.0+cu128）。
4. 安装其余Python依赖

   pip install -r requirements.txt
   # 许多模型推荐（可选）
   pip install flash-attn
   

5. （如果需要Cosmos视频） 使用README中说明的环境变量技巧安装transformer_engine。

运行训练任务

• 编辑examples/中的示例TOML配置文件（主要是examples/main_example.toml）。调整数据集路径、输出目录、学习率等。
• 使用DeepSpeed启动，例如：

  NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 \
  deepspeed --num_gpus=4 train.py --deepspeed --config examples/hunyuan_video.toml
  

  --num_gpus告诉DeepSpeed要使用多少GPU；配置中的pipeline_stages控制每个模型被分割成多少个管道分区。
• 要恢复，添加--resume_from_checkpoint <文件夹>并重用完全相同的配置文件。

数据准备

• 将媒体文件（图像或视频）组织到一个或多个目录中。
• 为每个媒体文件创建一个同名的.txt字幕文件。
• 图像可以是任何Pillow兼容格式；视频可以是任何ImageIO兼容格式（WebP视频不支持）。
• 脚本会将VAE潜在变量和CLIP风格的文本嵌入缓存到每个数据集目录下的cache/子文件夹中。

输出布局

• 每次运行在output_dir下创建一个新的子文件夹。
• 其中包含：
  • DeepSpeed检查点文件夹（global_stepXXXX）。
  • 包含safetensors权重和PEFT适配器JSON的模型快照（epoch1/、epoch2/等）。
  • 用于可复现性的原始diffusion-pipe TOML配置。
  • 用于损失/验证曲线的TensorBoard事件文件。

平台注意事项

• 由于DeepSpeed对Windows的支持有限，Windows原生训练可能无法工作。如果必须使用Windows主机，请使用WSL 2。

相关项目

• 作者还维护qlora-pipe，一个类似的大语言模型管道并行训练器。

---

总结：diffusion-pipe是一个实用、文档完善的工具包，适合需要训练超出单GPU内存预算的最先进扩散模型的研究人员和工程师。它处理了管道并行、缓存和检查点管理的繁重工作，让您可以专注于数据和超参数。