DistriFusion: 고해상도 디퓨전 모델을 위한 분산 병렬 추론

DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models

February 29, 2024

저자: Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li, Song Han

cs.AI

초록

디퓨전 모델은 고품질 이미지 합성에서 큰 성공을 거두었습니다. 그러나 디퓨전 모델을 사용하여 고해상도 이미지를 생성하는 것은 여전히 막대한 계산 비용으로 인해 어려운 과제이며, 이는 인터랙티브 애플리케이션에서 사용하기에는 지연 시간이 너무 길게 만듭니다. 본 논문에서는 이 문제를 해결하기 위해 다중 GPU 간의 병렬 처리를 활용한 DistriFusion을 제안합니다. 우리의 방법은 모델 입력을 여러 패치로 분할하고 각 패치를 GPU에 할당합니다. 그러나 이러한 알고리즘을 단순히 구현하면 패치 간의 상호작용이 끊어져 충실도가 떨어지고, 이러한 상호작용을 포함시키면 엄청난 통신 오버헤드가 발생합니다. 이러한 딜레마를 극복하기 위해, 우리는 인접한 디퓨전 단계 간의 입력이 매우 유사하다는 점을 관찰하고, 이전 타임스텝에서 미리 계산된 특징 맵을 재사용하여 현재 단계에 컨텍스트를 제공하는 displaced patch parallelism을 제안합니다. 따라서 우리의 방법은 비동기 통신을 지원하며, 이를 계산과 파이프라인으로 처리할 수 있습니다. 광범위한 실험을 통해 우리의 방법이 최신 Stable Diffusion XL에 적용 가능하며 품질 저하 없이 8개의 NVIDIA A100에서 단일 GPU 대비 최대 6.1배의 속도 향상을 달성할 수 있음을 보여줍니다. 우리의 코드는 https://github.com/mit-han-lab/distrifuser에서 공개되어 있습니다.

English

Diffusion models have achieved great success in synthesizing high-quality images. However, generating high-resolution images with diffusion models is still challenging due to the enormous computational costs, resulting in a prohibitive latency for interactive applications. In this paper, we propose DistriFusion to tackle this problem by leveraging parallelism across multiple GPUs. Our method splits the model input into multiple patches and assigns each patch to a GPU. However, na\"{\i}vely implementing such an algorithm breaks the interaction between patches and loses fidelity, while incorporating such an interaction will incur tremendous communication overhead. To overcome this dilemma, we observe the high similarity between the input from adjacent diffusion steps and propose displaced patch parallelism, which takes advantage of the sequential nature of the diffusion process by reusing the pre-computed feature maps from the previous timestep to provide context for the current step. Therefore, our method supports asynchronous communication, which can be pipelined by computation. Extensive experiments show that our method can be applied to recent Stable Diffusion XL with no quality degradation and achieve up to a 6.1times speedup on eight NVIDIA A100s compared to one. Our code is publicly available at https://github.com/mit-han-lab/distrifuser.