DistriFusion: Inferencia Paralela Distribuida para Modelos de Difusión de Alta Resolución
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
Resumen
Los modelos de difusión han logrado un gran éxito en la síntesis de imágenes de alta calidad. Sin embargo, generar imágenes de alta resolución con modelos de difusión sigue siendo un desafío debido a los enormes costos computacionales, lo que resulta en una latencia prohibitiva para aplicaciones interactivas. En este artículo, proponemos DistriFusion para abordar este problema aprovechando el paralelismo en múltiples GPUs. Nuestro método divide la entrada del modelo en múltiples parches y asigna cada parche a una GPU. Sin embargo, implementar este algoritmo de manera ingenua rompe la interacción entre los parches y reduce la fidelidad, mientras que incorporar dicha interacción generaría una sobrecarga de comunicación considerable. Para superar este dilema, observamos la alta similitud entre las entradas de pasos de difusión adyacentes y proponemos el paralelismo de parches desplazados, que aprovecha la naturaleza secuencial del proceso de difusión reutilizando los mapas de características precalculados del paso de tiempo anterior para proporcionar contexto al paso actual. Por lo tanto, nuestro método permite comunicación asíncrona, que puede ser encadenada con el cálculo. Experimentos extensivos muestran que nuestro método puede aplicarse al reciente Stable Diffusion XL sin degradación de calidad y lograr una aceleración de hasta 6.1 veces en ocho NVIDIA A100s en comparación con una. Nuestro código está disponible públicamente en 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.