DistriFusion : Inférence parallèle distribuée pour les modèles de diffusion haute résolution
DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models
February 29, 2024
Auteurs: Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li, Song Han
cs.AI
Résumé
Les modèles de diffusion ont obtenu un grand succès dans la synthèse d'images de haute qualité. Cependant, la génération d'images haute résolution avec ces modèles reste un défi en raison des énormes coûts de calcul, entraînant une latence prohibitive pour les applications interactives. Dans cet article, nous proposons DistriFusion pour résoudre ce problème en exploitant le parallélisme sur plusieurs GPU. Notre méthode divise l'entrée du modèle en plusieurs patches et attribue chaque patch à un GPU. Cependant, une implémentation naïve d'un tel algorithme rompt l'interaction entre les patches et entraîne une perte de fidélité, tandis que l'intégration de cette interaction engendre un surcoût de communication considérable. Pour surmonter ce dilemme, nous observons la forte similarité entre les entrées des étapes de diffusion adjacentes et proposons un parallélisme de patches décalés, qui tire parti de la nature séquentielle du processus de diffusion en réutilisant les cartes de caractéristiques précalculées de l'étape précédente pour fournir un contexte à l'étape actuelle. Ainsi, notre méthode prend en charge une communication asynchrone, qui peut être pipelinée par le calcul. Des expériences approfondies montrent que notre méthode peut être appliquée au récent Stable Diffusion XL sans dégradation de qualité et atteindre jusqu'à une accélération de 6,1 fois sur huit NVIDIA A100 par rapport à un seul. Notre code est disponible publiquement à l'adresse 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.