Ironwood TPU を使用したトレーニングに関するデベロッパー ガイド

※この投稿は米国時間 2026 年 3 月 24 日に、Google Cloud blog に投稿されたものの抄訳です。

数兆単位のパラメータを扱う AI モデルへの移行により、演算リソースの需要が急激に高まり、従来のインフラストラクチャの限界が試されています。第 7 世代の Ironwood TPU は、Google がカスタム設計した AI インフラストラクチャです。チップ間相互接続（ICI）、光回路スイッチ（OCS）、データセンター ネットワーク（DCN）、および大規模な集約型高帯域幅メモリ（HBM）容量を組み合わせることで、最大 9,216 個のチップを格納できる Pod に対応する包括的なシステムとしてスケールできるように設計されています。さらに、Ironwood はハードウェア アーキテクチャとソフトウェアの統合された共同設計を特徴としており、コンパイラ中心の XLA、および Pallas や Mosaic などの Python ネイティブ カーネルといったイノベーションが導入されています。組織はこれらの機能を組み合わせることで、高度なフロンティア モデルをトレーニングおよび提供する能力を大幅に高め、AI ライフサイクル全体を最適化し、高いパフォーマンスを維持できます。

この技術概要では、Ironwood ハードウェア上でのトレーニング効率の向上と、卓越したパフォーマンスの実現を目指して設計された、JAX および MaxText エコシステムにおける具体的な手法とツールについて説明します。

Ironwood の主な最適化戦略

1. MaxText によるネイティブ FP8 の活用

Ironwood は、行列乗算ユニット（MXU）で 8 ビット浮動小数点（FP8）をネイティブにサポートする最初の TPU 世代です。重み、アクティベーション、勾配に FP8 精度を利用することで、ユーザーは理論上、スループットを Brain Floating Point 16（BF16）の 2 倍に高められます。FP8 レシピを正しく構成すると、モデルの品質を損なうことなく効率を向上させることができます。

これらの FP8 トレーニング レシピを実装するには、Qwix ライブラリから始めます。この機能は、MaxText 構成内で関連するフラグを指定すると有効になります。,  

詳しくは、Google デベロッパー フォーラムのブログ投稿 Ironwood での FP8 トレーニングの最適化についてをご覧ください。

2. Tokamax カーネルによる加速

Tokamax は、TPU 向けに最適化された高パフォーマンスの JAX カーネルのライブラリです。これらのカーネルは、次のメカニズムを通じて特定のボトルネックを軽減するように設計されています。

1. Splash Attention: このメカニズムは、標準的なアテンション プロセスに内在する I/O の制限に対処します。オンチップ SRAM 内で計算を維持することで、メモリ帯域幅が制約になることが多い長いコンテキストの処理に特に効果を発揮します。

2. Megablox グループ化行列乗算（GMM）: これは、混合エキスパート（MoE）モデルでよく見られる「不規則な」なテンソルを管理します。GMM を利用すると、システムは非効率的なパディングを回避し、MXU の使用率を高められます。

3. カーネル チューニング: Tokamax ライブラリには、ハイパーパラメータを最適化するためのユーティリティが含まれています。これらのツールを使用すると、Ironwood TPU の特定のメモリ階層に合わせて、タイルサイズやその他の構成を調整できます。

3. SparseCore への集団のオフロード

Ironwood の第 4 世代 SparseCore は、不規則なメモリアクセス パターンを管理するために特別に設計されたプロセッサです。ユーザーは、特定の XLA フラグを使用して、All-Gather や Reduce-Scatter などの集団通信演算を SparseCore に直接オフロードできます。

このオフロード メカニズムにより、TensorCore を主要なモデル計算に専念させながら、通信タスクを並行して実行できます。このような機能の重複は、通信のレイテンシを隠し、MXU へのデータ スループットを一定に保つための重要な戦略です。

4. VMEM 上でのメモリ パイプラインのファインチューニング

TPU メモリ アーキテクチャの重要な部分である VMEM は、カーネルのパフォーマンスを最適化するように設計された高速なオンチップ SRAM です。現在の演算と将来の重みのプリフェッチの間で VMEM の割り当てを調整することで、実行速度を全体的に向上させることができます。たとえば、現在のスコープ用に予約されている VMEM を増やすと、カーネルで使用されるタイルサイズを大きくすることができます。これにより、潜在的なメモリストールが解消され、カーネルのパフォーマンスが向上します。

TPU メモリ アーキテクチャの詳細については、TPU パイプラインをご覧ください。

5. 最適なシャーディング戦略の選択

最後に、MaxText は、すべての TPU で利用できるさまざまな並列処理手法をサポートしています。最適な選択は、モデルサイズ、アーキテクチャ（Dense や MoE）、シーケンス長によって異なります。適切なシャーディング戦略を選択すると、モデルのパフォーマンスを高められます。

• 完全にシャーディングされたデータ並列処理（FSDP）: これは、単一チップのメモリ容量を超える大規模モデルをトレーニングする場合に推奨される戦略です。FSDP は、モデルの重み、勾配、オプティマイザの状態を複数のチップにシャーディングします。デバイスごとのバッチサイズを増やし、より多くの演算を導入することで、All-Gather 演算のレイテンシを隠し、効率を向上させることができます。

• テンソル並列処理（TP）: 個々のテンソルをシャーディングします。Ironwood は演算密度が高いため、モデルの次元が極めて大きい場合に TP が最大の効果を発揮します。TP を 2 分割して活用すると、Ironwood のデュアル チップレット設計における高速なダイ間相互接続を利用できます。

• エキスパート並列処理（EP）: MoE モデルでエキスパートをデバイス間で分散するのに役立ちます。

• コンテキスト並列処理（CP）: 非常に長いシーケンスに必要で、シーケンスの次元に沿ってアクティベーションをシャーディングします。

• ハイブリッド アプローチ: 大規模な実行で演算、メモリ、通信のバランスを取るには、戦略の組み合わせが必要になる場合が多いです。

上述の 2～5 の手法について詳しくは、デベロッパー フォーラムの投稿 Optimizing Frontier Model Training on TPU v7x Ironwood（TPU v7x Ironwood でのフロンティア モデル トレーニングの最適化）をご覧ください。

Ironwood のメリット: システムレベルのパフォーマンス

これらの最適化手法と、高速の 3D トーラス チップ間相互接続（ICI）や大容量 HBM などの Ironwood のアーキテクチャ上の強みを組み合わせることで、フロンティア モデルのトレーニング向け高性能プラットフォームが実現します。ハードウェア、コンパイラ（XLA）、フレームワーク（JAX、MaxText）間の緊密な共同設計により、AI インフラストラクチャから最大限のパフォーマンスを引き出すことができます。

AI の取り組みを加速させる準備は整いましたか？以下のリソースで、各最適化手法について詳しく確認できます。

関連情報

• Ironwood での FP8 トレーニングの最適化について

• TPU v7x Ironwood でのフロンティア モデル トレーニングの最適化

このブログ投稿に協力してくれた Hina Jajoo と Amanda Liang に感謝します。

- プロダクト戦略およびオペレーション担当、Lillian Yu

- Google TPU 担当プロダクト マネージャー、Liat Berry