diffusion-pipe: 大規模な画像およびビデオ拡散モデルのためのパイプライン並列トレーニングフレームワーク

何を解決するか

パイプライン並列を利用することで、単一のGPUに収まらないほど巨大な拡散モデルのトレーニングを可能にします。画像およびビデオ生成モデルの両方をトレーニングするための統一フレームワークを提供し、異なるアーキテクチャを管理する複雑さを軽減します。

仕組み

このプロジェクトは、DeepSpeed を介したデータ並列とパイプライン並列のハイブリッドを使用し、モデルを複数のGPUに分散します。VRAM使用量を最適化するため、ディスクへの latents とテキスト埋め込みのプリキャッシュを実装し、実際のトレーニングプロセス中に VAE とテキストエンコーダーをアンロードできるようにします。さらに、8ビット最適化器、ブロック スワッピング、アクティベーション チェックポイントなどのメモリ節約技術もサポートします。

対象者

Flux、SDXL、または HunyuanVideo などの大規模な画像およびビデオ拡散モデルをファインチューニングまたはトレーニングする必要がある AI リサーチャーおよび開発者で、単一の巨大な GPU の VRAM が不足している方々。

ハイライト

• 広範なモデルサポート: SDXL、Flux、LTX-Video、HunyuanVideo、Cosmos、Wan2.1 などの幅広い最新モデルと互換性があります。
• パイプライン並列: 単一のGPUのメモリを超えるモデルを扱うために特別に設計されています。
• 統合された画像/ビデオワークフロー: 単一のトレーニングスクリプトで両方のモダリティをサポートします。
• 効率的なキャッシュ: トレーニングを高速化しメモリを節約するため、マルチプロセス、マルチGPU でのエンベディングと latents のプリキャッシュを行います。
• トレーニング管理: Tensorboard ロギング、保持セットでの評価、トレーニング再開のための状態チェックポイントを含みます。