Cloud TPU の概要

Tensor Processing Unit（TPU）は、Google 独自に開発された特定用途向け集積回路（ASIC）であり、機械学習ワークロードの高速化に使用されます。TPU ハードウェアの詳細については、TPU アーキテクチャをご覧ください。Cloud TPU は、TPU を Google Cloudでスケーラブルなコンピューティング リソースとして利用できるようにするウェブサービスです。

TPU は、機械学習アルゴリズムでよく見られる大規模な行列演算を実行するために設計されたハードウェアを使用して、モデルを効率的にトレーニングします。TPU にはオンチップの高帯域幅メモリ（HBM）が搭載されており、大きいモデルやバッチサイズを使用できます。TPU はスライスと呼ばれるグループに接続でき、コードをほとんど変更することなくワークロードをスケールアップできます。

TPU で実行されるコードは、アクセラレータ線形代数（XLA）コンパイラでコンパイルする必要があります。XLA は、ML フレームワーク アプリケーションによって出力されたグラフを受け取り、グラフの線形代数、損失、勾配のコンポーネントを TPU マシンコードにコンパイルするジャストインタイム コンパイラです。プログラムの残りの部分は TPU ホストマシンで実行されます。XLA コンパイラは、TPU ホストマシンで実行される TPU VM イメージの一部です。

TPU の用途

Cloud TPU は特定のワークロード用に最適化されています。状況によっては、Compute Engine インスタンスで GPU または CPU を使用して機械学習ワークロードを実行する場合があります。一般に、次のガイドラインに基づいてワークロードに最適なハードウェアを決定できます。

CPU

  • 最大限の柔軟性を必要とする迅速なプロトタイピング
  • トレーニングに時間がかからない単純なモデル
  • 実際のバッチサイズが小さい小規模なモデル
  • C++ で記述されたカスタムの TensorFlow オペレーションが多く含まれるモデル
  • ホストシステムの使用可能な I/O またはネットワーク帯域幅によって制限が課せられるモデル

GPU

  • CPU 上で少なくとも部分的に実行しなければならない多数のカスタム PyTorch/JAX 演算を使用するモデル
  • Cloud TPU で利用できない TensorFlow 演算を使用するモデル（利用可能な TensorFlow 演算のリストをご覧ください）
  • 実際のバッチサイズが大きい中～大規模なモデル

TPU

  • 行列計算が多くを占めるモデル
  • メインのトレーニング ループ内にカスタムの PyTorch/JAX 演算がないモデル
  • トレーニングに数週間または数か月かかるモデル
  • 実際のバッチサイズが大きい大規模なモデル
  • 高度なランキングとレコメンデーションのワークロードでよく見られる極めて大規模なエンベディングを使用するモデル

Cloud TPU は、以下のワークロードには適していません。

  • 頻繁な分岐を必要とするか、要素ごとの代数演算を多く含む線形代数プログラム
  • 高精度の演算を必要とするワークロード
  • メインのトレーニング ループにカスタム演算が含まれるニューラル ネットワーク ワークロード

TPU Google Cloud

TPU は、Cloud TPU VM、Google Kubernetes Engine、Vertex AI で使用できます。次の表に、各 Google Cloudサービスのリソースを示します。

Google Cloud サービス リソース
Cloud TPU Cloud TPU VM を使ってみる
Google Kubernetes Engine

GKE の TPU について

TPU を使用して GKE で Ray を実行する
Vertex AI

TPU を使用した Vertex AI でのトレーニング

Vertex AI でオンライン予測に TPU を使用する

モデル開発のベスト プラクティス

計算の多くを非行列演算（加算、再形成、連結など）が占めるプログラムは、MXU 使用率が高くならない可能性があります。次に、Cloud TPU に適したモデルを選択して作成する際に役立つガイドラインを示します。

レイアウト

XLA コンパイラでは、マトリックス ユニット（MXU）で効率的に計算を実行するために、行列乗算をより小さなブロックにタイリングするなどのコード変換が行われます。MXU ハードウェアの構造、128x128 のシストリック アレイ、8 の倍数の次元を基本とする TPU のメモリ サブシステムの設計が、タイリングの効率化のために XLA コンパイラによって使用されます。その結果、特定のレイアウトはタイリングに役立ちますが、他のレイアウトではタイリングの前に再形成が必要になります。再形成オペレーションは、多くの場合、Cloud TPU でメモリバウンドです。

図形

XLA コンパイラによって、最初のバッチに合わせて ML グラフがコンパイルされます。後続のバッチに異なる形状がある場合、モデルは機能しません。（形状が変化するたびにグラフを再コンパイルするのでは遅すぎます）。したがって、動的形状を持つテンソルがあるモデルは、TPU にはあまり適していません。

パディング

高パフォーマンスの Cloud TPU プログラムは、密な計算を 128x128 のチャンクにタイリングできるプログラムです。MXU 全体が行列計算で占有されない場合、コンパイラによってテンソルはゼロでパディングされます。パディングには次の 2 つの欠点があります。

  • ゼロでパディングされたテンソルでは、TPU コアの利用が不十分になります。
  • パディングによって、テンソルに必要なオンチップ メモリ ストレージ量が増加し、極端な場合にはメモリ不足エラーにつながる可能性があります。

パディングは必要に応じて XLA コンパイラによって自動的に行われますが、op_profile ツールを使用して、行われるパディングの量を決定できます。TPU に適したテンソル次元を選択することで、パディングを回避できます。

次元

適切なテンソル次元を選択すると、TPU ハードウェア、特に MXU から最大のパフォーマンスを引き出すために役立ちます。XLA コンパイラは、バッチサイズまたは特徴ディメンションのいずれかを使用して MXU を最大限に利用しようとします。したがって、これらのいずれかは 128 の倍数である必要があります。そうでない場合、コンパイラによっていずれかが 128 までパディングされます。理想的には、バッチサイズとフィーチャ次元は 8 の倍数である必要があり、これによりメモリ サブシステムから高パフォーマンスを引き出すことができます。

Cloud TPU のスタートガイド

サポートの依頼

Cloud TPU サポートにお問い合わせください。有効な Google Cloud プロジェクトがある場合は、次の情報を提供できるようにしてください。

  • Google Cloud プロジェクト ID
  • TPU 名（存在する場合）
  • ご提供いただく必要があるその他の情報

次のステップ

Cloud TPU の詳細を調べるには、次のリソースも役立ちます。