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

扩散模型在合成高质量图像方面取得了巨大成功。然而，由于巨大的计算成本，使用扩散模型生成高分辨率图像仍然是一个挑战，这导致交互式应用程序的延迟过高。在本文中，我们提出了DistriFusion来解决这个问题，通过利用多GPU之间的并行性。我们的方法将模型输入分割成多个补丁，并将每个补丁分配给一个GPU。然而，简单地实现这样的算法会破坏补丁之间的交互并且失去保真度，而纳入这样的交互将会带来巨大的通信开销。为了克服这个困境，我们观察到相邻扩散步骤输入之间的高相似性，并提出了位移补丁并行性，这利用了扩散过程的顺序性质，通过复用前一个时间步骤的预计算特征图为当前步骤提供上下文。因此，我们的方法支持异步通信，可以由计算来流水线处理。广泛的实验表明，我们的方法可以应用于最近的Stable Diffusion XL而不降低质量，并且与单个NVIDIA A100相比，实现了高达6.1倍的加速。我们的代码在 https://github.com/mit-han-lab/distrifuser 上公开可用。