CVPR 2024 | MIT发布DistriFusion技术，图像生成实现6倍加速！

扩散模型能够生成高质量图像，然而，由于计算成本巨大，导致交互应用的延迟过长，因此使用扩散模型生成高分辨率图像仍然是一个挑战。

MIT提出了DistriFusion是能在多个设备上并行运行扩散模型的方法，以减少单样本生成的延迟，而不会损害图像质量。方法基于patch parallelism，它将图像分成多个patch，每个patch分配给不同的设备。是一种无需训练的算法，可利用多个 GPU 来加速扩散模型推理，而不会牺牲图像质量。

设置：SDXL 配备50 Euler拉采样器，分辨率为1280×1920。延迟是在 A100上测量的。

相关链接

论文地址：https://arxiv.org/pdf/2402.19481

代码地址：https://github.com/mit-han-lab/distrifuser

项目地址：https://hanlab.mit.edu/projects/distrifusion

论文阅读

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

摘要

扩散模型在合成高质量图像方面取得了巨大的成功。然而，由于巨大的计算成本，使用扩散模型生成高分辨率图像仍然具有挑战性，导致交互应用程序的延迟令人望而却步。在本文中，我们提出了利用多个gpu的并行性来解决这个问题的DistriFusion。我们的方法将模型输入分割成多个补丁，并将每个补丁分配给GPU。然而，naïvely实现这样的算法会破坏补丁之间的交互并失去保真度，而合并这样的交互将产生巨大的通信开销。为了克服这一困境，我们观察到相邻扩散步骤的输入之间的高度相似性，并提出了位移补丁并行性，该方法通过重用前一个时间步骤预先计算的特征映射来利用扩散过程的顺序性，为当前步骤提供上下文。因此，我们的方法支持异步通信，可以通过计算实现流水线化。大量的实验表明，我们的方法可以应用于最近的stable Diffusion XL，没有质量下降，并且在8个A100 gpu上实现了高达6.1倍的加速。

方法

DistriFusion概述。为简单起见省略t和c的输入，并使用N = 2个设备作为示例。上标(1)和(2)分别表示第一个和第二个patch。前一步的陈旧激活将变暗。在每一步t中，我们首先将输入xt分成N个patch x(1)t，…, x t (N)。对于每一层l和设备i，在获得输入激活补丁Al，(i)t后，两个操作将异步进行:首先，在设备i, Al，(i)上，它从上一步分散到陈旧的激活Alt+1中。然后输入这个分散操作的输出。我们对每一层重复这个过程。然后将最终输出聚合在一起以近似ϵθ(xt)，用于计算xt−1。用于预测ϵθ(xt)的每个设备的时间轴可视化下图所示：

时间轴在每个设备上可视化时predictingϵθ(xt)。Comm.是指通信，它与计算是异步的。AllGather开销完全隐藏在计算中。