Google Cloud の TPU について

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 のライフサイクル全体を管理できます。Google では、レガシーの Cloud TPU API ではなく Compute Engine を使用して TPU リソースをプロビジョニングすることをおすすめします。

詳細については、Compute Engine の Cloud TPU リソースをご覧ください。

Google Kubernetes Engine

Google Kubernetes Engine（GKE）は、大規模な AI ワークロードをオーケストレートするためのフルマネージドのマルチテナント Kubernetes 環境を提供します。GKE は、TPU VM の作成、構成、削除など、TPU ノードとノードプールのライフサイクル管理をサポートしています。

詳細については、GKE での TPU についてをご覧ください。

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 マシンをご覧ください。

次のステップ