Cloud TPU v5e上のJAXを使用したStable Diffusion XL推論の高速化

Hugging Faceは、DiffusersライブラリにJAXサポートを統合し、Cloud TPU v5e上でStable Diffusion XL（SDXL）の高パフォーマンスかつコスト効率の良い推論を可能にしました。この統合は、前身より約3倍大きいUNetを持つSDXLの計算上の課題に対処し、JAXのJust-In-Time（JIT）コンパイルとXLA駆動の並列処理を活用します。

JAXとTPU v5eによる技術的最適化

Cloud TPU v5e上でSDXLを提供することにより、2つの主要なソフトウェアおよびハードウェアメカニズムであるJITコンパイルとSPMD並列処理を通じて高い効率が達成されます。

静的形状のためのJITコンパイル

JAXは、初回実行時にコードをトレースし、その後の呼び出しのために最適化されたTPUバイナリを生成するJust-In-Time（JIT）コンパイルを利用します。このプロセスでは、静的な入力、中間、出力の形状が必要です。SDXLは次の理由からJITコンパイルと高い互換性があります：

一定の出力形状: 画像生成では通常、固定された画像数と一貫したサイズが使用されます。
固定形状の埋め込み: Stable DiffusionおよびSDXLは、テキストプロンプトに対してパディングを含む固定形状の埋め込みベクトルを使用します。

初回のコンパイルには数分かかります（提供された例では約3分）、その後の推論呼び出しは大幅に高速化されます。

XLA並列処理とスループット

JAXのpmapは、単一プログラム複数データ（SPMD）実行を可能にし、複数のXLAデバイスにわたってワークロードをスケールできます。これにより、画像生成の線形スケーリングが可能になります：たとえば、8チップのTPUは、1チップが1枚の画像を生成する時間と同じ時間で8枚の画像を生成できます。Cloud TPU v5eインスタンスは、1〜256チップのさまざまな構成で利用可能であり、超高速ICIリンクで接続されているため、ユーザーは特定のスループットニーズに基づいてスケールできます。

JAXにおける実装パイプライン

JAXを使用したSDXL推論の実行は、モデルパラメータをパイプラインとは別に扱う関数型アプローチを含みます。主な実装手順は次のとおりです：

モデルの読み込み: FlaxStableDiffusionXLPipeline.from_pretrained を使用して、SDXL 1.0 ベースモデルを読み込みます。
精度管理: モデルパラメータをbfloat16に変換してメモリ使用量を削減し、速度を向上させながら、スケジューラの状態をfloat32のまま保持し、低品質または黒い画像を引き起こす精度エラーを防ぎます。
入力の準備: prepare_inputs を使用して、呼び出し間でプロンプトの次元が一貫していることを確認します。これはJITコンパイルに必要です。
デバイスレプリケーション: 利用可能なTPUチップ間でパラメータと入力をレプリケートします（例：TPU v5e-4に対してreplicateを使用）。各チップに一意のランダムシードを割り当て、多様な画像出力を確保します。
実行: jit=True を指定してパイプラインを呼び出し、XLAコンパイルプロセスをトリガーします。

パフォーマンスベンチマーク

Euler Discrete スケジューラを使用して20ステップで実施されたSDXL 1.0ベースのベンチマークでは、TPU v5eがTPU v4よりも優れたコスト効率を示しています。

ハードウェア	バッチサイズ	レイテンシ	ドルあたりの性能
TPU v5e-4 (JAX)	4	2.33s	21.46
TPU v5e-4 (JAX)	8	4.99s	20.04
TPU v4-8 (JAX)	4	2.16s	9.05
TPU v4-8 (JAX)	8	4.17s	8.98

TPU v5eは、TPU v4と比較してドルあたりの性能が最大2.4倍向上します。性能は、スループット（チップあたりのレイテンシで割ったバッチサイズ）を計算し、その結果をハードウェアのリスト価格で割ることによって測定されます。

デプロイメントアーキテクチャ

現在の実装では、ロードバランシングサーバーがユーザーリクエストをランダムにプリアルokされたCloud TPU v5e-4インスタンス上で動作するバックエンドサーバーにルーティングします。各インスタンスは、約4秒（フロントエンド処理および通信を含む）で1024×1024ピクセルの画像を4枚生成し、実際の生成時間は約2.3秒です。