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）