PaperlayerView exact revision on arXiv

Source-linked AI summary

PipeFusion: Patch-level Pipeline Parallelism for Diffusion Transformers Inference

Jiarui Fang, Jinzhe Pan, Aoyu Li, Xibo Sun, Jiannan Wang

arXiv:2405.14430v4cs.CVcs.AIcs.PF

TL;DR

Long-sequence DiT inference faces quadratic latency growth, and existing parallel methods can incur high communication costs. PipeFusion pipelines image patches across GPUs and reuses stale activations from previous steps; it improves latency across Pixart, Stable Diffusion 3, and Flux.1, reaching up to 1.5× faster inference on commodity multi-GPU clusters.

  • Problem

    DiT attention causes quadratic latency growth for long sequences, while tensor parallelism can incur communication costs that outweigh its parallel-computation benefits.

    Source: §1 · p. 1

  • Method

    PipeFusion partitions DiT layers and image patches across GPUs, using stale activations from prior timesteps in a patch-level pipeline to reduce communication and memory costs.

    Source: §1 · p. 2

  • Results

    PipeFusion improves latency across Pixart, Stable Diffusion 3, and Flux.1, achieving up to 1.5× faster inference on commodity multi-GPU clusters.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Temporal redundancy offers a parallelism strategy that composes with sequence or tensor parallelism as a path toward scalable large-cluster deployment.

    Source: §6 · p. 10

    PipeFusion is less effective for one-step or few-step diffusion models, where temporal redundancy is minimal.

    Source: §5 · p. 10

Abstract

from arXiv · show

This paper presents PipeFusion, an innovative parallel methodology to tackle the high latency issues associated with generating high-resolution images using diffusion transformers (DiTs) models. PipeFusion partitions images into patches and the model layers across multiple GPUs. It employs a patch-level pipeline parallel strategy to orchestrate communication and computation efficiently. By capitalizing on the high similarity between inputs from successive diffusion steps, PipeFusion reuses one-step stale feature maps to provide context for the current pipeline step. This approach notably reduces communication costs compared to existing DiTs inference parallelism, including tensor parallel, sequence parallel and DistriFusion. PipeFusion enhances memory efficiency through parameter distribution across devices, ideal for large DiTs like Flux.1. Experimental results demonstrate that PipeFusion achieves state-of-the-art performance on 8$\times$L40 PCIe GPUs for Pixart, Stable-Diffusion 3, and Flux.1 models. Our source code is available at https://github.com/xdit-project/xDiT.

1 Introduction

DiTs face rising inference latency on long sequences, while existing parallel methods incur substantial communication or memory costs. PipeFusion uses patch-level pipeline parallelism and temporal redundancy to reduce those costs while preserving image quality.

1 Introduction

  • DiT attention causes quadratic latency growth for long sequences, and tensor parallelism can lose its benefit when activation communication costs outweigh parallel computation.The paper identifies single-GPU latency as insufficient for practical requirements.Source: §1 · p. 1
  • Successive diffusion steps have similar inputs and activations, enabling stale-feature reuse; DistriFusion applies this idea but retains full-spatial attention keys and values across layers.That requirement makes DistriFusion’s memory costs scale linearly with sequence length.Source: §1 · p. 2
  • PipeFusion pipelines image patches across GPUs, reducing communication and attention K,V buffers while distributing parameters and retaining higher image-generation accuracy than DistriFusion.It transfers only the initial-layer input and final-layer output per device, and uses 1/N of DistriFusion’s attention K,V memory buffers.Source: §1 · p. 2
  • On 8×L40 PCIe GPUs, PipeFusion achieves the lowest DiT inference latency among the compared sequence-parallel, DistriFusion, and tensor-parallel methods.The paper also identifies 12B Flux.1 as a model suited to PipeFusion’s distributed-parameter memory approach.Source: §1 · p. 2

2 Background & Related Works

Diffusion models iteratively denoise latent inputs, while DiTs use transformer blocks to process latent patches. Related work exploits temporal redundancy by caching or reusing activations and attention outputs across diffusion steps.

2 Background & Related Works

  • Diffusion models iteratively predict noise from a noisy latent and use a sampler update to produce the previous latent in the denoising sequence.The model takes the noisy latent, timestep, and a condition such as text or an image as input.Source: §2 · p. 2
  • DiTs divide noisy latent representations into patches, embed them as tokens, and process them with transformer blocks that use self-attention.The paper contrasts this with U-Nets, which use convolutional layers to capture spatial hierarchies.Source: §2 · p. 2
  • Prior work exploits similarity across successive diffusion steps by caching or reusing activations, attention outputs, or selected blocks.Examples include DeepCache, TGATE, Δ-DiT, PAB, and DiTFastAttn, which use different features of temporal or other redundancy.Source: §2 · pp. 2–3

3 DiT Parallel Methods

The section compares tensor, sequence, and asynchronous sequence parallelism for DiT inference before introducing PipeFusion’s patch-level pipeline. PipeFusion partitions both image patches and model layers across GPUs and reuses stale activations to support pipelined computation.

3.1 Tensor Parallelism & Sequence Parallelism

  • Tensor parallelism uses two all-reduces per layer, while sequence parallelism partitions images into patches and exchanges partial Q, K, and V for global attention.Ulysses uses All2All to switch partitioning, while Ring-Attention transfers K/V subblocks point to point.Source: §3.1 · p. 3, Figure 2 · p. 3

3.2 DistriFusion: Asynchronous Sequence Parallelism

  • DistriFusion uses asynchronous all-gather to collect remote K and V activations, computing with stale activations from a previous timestep while current-step communication proceeds.Its communication-computation overlap comes with increased memory use for full-spatial K and V buffers on each device.Source: §3.2 · p. 4

3.3 PipeFusion: Patch-level Pipeline Parallel

  • PipeFusion assigns consecutive DiT layers to GPU stages and partitions the input into patches, using previous-timestep activations as context while stages process patches asynchronously.Asynchronous point-to-point transfers overlap communication with computation.Source: §3.3 · pp. 4–5
  • 98.5% is the effective computation ratio when M = N = 4 and S = 50; the authors recommend M = N, and M > N avoids bubbles after startup.The pipeline has no waiting after initialization; bubbles occur at its beginning.Source: §3.3 · p. 5
  • PipeFusion’s fresh activation area grows across micro-steps, whereas DistriFusion keeps one of M patches fresh, making PipeFusion closer to the original diffusion method.Figure 4 distinguishes fresh activations from stable activations.Source: Figure 4 · p. 5, §3.3 · p. 5

3.4 Comparison Between Different Parallelisms

  • PipeFusion has the lowest communication cost when N < 2L, with a total cost of 2O(p × hs) independent of L and communication overlapped with computation.The comparison is for one diffusion timestep.Source: §3.4 · pp. 5–6, Table 1 · p. 6
  • PipeFusion and tensor parallelism reduce parameter-memory cost as GPU count grows; PipeFusion’s KV buffer also shrinks with N, unlike DistriFusion’s.Both PipeFusion and DistriFusion maintain KV buffers for transformer layers, creating activation-memory overhead, especially for long sequences.Source: §3.4 · p. 6

4 Experiments

Experiments compare parallel inference on 8×L40 GPUs across Pixart, SD3-medium, and Flux.1-dev, evaluating latency, memory use, and image quality. PipeFusion reports faster inference than sequence-parallel alternatives and lower memory use in the reported settings.

4 Experiments

  • The evaluation used an 8×L40-48GB PCIe Gen4x16 cluster and covered 0.6B-parameter Pixart, 2B SD3-medium, and 12B Flux.1-dev.Latency comparisons exclude VAE decoding; reported figures average five runs, with the patch count selected by searching M from 2, 4, 8, 16, and 32.Source: §4 · p. 6, §4.1 · pp. 6–7

4.1.1 Latency Evaluation

  • 1.23× and 1.25× speedups over the best sequence-parallel configuration on Pixart’s two evaluated cases.The best sequence-parallel configuration was USP with ulysses=4 and ring=2.Source: Figure 5 · p. 7, §4.1.1 · p. 8
  • 1.57× and 1.30× speedups over the best sequence-parallel configuration on SD3-medium’s two evaluated cases.The model cannot generate 4096px images because of positional-encoding limitations.Source: Figure 6 · p. 7, §4.1.1 · p. 8
  • 1.48×, 1.55×, and 1.16× faster than the second-best parallel approaches on Flux.1-dev’s 1024px, 2048px, and 4096px tasks, respectively.With CFG parallel (cfg=2, pipefusion=4), speedups versus the 1-GPU baseline were 5.09x, 7.91x, and 8.59x.Source: Figure 7 · p. 7, §4.1.1 · pp. 7–8
  • At 4096px, PipeFusion’s end-to-end latency was 32.1s versus SP’s 37.3s, while its communication share was 4.6% versus SP’s 17.9%.DistriFusion had no valid 4096px data because of out-of-memory issues; PipeFusion’s communication cost was 1.49s versus SP’s 6.69s, a 78% reduction.Source: §4.1.1 · pp. 8–9

4.1.2 Memory Efficiency

  • On Flux.1-dev, PipeFusion’s overall memory use was 32% and 36% of SP’s at 1024px and 2048px, respectively.Its activation use increased because of the KV Buffer, but this increase was relatively minor compared with parameter memory.Source: Figure 9 · p. 9, §4.1.2 · p. 9

4.2 Quality Results

  • PipeFusion achieved lower FID scores than DistriFusion across identical device configurations on the Pixart and Flux.1 datasets.The accompanying visual comparison reports generated images as nearly indistinguishable from the original images to the human eye.Source: §4.2 · pp. 9–10, Table 2 · p. 10

5 Limitations & Discussion

PipeFusion is not a universal solution: its effectiveness depends on temporal redundancy, while video generation is a promising extension. A hybrid can use sequence parallelism on high-bandwidth links and PipeFusion across low-bandwidth inter-node links.

5 Limitations & Discussion

  • For cross-node deployment, sequence parallelism suits high-bandwidth NVLink settings, while PipeFusion is particularly effective across low-bandwidth inter-node links.The authors describe combining the approaches as a practical direction, rather than treating PipeFusion as a universal solution.Source: §5 · p. 10
  • PipeFusion is less effective for few-step or one-step diffusion methods because they provide minimal temporal redundancy.The paper notes this is not a major practical constraint because state-of-the-art DiTs generally require many steps for competitive quality.Source: §5 · p. 10
  • Video generation is a promising extension because PipeFusion partitions along the sequence dimension, which accommodates video models that flatten spatial and temporal dimensions.The paper also argues video models’ larger size and longer sequences amplify sequence-parallel communication overhead, which PipeFusion mitigates through parameter partitioning and pipelined execution.Source: §5 · p. 10

6 Conclusion

PipeFusion reuses temporally redundant activations and overlaps communication with computation to reduce bandwidth demands and memory footprint while preserving image quality. Experiments on Pixart, Stable Diffusion 3, and Flux.1 show consistent latency improvements over strong baselines, up to 1.5× faster inference on commodity multi-GPU clusters.

6 Conclusion

  • Up to 1.5× faster inference on commodity multi-GPU clusters, with consistent latency improvements over strong baselines on Pixart, Stable Diffusion 3, and Flux.1.PipeFusion reduces bandwidth demands and memory footprint while preserving image quality.Source: §6 · p. 10
  • The study focuses on single-node, multi-GPU inference, and describes composing PipeFusion with sequence or tensor parallelism as a path toward scalable large-cluster deployment.The authors identify temporal redundancy as a systems opportunity for designing new forms of parallelism.Source: §6 · p. 10

NeurIPS Paper Checklist

The checklist describes how the paper’s claims align with PipeFusion’s contributions and experimental results, documents limitations and reproducibility information—including areas not addressed—and notes that the paper contains no important, original, or non-standard LLM component.

Claims and limitations

  • The checklist says PipeFusion’s claims accurately reflect its patch-level pipeline-parallel contribution and are supported by experiments, while noting limited-step models and extremely long sequences as limitations.It also says the method is not well-suited to distilled or one-step diffusion models.Source: NeurIPS Paper Checklist · p. 13

Theory assumptions and proofs

  • The checklist states that the paper reports no formal theoretical results, theorems, or proofs because it focuses on proposing and experimentally evaluating PipeFusion.Source: NeurIPS Paper Checklist · p. 14

Experimental reproducibility

  • The paper provides experimental setup, model, software, and hardware details that the checklist says are sufficient to reproduce the main results, although code and data are not explicitly provided there.The listed setup includes 8×L40 PCIe GPUs and models Pixart, Stable-Diffusion 3, and Flux.1.Source: NeurIPS Paper Checklist · p. 14

Experiment statistical significance

  • The experiments report average latency and FID scores across multiple runs but provide no error bars, confidence intervals, or statistical significance tests.The checklist attributes this to results being stable among tests, without reporting variability estimates.Source: NeurIPS Paper Checklist · p. 16

Experiments compute resources

  • The checklist says the paper specifies compute resources, including 8×L40 PCIe GPUs with 48GB memory each and PyTorch 2.4.1 with CUDA Runtime 12.1.105.Source: NeurIPS Paper Checklist · p. 17

Societal impact and ethics

  • The checklist says the paper does not discuss broader societal impacts, while reporting that it involves no human subjects or sensitive data requiring special ethical considerations.Source: NeurIPS Paper Checklist · p. 17

NeurIPS Paper Checklist

  • The authors answer [No] on LLM use, saying PipeFusion focuses on parallelizing diffusion transformers rather than using LLMs as a core component.The checklist notes that NA applies when core method development does not involve LLMs as important, original, or non-standard components.Source: NeurIPS Paper Checklist · p. 20

A Image Quality

The visual comparison presents PipeFusion, DistriFusion, and original-model samples; PipeFusion outputs are nearly indistinguishable from the originals to the human eye across configurations.

A Image Quality

  • PipeFusion images are nearly indistinguishable from original-model images to the human eye across different patch and device configurations.Figure 10 uses sample images generated from prompts from DistriFusion; the paper states that FID results are independent of those demonstration prompts.Source: A Image Quality · p. 20
  • Figure 10 uses sample prompts from DistriFusion to support visual assessment of generation quality.The paper distinguishes this visual demonstration from FID evaluation, which is independent of the specific prompts shown.Source: A Image Quality · p. 20
  • FID evaluation results are independent of the specific prompts used for the visual demonstrations.Source: A Image Quality · p. 20

B Ablation Study

PipeFusion’s latency depends on patch and warmup settings, while its reported advantage persists on NVLink GPUs. The section also reports image-quality and warmup optimizations.

Patch Number

  • Latency is lowest when the patch number M equals the number of GPUs N.Larger M can reduce operator efficiency, while smaller M reduces communication-computation overlap; M does not affect communication cost.Source: B Ablation Study · p. 20, Figure 11 · p. 22

Warmup Step

  • More warmup steps can increase latency, and their overhead is lessened when they form a smaller share of the sampling steps.The experiments used 20 or 28 sampling steps; related work commonly uses 50 or 100, which helps alleviate warmup overhead.Source: B Ablation Study · p. 20

Warmup optimization

  • A proposed optimization separates warmup from working steps, using sequence parallelism during warmup and transferring its output feature maps to working devices.The passages describe this as a way to mitigate warmup-related performance degradation.Source: B Ablation Study · pp. 20, 22

Image quality

  • Generated-image quality closely resembles the original across four or eight devices and varying patch numbers.FID evaluation used COCO Captions 2014 images and Clean-FID; the passage gives no FID values for these comparisons.Source: Figure 10 · p. 21

Warmup optimization

  • PixArt can produce satisfactory images without warmup when its KV buffer is initialized to zero; automatic warmup selection is proposed for future work.The proposed detection uses the difference between the latent-space input and the previous diffusion step.Source: B Ablation Study · p. 22

Evaluation on NVLink connected GPUs

  • PipeFusion consistently achieves the lowest PixArt latency across resolutions on 8×A100 GPUs with NVLink.The reported performance gap narrows at higher resolutions, where computation costs dominate more of the workload.Source: B Ablation Study · p. 22, Table 4 · p. 22

Related papers

Loading 2405.14430v4…