$ show tdrussell/diffusion-pipe
This repository is a training framework for diffusion models (image and video generation) using DeepSpeed pipeline parallelism. It supports over 25 model architectures including SDXL, Flux, HunyuanVideo, Cosmos, Wan2.1, and many others. The framework is designed for efficient training on limited GPU memory through techniques like block swapping (offloading intermediate layers to CPU), activation checkpointing, and fp8/quantization. It provides a unified pipeline interface with custom optimizers (Automagic, AdamW8bitKahan, GenericOptim with multiple advanced modes) and a comprehensive caching system for precomputed latents and text embeddings. The architecture is centered around the `train.py` entry point, which expects a TOML configuration file and throws pipeline layers across GPUs. Each model is implemented as a subclass of `BasePipeline` or `ComfyPipeline`, defining `to_layers()` for sequential partitioning and `prepare_inputs()` for data preprocessing. The system integrates deeply with PEFT for LoRA/LoKr adapter training and supports multi-modal data (images, videos, audio) via a unified media preprocessor. A sophisticated caching layer using SQLite metadata and binary shards avoids duplicate computations across multi-GPU runs. While primarily an application training pipeline, the repository contains research-oriented components such as custom attention wrappers, advanced optimizers with subspace momentum and polar-express updates, and experimental projection methods (GaLore, SRHT, top-k norm). The code heavily patches external libraries (ComfyUI, DeepSpeed, PEFT) to enable training features these libraries do not natively support, which introduces maintenance risk but demonstrates deep engineering skill. The framework is actively maintained with frequent commits, a detailed model support matrix, and example configurations for each supported architecture.
The LLM identified 4 files worth highlighting for reuse:
Implements a block-swapping mechanism (ModelOffloader) that offloads intermediate transformer blocks to CPU during forward and backward passes, reducing peak GPU memory. Supports re-materialization via hooks and is designed for pipeline-parallel training of large models. This isolated pattern can be adapted to any sequential model training where memory is a bottleneck.
Provides a distributed-safe caching system for precomputed latents and text embeddings. Uses SQLite for metadata management and binary shard files for data storage, with fingerprint-based invalidation to detect changes in source images or model configurations. Designed for multi-GPU training to avoid duplicate computation and can be reused as a general-purpose caching layer for any data processing pipeline.
Unified checkpoint saving and loading handler that manages full model saves, LoRA adapter saves in ComfyUI-compatible format, and DeepSpeed pipeline checkpoints. Supports signal files (save, save_quit) for graceful interruption and resumption of training. Demonstrates non-obvious engineering by handling multiple save formats and distributed state in a single, lightweight abstraction.
Wraps Flash Attention v2/v3 with a clean fallback to PyTorch's scaled_dot_product_attention and manual kernels for varlen and windowed attention. Provides a single `attention` function that is used across the Wan pipeline. This pattern is reusable in any model requiring batched attention with optional flash kernel support and demonstrates careful handling of API compatibility.