TPUの詳細分析

 (henryhmko.github.io)
13 ポイント 投稿者 GN⁺ 2025-06-23 | まだコメントはありません。 | WhatsAppで共有
  • TPU は、Google が開発した大規模AIの学習および推論向けカスタムチップで、GPUとは異なる 設計思想 を持つ
  • スケーラビリティ と エネルギー効率 を重視し、ハードウェア（例: SoC構成、大容量オンチップメモリ）とソフトウェア（XLAコンパイラ）を一体で設計している
  • 中核構造は シストリックアレイ、パイプライニング、事前コンパイル方式で、ほとんどのディープラーニング演算（特に行列積）に最適化されている
  • OCI および OCS 技術により、柔軟なノード構成、高性能な並列処理、多様なトポロジ選択が可能
  • TPU システムは単一チップから超大規模マルチポッドまで階層的に拡張でき、大規模AIモデルの学習と効率的なリソース活用を実現する

概要と背景

  • TPU は、GoogleがAIサービスの大規模拡張に対応するために開発した ASIC ベースのチップ
  • 当初はGPU、FPGA、ASICの導入を比較検討していたが、2013年に音声検索などでディープラーニング活用が拡大したことで、専用ハードウェアの必要性が高まった
  • 現在は Gemini、Veo など主要AIサービスで不可欠な役割を担い、推薦モデル（DLRM）などにも広く利用されている

TPU単一チップ構造

基本構成

  • TPUv4 では、単一チップに2基の TensorCore が存在する（推論特化TPUでは1基）
  • 各 TensorCore は、CMEM（128MiB）、HBM（32GiB）のメモリ装置に接続される

TensorCore内部構成

  • Matrix Multiply Unit (MXU) : 128x128シストリックアレイベースで、行列積を専任で処理
  • Vector Unit (VPU) : 一般的な要素単位演算を実行
  • Vector Memory (VMEM; 32MiB) : HBMからデータをコピーし、作業準備用メモリとして活用
  • Scalar Unit + Scalar Memory (SMEM; 10MiB) : 制御フロー、スカラー演算、メモリアドレス管理を担当

GPUとの構造的な違い

  • TPU はオンチップメモリ（CMEM、VMEM、SMEM）がGPUよりはるかに大きい
  • HBM 容量はGPUの方が大きく、演算コア数もGPUの方がはるかに多い
  • TPUv5p では、チップあたり500 TFLOPs/sec、ポッド全体（8960チップ）で4.45 ExaFLOPs/sec級の高性能を実現している

TPUの設計思想

1. シストリックアレイとパイプライニング

  • シストリックアレイ は、処理要素（PE）が配列され、隣接要素へ演算結果を受け渡す構造
  • データ入力後は追加の制御なしに連続演算が可能で、メモリの読み書きは入出力時にのみ発生する
  • 主に 行列積、畳み込み に最適化されている
  • パイプライン処理により、演算とデータ移動を同時進行させてスループットを最適化する

シストリックアレイの欠点 - スパース性(sparsity)

  • すべての演算要素が常に有効化されるため、疎行列 には適していない
  • 今後DLモデルが不規則なスパース性を志向する場合、制約として作用する可能性がある

2. 事前コンパイル(AoT)とキャッシュ最小化

  • TPU-XLAコデザイン により、不規則なメモリアクセスを必要とするキャッシュ依存を最小化し、エネルギー消費を削減
  • XLAコンパイラが演算グラフの解析を通じてメモリアクセスパターンを事前に算出し、キャッシュの代わりにスクラッチパッドメモリ中心で運用する
  • JAX の @jit はJITとAoTの中間形態で、初回実行時に静的グラフを生成し、その後XLAでAoTコンパイルを行う
  • 入力形状が変わると再コンパイルが必要で、動的パディングや反復処理には非効率な面がある
  • 高いエネルギー効率を達成する一方で、柔軟性の不足が欠点

TPUv4のエネルギー効率

  • 現代のチップは HBM3 メモリを使用し、エネルギー消費を低減している
  • 演算に比べてメモリ処理のエネルギー消費は数十倍から数百倍に達するため、メモリアクセスを最小化することで大幅な効率向上が可能

TPUマルチチップ構造

トレイレベル(Tray/Board; 4チップ)

  • 1トレイは4個のTPUチップ（8個のTensorCore）とCPUホストで構成される
  • Host↔Chip接続はPCIe、Chip↔Chip接続は Inter-Core Interconnect(ICI) で、より高い帯域幅を提供する

ラックレベル(Rack; 4x4x4=64チップ)

  • 1ラックは64個のTPUチップで構成され、3Dトーラス(4x4x4) でICI、OCS(Optical Circuit Switching)により接続される
  • Googleではラック(Rack)、ポッド(Pod)、スライス(Slice)を区別している
    • ラック: 64チップの物理単位（=キューブ）
    • ポッド: ICIおよびOCSで接続可能な最大単位（例: TPUv4=4096チップ=64ラック）
    • スライス: 4チップ〜Superpodの間で任意に構成できる抽象単位

OCSの利点

  1. ラップアラウンド(Wraparound) : 各軸をリング（1Dトーラス）化し、ノード間の最悪移動ホップ数を削減
  2. 柔軟な非連続スライス: OCSスイッチング構造により、物理的に離れたノードも1つのスライスとして構成可能で、リソース活用や保守が容易
  3. ツイストトポロジ: 同一チップ数（x, y, z固定）でも接続構造を変更し、特定の演算パターンに対して速度を最適化できる（例: twisted torus）

トポロジ活用事例

  • キューブ型: データ並列・テンソル並列に適する（最大帯域幅）
  • 直線型（葉巻型） : パイプライン並列に有利
  • ツイストトーラス: all-to-all通信が必要な場合（例: tensor parallel）に高速化

スーパーポッド(Full Pod/Superpod; TPUv4: 4096チップ/64ラック)

  • 複数のラックをICIおよびOCSで接続し、超大規模システムを構築
  • スライストポロジ の種類によって通信帯域幅や並列処理性能が変化する
  • OCSにより、非連続スライスやツイストトポロジなどの柔軟性を確保

マルチポッドレベル(Multi-pod/Multislice; TPUv4: 4096個超)

  • 複数のポッドを データセンターネットワーク(DCN) で接続し、大規模学習インフラを提供するが、帯域幅はICIより低い
  • PaLM の学習では2ポッド（6144 TPUv4）を活用し、合計6ポッドのリソース管理で運用された
  • 大規模モデル開発では、研究者の設定値（並列次元など）に応じて XLA コンパイラが通信パターンを最適化する
  • XLA が各スライス・ポッド間の通信演算を挿入し、最小限のコード変更で大規模分散学習を実現する

実機ハードウェアと図示例

  • TPUラック: 4x4x4の3Dトーラスが1ユニットで、各行ごとに2トレイ（8チップ）
  • TPUv4トレイ: 実際には4つのPCIeポート（各TPUに1つずつ）
  • TPUv4チップ: 中央にASIC、周囲に4つのHBMスタック（2 TensorCore基準）
  • TPUv4i(推論型) チップフロアプラン: 1 TensorCoreで、広いCMEM面積を占める

まとめ

  • Google TPU Research Cloud(TRC) に研究支援への謝意を示す

参考資料

関連記事

まだコメントはありません。

まだコメントはありません。