DistriFusion: Gedistribueerde Parallelle Inferentie voor Hoogresolutie Diffusiemodellen

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

Samenvatting

Diffusiemodellen hebben grote successen geboekt bij het synthetiseren van hoogwaardige afbeeldingen. Het genereren van afbeeldingen met hoge resolutie met diffusiemodellen blijft echter een uitdaging vanwege de enorme rekenkosten, wat resulteert in een onaanvaardbare latentie voor interactieve toepassingen. In dit artikel stellen we DistriFusion voor om dit probleem aan te pakken door parallellisme over meerdere GPU's te benutten. Onze methode splitst de modelinvoer in meerdere patches en wijst elke patch toe aan een GPU. Een naïeve implementatie van een dergelijk algoritme verbreekt echter de interactie tussen patches en verliest aan kwaliteit, terwijl het opnemen van een dergelijke interactie enorme communicatie-overhead met zich meebrengt. Om dit dilemma te overwinnen, observeren we de hoge gelijkenis tussen de invoer van aangrenzende diffusiestappen en stellen we displaced patch parallellisme voor, dat gebruikmaakt van het sequentiële karakter van het diffusieproces door vooraf berekende kenmerkkaarten van de vorige tijdstap te hergebruiken om context te bieden voor de huidige stap. Onze methode ondersteunt daarom asynchrone communicatie, die kan worden gepipelineerd door berekening. Uitgebreide experimenten tonen aan dat onze methode kan worden toegepast op de recente Stable Diffusion XL zonder kwaliteitsverlies en een snelheidsverbetering tot 6,1 keer kan bereiken op acht NVIDIA A100's in vergelijking met één. Onze code is openbaar beschikbaar op 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.