DistriFusion: Inferência Paralela Distribuída para Modelos de Difusão de Alta Resolução

DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models

February 29, 2024

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

cs.AI

Resumo

Modelos de difusão têm alcançado grande sucesso na síntese de imagens de alta qualidade. No entanto, gerar imagens de alta resolução com modelos de difusão ainda é desafiador devido aos enormes custos computacionais, resultando em uma latência proibitiva para aplicações interativas. Neste artigo, propomos o DistriFusion para enfrentar esse problema aproveitando o paralelismo em múltiplas GPUs. Nosso método divide a entrada do modelo em vários patches e atribui cada patch a uma GPU. No entanto, implementar tal algoritmo de forma ingênua quebra a interação entre os patches e resulta em perda de fidelidade, enquanto incorporar essa interação acarretaria um enorme overhead de comunicação. Para superar esse dilema, observamos a alta similaridade entre as entradas de passos de difusão adjacentes e propomos o paralelismo de patches deslocados, que aproveita a natureza sequencial do processo de difusão reutilizando mapas de características pré-computados do passo anterior para fornecer contexto para o passo atual. Assim, nosso método suporta comunicação assíncrona, que pode ser encadeada com a computação. Experimentos extensivos mostram que nosso método pode ser aplicado ao recente Stable Diffusion XL sem degradação de qualidade e alcançar um speedup de até 6,1 vezes em oito NVIDIA A100s em comparação com uma. Nosso código está publicamente disponível em 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.