tdrussell/diffusion-pipe
A pipeline parallel training script for diffusion models.
What it solves
它通过管线并行，使得过于庞大而无法在单个 GPU 上容纳的扩散模型得以训练。提供一个统一的框架，同时支持图像和视频生成模型，降低管理不同架构的复杂度。
How it works
该项目使用 DeepSpeed 的数据与管线并行混合方式，将模型分布到多个 GPU。为优化 VRAM 使用，实现了将潜在向量和文本嵌入预先缓存到磁盘的机制，使 VAE 与文本编码器在实际训练时可以卸载。还支持 8 位优化器、块交换、激活检查点等省显存技术。
Who it’s for
需要微调或训练大规模图像和视频扩散模型（如 Flux、SDXL、或 HunyuanVideo）的 AI 研究者和开发者，且单个大型 GPU 的显存不足。
Highlights
- Broad Model Support: 兼容众多现代模型，包括 SDXL、Flux、LTX-Video、HunyuanVideo、Cosmos 与 Wan2.1。
- Pipeline Parallelism: 专为处理超过单个 GPU 内存上限的模型而设计。
- Unified Image/Video Workflow: 在同一训练脚本中同时支持图像与视频两种模式。
- Efficient Caching: 多进程、多 GPU 预缓存嵌入和潜在向量，加速训练并节省显存。
- Training Management: 包含 Tensorboard 日志、在保留测试集上的评估，以及可用于恢复训练的状态检查点。