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

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