DiffusionPipe: Training Large Diffusion Models with Efficient Pipelines

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Tian, Ye, Jia, Zhen, Luo, Ziyue, Wang, Yida, Wu, Chuan
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909188072931328
author Tian, Ye
Jia, Zhen
Luo, Ziyue
Wang, Yida
Wu, Chuan
author_facet Tian, Ye
Jia, Zhen
Luo, Ziyue
Wang, Yida
Wu, Chuan
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2405_01248
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle DiffusionPipe: Training Large Diffusion Models with Efficient Pipelines
Tian, Ye
Jia, Zhen
Luo, Ziyue
Wang, Yida
Wu, Chuan
Distributed, Parallel, and Cluster Computing
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.
title DiffusionPipe: Training Large Diffusion Models with Efficient Pipelines
topic Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2405.01248