分配器

其他源码

v0.0.1beta0

下载

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

纸|项目|博客 |幻灯片| Youtube |海报

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

【2024年4月4日】 DistriFusion入选CVPR 2024亮点海报！

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

预告片我们引入了 DistriFusion，这是一种免训练算法，可利用多个 GPU 来加速扩散模型推理，而不会牺牲图像质量。 Naïve Patch（概述 (b)）由于缺乏补丁交互而存在碎片问题。所提供的示例是使用 SDXL 使用 50 步欧拉采样器以 1280×1920 分辨率生成的，并在 A100 GPU 上测量延迟。

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

概述

(a)在单个设备上运行的原始扩散模型。 (b)由于补丁之间缺乏交互，将图像天真地分成跨 2 个 GPU 的 2 个补丁，因此边界处有明显的接缝。 (c)我们的 DistriFusion 在第一步采用同步通信进行补丁交互。之后，我们通过异步通信重用上一步中的激活。通过这种方式，通信开销可以隐藏到计算管道中。

表现

加速

使用 50 步 DDIM 采样器测量 DistriFusion with SDXL 的总延迟，以便在 NVIDIA A100 GPU 上生成单个图像。提高分辨率时，可以更好地利用 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 [email protected]: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 = "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" )

具体来说，我们的distrifuser与扩散器共享相同的 API，并且可以以类似的方式使用。您只需要定义一个DistriConfig并使用我们包装的DistriSDXLPipeline来加载预训练的 SDXL 模型。然后，我们可以生成像扩散器中的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 和扩散器 0.24.0。首先，您可能需要安装一些额外的依赖项：

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

可可品质

您可以使用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 }
}