tdrussell/diffusion-pipe

diffusion‑pipe – Pipeline‑parallel training for diffusion models

What it is

• An open‑source training script that lets you train very large diffusion models (image‑to‑image, text‑to‑image, text‑to‑video, etc.) across multiple GPUs using DeepSpeed’s pipeline‑parallelism.
• Supports a long list of recent diffusion families (SDXL, Flux, Stable Diffusion 3, Qwen‑Image, Cosmos, Hunyuan, MiniMax H3, etc.) and works for both image and video data.

Key capabilities

Feature	Why it matters
Pipeline parallelism	Splits a single model across several GPUs so models that don’t fit on one card can still be trained.
Hybrid data‑ + pipeline‑parallelism	Uses the remaining GPUs for data parallelism, automatically scaling the batch size.
TensorBoard logging & evaluation metrics	Lets you monitor loss, VRAM usage, and validation scores while training.
Checkpointing & resume	Saves full training state (weights, optimizer, dataloader) and can restart from any checkpoint without re‑running the config.
Pre‑caching of latents & text embeddings	Moves the heavy VAE and text‑encoder work to a preprocessing step, freeing VRAM for the main model.
Unified image/video pipeline	The same code path handles still images and video frames, with optional audio handling for video models.
Extensible model interface	Adding a new diffusion model only requires a small subclass implementation.
Low‑VRAM tricks	Example configs show how to use 8‑bit AdamW, block swapping, activation checkpointing, and CUDA memory‑allocator tricks to train 14‑B‑parameter models on a single 24 GB GPU.

How to get started

1. Clone the repo (including submodules)

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

2. Create a Conda environment (Python 3.12 is required)

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

3. Install PyTorch that matches your GPU/CUDA version (e.g. torch==2.9.0+cu128 for an RTX 4090).
4. Install the remaining Python deps

   pip install -r requirements.txt
   # optional but recommended for many models
   pip install flash-attn
   

5. (If you need Cosmos video) install transformer_engine with the environment‑variable hack shown in the README.

Running a training job

• Edit one of the example TOML configs in examples/ (the main one is examples/main_example.toml). Adjust dataset paths, output directory, learning‑rate, etc.
• Launch with DeepSpeed, e.g.:

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

  --num_gpus tells DeepSpeed how many GPUs to use; pipeline_stages inside the config controls how many pipeline partitions each model gets.
• To resume, add --resume_from_checkpoint <folder> and reuse the exact same config file.

Data preparation

• Organise media files (images or videos) in one or more directories.
• For each media file create a .txt caption file with the same base name.
• Images can be any Pillow‑compatible format; videos any ImageIO‑compatible format (WebP videos are not supported).
• The script will cache VAE latents and CLIP‑style text embeddings under a cache/ sub‑folder inside each dataset directory.

Output layout

• A new sub‑folder under the output_dir is created per run.
• Inside you’ll find:
  • DeepSpeed checkpoint folders (global_stepXXXX).
  • Model snapshots (epoch1/, epoch2/ …) containing safetensors weights and PEFT adapter JSON.
  • The original diffusion‑pipe TOML config for reproducibility.
  • TensorBoard event files for loss/validation curves.

Platform notes

• Native Windows training is unlikely to work because DeepSpeed’s Windows support is limited. Use WSL 2 if you must stay on a Windows host.

Related work

• The author also maintains qlora‑pipe, a similar pipeline‑parallel trainer for large language models.

---

Bottom line: diffusion-pipe is a practical, well‑documented toolkit for researchers and engineers who need to train state‑of‑the‑art diffusion models that exceed a single‑GPU memory budget. It handles the heavy lifting of pipeline parallelism, caching, and checkpoint management, letting you focus on data and hyper‑parameters.