Tensor Processing Unit（TPU）は、Google 独自に開発された特定用途向け集積回路（ASIC）であり、ML ワークロードと AI ワークロードの高速化に使用されます。複雑な基盤モデルを数週間トレーニングする場合でも、大規模な推論を実行する場合でも、TPU は JAX や PyTorch などのフレームワークに最適化された、スケーラブルで専用のコンピューティング リソースを提供します。
Cloud TPU は、最も要求の厳しい AI ワークロードに対応するように設計されています。主なメリットは次のとおりです。
行列計算向けに最適化: TPU は、ML アルゴリズムの基本となる大規模な行列演算を非常に効率的に実行するために、行列乗算ユニット（MXU）を使用して特別に設計されています。
高帯域幅メモリ（HBM）: オンチップの高帯域幅メモリにより、より大きなモデルのトレーニングとサービングが可能になり、より大きなバッチサイズを効果的に利用できます。
スライスによる大規模なスケーラビリティ: TPU チップは、スライスと呼ばれるグループに接続できます。スライスを使用すると、ワークロードを最大数千の TPU チップにスケールアップして、大規模なトレーニング ジョブを実行できます。
TPU を使用するケース
TPU は、次のような特定のワークロード用に最適化されています。
- 行列計算が多くを占めるモデル
- メインのトレーニング ループ内にカスタムの PyTorch / JAX 演算がないモデル
- トレーニングに数週間または数か月かかるモデル
- 実際のバッチサイズが大きい大規模なモデル
- 高度なランキングやレコメンデーションのワークロードでよくある極めて大規模なエンベディングを含むモデル
TPU は、以下のワークロードには適していません。
- 頻繁な分岐を必要とするか、要素ごとの代数演算を多く含む線形代数プログラム
- 高精度の演算を必要とするワークロード
- メインのトレーニング ループにカスタム演算が含まれるニューラル ネットワーク ワークロード
Google Cloudのプロビジョニング オプション
運用上のニーズに応じて、次の Google Cloud プロダクトを使用して TPU にアクセスし、プロビジョニングできます。
Compute Engine
Compute Engine を使用すると、個々の TPU VM またはスライスを作成して管理できます。これにより、TPU VM のライフサイクル全体を管理できます。TPU リソースをプロビジョニングする場合は、以前の Cloud TPU API ではなく Compute Engine を使用することをおすすめします。
詳細については、Compute Engine の Cloud TPU リソースをご覧ください。
Google Kubernetes Engine
Google Kubernetes Engine（GKE）は、大規模な AI ワークロードをオーケストレートするためのフルマネージドのマルチテナント Kubernetes 環境を提供します。GKE は、TPU VM の作成、構成、削除など、TPU ノードとノードプールのライフサイクル管理をサポートしています。
詳しくは、GKE での TPU についてをご覧ください。
バッチ
Batch は、Compute Engine TPU VM で実行されるコンテナ化されたワークロードまたはスクリプトベースのワークロードのスケジューリング、キューイング、実行を行えるフルマネージド サービスです。
詳細については、TPU を使用する Batch ジョブを作成して実行するをご覧ください。
Cloud TPU
Google Cloud CLI や Cloud TPU 用の Cloud クライアント ライブラリを含む Cloud TPU API は、開発が終了しています。TPU リソースのプロビジョニングと管理には、オーケストレーションとワークロードのニーズに基づいて、Compute Engine または GKE を使用することをおすすめします。
詳細については、Cloud TPU API から移行するをご覧ください。
Compute Engine でサポートされている TPU バージョン
Compute Engine は、次の TPU バージョンをサポートしています。
- TPU7x（Ironwood）
- TPU v6e（Trillium）
- TPU v5p
各 TPU バージョンの詳細については、TPU マシンをご覧ください。
次のステップ
- Compute Engine の Cloud TPU リソースについて学習する。
- TPU ハードウェアについて確認する