DistriFusion: Распределенное параллельное вывод для моделей диффузии высокого разрешения

DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models

February 29, 2024

Авторы: Muyang Li, Tianle Cai, Jiaxin Cao, Qinsheng Zhang, Han Cai, Junjie Bai, Yangqing Jia, Ming-Yu Liu, Kai Li, Song Han

cs.AI

Аннотация

Модели диффузии достигли большого успеха в синтезе изображений высокого качества. Однако генерация изображений высокого разрешения с помощью моделей диффузии по-прежнему представляет собой сложную задачу из-за огромных вычислительных затрат, что приводит к запретительной задержке для интерактивных приложений. В данной статье мы предлагаем DistriFusion для решения этой проблемы путем использования параллелизма на нескольких графических процессорах (GPU). Наш метод разбивает входную модель на несколько участков и назначает каждый участок одному GPU. Однако наивная реализация такого алгоритма нарушает взаимодействие между участками и теряет точность, в то время как включение такого взаимодействия повлечет за собой огромные накладные расходы на коммуникацию. Чтобы преодолеть это дилемму, мы замечаем высокую схожесть между входами смежных шагов диффузии и предлагаем параллелизм с смещенными участками, который использует последовательную природу процесса диффузии путем повторного использования предварительно вычисленных карт признаков с предыдущего временного шага для обеспечения контекста текущего шага. Таким образом, наш метод поддерживает асинхронную коммуникацию, которая может быть организована в виде конвейера вычислений. Обширные эксперименты показывают, что наш метод может быть применен к недавней модели Stable Diffusion XL без ухудшения качества и достигает ускорения до 6,1 раза на восьми графических процессорах NVIDIA A100 по сравнению с одним. Наш код доступен по адресу 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.