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

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

[2024年12月1日] DistriFusion已被集成到NVIDIA的TensorRT-LLM中，用于高分辨率图像生成的分布式推理。

[2024年7月29日] DistriFusion现已支持ColossalAI！

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

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

我们提出了DistriFusion，这是一种无需训练的算法，能够利用多张GPU加速扩散模型的推理，同时不牺牲图像质量。朴素的分块方法（概述(b)）由于缺乏补丁之间的交互而存在碎片化问题。所展示的例子是使用SDXL和50步欧拉采样器在1280×1920分辨率下生成的，延迟是在A100 GPU上测量的。

DistriFusion：高分辨率扩散模型的分布式并行推理
Muyang Li*, Tianle Cai*, Jiaxin Cao，Qinsheng Zhang，Han Cai，Junjie Bai，Yangqing Jia，Ming-Yu Liu，Kai Li，以及Song Han
麻省理工学院、普林斯顿大学、Lepton AI和NVIDIA
发表于CVPR 2024。

概述

(a) 原始扩散模型在单个设备上运行。(b) 将图像简单地分成2个补丁并在2张GPU上处理时，由于补丁之间缺乏交互，边界处会出现明显的接缝。(c) 我们的DistriFusion在第一步中采用同步通信来实现补丁间的交互。之后，我们通过异步通信复用前一步的激活值。这样，通信开销可以隐藏在计算流水线中。

性能

加速比

使用50步DDIM采样器，在NVIDIA A100 GPU上使用SDXL生成单张图像时，DistriFusion的总延迟测量结果。随着分辨率的提高，GPU设备的利用率也会更高。值得注意的是，在生成3840×3840图像时，使用2、4和8张A100 GPU时，DistriFusion分别实现了1.8倍、3.4倍和6.1倍的加速。

质量

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

参考文献：

• 去噪扩散隐式模型（DDIM），Song et al.，ICLR 2021
• 阐明基于扩散的生成模型的设计空间，Karras et al.，NeurIPS 2022
• 扩散模型的并行采样，Shih et al.，NeurIPS 2023
• SDXL：改进潜在扩散模型以用于高分辨率图像合成，Podell et al.，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，并且可以用类似的方式使用。您只需要定义一个DistriConfig，然后使用我们封装的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 使用 欧拉采样器，或选择 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: 高分辨率扩散模型的分布式并行推理},
  author={李牧阳、蔡天乐、曹嘉欣、张钦盛、蔡涵、白俊杰、贾扬青、刘明宇、李凯、韩松},
  booktitle={IEEE/CVF 计算机视觉与模式识别会议（CVPR）论文集},
  year={2024}
}

致谢

我们的代码基于 huggingface/diffusers 和 lmxyy/sige 开发。我们感谢 torchprofile 提供的 MACs 测量工具，clean-fid 提供的 FID 计算功能，以及 Lightning-AI/torchmetrics 提供的 PSNR 和 LPIPS 工具。

我们感谢 Jun-Yan Zhu 和 Ligeng Zhu 的有益讨论和宝贵反馈。本项目得到了 MIT-IBM Watson AI 实验室、亚马逊公司、MIT 科学中心以及美国国家科学基金会的支持。

DistriFusion 快速上手指南

DistriFusion 是一种无需训练的分布式并行推理算法，旨在利用多张 GPU 加速高分辨率扩散模型（如 SDXL）的生成过程，同时保持图像质量。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux
• Python: Python 3
• GPU: NVIDIA GPU
• CUDA: 版本 >= 12.0
• CuDNN: 与 CUDA 版本对应
• PyTorch: 版本必须为 2.2 (torch==2.2)

注意：请确保已正确安装与 CUDA 12.0+ 匹配的 PyTorch 2.2 版本。

安装步骤

您可以选择通过 PyPI 直接安装，或从 GitHub 源码安装。

方式一：通过 PyPI 安装（推荐）

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 .

基本使用

distrifuser 的 API 设计与 Hugging Face diffusers 完全兼容。只需定义 DistriConfig 并使用封装好的 DistriSDXLPipeline 即可加载模型并进行分布式推理。

以下是一个基于 SDXL 模型的最小化运行示例（保存为 sdxl_example.py）：

import torch

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

# 配置分布式参数，例如生成 1024x1024 图像，预热步数为 4
distri_config = DistriConfig(height=1024, width=1024, warmup_steps=4)

# 加载预训练的 SDXL 模型
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,
)

# 仅在主进程 (rank 0) 显示进度条
pipeline.set_progress_bar_config(disable=distri_config.rank != 0)

# 执行生成
image = pipeline(
    prompt="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    generator=torch.Generator(device="cuda").manual_seed(233),
).images[0]

# 仅在主进程保存图像
if distri_config.rank == 0:
    image.save("astronaut.png")

运行命令

使用 torchrun 启动多卡推理，将 $N_GPUS 替换为您想要使用的 GPU 数量：

torchrun --nproc_per_node=$N_GPUS sdxl_example.py

该工具同样支持 SD 1.4/2.0 模型，使用方法类似（参考项目中的 scripts/sd_example.py）。