DistriFusion: Inferenza Parallela Distribuita per Modelli di Diffusione ad Alta Risoluzione

DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models

February 29, 2024

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

cs.AI

Abstract

I modelli di diffusione hanno ottenuto un grande successo nella sintesi di immagini di alta qualità. Tuttavia, generare immagini ad alta risoluzione con i modelli di diffusione rimane una sfida a causa degli enormi costi computazionali, che si traducono in una latenza proibitiva per applicazioni interattive. In questo articolo, proponiamo DistriFusion per affrontare questo problema sfruttando il parallelismo su più GPU. Il nostro metodo divide l'input del modello in più patch e assegna ciascuna patch a una GPU. Tuttavia, implementare in modo ingenuo un tale algoritmo rompe l'interazione tra le patch e compromette la fedeltà, mentre incorporare tale interazione comporterebbe un enorme sovraccarico di comunicazione. Per superare questo dilemma, osserviamo l'elevata somiglianza tra l'input di passaggi di diffusione adiacenti e proponiamo il parallelismo a patch spostate, che sfrutta la natura sequenziale del processo di diffusione riutilizzando le mappe di caratteristiche pre-calcolate dal passaggio temporale precedente per fornire contesto al passaggio corrente. Pertanto, il nostro metodo supporta la comunicazione asincrona, che può essere eseguita in pipeline con il calcolo. Esperimenti estensivi dimostrano che il nostro metodo può essere applicato al recente Stable Diffusion XL senza degradazione della qualità e ottenere un accelerazione fino a 6,1 volte su otto NVIDIA A100 rispetto a una. Il nostro codice è disponibile pubblicamente all'indirizzo 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.