Googleは第7世代Ironwood TPUの開発者向けトレーニングガイドを公開、システムレベルのパフォーマンス最適化を詳述

ME News ニュース、4月2日（UTC+8）、Google公式は最近、第7世代Ironwood TPU向けの開発者トレーニングガイドを公開しました。このガイドは、開発者がIronwood TPUのシステムレベルの性能を最大限に活用し、最先端のAIモデルの効率的なトレーニングと展開を支援することを目的としています。Ironwood TPUは、兆のパラメータを持つモデルの計算能力ニーズに応えるために設計されたカスタムAIインフラストラクチャであり、チップ間インターコネクション（ICI）、光路スイッチ（OCS）、データセンターネットワーク（DCN）、高帯域幅メモリ（HBM）などの技術を用いて、最大9,216チップをサポートする完全なシステムを構築しています。本稿では、このハードウェアに対する複数の重要な最適化戦略について詳述しています。具体的には、行列乗算ユニット（MXU）がネイティブにFP8トレーニングをサポートしスループットを向上させること、TPU最適化専用のJAXコアライブラリTokamaxを採用し、「スプラッシュアテンション」や「Megabloxグループ行列乗算」を用いて長いコンテキストや混合エキスパートモデルの不規則なテンソルを処理すること、第四世代の疎結合コア（SparseCore）を利用して集団通信操作をオフロードし遅延を隠すこと、TPUの高速オンチップSRAM（VMEM）の割り当てを微調整してメモリ待ちを削減すること、そしてモデルの規模、アーキテクチャ、シーケンス長に応じて最適なパーティショニング戦略（FSDP、TP、EPなど）を選択することを含みます。（出典：InFoQ）