ハクソク

世界を動かす技術を、日本語で。

TPUの深層解析

338日前原文(henryhmko.github.io)

概要

  • TPU はGoogleが開発した AI専用ASIC で、 行列演算性能とエネルギー効率 が特徴。
  • ハードウェアとソフトウェアの協調設計 （例：XLAコンパイラ）により 高いスケーラビリティ を実現。
  • シストリックアレイ や パイプライン処理、 AOTコンパイル などが設計の要。
  • ラックやポッド単位での拡張性 に優れ、GoogleのAIサービス基盤を支える。
  • メモリアクセスの最適化 と エネルギー効率の向上 が設計思想の中核。

TPUの設計思想と概要

  • TPU（Tensor Processing Unit） は、GoogleがAIワークロード向けに開発した 専用ASIC。
  • 主な強みは 行列演算（matmul）スループット と エネルギー効率 の両立。
  • 2006年 に設計検討が始まり、当初は CPU の余剰計算力で十分だったが、 2013年 の音声検索AI導入で 計算需要が急増。
  • 現在は Gemini や Veo などのAIモデルの 学習・推論、 DLRM などのレコメンデーションモデルにも利用。
  • TPUのスケーラビリティ は、 ハードウェア （省電力・モジュール性）と ソフトウェア （XLAコンパイラ）の 協調設計 によって実現。

TPU単一チップ構成（TPUv4を例に）

  • 1チップに 2つのTensorCore （推論専用TPUは1つのみ）。
  • 各TensorCoreは CMEM（128MiB） と HBM（32GiB） を共有。
  • TensorCore内部 には以下のユニットが存在：
    • Matrix Multiply Unit（MXU） ：128×128のシストリックアレイ、行列演算を担当。
    • Vector Unit（VPU） ：ReLUや加算・乗算など要素単位演算。
    • Vector Memory（VMEM; 32MiB） ：HBMからデータをバッファリング。
    • Scalar Unit + Scalar Memory（SMEM; 10MiB） ：制御・スカラー演算・アドレス生成。
  • TPUのオンチップメモリ（CMEM, VMEM, SMEM）はGPUのL1/L2キャッシュより大容量。
  • HBM容量はGPUより小さい が、 計算コア数は少数精鋭。
  • 高スループット ：TPUv5pは 500TFLOPs/秒/チップ、ポッド全体で 4.45ExaFLOPs/秒 （8960チップ）。
  • 最新のTPUv7（Ironwood） は 42.5ExaFLOPS/秒/ポッド （9216チップ）。

TPU設計思想の要点

  • シストリックアレイ＋パイプライン、 Ahead-of-Time（AOT）コンパイル、 演算パターンの予測容易性 が三本柱。
  • 行列演算・畳み込み のような データフローが固定的な処理 に特化。
  • シストリックアレイ は、隣接する処理素子間でデータを流しながら計算を進める構造。
    • 制御ロジック不要、 入出力以外のメモリアクセス削減。
    • パイプライン処理 で計算とデータ移動を重ね合わせ。
  • 欠点 ：スパース（疎）行列には不向き、全PEが常時稼働してしまう。

AOTコンパイルとキャッシュレス設計

  • 従来のキャッシュ は 不規則なメモリアクセス に対応するためのもの。
  • TPUは用途が限定的でアクセスパターンが予測可能 なため、 キャッシュを排除しスクラッチパッドメモリで対応。
  • XLAコンパイラ が事前解析し、最適なメモリアクセスと演算スケジューリングを生成。
  • JAX は @jit による JIT＋AOTハイブリッド。最初の呼び出し時に計算グラフを静的に生成しXLAへ渡す。
    • 入力形状が変わるたびに再コンパイルが必要、動的パディングやループには弱い。
  • 柔軟性の欠如 や コンパイラ依存 がデメリットだが、 エネルギー効率 と 高スループット を優先。

TPUのエネルギー効率（TPUv4例）

  • メモリアクセスが演算よりも桁違いにエネルギーを消費。
  • HBM3 など最新メモリは従来DRAMより省電力。
  • メモリアクセス削減 が プログラムの高速化と省電力化 の両方に寄与。
  • FLOPS増加とメモリアクセス削減のトレードオフ が現代スケーリング則の要点。

TPUのマルチチップ構成

  • トレイ（Tray/Board） ：4チップ（8 TensorCore）、各トレイにCPUホスト。
    • ホスト⇔チップ はPCIe、 チップ間 は Inter-Core Interconnect（ICI） で高速通信。
  • ラック（Rack） ：64チップ（4×4×4）、3Dトーラス構成で拡張性重視。
  • ポッド（Pod） ：最大単位、例：TPUv4は4096チップ（64ラック）。
  • スライス（Slice） ：4チップ～ポッドまでの任意構成、物理単位ではなく抽象単位。

OCS（Optical Circuit Switching）の利点

  • OCSは3Dトーラスの面に配置、ノード間のラップアラウンド通信を実現。
  • 最悪ケースのホップ数を半減、 高並列化時の通信遅延を大幅削減。
  • 拡張性が高く、スーパーコンピュータ級のシステム構築が容易。

トポロジーの違い

  • 古いTPU世代（v2, v3）や一部推論用TPU（v5e, v6e）は2Dトーラス。
  • TPUv7（Ironwood）は3Dトーラス とされるが詳細は未公開（2025年6月時点）。

まとめ

  • TPUはAI専用に特化した設計思想 により、 高スループット・省電力・優れたスケーラビリティ を実現。
  • シストリックアレイやAOTコンパイル など独自アプローチで、 GPUとは異なる設計哲学 を持つ。
  • GoogleのAIサービスの根幹インフラ として今後も発展が期待される。

Hackerたちの意見

要するに、キャッシュはハードウェアを柔軟にして、さまざまなアプリケーションに適応できるようにするんだ。これが、GPUがとても柔軟なハードウェアである大きな理由の一つだよ（TPUと比べてね）。これは正しいけど、言い方がちょっと違うね。エネルギーを消費するのはキャッシュそのものじゃなくて、「ページフォールト」の時に自動でキャッシュにロード/フェッチ/ストアするMMUなんだ。TPUにはMMUがないし、さらにプッシュアーキテクチャだからね（プルとは違う）。

TPUが使っているレギュラーグリッドのシストリックアレイにうまくマッピングされるアルゴリズムの良い参考文献を教えてくれない？ この記事では行列の掛け算や畳み込みが良いって言ってるけど、他に何かある？ 固有分解？ SVD？ 行列の指数関数？ Ax = bやAX = Bを解くこと？ コレスキー分解？

└

128x128（できればもっと大きい）な密行列の掛け算として表現できるもので、他には何もいらない。

└

https://jax-ml.github.io/scaling-book/ は、読むべき最高の参考文献の一つだと思うよ。単一デバイスと分散計算がTPUのハードウェア機能にどうマッピングされるかが詳しく説明されてる。特に、トランスフォーマーの計算を前向きと後ろ向きの両方でマッピングすることに重点が置かれてるから、トランスフォーマーネットワークの構造についての知識が少し必要だね。

└

SVDや固有分解は、よく多くの行列積を作ることになる（例えば、Krylovベースの手法を使うとき、アーノルディやKrylov-Schurなど）。だから、TPUはそこでうまく機能すると思うよ。Ax = bを解くための一つの方法であるGMRESもアーノルディ分解に基づいてるし。

その冷却チャンネルにはポンプかメーターバルブとしてNEMAステッパーが付いてるの？[0] もしそうなら、すごいね。それはやりすぎな気がする。[0]: https://henryhmko.github.io/posts/tpu/images/tpu_tray.png

└

確実にクローズドループだね、サーボかもしれない。

どうしてGoogleで働いてないのに、こんなにTPUについて詳しい人がいるの？

└

最後の謝辞から推測するに、著者はhttps://sites.research.google/trc/about/を通じてTPUにアクセスできるみたいだね。でも、これが唯一の方法じゃないよ。TPUはGCPで運営している企業にも利用可能で、GPUとは異なる価格/性能のポイントがあるから、TPUを実際に使ってみる別の方法になるよ。

└

ブログ記事に書いてあることは、基本的にみんな知ってることだよね。GoogleはTPUについて論文を出したり、講演をしたりしてるし。ただ、詳細が足りない部分もあって、いくつかの仮定や推測が必要だね。JaxやXLAは（部分的に）オープンソースで、TPUがどう動いてるかの手がかりも与えてくれるよ。 https://arxiv.org/abs/2304.01433 https://jax-ml.github.io/scaling-book/

いい記事だね！

ELI5: GPUとTPUの最適化がLLMの決定論に具体的にどう影響するの？それともこれは全部神話なの？

└

LLMは一般的に決定論的だよ。トークンのサンプリングステップは、結果を良くするためにある程度ランダム化されることが多い（創造性を高めたり、ループを避けたりするため）。でも、それをオフにすることもできるよ（単純なサンプラーの場合は温度をゼロにすればいい）。

└

結果の決定論には影響しないけど、異なるアーキテクチャはパフォーマンスに関して異なる決定論の保証があるよ。スケジューリングや他の要因の結果としてね。TPUはGroqのTPUアクセラレーターと似た系譜を持ってる（ちなみに、私はGroqで働いてる）。実際、GroqのTPUは完全に決定論的で、決定論的な出力を得られるだけでなく、決定論的なサイクル数で得られるんだ。ただし、ハードウェアを決定論的にするには、HWレベルのスケジューリングや他の非決定論的な要素を諦める必要がある。これにより、アーキテクチャは「十分に賢いコンパイラ」に非常に依存することになる。TPUやそれに似たプロセッサは一般的にVLIWと見なされていて、すべて同様にコンパイラが賢いスケジューリングの決定を前もって行うことに依存してるから、良い計算/IOのオーバーラップを確保したり、パイプラインバブルを排除したりするのが重要なんだ。一方、GPUはチップ自体に非常に洗練されたスケジューリングシステムを持っていて、カーネルのスワッピングなどもあって、もっと柔軟で、コンパイラにあまり依存せず、あまり労力をかけずにプロセッサの高い利用率を達成しやすいんだ。要するに、TPUは決定論的なサイクル保証があるかもしれないけど、現在の世代/アーキテクチャのGPUは非決定論的なスケジューリングやメモリアクセスパターンを使ってるから無理だよ。でも、どちらも決定論的なプログラムに対しては決定論的な出力を生成するよ。

誰か、以下のことがどうして真実になり得るのか教えてくれない？ 1. TPUはNvidiaチップの本格的な競争相手。 2. 最高のチップを持つチップメーカーは1〜3.5兆ドルの価値がある。 3. Googleの時価総額は2兆ドル。 4. GoogleがTPUを売らないのは正しい。レンタルの方がいいって話は聞いたけど、実際に良いものなら、売るのも会社の他の部分と同じくらい良いビジネスになるんじゃない？

└

売るためには、それをサポートするためのインフラをかなり構築しなきゃいけないよ。レンタルで高いマージンを得てるのに、わざわざその手間をかける理由はないよね。それに、もし本当に良いものなら、競合とその技術を共有して平等にするのは避けた方がいいし。「最高のチップを持つチップメーカー」って言ったらNvidiaのことだし、他にはあまりないよ。それに、AlphabetはTPUを生産するだけじゃないからね。