Cloud TPU

解決する課題

• 大規模な深層学習モデルの学習に時間がかかりすぎるため、行列演算を高速化したい
• 学習・推論のスループットあたりのコストを、汎用 GPU 以外の選択肢でも最適化したい
• モデルが1チップに収まらないほど大きく、多数のアクセラレーターを高速に相互接続して分散学習したい
• アクセラレーターのプロビジョニングや相互接続の構築を自分で組まずにマネージドに使いたい

これらに対し、Google が機械学習向けに自社設計した専用プロセッサが Cloud TPU です。汎用計算ではなく、ニューラルネットワークの中心となる大規模な行列演算に特化することで、対象ワークロードでの効率を高めます。

主要概念と用語

• TPU（Tensor Processing Unit）: Google が機械学習向けに設計した専用プロセッサ（ASIC）。行列演算を集中的に処理する
• チップ / コア: 演算の物理単位。1チップに高速メモリ（HBM）が直結し、複数のコアを持つ世代もある
• TPU Pod / スライス: 多数のチップを専用の高速相互接続で束ねた構成が Pod。その一部を切り出した単位がスライスで、必要な規模だけ確保できる
• ホスト VM: TPU に接続される Compute Engine ベースの VM。ここで学習スクリプトが動き、TPU に演算をオフロードする
• XLA: 計算グラフを TPU 向けに最適化・コンパイルするコンパイラ。JAX/TensorFlow/PyTorch などのフロントエンドを TPU 上で動かす土台
• 世代（v シリーズ）: TPU には世代があり、新しい世代ほど演算性能やメモリ、相互接続が強化される。推論特化の派生もある
• データ並列 / モデル並列: 同じモデルを多数のチップで分担して学習する方式。前者はデータを、後者はモデル自体を分割する

仕様・制限・クォータ

• TPU は特定のリージョン・ゾーンで提供され、世代ごとに利用可能な場所が異なる
• 確保の単位は**チップ数（スライスサイズ）**で、小規模なスライスから大規模な Pod まで選べる
• 課金はおおむね**チップ時間（使った規模 × 時間）**ベースで、確保している間は課金が続く
• TPU の数（チップ数）にはクォータがあり、世代・リージョンごとに上限が設定される。大規模利用は引き上げ申請や予約が前提になる
• 行列演算に特化するため、分岐の多い汎用処理や前後処理は不向きで、その部分はホスト VM 側や他サービスに任せる
• 具体的な世代名・性能値・単価・クォータ値は変動するため、本番設計では公式ドキュメントと自プロジェクトの割り当てを確認する

内部の仕組み

TPU の中核は、行列の積和演算を大量に並列実行するシストリックアレイと呼ばれる演算器です。データを格子状の演算ユニットへ流し込み、途中結果をメモリに書き戻さずユニット間で受け渡すことで、行列演算のスループットを高めます。各チップには高速な**HBM（高帯域メモリ）**が直結し、重みや中間データの供給律速を抑えます。

利用者は JAX・TensorFlow・PyTorch といったフロントエンドでモデルを記述し、XLA が計算グラフを TPU 向けに最適化してコンパイルします。実際の学習スクリプトは TPU に接続されたホスト VM上で動き、行列演算の重い部分が TPU にオフロードされる、という分担になります。

1チップに収まらない大規模モデルでは、多数のチップを専用の高速相互接続で束ねた Pod を使い、データ並列やモデル並列でチップ間に処理を分散します。チップ間通信が速いほど分散学習のスケール効率が上がるため、相互接続は TPU の重要な設計要素です。

設計パターン / ベストプラクティス

• まずフレームワークの対応を確認する。JAX や TensorFlow、PyTorch/XLA など TPU 対応の経路でモデルを組む
• 入力パイプラインを TPU の処理速度に追従させる。前処理や I/O がボトルネックだと TPU が遊ぶため、データ供給を並列化・先読みする
• 大規模モデルはスライスサイズを段階的に拡大し、スケール効率（チップを増やしたぶん速くなるか）を測りながら広げる
• 中断に強い学習にするため、チェックポイントを Cloud Storage に定期保存し、再開できるようにする
• 推論主体なら推論に適した世代・構成を選び、学習用の大規模 Pod をそのまま使わない
• バッチサイズやデータ型（混合精度など）をTPU が効率を出せる形に合わせる

運用・監視

• ジョブやホスト VM のメトリクス・ログは Cloud Monitoring / Cloud Logging で監視し、TPU の利用率や入力パイプラインの詰まりを把握する
• 学習が遅いときは、まずTPU が演算で律速か、入力パイプラインで律速かを切り分ける。後者ならデータ供給側を改善する
• 大規模ジョブはプロファイリングツールでステップ時間の内訳を確認し、通信・演算・I/O のどこが支配的かを見る
• 長時間ジョブはチェックポイントからの再開を前提に、障害・中断時のリカバリ手順を用意する
• 大規模・継続利用では予約を活用し、必要なときに容量を確保できるようにする

コスト

Cloud TPU の課金は、おおむね確保したチップの規模（スライスサイズ）× 利用時間に対して発生します。したがってコスト最適化の鍵は、適切な規模を必要な時間だけ使い、TPU を遊ばせないことにあります。

セキュリティ

• TPU を使うジョブには専用のサービスアカウントを割り当て、学習データのバケットや必要な API だけに権限を絞る（最小権限）
• 学習データやチェックポイントを置く Cloud Storage は、暗号化と IAM によるアクセス制御を前提にする
• ホスト VM はVPC・サブネットに配置し、ネットワークの到達範囲を最小化する
• TPU リソースの作成・利用権限は IAM のロールで制御し、誰がどの規模を確保できるかを管理する
• 認証情報や API キーをコードに直書きせず、Secret Manager などから取得する

関連サービス・比較

機械学習アクセラレーターとして最も比較されるのは、Compute Engine などで使える GPU です。両者は競合というより、フレームワーク対応やワークロード特性に応じて使い分ける関係にあります。

TPU は単体で使うほか、Vertex AI のマネージドな学習・推論や GKE 上のワークロードからも利用でき、データは Cloud Storage を介して受け渡すのが一般的です。汎用性と既存資産を重視するなら GPU、対応フレームワークで大規模学習のコスト効率を狙うなら TPU、という整理になります。

ハンズオン / CLI例

# TPU VM（ホスト VM と TPU が一体の構成）を作成する例
# accelerator-type で世代とスライスサイズ、version でソフトウェアイメージを指定する
gcloud compute tpus tpu-vm create demo-tpu \
  --zone=asia-northeast1-b \
  --accelerator-type=<利用する世代とスライス> \
  --version=<対応するランタイムイメージ>

# 作成した TPU VM の一覧を確認
gcloud compute tpus tpu-vm list \
  --zone=asia-northeast1-b

# TPU VM に SSH 接続して学習スクリプトを実行
gcloud compute tpus tpu-vm ssh demo-tpu \
  --zone=asia-northeast1-b

# 使い終わったら削除して課金を止める
gcloud compute tpus tpu-vm delete demo-tpu \
  --zone=asia-northeast1-b