tdrussell/diffusion-pipe

A pipeline parallel training script for diffusion models.

What it solves

パイプライン並列化を利用することで、単一 GPU に収まらないほど巨大な拡散モデルの訓練を可能にします。画像生成モデルと動画生成モデルの両方を統一されたフレームワークで扱えるため、異なるアーキテクチャの管理複雑性が低減されます。

How it works

本プロジェクトは DeepSpeed を用いたデータ並列とパイプライン並列のハイブリッドで、モデルを複数 GPU に分散させます。VRAM 使用率を最適化するため、潜在表現とテキスト埋め込みをディスクに事前キャッシュし、実際の訓練中に VAE とテキストエンコーダをアンロードできるようにしています。また、8 ビットオプティマイザ、ブロックスワップ、アクティベーションチェックポイントなどのメモリ節約技術もサポートしています。

Who it’s for

Flux、SDXL、または HunyuanVideo などの大規模画像・動画拡散モデルをファインチューニングまたは訓練したいが、単一の大容量 GPU の VRAM が足りない AI 研究者・開発者向けです。

Highlights

  • Broad Model Support: SDXL、Flux、LTX-Video、HunyuanVideo、Cosmos、Wan2.1 など、幅広い最新モデルに対応。
  • Pipeline Parallelism: 単一 GPU のメモリを超えるモデルを扱えるよう特別に設計。
  • Unified Image/Video Workflow: 1 つの訓練スクリプトで画像と動画の両モダリティをサポート。
  • Efficient Caching: マルチプロセス・マルチ GPU による埋め込みと潜在表現の事前キャッシュで訓練を高速化し、メモリ使用を削減。
  • Training Management: Tensorboard ロギング、保持テストセットでの評価、再開可能なチェックポイント機能を搭載。