JAXとCloud TPU v5eでStable Diffusion XL推論を高速化
🧨 Accelerating Stable Diffusion XL Inference with JAX on Cloud TPU v5e
Google Cloud TPU v5eは、SDXLのような大規模生成AIモデルの推論とトレーニングに最適化されたカスタムAIアクセラレータです。TPU v4と比較して半額以下で、コスト効率とパフォーマンスを両立させています。Diffusers JAX統合により、XLA経由でTPU上でSDXLを実行でき、デモではTPU v5e-4インスタンス(各4チップ)を使用し、約4秒で4枚の1024x1024画像を生成しました。JAXのjitコンパイルとXLAコンパイラによる並列化(pmap)が、静的な形状を持つ画像生成ワークロードの最適化に貢献します。TPU v5eは、1チップから256チップまで多様な形状があり、JAXの並列処理を活用できます。bfloat16へのパラメータ変換によりメモリ削減と高速化を実現しつつ、スケジューラ状態はfloat32を維持します。TPU v5eはTPU v4と比較して最大2.4倍のパフォーマンス/ドルを達成し、コスト効率の高さを示しています。デモコードはHugging Face Diffusersで公開されています。
- Google CloudがTPU v5e上でJAXを用いたStable Diffusion XL推論の高速化を実現したと発表
GoogleのTPU v5eとJAXを用いたStable Diffusion XLの推論高速化は、クラウド上の生成AI推論のコスト効率が大幅に向上していることを示す。特にTPU v4比で最大2.4倍のパフォーマンス/ドルを達成し、4秒で4枚の高解像度画像を生成できる点は、AI画像生成の実用化・商用化を加速させる。静的形状ワークロードにおけるJAXのjit/pmap並列化の優位性が明確であり、AWS/NVIDIA GPUに対するGoogle Cloudの差別化要因としてTPUエコシステムの競争力が増している。クラウドAIインフラ市場の競争構造や、GPU代替としてのTPU需要拡大に注目すべき。