HeyGenとGoogle Cloud：Avatar IVをTPUにもたらす

HeyGenがGoogle Cloud TPUs（v6e Trillium）上で18B以上のパラメータを持つ動画生成モデル「Avatar IV」を動作させることに成功。torchax、XLA、FSDP、Ulyssesなどを活用し、8チップ構成でリアルタイムストリーミングを1.86倍高速化。カスタムPallasカーネルや最適化手法が貢献した。

• HeyGenはAI動画生成プラットフォーム「Avatar IV」をGoogle CloudのTPU上で動作させ、1.86倍高速化に成功した。
• Avatar IVは180億以上のパラメータを持つディフュージョンモデルで、話す顔の動画生成を担う。
• Google CloudのAIインフラ最適化チームと連携し、TPU v6eホスト上で性能向上を実現した。
• 速度向上のための最適化は、メッシュのコレクティブ通信、スパースアテンション、ソフトマックス内部ループの並列化など多岐にわたる。
• GPU向けに開発されたモデルコードをJAXのPyTorchフロントエンドであるtorch.jax経由でTPU上で動作させることに成功した。