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グループ化行列乗算」を用いて長いコンテキストやハイブリッドエキスパートモデルの不規則なテンソルを処理すること；第4世代の疎結合コア（SparseCore）を利用して集団通信操作をオフロードし遅延を隠すこと；TPUの高速オンチップSRAM（VMEM）の割り当てを細かく調整しメモリ待ちを削減すること；モデルの規模、アーキテクチャ、シーケンス長に応じて最適なパーティショニング戦略（FSDP、TP、EPなど）を選択すること。
（出典：InFoQ）