diffusion-pipe：拡散モデル向けパイプライン並列トレーニングスクリプト

diffusion-pipeでこのスクリプトが訓練するものとモデル分割の仕組み

README の冒頭では、diffusion-pipe は拡散モデル向けのパイプライン並列トレーニングスクリプトであると説明されている。トレーニングコードは DeepSpeed のパイプライン並列を基盤としており、単一 GPU に収まらないモデルを訓練するための仕組みになっている。README のモデル一覧には、SDXL、Flux、LTX-Video、HunyuanVideo（t2v）、Cosmos、Lumina Image 2.0、Wan2.1（t2v と i2v）、Chroma、HiDream、Stable Diffusion 3、Cosmos-Predict2、OmniGen2、Flux Kontext、Wan2.2、Qwen-Image、Qwen-Image-Edit、HunyuanImage-2.1、AuraFlow、Z-Image、HunyuanVideo-1.5、Flux 2（Dev と Klein）、Anima、Ernie-Image、LTX 2.3、Ideogram4、Krea 2 が含まれる。README にはこれらのモデルの訓練ベンチマークや性能比較は含まれておらず、この記事でも報告しない。

diffusion-pipeでインストールと README が利用者に委ねている依存関係

README のインストール手順は、再帰的クローンから始まる：git clone --recurse-submodules https://github.com/tdrussell/diffusion-pipe。リポジトリには ComfyUI サブモジュールが含まれているためである。最初のクローンで --recurse-submodules を付けなかった場合は、git submodule init と git submodule update を実行する必要がある。環境は Miniconda と Python 3.12 で作成し、PyTorch を最初にインストールする。PyTorch は意図的に requirements.txt に含まれておらず、GPU によって適切な PyTorch や CUDA のバージョンが異なるためである。README の記述（2025年10月26日時点）では、著者の 4090 では PyTorch 2.9.0 と CUDA 12.8 が動作したとある。Flash Attention はオプションで、一部のモデルで必要になる。NVIDIA Cosmos には追加で TransformerEngine が必要で、README には著者の conda パスに依存した環境変数付きインストールコマンドが記載されている。どのモデルに Flash Attention が必要かは README に明記されておらず、モデルごとの確認事項となる。

diffusion-pipeでデータセットの構成とディスクキャッシュ

データセットは、画像または動画ファイルと対応するキャプションファイルを含む1つ以上のディレクトリで構成される。キャプションはメディアファイルと同じベース名の .txt ファイルであり、image1.png には同じディレクトリに image1.txt が必要となる。キャプションファイルがない場合は警告が出力され、空のキャプションでトレーニングが続行される。画像形式は Pillow が読み込めるもの、動画形式は ImageIO が読み込めるものが想定されている。WebP 動画は ImageIO が複数フレームの WebP を読み込めないためサポートされない。latent とテキスト埋め込みはトレーニング前にディスクへキャッシュされ、各データセットディレクトリ内の cache ディレクトリに保存される。キャッシュには Huggingface Datasets ライブラリが使われる。キャッシュは複数回のトレーニング実行間で再利用される。関連するフラグは3つある：--cache_only はキャッシュ処理だけを行って終了する。--regenerate_cache はキャッシュの再生成を強制する。--trust_cache はデータファイルが変更されたかどうかを指紋で確認せずにキャッシュメタデータを読み込むもので、README によると 100,000 枚以上の大規模データセットで読み込みを高速化できる。テキストエンコーダの出力もキャッシュされるため、README は現在テキストエンコーダ用の LoRA トレーニングをサポートしていないと明記している。

diffusion-pipeで設定ファイル、起動コマンド、再開方法

README はトレーニング前に examples ディレクトリを確認するよう指示している。main_example.toml にはほとんどの設定を説明するコメントが含まれ、dataset.toml にはデータセットオプションのドキュメントがある。設定ファイルをコピーして編集し、少なくともすべてのパスを変更する必要がある。README に示された起動コマンドは次のとおり：NCCL_P2P_DISABLE="1" NCCL_IB_DISABLE="1" deepspeed --num_gpus=1 train.py --deepspeed --config examples/hunyuan_video.toml。RTX 4000 シリーズではこの2つの環境変数が必要で、他の GPU では不要な場合がある。チェックポイントが有効な場合、同じコマンドに --resume_from_checkpoint フラグを付けて再実行すると、最新のチェックポイントから再開される。フラグの後にチェックポイントフォルダ名を指定することもできる。再開時にはコマンドラインで渡した設定ファイルが使われ、出力ディレクトリに保存された設定ファイルは使われない。ユーザーは以前の設定と一致させる責任がある。

diffusion-pipeで出力ディレクトリとチェックポイントの内容

各トレーニング実行は output_dir の下に新しいディレクトリを作成する。このディレクトリにはチェックポイント、保存されたモデル、Tensorboard のメトリクスが含まれる。保存されたモデルや LoRA は epoch1、epoch2 などのディレクトリに配置される。DeepSpeed のチェックポイントは global_step1234 のようなディレクトリに配置される。README によると、チェックポイントには重み、オプティマイザ状態、データローダ状態などすべてのトレーニング状態が含まれるが、推論には直接使用できない。保存されたモデルディレクトリには safetensors の重み、PEFT adapter の設定 JSON、diffusion-pipe の設定ファイルが含まれる。Tensorboard のメトリクス名やチェックポイントのフォーマットバージョンは README に明記されておらず、この記事でも扱わない。

diffusion-pipeでVRAM 削減設定と並列戦略

wan_14b_min_vram.toml のサンプルファイルに、VRAM 削減のための設定がまとめられている。AdamW8BitKahan オプティマイザ、blocks_to_swap = 32 によるブロックスワッピング、activation_checkpointing = 'unsloth' による unsloth アクティベーションチェックポイント、PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True による CUDA アロケータの設定などが含まれる。README によると、expandable_segments は動画トレーニングや複数のアスペクト比バケットで効果があったが、著者のシステムではランダムな CUDA エラーを引き起こすこともあり、数ステップを通過すれば通常は問題なくトレーニングが続くという。並列戦略はデータ並列とパイプライン並列のハイブリッドである。--num_gpus フラグで GPU 数を指定し、設定ファイルの pipeline_stages でパイプライン並列度を制御する。データ並列度は自動的に決定され、全 GPU が使用されるため、GPU 数は pipeline_stages で割り切れる必要がある。4 GPU で pipeline_stages=2 の場合、2つのモデルインスタンスが実行され、各インスタンスが2 GPU に分割される。

diffusion-pipeでGPL-3.0 と ComfyUI ベースのモデル追加

このリポジトリは GPL-3.0 ライセンスで公開されている。2025年11月29日の変更履歴には、ComfyUI コードを利用するためにライセンスを GPL-3 に変更したこと、今後のモデル実装は可能な限り ComfyUI のバックエンドコードを使用することが記されている。ライセンス本文は、GPL-3.0 がコピーレフトライセンスであり、ソフトウェアを共有・変更する自由を認めるものであること、このフリーソフトウェアには保証がないことを定めている。セキュリティ保証、サポート、保証の詳細についてはライセンス本文も README も言及していない。README の Features では、新しいモデルは単一のサブクラスを実装することで追加できると説明されており、これが文書化された拡張方法である。本番利用の可否について README もライセンス抜粋も触れていないため、利用者側で確認すべき項目となる。

編集部の結論

diffusion-pipe は DeepSpeed のパイプライン並列を中核とする拡散モデル用トレーニングスクリプトであり、多くの画像・動画モデルをサポートする。設定ファイル、レイテントとテキスト埋め込みのディスクキャッシュ、チェックポイント再開を中心にワークフローが構成されている。README にはベンチマークや本番利用の結果は含まれていないため、それらは各ユーザーが自分のハードウェアとデータセットで確認する必要がある。

公式情報源

1. Official README
2. Project repository