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

要旨

拡散モデルは高品質な画像合成において大きな成功を収めています。しかし、拡散モデルを用いて高解像度の画像を生成することは依然として課題であり、膨大な計算コストがかかるため、インタラクティブなアプリケーションでは許容できないレイテンシが発生します。本論文では、この問題を解決するために、複数のGPUを活用した並列処理手法であるDistriFusionを提案します。本手法では、モデルの入力を複数のパッチに分割し、各パッチをGPUに割り当てます。しかし、このようなアルゴリズムを単純に実装すると、パッチ間の相互作用が失われ、忠実度が低下します。一方で、この相互作用を取り入れると、膨大な通信オーバーヘッドが発生します。このジレンマを克服するため、我々は隣接する拡散ステップ間の入力の高い類似性に着目し、拡散プロセスの逐次的な性質を活用した「displaced patch parallelism」を提案します。この手法では、前のタイムステップで計算済みの特徴マップを再利用して、現在のステップのコンテキストを提供します。これにより、非同期通信をサポートし、計算とパイプライン化することが可能になります。大規模な実験により、本手法を最新のStable Diffusion XLに適用しても品質の劣化がなく、8台のNVIDIA A100を使用した場合、1台と比較して最大6.1倍の高速化を達成できることが示されました。本手法のコードは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.