DistriFusion：高分辨率扩散模型的分布式并行推理

论文 | 项目 | 博客 | 幻灯片 | YouTube | 海报

[2024年7月29日] DistriFusion 已在 ColossalAI 中得到支持！

[2024年4月4日] DistriFusion 被选为 CVPR 2024 亮点海报展示！

[2024年2月29日] DistriFusion 被 CVPR 2024 接收！我们的代码已公开发布！

概览图我们提出了 DistriFusion，这是一种免训练的算法，可利用多个 GPU 加速扩散模型推理，同时不牺牲图像质量。简单的图像块处理（概览(b)）由于缺乏图像块之间的交互而存在分割问题。展示的示例是使用 SDXL 在 1280×1920 分辨率下生成的，采用 50 步 Euler 采样器，延迟在 A100 GPU 上测量。

DistriFusion：高分辨率扩散模型的分布式并行推理 李牧洋*，蔡天乐*，曹嘉欣，张钦胜，蔡涵，白俊杰，贾扬清，刘明宇，李凯，以及宋涵 麻省理工学院、普林斯顿大学、Lepton AI 和 NVIDIA 发表于 CVPR 2024。

概述

(a) 在单个设备上运行的原始扩散模型。(b) 简单地将图像分割成 2 个块并分布在 2 个 GPU 上，由于块之间缺乏交互，边界处会出现明显的接缝。(c) 我们的 DistriFusion 在第一步使用同步通信进行块交互。之后，我们通过异步通信重用前一步的激活。通过这种方式，通信开销可以隐藏在计算流水线中。

性能

加速效果

<img src="https://yellow-cdn.veclightyear.com/835a84d5/43cdba14-514d-4107-a9dc-2007345a78fc.jpg" width="80%"/> 在 NVIDIA A100 GPU 上使用 SDXL 和 50 步 DDIM 采样器生成单张图像时，DistriFusion 的总延迟测量结果。随着分辨率的提高，GPU 设备得到更好的利用。值得注意的是，在生成 3840×3840 图像时，DistriFusion 分别在 2、4 和 8 个 A100 上实现了 1.8 倍、3.4 倍和 6.1 倍的加速。

质量

SDXL 的定性结果。FID 是与真实图像相比计算得出的。我们的 DistriFusion 可以根据使用的设备数量减少延迟，同时保持视觉保真度。

参考文献：

去噪扩散隐式模型（DDIM），Song 等人，ICLR 2021
阐明基于扩散的生成模型的设计空间，Karras 等人，NeurIPS 2022
扩散模型的并行采样，Shih 等人，NeurIPS 2023
SDXL：改进高分辨率图像合成的潜在扩散模型，Podell 等人，ICLR 2024

先决条件

Python3
NVIDIA GPU + CUDA >= 12.0 及相应的 CuDNN
PyTorch = 2.2

入门指南

安装

安装 PyTorch 后，您应该能够通过 PyPI 安装 distrifuser

pip install distrifuser

或通过 GitHub：

pip install git+https://github.com/mit-han-lab/distrifuser.git

或在本地进行开发安装

git clone git@github.com:mit-han-lab/distrifuser.git
cd distrifuser
pip install -e .

使用示例

在 scripts/sdxl_example.py 中，我们提供了一个使用 DistriFusion 运行 SDXL 的最小脚本。

import torch

from distrifuser.pipelines import DistriSDXLPipeline
from distrifuser.utils import DistriConfig

distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)
pipeline = DistriSDXLPipeline.from_pretrained(
    distri_config=distri_config,
    pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0",
    variant="fp16",
    use_safetensors=True,
)

pipeline.set_progress_bar_config(disable=distri_config.rank != 0)
image = pipeline(
    prompt="宇航员在丛林中，冷色调，柔和的颜色，细节丰富，8k",
    generator=torch.Generator(device="cuda").manual_seed(233),
).images[0]
if distri_config.rank == 0:
    image.save("astronaut.png")

具体来说，我们的 distrifuser 与 diffusers 共享相同的 API，可以以类似的方式使用。您只需定义一个 DistriFusion 并使用我们封装的 DistriSDXLPipeline 来加载预训练的 SDXL 模型。然后，我们可以像 diffusers 中的 StableDiffusionXLPipeline 一样生成图像。运行命令是

torchrun --nproc_per_node=$N_GPUS scripts/sdxl_example.py

其中 $N_GPUS 是您想要使用的 GPU 数量。

我们还在 scripts/sd_example.py 中提供了一个使用 DistriFusion 运行 SD1.4/2 的最小脚本。使用方法相同。

基准测试

我们的基准测试结果使用 PyTorch 2.2 和 diffusers 0.24.0。首先，您可能需要安装一些额外的依赖项：

pip install git+https://github.com/zhijian-liu/torchprofile datasets torchmetrics dominate clean-fid

COCO 质量

您可以使用 scripts/generate_coco.py 生成带有 COCO 描述的图像。命令是

torchrun --nproc_per_node=$N_GPUS scripts/generate_coco.py --no_split_batch

其中$N_GPUS是您想使用的GPU数量。默认情况下，生成的结果将存储在results/coco中。您也可以使用--output_root自定义它。一些您可能想调整的额外参数：

--num_inference_steps：推理步骤数。我们默认使用50步。
--guidance_scale：无分类器引导尺度。我们默认使用5。
--scheduler：扩散采样器。我们默认使用DDIM采样器。您也可以使用euler表示Euler采样器和dpm-solver表示DPM求解器。
--warmup_steps：额外预热步骤数（默认为4）。
--sync_mode：不同的GroupNorm同步模式。默认使用我们修正的异步GroupNorm。
--parallelism：您使用的并行范式。默认为补丁并行。您可以使用tensor表示张量并行和naive_patch表示朴素补丁。

生成所有图像后，您可以使用我们的脚本scripts/compute_metrics.py来计算PSNR、LPIPS和FID。用法如下：

python scripts/compute_metrics.py --input_root0 $IMAGE_ROOT0 --input_root1 $IMAGE_ROOT1

其中$IMAGE_ROOT0和$IMAGE_ROOT1是您想比较的图像文件夹路径。如果IMAGE_ROOT0是真实值文件夹，请添加--is_gt标志以进行调整大小。我们还提供了一个脚本scripts/dump_coco.py来导出真实值图像。

延迟

您可以使用scripts/run_sdxl.py来对我们不同方法的延迟进行基准测试。命令如下：

torchrun --nproc_per_node=$N_GPUS scripts/run_sdxl.py --mode benchmark --output_type latent

其中$N_GPUS是您想使用的GPU数量。与scripts/generate_coco.py类似，您也可以更改一些参数：

--num_inference_steps：推理步骤数。我们默认使用50步。
--image_size：生成的图像大小。默认为1024×1024。
--no_split_batch：禁用无分类器引导的批次分割。
--warmup_steps：额外预热步骤数（默认为4）。
--sync_mode：不同的GroupNorm同步模式。默认使用我们修正的异步GroupNorm。
--parallelism：您使用的并行范式。默认为补丁并行。您可以使用tensor表示张量并行和naive_patch表示朴素补丁。
--warmup_times/--test_times：预热/测试运行次数。默认分别为5和20次。

引用

如果您在研究中使用此代码，请引用我们的论文。

@inproceedings{li2023distrifusion,
  title={DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models},
  author={Li, Muyang and Cai, Tianle and Cao, Jiaxin and Zhang, Qinsheng and Cai, Han and Bai, Junjie and Jia, Yangqing and Liu, Ming-Yu and Li, Kai and Han, Song},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2024}
}