PaperlayerView exact revision on arXiv

Source-linked AI summary

DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models

Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li, Song Han

arXiv:2402.19481v4cs.CV

TL;DR

High-resolution diffusion generation is computationally costly and too slow for interactive applications. DistriFusion splits inputs across GPUs and reuses activations from adjacent denoising steps to preserve patch interaction while overlapping communication. On SDXL, it reports up to 6.1× speedup on eight A100 GPUs without quality degradation.

  • Problem

    High-resolution image generation with diffusion models requires substantial computation, creating latency that impedes interactive applications.

    Source: §1 · p. 2

  • Method

    DistriFusion assigns image patches to GPUs and reuses prior-step activations as context for patch interactions, with asynchronous communication overlapped by computation.

    Source: §1 · pp. 2–3

  • Results

    6.1× SDXL U-Net speedup on 8 A100 GPUs is reported while preserving visual fidelity.

    Source: §1 · p. 3

  • Takeaways & Limitations

    DistriFusion provides a multi-GPU approach for reducing single-image diffusion latency without reported quality degradation.

    Source: §6 · p. 9

    NVLink is essential to fully hide communication overhead and maximize DistriFusion's speedup.

    Source: §6 · p. 9

Abstract

from arXiv · show

Diffusion models have achieved great success in synthesizing high-quality images. However, generating high-resolution images with diffusion models is still challenging due to the enormous computational costs, resulting in a prohibitive latency for interactive applications. In this paper, we propose DistriFusion to tackle this problem by leveraging parallelism across multiple GPUs. Our method splits the model input into multiple patches and assigns each patch to a GPU. However, naively implementing such an algorithm breaks the interaction between patches and loses fidelity, while incorporating such an interaction will incur tremendous communication overhead. To overcome this dilemma, we observe the high similarity between the input from adjacent diffusion steps and propose displaced patch parallelism, which takes advantage of the sequential nature of the diffusion process by reusing the pre-computed feature maps from the previous timestep to provide context for the current step. Therefore, our method supports asynchronous communication, which can be pipelined by computation. Extensive experiments show that our method can be applied to recent Stable Diffusion XL with no quality degradation and achieve up to a 6.1$\times$ speedup on eight NVIDIA A100s compared to one. Our code is publicly available at https://github.com/mit-han-lab/distrifuser.

1. Introduction

High-resolution diffusion generation raises latency that limits interactive applications, while available multi-GPU methods do not effectively accelerate single-image generation. DistriFusion uses patch parallelism and prior-step activations to preserve patch interaction while reducing latency; experiments report up to 6.1× speedup on eight A100 GPUs without loss of visual fidelity.

1. Introduction

  • High-resolution diffusion increases computation and latency, making real-time image generation challenging.The introduction contrasts Stable Diffusion's 512 × 512 images with SDXL's 1024 × 1024 images and Sora's 1080 × 1920 video.Source: §1 · p. 2
  • For single-image generation, tensor-parallel communication costs can outweigh distributed-computation savings, leaving multiple GPUs underused.The introduction notes that multiple GPUs are typically used for batch inference, while a single image usually runs on one GPU.Source: §1 · p. 2
  • Naïvely assigning image patches to separate GPUs creates visible seams, while synchronizing patch interactions can erase parallelism benefits.Figure 2 contrasts the seam-prone split with DistriFusion's communication approach.Source: §1 · p. 2, Figure 2 · p. 2
  • DistriFusion reuses activations from the previous denoising step after synchronous first-step communication, providing context for patch interactions.The method relies on similarity between inputs at adjacent denoising steps.Source: §1 · p. 2
  • Asynchronous communication is pipelined with computation, while convolutional and attention layers run only on assigned regions to reduce per-device computation.The paper names this parallelization approach displaced patch parallelism.Source: §1 · p. 3
  • 6.1× SDXL U-Net speedup on 8 A100 GPUs, with visual fidelity preserved, is reported for single-image generation.With batch splitting for classifier-free guidance, the paper reports 6.6× speedup on 8 GPUs for 3840 × 3840 images.Source: §1 · p. 3

2. Related Work

Prior diffusion acceleration work reduces sampling steps, optimizes inference, or distills faster generators, while established multi-GPU strategies face communication or quality trade-offs for diffusion models. DistriFusion instead partitions image inputs across devices and uses prior-step features to support patch interaction with lower communication demands.

2. Related Work

  • Diffusion acceleration includes more efficient denoising, training-free sampling, distillation, and neural-inference optimization, but expedited generators can retain a quality gap.The cited sampling approaches use differential-equation methods and exponential integrators to reduce steps while maintaining numerical accuracy.Source: §2 · p. 3
  • Tensor parallelism's communication burden is substantial for diffusion's large spatial activations, whereas data parallelism provides no latency improvement for serving.ParaDiGMS parallelizes denoising steps using Picard iterations but tends to waste computation and deviate from original-model results.Source: §2 · p. 3
  • Displaced patch parallelism reuses previous-step features as context and favors AllGather over AllReduce to reduce communication demands.The paper presents it as a parallelism strategy tailored to diffusion's sequential characteristics.Source: §2 · p. 3
  • Earlier sparse-computation work includes activation sparsity and structured approaches such as spatial masks and edit-driven regions.The related-work discussion also describes patch-based inference for reducing memory use in classification and detection.Source: §2 · p. 3
  • DistriFusion partitions image-generation inputs into patches processed by separate devices, reducing per-device computation by limiting work to assigned regions.The paper distinguishes this latency-focused image-generation use from patch-based work on classification and detection.Source: §2 · p. 4

3. Background

Diffusion models iteratively predict noise and update a noisy image toward a clean output, with model forward passes dominating latency. Large high-resolution activations make single-image multi-GPU acceleration difficult, while speculative-step parallelism and tensor parallelism have practical drawbacks.

3. Background

  • At each denoising step, a noise-prediction model takes the noisy image, time step, and condition as input and predicts the noise used to derive the next image state.The process starts from Gaussian noise and iterates toward the final clean image.Source: §3 · p. 4
  • Model forward passes dominate latency, and SDXL requires 6,763 GMACs per step for a 1024 × 1024 image.The passage says computation grows more than quadratically as resolution increases.Source: §3 · p. 4
  • Speculative denoising-step parallelism can waste computation and deviate from original-model quality, while tensor-parallel communication costs can be intolerable.The passage frames alternative workload distribution as an open challenge for accelerating single-image generation.Source: §3 · p. 4

4. Method

DistriFusion splits the image into patches and reuses slightly outdated activations from the previous diffusion step to preserve patch interactions while overlapping communication with computation. It also limits layer computation to fresh regions and adds strategies for normalization and warm-up steps.

Displaced patch parallelism

  • Reusing stale activations from the previous diffusion step addresses the seam and communication costs of independent patches and synchronous cross-patch exchange.The method calls this activation displacement and uses it to overlap communication and computation.Source: §4 · p. 4
  • Each device processes one patch; its fresh activation patch is scattered into the matching spatial location of stale activations before layer computation.The layer operation is applied selectively to fresh regions, and outputs are synchronized to approximate ϵθ(xt).Source: §4 · pp. 4–5

Communication

  • After the first step, devices broadcast activations with asynchronous AllGather so communication can overlap computation; the first step instead uses the standard synchronous process.The timeline caption says AllGather overhead is fully hidden within computation.Source: §4 · p. 4, Figure 4 · p. 5

Sparse operations

  • Sparse operators compute only on fresh regions; self-attention retains full-feature-map keys and values while preserving queries only for fresh-area tokens.The reported operator cost is proportional to the fresh-area size.Source: §4 · p. 5

Corrected asynchronous GroupNorm

  • For GroupNorm, a correction to cached stale statistics approximates current global statistics while precise statistics are aggregated asynchronously.The method also approximates variance; for rare negative approximated variances, it falls back to the fresh patch’s local variance.Source: §4 · p. 5

Warm-up steps

  • Several synchronous patch-parallel iterations serve as warm-up before displaced patch parallelism, improving image quality especially with reduced-step samplers.The motivation is that early sampling shapes low-frequency layout and semantics, while later sampling recovers high-frequency detail.Source: §4 · p. 6

5. Experiments

Experiments benchmark DistriFusion on Stable Diffusion XL using COCO Captions and assess image quality, latency, communication, and design choices. Results report quality comparable to the original model alongside lower latency, with larger speedups at higher resolutions and ablations examining warm-up, normalization, and communication.

5.1. Setups

  • Experiments use pre-trained Stable Diffusion XL, a 5K-image COCO Captions validation subset, and default 50-step DDIM sampling with a four-step synchronous warm-up after the first step.Latency is benchmarked in PyTorch after three warm-up runs and averaged over ten runs, excluding the fastest and slowest.Source: §5.1 · p. 6

5.2. Main Results

  • With 8 devices, DistriFusion achieves PSNR, LPIPS, and FID scores comparable to the original model while maintaining total MACs.The evaluation uses 1024 × 1024 images; Table 1 reports cumulative computation across devices and latency on NVIDIA A100 GPUs.Source: §5.2 · p. 6, Table 1 · p. 8

FID: 24.3

  • In the 25-step setting, DistriFusion mirrors original quality and accelerates the model by 2.7×, while ParaDiGMS achieves only 1.3×.In the 50-step setting, ParaDiGMS reaches 2.8× speedup at the cost of compromised image quality.Source: FID: 24.3 · p. 7
  • At 3840 × 3840, DistriFusion achieves 1.8×, 3.4×, and 6.1× speedups with 2, 4, and 8 A100s, respectively.With batch splitting for classifier-free guidance, the speedups become 3.6× and 6.6× with 4 and 8 A100s.Source: FID: 24.3 · pp. 7–8, Figure 6 · p. 8

5.3. Ablation Study

  • Across consecutive steps, average input difference is 0.02 within a [−4, 4] range, and DistriFusion cuts synchronous patch-parallel communication overhead by over 50%.The reported difference is about 0.3%; the remaining overhead mainly comes from NCCL kernels using GPU SMs during asynchronous communication.Source: §5.3 · p. 8
  • For 10-step DPM-Solver, naïve DistriFusion struggles to preserve image quality, while a two-step warm-up significantly recovers performance with only slightly increased latency.The paper identifies this 10-step configuration as the threshold for training-free samplers to maintain image quality.Source: §5.3 · p. 9
  • Corrected asynchronous GroupNorm matches synchronized GroupNorm image quality without synchronous communication overhead.Separate GroupNorm is faster but lowers fidelity; stale statistics can cause fog-like noise, while the correction term addresses the distribution gap between stale and fresh statistics.Source: §5.3 · p. 9

6. Conclusion & Discussion

DistriFusion parallelizes image patches across GPUs and reuses prior-step activations to maintain patch interactions. On SDXL, it achieves up to a 6.1× speedup on eight NVIDIA A100s, subject to hardware, resolution, and sampler limits.

6. Conclusion & Discussion

  • Up to 6.1× speedup on eight NVIDIA A100s for Stable Diffusion XL.The method reuses pre-computed activations from previous steps to maintain patch interactions.Source: §6 · p. 9
  • DistriFusion divides images into patches assigned to separate GPUs and reuses pre-computed activations from previous steps to maintain patch interactions.Source: §6 · p. 9
  • Maximum speedup depends on NVLink; speedups are limited for low-resolution images, and extremely few-step methods may not work, though preliminary results suggest 10 steps can retain quality.Source: §6 · p. 9

Related papers

Loading 2402.19481v4…