Decentralized Diffusion Models

要約

大規模な AI モデルのトレーニングでは、数千の GPU に作業を分割し、各ステップでそれらの GPU 間の勾配を同期します。
これにより、集中型のモノリシック クラスターのみがサポートできる重大なネットワーク負荷が発生し、インフラストラクチャのコストが上昇し、電力システムに負担がかかります。
私たちは、集中型の高帯域幅ネットワーキング ファブリックへの依存を排除​​することで、独立したクラスターまたはデータセンター全体に拡散モデルのトレーニングを分散するためのスケーラブルなフレームワークである分散型拡散モデルを提案します。
私たちの方法では、データセットのパーティション全体にわたって、それぞれが互いに完全に分離された一連のエキスパート拡散モデルをトレーニングします。
推論時には、専門家が軽量ルーターを介してアンサンブルします。
データセット全体でトレーニングされた単一のモデルと同じ目的をアンサンブルが集合的に最適化することを示します。
これは、トレーニングの負担を多数の「コンピューティング アイランド」に分散できることを意味し、インフラストラクチャのコストを削減し、局所的な GPU 障害に対する回復力を向上させることができます。
分散型拡散モデルにより、研究者は中央の統合システムではなく、オンデマンド GPU ノードのような、より小型でコスト効率が高く、すぐに利用できるコンピューティングを活用できるようになります。
私たちは ImageNet と LAION Aesthetics について広範な実験を実施し、分散型拡散モデル FLOP-for-FLOP が標準拡散モデルよりも優れていることを示しています。
最終的にアプローチを 240 億パラメータに拡張し、高品質の拡散モデルが 8 つの個別の GPU ノードで 1 週間以内にトレーニングできることを実証しました。

要約(オリジナル)

Large-scale AI model training divides work across thousands of GPUs, then synchronizes gradients across them at each step. This incurs a significant network burden that only centralized, monolithic clusters can support, driving up infrastructure costs and straining power systems. We propose Decentralized Diffusion Models, a scalable framework for distributing diffusion model training across independent clusters or datacenters by eliminating the dependence on a centralized, high-bandwidth networking fabric. Our method trains a set of expert diffusion models over partitions of the dataset, each in full isolation from one another. At inference time, the experts ensemble through a lightweight router. We show that the ensemble collectively optimizes the same objective as a single model trained over the whole dataset. This means we can divide the training burden among a number of ‘compute islands,’ lowering infrastructure costs and improving resilience to localized GPU failures. Decentralized diffusion models empower researchers to take advantage of smaller, more cost-effective and more readily available compute like on-demand GPU nodes rather than central integrated systems. We conduct extensive experiments on ImageNet and LAION Aesthetics, showing that decentralized diffusion models FLOP-for-FLOP outperform standard diffusion models. We finally scale our approach to 24 billion parameters, demonstrating that high-quality diffusion models can now be trained with just eight individual GPU nodes in less than a week.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google