diffusion-pipe：把扩散模型训练拆到多卡上的管线并行方案

它解决的是单卡放不下模型的问题

训练扩散模型时，显存瓶颈往往不在数据，而在模型本身。SDXL 或 Flux 这类模型，单张 24GB 的显卡勉强能跑，但 HunyuanVideo 或 Wan2.1 这种视频模型，参数量和中间激活值会轻易超出单卡显存。diffusion-pipe 的核心思路是管线并行，把模型的不同层分配到多张 GPU 上，每张卡只负责一部分计算。它面向的是已经有多个 GPU、但不想写分布式训练代码的研究者或爱好者。README 明确说，整个脚本围绕 DeepSpeed 的管线并行构建，这意味着它的适用边界从一开始就由 DeepSpeed 决定。

管线并行与统一模型接口

diffusion-pipe 的工作方式可以拆成两层。底层是 DeepSpeed 的 pipeline parallelism，它把模型按层切分，数据像流水线一样依次经过各卡。上层是脚本自己实现的模型抽象，README 说添加新模型只需要实现一个子类。这个抽象层把图像模型和视频模型统一起来，训练脚本本身不关心你喂给它的是 PNG 还是 MP4。数据集处理上，它支持多进程、多 GPU 预缓存 latents 和文本嵌入，这一步能显著减少训练时的重复计算。但缓存也带来了麻烦，后面会专门讲。

安装流程：PyTorch 和 CUDA 的坑

安装不是一条 pip install 命令搞定的事。README 要求先装 Miniconda，创建 Python 3.12 环境，然后手动安装 PyTorch。PyTorch 不在 requirements.txt 里，因为不同 GPU 需要不同版本，作者说 PyTorch 2.9.0 配 CUDA 12.8 在他自己的 4090 上工作正常。接着要装 nvcc，版本尽量匹配 PyTorch 的 CUDA 版本。然后是 requirements.txt 里的其余依赖。Flash Attention 也是可选的，但某些模型必须要它。如果你要训练原版 Cosmos 视频模型，还得额外装 TransformerEngine，README 给了一条带环境变量的安装命令。整个流程对 Linux 用户友好，但 Windows 原生环境基本没戏，因为 DeepSpeed 只部分支持 Windows，作者建议用 WSL 2。

配置与启动：TOML 文件是核心

训练前必须读 examples 目录里的配置文件。主示例文件 examples/main_example.toml 几乎每个设置都有注释，数据集配置在 examples/dataset.toml。你需要复制一份，改掉所有路径，包括数据集配置里的路径。启动命令是：NCCL_P2P_DISABLE="1" NCCL_IB_DISABLE="1" deepspeed --num_gpus=1 train.py --deepspeed --config examples/hunyuan_vi。注意这里 --num_gpus=1 看起来矛盾，但这是 DeepSpeed 的启动方式，实际使用的 GPU 数量在 TOML 配置里指定。数据集格式很简单：图片或视频文件配同名 .txt 文件作为说明文字，没有匹配的说明文件会打印警告，但训练会继续，用空说明代替。

缓存陷阱：不删缓存就会报 Tensor shape 错误

diffusion-pipe 最容易被忽视的问题是缓存。README 在多个版本更新里反复警告：如果模型代码或预处理逻辑变了，旧缓存文件会导致 Tensor shape 错误。比如 2026-06-07 的更新中，Flux2 去掉了 attention masking，作者明确说要么删掉 cache 文件夹，要么用 --regenerate_cache 参数，否则会报错。2026-08-06 的更新也提到，Z-Image、Flux2、Ernie-Image 的 latent scaling 逻辑改到了缓存阶段，训练这些模型时必须 --regenerate_cache。这是一个真实的维护成本，每次更新脚本或切换模型时，你都得记得处理缓存，否则训练会莫名其妙地失败。

模型支持范围广，但各有局限

README 列出的模型支持列表很长：SDXL、Flux、LTX-Video、HunyuanVideo、Cosmos、Wan2.1、Chroma、HiDream、Stable Diffusion 3、Qwen-Image、Flux 2、MiniMax H3 等等。但每个模型的支持程度不同。比如 LTX 2.3 只支持 T2I 和 T2V 训练，不支持音频。HunyuanVideo-1.5 目前也只支持 T2I 和 T2V。MiniMax H3 只支持 T2I 和 T2VA。这意味着你不能假设列表里的模型都支持全部功能，使用前必须查 docs/supported_models.md 确认具体选项和限制。另外，视频格式上，WebP 视频不被支持，因为 ImageIO 无法加载多帧 WebP，这是数据准备阶段容易踩的坑。

与同类工具的差异：DeepSpeed 绑定是双刃剑

扩散模型训练脚本有不少替代品，比如 Hugging Face 的 diffusers 训练脚本，或者 SimpleTuner 这类项目。diffusers 的训练脚本通常基于 Accelerate 做数据并行，模型必须能放进单卡，或者依赖模型并行插件。diffusion-pipe 的不同之处在于它完全围绕 DeepSpeed 的管线并行构建，这使它能够处理超过单卡显存的模型，这是数据并行做不到的。但代价是它绑死了 DeepSpeed，而 DeepSpeed 对 Windows 支持不完整，且安装配置复杂。如果你只是想在单卡上微调一个小模型，diffusers 的脚本会更轻量。diffusion-pipe 的取舍很明确：为了管线并行，接受 DeepSpeed 的复杂度。

维护与升级：依赖不锁定，更新需谨慎

requirements.txt 里的依赖大多故意不锁版本，README 建议用 pip install -r requirements.txt -U 更新。这给了灵活性，但也意味着升级可能引入不兼容。更新脚本本身要 git pull 和 git submodule update，作者特别强调后者，因为子模块可能更新到新 commit。许可证是 GPL-3.0，这意味着如果你修改并分发代码，衍生作品也必须开源。对个人训练使用影响不大，但如果你打算把它集成到商业产品里，GPL 会是一个需要认真考虑的因素。最近更新记录显示项目活跃，2026 年 8 月还在加新模型支持，但每次更新都可能伴随缓存或配置格式的变化，升级前最好先备份配置和缓存。

编辑结论

diffusion-pipe 适合已有 DeepSpeed 和 CUDA 环境、需要训练超过单卡显存上限的扩散模型的个人或小团队。它不适合刚入门、没有 Linux 经验、或依赖 Windows 原生环境的用户。采用前先确认三件事：你的 GPU 能否装对 PyTorch 和 CUDA 组合，你的数据集是否包含 WebP 视频（不支持），以及你是否接受 GPL-3.0 协议对衍生代码的约束。若这些都能接受，diffusion-pipe 是目前少有的、把管线并行和多种新模型支持打包在一起的实用脚本。

官方来源

1. Official README
2. Project repository