DistriFusion: Distributed Inference

• DistriFusion is a distributed framework that leverages displaced patch parallelism and asynchronous communication to efficiently perform high-resolution diffusion model inference and Bayesian posterior fusion.
• It reduces per-step latency by overlapping computation with communication, ensuring high-fidelity outputs and eliminating visible seam artifacts in patch-wise generative inference.
• The methodology extends to a scalable Sequential Monte Carlo fusion approach, offering robust and exact product-pooling of sub-posteriors in large-scale Bayesian computing.

DistriFusion encompasses a family of distributed strategies and frameworks designed for efficient parallel inference, estimation, and sampling in systems led by either probabilistic inference, generative models (such as diffusion models), or distributed Bayesian computation. Initially formulated in the context of distributed diffusion model inference, with further generalizations in distributed posterior fusion, DistriFusion now refers to both a practical, high-throughput patch-parallel inference system for high-resolution diffusion models and a rigorous Monte Carlo methodology for exact product-pooling of sub-posteriors in big data Bayesian problems.

1. Motivation and Core Problem Domains

DistriFusion addresses multi-node or multi-device parallelization challenges endemic to two computational domains:

• High-resolution diffusion model inference: Modern diffusion models, such as Stable Diffusion XL (SDXL), possess large activation maps and require iterative Markovian denoising steps, resulting in prohibitive latency and memory demands for single GPU inference at $1024\times1024$ and above. Patch-wise data parallelization offers computational relief, but naive approaches yield visible seams at patch boundaries due to the loss of spatial interaction between patches. Full synchronization of activations at each layer is bandwidth-prohibitive, motivating hybrid schemes (Li et al., 2024, Zhang et al., 2024).
• Distributed Bayesian inference: In distributed data settings or multi-party privacy-constrained computation, posterior fusion from sub-posteriors (arising from data splits) is necessary. Standard methods may fail when sub-posteriors are non-Gaussian or poorly approximated, and communication costs for centralized "one-shot" fusion scale poorly in large C (number of partitions) regimes (Chan et al., 2021).

The central theme is the design of communication-efficient, provably correct, and high-fidelity algorithms for distributed fusion—either of neural generative outputs or of probabilistic estimates.

2. Displaced Patch Parallelism for Diffusion Inference

The flagship DistriFusion method for high-resolution generative diffusion models is based on "displaced patch parallelism" (Li et al., 2024):

• The noisy latent input $x_t$ is partitioned into $N$ non-overlapping patches $\{x_t^{(i)}\}$, each assigned to a GPU running an identical replica of the diffusion U-Net.
• At the initial (noisiest) step ($T$), all GPUs synchronize to gather the full activation (AllGather) for exact cross-patch interaction.
• For later steps ($t<T$), each device uses the assembled (but stale) global activation $A_{t+1}^l$ as a spatial canvas, updating only its assigned patch region with the freshly computed local activation $A_t^{l,(i)}$. This updated activation $\tilde{A}_t^l$ is then used for local computation, while a non-blocking AllGather of the fresh patches is launched to assemble $A_t^l$ asynchronously, pipelined with computation (see (Li et al., 2024), Figure 1).
• The sequential denoising nature of diffusion yields very small differences between $x_t$0 and $x_t$1 (average $x_t$2), so this staleness incurs negligible fidelity loss.
• Because communication is overlapped with patchwise computation, the per-step latency is reduced to $x_t$3, where $x_t$4 is the effective comm/compute overlap.

This architecture eliminates the large-scale AllReduce patterns conventional in tensor or sequence parallelism, instead relying on spatial decomposition and asynchronous collective communication.

3. Implementation Primitives and Scheduling

DistriFusion's computational pipeline is characterized by:

• Fully replicated model weights: Each GPU holds the entire U-Net weight set, obviating the need for parameter synchronization.
• Patch partitioning and scatter-gather: Each per-layer activation is AllGathered across devices, but only asynchronously after the first step. Patches are scattered into stale global canvases and locally updated.
• Asynchronous overlap: Communication and computation are aggressively overlapped. The scheduling at each U-Net layer proceeds as

1. Wait for the stale activation from the previous step (non-blocking).
2. Compute the local patch forward pass.
3. Scatter new patch activations into the stale global activation.
4. Launch a non-blocking AllGather for the current patch.
5. Proceed to the next layer, overlapping steps 3-4 of layer $x_t$5 with step 1 of layer $x_t$6.

• Classifier-free guidance adaptation: In conditional generation, patch groups are statically split between unconditional and conditional passes, with guidance fusion performed at the end.

This design supports nearly $x_t$7-fold ideal compute speedup until communication volume or patch staleness saturate improvements at high $x_t$8.

4. Performance, Complexity, and Comparison

Quantitative evaluation (Li et al., 2024, Zhang et al., 2024) highlights:

• Compute and communication analysis:
  • Per-GPU compute scales as $x_t$9, communication as $N$0.
  • Typical communication cost per denoising step (Stable Diffusion XL, 8 GPUs, 1024×1024): $N$1 GB per step.
• Empirical speedup:
  • 4 GPUs: $N$2 for $N$3, up to $N$4 for $N$5.
  • 8 GPUs: similar performance, as guidance halves effective patch utilization.
• Comparison to Partially Conditioned Patch Parallelism (PCPP):

| Method | Comm./step (GB, $N$6 GPUs, $N$7) | Speedup ($N$8) | FID (4/8 patch GPUs) | |-----------------|--------------------------------------|-----------------------------|----------------------| | DistriFusion | 1.53 | $N$9 | 20.8 / 24.1 | | PCPP | 0.48 (–69%) | $\{x_t^{(i)}\}$0 | 38.4 / 42.2 |

PCPP sharply reduces communication by conditioning only on partial neighbor patches (not full AllGather), at the expense of higher perceptual error and FID. DistriFusion maximizes quality but is communication-bound in extreme settings (Zhang et al., 2024).

• Qualitative metrics: DistriFusion's FID matches original SDXL (24.4 at $\{x_t^{(i)}\}$1 GPUs, cf. $\{x_t^{(i)}\}$2 baseline), while maintaining comparable PSNR/LPIPS and showing no seam artifacts. Naive patching (no cross-patch sync) causes FID to rise significantly (to $\{x_t^{(i)}\}$3).
• Memory: Each GPU holds full weights and one full stale activation per layer, but does not require extra AllReduce buffers as in tensor parallelism.

5. Algorithmic Generalizations: Distributed Bayesian Fusion

DistriFusion is also used to denote a divide-and-conquer Sequential Monte Carlo (SMC) methodology for exact Monte Carlo product-pooling of sub-posteriors in distributed Bayesian inference (Chan et al., 2021):

• Recursive fusion tree: Sub-posteriors $\{x_t^{(i)}\}$4 are organized into a binary tree, with fusion at each node achieved by sampling from a Gaussian proposal derived from child importance-weighted samples and their empirical covariances.
• Incremental fusion: Each fusion step produces weighted particle approximations to the product density; at each node, resampling is triggered if the effective sample size falls below a threshold to control degeneracy.
• Communication efficiency: Per-level communication requires only $\{x_t^{(i)}\}$5 particles, and total per-core bandwidth is $\{x_t^{(i)}\}$6 versus $\{x_t^{(i)}\}$7 for centralized fusion.
• Cost and robustness: Computational and bandwidth costs scale logarithmically with the number of partitions $\{x_t^{(i)}\}$8. The method extends to high-dimensional (up to $\{x_t^{(i)}\}$9) posteriors and hundreds of clients.
• Empirical superiority: In simulated and real-data settings, this SMC-based DistriFusion yields lower integrated absolute errors (vs. exact full-data posteriors) compared to Consensus Monte Carlo, Weierstrass, and KDE-MC, particularly as $T$0 grows (Chan et al., 2021).

6. Variants and Extensions

Additional DistriFusion-inspired adaptations include:

• Hybrid diversity distillation in fast generative inference: In the context of diffusion model distillation (Gandikota et al., 13 Mar 2025), DistriFusion refers to a hybrid inference recipe: the first denoising step is run using the full ("base") model, the remainder with the compressed ("distilled") model. This hybrid disables the diversity collapse seen in pure distilled inference and restores (or exceeds) base-level sample diversity at nearly the cost of the compressed model. Both low-rank concept adapters (LoRA) and attribute control modules are compatible across the hybrid steps. This approach requires both models in memory and can be extended for prompt-adaptive cutoff point selection.
• Patch parallel quantization and ablation: Warm-up steps with full synchronization, and strategies for correcting local statistics (e.g., GroupNorm correction via displaced activations) are essential to maintaining accuracy in extreme-step or low-data regimes (Li et al., 2024, Zhang et al., 2024).

7. Impact, Limitations, and Outlook

DistriFusion methodologies have established new standards for high-resolution multi-GPU diffusion inference fidelity-vs-throughput tradeoff and for scalable exact posterior fusion. Their strengths include:

• Near-linear multi-GPU scaling for high-resolution image generation, with up to $T$1 speedup and no retraining requirement for generative models (Li et al., 2024).
• Robustness to communication bottlenecks by aggressive comm/compute overlap and minimized synchronization.
• Theoretically sound SMC-based product fusion scalable to hundreds of clients in big-data Bayesian scenarios (Chan et al., 2021).
• Preservation of fine-grained control and sample diversity in generative distillation kits (Gandikota et al., 13 Mar 2025).

Limitations stem from:

• Full per-GPU model replication—potentially infeasible with extreme model sizes compared to tensor parallelism.
• Communication cost eventually dominating at extremes of image size or GPU count; PCPP addresses this via neighbor-limited context with a perceptual quality tradeoff.
• Memory cost for hybrid inference (distilled/base coexistence).
• Over-partitioning (very large $T$2) can erode fidelity as the approximation from stale activations degrades; warm-up steps (synchronous exchange) can partially mitigate this.

Future advancements may focus on further tradeoff tuning in neighbor conditioning, prompt-adaptive switching in hybrid inference, consolidating hybrid policies into a single model, and extending SMC-based fusion to nonparametric or hierarchical frameworks. The codebase for DistriFusion's patch-parallel inference is available at https://github.com/mit-han-lab/distrifuser.

---

References:

• "DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models" (Li et al., 2024)
• "Partially Conditioned Patch Parallelism for Accelerated Diffusion Model Inference" (Zhang et al., 2024)
• "Divide-and-Conquer Fusion" (Chan et al., 2021)
• "Distilling Diversity and Control in Diffusion Models" (Gandikota et al., 13 Mar 2025)