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

MEニュース メッセージ、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）