技術サマリー：PipeFusion

問題定義

拡散トランスフォーマー（DiT）は、高解像度の画像および動画合成における主要なアーキテクチャとして台頭し、従来の U-Net モデルを上回っています。しかし、DiT は推論時に、特に高解像度画像や長いシーケンスを生成する際に、顕著なレイテンシ課題に直面しています。このレイテンシは、シーケンス長に対するアテンション計算の二次的な成長に起因します。

既存の並列推論戦略は、マルチ GPU 環境においてこれらの課題を効率的に解決することに苦慮しています：

• テンソル並列化（TP）： 大規模言語モデル（LLM）では効果的ですが、DiT では大きな活性化量により非効率です。通信オーバーヘッド（特に All-Reduce 操作）は、特に層数の多いモデルにおいて、並列計算の利点を上回ることがあります。
• シーケンス並列化（SP）： DeepSpeed-Ulysses や Ring-Attention などの手法は、入力シーケンス（画像パッチ）をデバイス間で分割します。メモリ圧力を軽減しますが、各層で広範なデータ交換（Query、Key、Value 活性化の送信）が必要となり、高い通信コストをもたらします。
• DistriFusion： 入力時間的冗長性（連続する拡散タイムステップ間での入力および活性化の類似性）を利用する非同期シーケンス並列化手法です。通信と計算をオーバーラップさせますが、メモリ効率の面で課題があります。DistriFusion は、すべてのデバイスで全層の Key（K）および Value（V）活性化の完全な空間次元を維持する必要があり、メモリコストがシーケンス長に比例して増加し、Flux.1 のような大規模モデルの効率的なスケーリングを妨げます。

手法：PipeFusion

本論文は、DiT 推論に特化した新しいパッチレベルのパイプライン並列化戦略であるPipeFusionを提案します。PipeFusion は、モデル層と入力画像パッチを複数の GPU 間で同時に分割することで、既存の手法の限界を克服します。

中核メカニズム

1. 二重分割：

   • 層分割： DiT のバックボーンをデータフロー方向に分割します。各 GPU には、連続したトランスフォーマー層のセットが割り当てられます。
   • パッチ分割： 入力画像は $M$ 個の非重複パッチに分割されます。各 GPU は、割り当てられた層を通じて 1 つのパッチをパイプライン方式で処理します。

2. 時間的冗長性の活用：

   • 拡散プロセスは、連続するタイムステップ間の入力および活性化に高い類似性を示します。PipeFusion はこれを活用し、直前の拡散タイムステップ（$T+1$）からの「古くなった（stale）」特徴マップを再利用して、現在のタイムステップ（$T$）のコンテキストを提供します。
   • 現在のタイムステップのために全デバイスから新しい活性化を待つ代わりに、各デバイスは前ステップからの古くなった活性化を使用して計算を開始できます。これにより、初期起動時を除き、パイプラインは最小限のアイドル時間（バブル）で継続的に実行されます。

3. 通信の最適化：

   • 各層の活性化を送信する SP や TP と異なり、PipeFusion は特定のステージに対して、初期層の入力と最終層の出力のみを送信します。
   • 非同期のポイントツーポイント（P2P）通信を利用して、マイクロステップのパッチ活性化を次のデバイスへ送信し、通信と計算を効果的にオーバーラップさせます。

4. メモリ効率：

   • PipeFusion はモデルパラメータをデバイス間で分散させ、パラメータメモリ使用量を $N$（デバイスの数）倍削減します。
   • 重要なのは、アテンションに必要な KV バッファに関して、PipeFusion はそのデバイスに割り当てられた層に対して、完全な空間 K および V テンソルの $1/N$ だけを格納すればよい点です。これに対し、DistriFusion はすべてのデバイスで全層の完全な空間 K および V テンソルを必要とします。これにより、PipeFusion は大規模モデルにおいて著しくメモリ効率的になります。

主要な貢献

• 新しい並列化戦略： DiT のフルアテンション機構に特化し、層ごとのパイプライン並列化とパッチレベルのシーケンス分割を組み合わせた PipeFusion の導入。
• 削減された通信コスト： PipeFusion は、層ごとの活性化ではなく、ステージの入力/出力のみを送信することで、デバイス間通信を削減します。理論的解析により、通信コストは層数（$L$）に依存せず、シーケンス長（$p$）と隠れ次元サイズ（$h_s$）にのみ比例してスケーリングすることが示されています。
• 向上したメモリスケーラビリティ： パラメータを分散させ、必要な KV バッファサイズを $1/N$ に削減することで、PipeFusion は他の手法が Out-of-Memory（OOM）エラーにより失敗する汎用ハードウェア上で、大規模モデル（例：12B の Flux.1）の推論を可能にします。
• 品質の維持： 理論的解析と実証結果は、PipeFusion が拡散プロセス中に「新しい（fresh）」活性化のより広い領域を利用するのに対し、DistriFusion は古くなったデータに大きく依存するため、PipeFusion が DistriFusion よりも高い生成品質を維持することを示しています。

実験結果

著者は、3 つの最先端 DiT モデル（PixArt（0.6B）、Stable Diffusion 3（2B）、Flux.1（12B））を使用して、8×L40 PCIe GPU クラスター上で PipeFusion を評価しました。

• レイテンシ性能：

  • PipeFusion は、テストされたすべての解像度（1024px から 4096px）およびモデルにおいて、最低の推論レイテンシを達成しました。
  • 8×L40 GPU において、PipeFusion は PixArt および SD3 に対して、最良のシーケンス並列化ベースライン（USP）よりも1.48 倍から 1.55 倍高速でした。
  • 12B の Flux.1 モデルにおいて、PipeFusion は最良のシーケンス並列化に対して1.23 倍から 1.25 倍の高速化を達成しました。
  • 高解像度タスク（4096px）では、DistriFusion は OOM 問題に直面しましたが、PipeFusion は総レイテンシのわずか 4.6%（SP の 17.9% に対して）という通信シェアで強力な性能を維持しました。

• メモリ効率：

  • PipeFusion は優れたメモリスケーラビリティを示しました。Flux.1 の 2048px において、PipeFusion の総メモリ使用量はシーケンス並列化が必要とするものの**32% から 36%**のみでした。
  • DistriFusion が失敗した PixArt における 8192px 解像度タスクを、PipeFusion は正常に処理しました。

• 画像品質：

  • Fréchet Inception Distance（FID）を用いた定量的評価により、PipeFusion が DistriFusion を一貫して上回り、元の単一デバイス実装に近いスコアを達成したことが示されました。
  • 視覚的比較により、PipeFusion によって生成された画像は元の出力とほとんど区別がつかないことが確認されました。

意義と主張

本論文は、PipeFusion が DiT 推論の並列化における根本的な転換点であると主張しています。入力時間的冗長性を体系的に活用することで、既存の手法が達成できない通信効率、メモリフットプリント、および生成品質のバランスを実現しています。

著者は、PipeFusion をすべての並列戦略の普遍的な代替手段としてではなく、ハイブリッド並列化スキームのための極めて効果的なコンポーネントとして位置づけています。PipeFusion は特に、低帯域幅のインターノード環境に適しており、シーケンス並列化は高帯域幅（NVLink）のイントラノード設定において有効であると提案しています。

この研究は、拡散プロセスにおける時間的冗長性が、新しい形態の並列化を設計するための実用的かつ未活用のリソースであることを浮き彫りにし、汎用マルチ GPU クラスターにおけるスケーラブルで高解像度の画像生成への道筋を示しています。また、著者はこのアプローチが、同様の MM-DiT アーキテクチャとより長いシーケンス長を共有する動画生成モデルと自然に互換性があることも指摘しています。