扩散模型能够生成高质量图像,然而,由于计算成本巨大,导致交互应用的延迟过长,因此使用扩散模型生成高分辨率图像仍然是一个挑战。
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开销完全隐藏在计算中。
(a)在单个设备上运行的原始扩散模型。(b) Naïvely将图像分成跨越2个gpu的2个补丁,由于补丁之间没有相互作用,在边界处有明显的接缝。(c) DistriFusion在第一步采用同步通信进行补丁交互。之后,我们通过异步通信重用前一步中的激活。通过这种方式,通信开销可以隐藏在计算管道中。
效果
定性的结果。FID是根据真实图像计算的。我们的DistriFusion可以根据使用设备的数量减少延迟,同时保持视觉保真度。
不同预热步骤下10步DPM-Solver的定性结果。LPIPS是针对整个COCO数据集上原始SDXL的样本计算的。Naïve没有预热步骤的DistriFusion有明显的质量下降。添加两步热身可以显著提高性能,同时避免高延迟上升。
8个a100的不同GN方案的定性结果。LPIPS是针对整个ecoco[5]数据集的原始样本计算的。单独的GN只使用设备补丁的统计信息。陈旧GN重用陈旧的统计信息。它们的质量下降了。同步。GN同步数据以确保准确的统计数据,但代价是额外的开销。我们的校正异步GN,通过校正陈旧的统计数据,避免了同步的需要,有效地恢复了质量。
结论
在本文中,我们引入了分布式融合来加速具有多个gpu的扩散模型以实现并行性。我们的方法将图像分成小块,并将每个小块分配给单独的GPU。我们重用前面步骤中预先计算的激活来维持补丁交互。在StableDiffusion XL上,我们的方法在8个NVIDIA a100上实现了高达6.1倍的加速。这一进步不仅提高了人工智能生成内容的创建效率,而且为未来人工智能应用的并行计算研究设定了新的基准。
局限性。完全隐藏通信开销在计算中,NVLink是使DistriFusion加速最大化的关键。然而,NVLink已经最近被广泛使用。此外,量化还可以减少我们方法的通信工作量。此外,对于低分辨率,DistriFusion的加速速度有限图像作为设备未得到充分利用。高级编译器将有助于利用这些设备并取得更好加速效果。由于环境的快速变化,我们的方法可能不适用于极少步骤的方法 去噪。然而,我们的初步实验表明稍微多一些步骤(例如,10个)对于DistriFusion来说就足够获得高质量的结果。
感谢你看到这里,也欢迎点击关注下方公众号并添加公众号小助手加入官方读者交流群,一个有趣有AI的AIGC公众号:关注AI、深度学习、计算机视觉、AIGC、Stable Diffusion、Sora等相关技术,欢迎一起交流学习💗~
没有评论:
发表评论