システム アーキテクチャ

Tensor Processing Unit(TPU)は、Google が設計した ML アクセラレータです。Cloud TPU では、TPU をスケーラブルな GPC クラウド リソースとして利用できます。TensorFlowPytorchJAX などの機械学習フレームワークを使用して、Cloud TPU で機械学習ワークロードを実行できます。

1 個の TPU デバイスには 4 つのチップが搭載されており、それぞれのチップに 2 つの TPU コアが含まれています。TPU コアには、1 つ以上の Matrix Multiply Unit(MXU)、Vector Processing Unit(VPU)、および Scalar Unit が含まれています。

MXU は、シストリック アレイ内の 128 x 128 の乗算アキュムレータで構成されています。MXU は、TPU チップで大部分の計算処理を行います。各 MXU により、bfloat16 の数値形式を使用して、サイクルごとに 16,000 の乗累算演算を処理できます。

VPU は、活性化、ソフトマックスなどの一般的な計算に使用されます。スカラー単位は、制御フロー、メモリアドレスの計算、その他のメンテナンス オペレーションに使用されます。

TPU のバージョン

TPU デバイスの正確なレイアウトは、使用する TPU のバージョンによって異なります。TPU v2 および TPU v3 のアーキテクチャの詳細とパフォーマンス特性については、ディープ ニューラル ネットワークをトレーニングするための一領域に特化したスーパーコンピュータをご覧ください。

TPU v2

TPU v2 ボードには、4 つの TPU チップと 16 GiB の HBM が搭載されています。各 TPU チップには、2 つのコアが含まれています。各コアには、MXU、ベクトル ユニット、スカラー ユニットがあります。

画像

TPU v3

TPU v3 ボードには、4 つの TPU チップと 32 GiB の HBM が搭載されています。各 TPU チップには、2 つのコアが含まれています。各コアには、MXU、ベクトル ユニット、スカラー ユニットがあります。

画像

Cloud TPU には、次の TPU 構成があります。

  • 単一の TPU デバイス
  • TPU Pod - 高速相互接続で接続された TPU デバイスのグループ
  • TPU スライス - TPU Pod のサブディビジョン

TPU v2 に対する TPU v3 のパフォーマンス上の利点

TPU v3 構成では、コアあたりの FLOPS とメモリ容量が増えているため、以下の方法でモデルのパフォーマンスを改善できます。

  • 計算依存型モデルの場合、TPU v3 構成を使用するコアごとのパフォーマンスが大幅に向上します。メモリ依存型モデルの場合、TPU v2 構成で TPU v3 構成と同等のパフォーマンスの向上は達成できません。

  • データが TPU v2 構成のメモリに収まらない場合、TPU v3 を使用することで、パフォーマンスが向上し、中間値の再計算(再実体化)も減少します。

  • TPU v3 構成では、TPU v2 構成に収まらないバッチサイズの新しいモデルを実行できます。たとえば、TPU v3 を使用すると、さらに深い ResNets を扱え、RetinaNet でより大きなイメージを処理できます。

TPU v2 で入力依存型(「インフィード」)に近いモデルは、トレーニング ステップで入力を待機するため、Cloud TPU v3 でも入力依存型になる可能性があります。インフィードの問題を解決する方法については、パイプライン パフォーマンス ガイドをご覧ください。

TPU の構成

TPU は、次の構成で使用できます。

  1. シングル TPU ボード
  2. TPU Pod
  3. TPU Pod スライス

シングル TPU ボード

シングルボード TPU 構成は、4 個の TPU チップ(8 個の TPU コア)を備えたスタンドアロン ボードで、他の TPU ボードへのネットワーク接続はありません。シングルボード TPU は TPU Pod の構成の一部ではなく、TPU Pod の一部分を占有することもありません。

TPU Pod とスライス

TPU Pod または TPU Pod スライスでは、TPU チップは高速相互接続を使用して接続され、各 TPU チップは TPU デバイス上の他のチップと直接通信します。TPU ランタイムは、Pod またはスライスの各 TPU コアへのデータの分散を自動的に処理します。Pod スライスは、32、128、512、1024、2048 のコアで使用できます。

Cloud TPU VM アーキテクチャ

TPU は行列演算のみを実行できるため、各 TPU ボードは CPU ベースのホストマシンに接続され、TPU では実行できない演算を実行できます。ホストマシンは、Cloud Storage からのデータの読み込み、データの前処理、TPU へのデータ送信を行います。

画像

TPU Pod では、TPU ボードごとに TPU ホストがあります。

画像

TPU ホスト(および TPU ボード)の操作方法は、使用する TPU VM アーキテクチャ(TPU ノードまたは TPU VM)によって異なります。

TPU ノード

TPU ノードは、TPU の本来の扱い方です。これには、gRPC を介して TPU ホストと通信する追加のユーザー VM が必要です。TPU ホストに直接アクセスすることはできません。

画像

TPU VM

TPU VM を使用する場合は、TPU デバイスに物理的に接続する Google Compute Engine VM に直接 SSH 接続します。VM への root アクセス権を取得すると、任意のコードを実行できます。コンパイラとランタイム デバッグログ、エラー メッセージにアクセスできます。

画像

サポートされているフレームワーク

JAX、PyTorch、TensorFlow などのフレームワークでは、すべての TPU VM で提供される libtpu という共有ライブラリを介して TPU にアクセスします。このライブラリには、TPU プログラムのコンパイルに使用される XLA コンパイラ、コンパイルされたプログラムの実行に使用される TPU ランタイム、TPU への低レベルアクセスにランタイムで使用される TPU ドライバが含まれます。

画像

TensorFlow

TPU VM

TPU VM を使用すると、ユーザー VM 上で実行する Python コードではなく、TPU ホストで直接実行できます。

画像

TensorFlow と Cloud TPU の詳細については、Cloud TPU での TensorFlow モデルの実行をご覧ください。

TPU ノード

Cloud TPU ノードシステムのアーキテクチャは、もともと TensorFlow 用に構築されています。TPU ホストはユーザーがアクセスできず、TensorFlow サーバーのヘッドレス コピーを実行します。Python や TensorFlow グラフとして表されていないユーザーコードは実行されません。ユーザーコードは、gRPC ネットワークを介して TPU ホストと通信する別のリモート VM で実行されます。

画像

PyTorch

TPU VM

TPU VM では、PyTorch コードは TPU ホスト上で直接実行されます。

画像

PyTorch と Cloud TPU の詳細については、Cloud TPU での PyTorch モデルの実行をご覧ください。

TPU ノード

PyTorch は、XRT と呼ばれるライブラリを使用して Cloud TPU ノード アーキテクチャ上で実行され、XLA グラフとランタイム命令を TensorFlow gRPC 接続を介して送信し、TensorFlow サーバーで実行できます。各 TPU ホストには、ユーザー VM が必要です。

画像

JAX

TPU VM

TPU VM では、コードを TPU ホストで直接実行できるため、ユーザー VM は必要ありません。

画像

Cloud TPU で JAX を実行する詳細については、JAX クイックスタートをご覧ください。

TPU ノード

Cloud TPU ノード上の JAX は、ホスト VM ごとに個別のユーザー VM が必要という点で PyTorch に似た形で動作します。

画像

次のステップ