Decentralized Diffusion Models

TL;DR

Decentralized Diffusion Models enable scalable, cost-effective training of diffusion models across independent clusters without centralized networking, maintaining performance and reducing infrastructure costs.

Contribution

We introduce a decentralized training framework for diffusion models that eliminates the need for centralized high-bandwidth networks, allowing training across independent clusters with ensemble inference.

Findings

Decentralized diffusion models outperform standard models on ImageNet and LAION datasets.

The approach scales to 24 billion parameters using only eight GPU nodes.

Training time is reduced to less than a week for large models.

Abstract

Large-scale AI model training divides work across thousands of GPUs, then synchronizes gradients across them at each step. This incurs a significant network burden that only centralized, monolithic clusters can support, driving up infrastructure costs and straining power systems. We propose Decentralized Diffusion Models, a scalable framework for distributing diffusion model training across independent clusters or datacenters by eliminating the dependence on a centralized, high-bandwidth networking fabric. Our method trains a set of expert diffusion models over partitions of the dataset, each in full isolation from one another. At inference time, the experts ensemble through a lightweight router. We show that the ensemble collectively optimizes the same objective as a single model trained over the whole dataset. This means we can divide the training burden among a number of "compute islands," lowering infrastructure costs and improving resilience to localized GPU failures. Decentralized diffusion models empower researchers to take advantage of smaller, more cost-effective and more readily available compute like on-demand GPU nodes rather than central integrated systems. We conduct extensive experiments on ImageNet and LAION Aesthetics, showing that decentralized diffusion models FLOP-for-FLOP outperform standard diffusion models. We finally scale our approach to 24 billion parameters, demonstrating that high-quality diffusion models can now be trained with just eight individual GPU nodes in less than a week.