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.