DiffusionPipe: Training Large Diffusion Models with Efficient Pipelines

TL;DR

DiffusionPipe tackles the bottleneck of training large diffusion models with pipeline parallelism by jointly optimizing backbone partitioning and non trainable component usage. It introduces a unified dynamic programming based backbone partitioning that minimizes the pipeline's critical path time $T^{max}$ under parameters $M$, $S$, and $D$, and extends this to cascaded backbones with bidirectional scheduling. A novel pipeline bubble filling strategy reuses non trainable computations and employs cross-iteration overlapping to hide latency, supported by a partial-batch processing design for long non trainable layers. Empirical results show up to 1.41x speedups over pipeline baselines and 1.28x over data parallel training, with pipeline bubbles reduced to under 5% of iteration time and the ability to use larger local batch sizes, thereby enabling more efficient training of diffusion models.

Abstract

Diffusion models have emerged as dominant performers for image generation. To support training large diffusion models, this paper studies pipeline parallel training of diffusion models and proposes DiffusionPipe, a synchronous pipeline training system that advocates innovative pipeline bubble filling technique, catering to structural characteristics of diffusion models. State-of-the-art diffusion models typically include trainable (the backbone) and non-trainable (e.g., frozen input encoders) parts. We first unify optimal stage partitioning and pipeline scheduling of single and multiple backbones in representative diffusion models with a dynamic programming approach. We then propose to fill the computation of non-trainable model parts into idle periods of the pipeline training of the backbones by an efficient greedy algorithm, thus achieving high training throughput. Extensive experiments show that DiffusionPipe can achieve up to 1.41x speedup over pipeline parallel methods and 1.28x speedup over data parallel training on popular diffusion models.

Figures (15)

• Figure 1: Training process of Stable Diffusion v2.1 rombach2022high and additional feedback of self-conditioning chen2022analog. Non-trainable components are marked in grey boxes.
• Figure 2: FIFO-1F1B schedule of a DNN. Gray blocks without numbers indicate pipeline bubbles. Potential critical paths are marked with a dashed line. Numbers indicate micro-batch index in both forward (blue) and backward (pink) steps.
• Figure 3: Bidirectional pipeline schedule of a DNN. Communication omitted. The same meaning of number and color with Fig. \ref{['fig:fifo_ofob_schedule']}. Micro-batch 0 to 3 pipeline from device 0 to 3 (down direction), while micro-batch 4 to 7 pipeline from device 3 to device 0 (up direction).
• Figure 4: Ratio of pipeline bubble time to iteration time (upper) and ratio of pipeline bubble time to non-trainable part execution time (lower) at batch size 64 using FIFO-1F1B scheduling.
• Figure 5: Execution time of non-trainable layers at batch size 64.