DistriFusion: Verteiltes paralleles Inferenz für hochauflösende Diffusionsmodelle
DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models
February 29, 2024
Autoren: Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li, Song Han
cs.AI
Zusammenfassung
Diffusionsmodelle haben große Erfolge bei der Synthese hochwertiger Bilder erzielt. Die Erzeugung hochauflösender Bilder mit Diffusionsmodellen bleibt jedoch aufgrund der enormen Rechenkosten eine Herausforderung, was zu einer prohibitiv hohen Latenz für interaktive Anwendungen führt. In diesem Artikel schlagen wir DistriFusion vor, um dieses Problem durch die Nutzung von Parallelität über mehrere GPUs zu lösen. Unser Verfahren teilt die Modelleingabe in mehrere Patches auf und weist jeden Patch einer GPU zu. Eine naive Implementierung eines solchen Algorithmus würde jedoch die Interaktion zwischen den Patches unterbrechen und die Bildqualität beeinträchtigen, während die Einbeziehung einer solchen Interaktion erhebliche Kommunikationskosten verursachen würde. Um dieses Dilemma zu überwinden, beobachten wir die hohe Ähnlichkeit zwischen den Eingaben benachbarter Diffusionsschritte und schlagen eine verschobene Patch-Parallelität vor, die den sequenziellen Charakter des Diffusionsprozesses nutzt, indem vorberechnete Feature-Maps aus dem vorherigen Zeitschritt wiederverwendet werden, um Kontext für den aktuellen Schritt bereitzustellen. Daher unterstützt unser Verfahren asynchrone Kommunikation, die durch Berechnungen gepipelt werden kann. Umfangreiche Experimente zeigen, dass unser Verfahren auf das aktuelle Stable Diffusion XL ohne Qualitätseinbußen angewendet werden kann und eine Beschleunigung von bis zu 6,1-fach auf acht NVIDIA A100s im Vergleich zu einer einzelnen GPU erreicht. Unser Code ist öffentlich unter https://github.com/mit-han-lab/distrifuser verfügbar.
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.