コレクションでコンテンツを整理 必要に応じて、コンテンツの保存と分類を行います。

システム アーキテクチャ

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

マトリックス乗算ユニット(MXU)は、シストリック アレイに含まれる 128 x 128 の乗算 / アキュムレータで構成されます。MXU は、TensorCore(TPU コア)の大部分の計算能力を提供します。各 MXU により、bfloat16 の数値形式を使用して、サイクルごとに 16,000 の乗累算演算を処理できます。

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

TPU のバージョン

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

TPU v4

最小の TPU v4 構成には、4 つの TPU チップ、32 GiB の HBM、128 MiB の共有共通メモリが含まれています。各 TPU チップには、2 つの TensorCore が含まれています。各 TensorCore には、4 つの MXU、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。次の図は、TPU v4 チップを示しています。

イメージ

TPU v3

最小の TPU v3 構成には、4 つの TPU チップと 32 GiB の HBM が含まれています。各 TPU チップには、2 つの TensorCore が含まれています。各 TensorCore には、2 つの MXU、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。次の図は、TPU v3 チップを示しています。

イメージ

TPU v2

最小の TPU v2 構成には、4 つの TPU チップと 16 GiB の HBM が含まれています。各 TPU チップには、2 つの TensorCore が含まれています。各 TensorCore には、1 つの MXU、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。次の図は、TPU v2 チップを示しています。

イメージ

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

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

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

メモリシステム:

  • v4 TPU チップには、チップ全体で統合された 32 GiB HBM メモリ空間があり、2 つのオンチップ TensorCore 間の調整を改善できます。
  • 最新のメモリ標準と速度を使用した HBM パフォーマンスの向上。
  • 512B の粒度での高パフォーマンス ストライプ化をネイティブにサポートすることで、DSP のパフォーマンス プロファイルを改善しました。

Tensor コア数:

  • 2 倍の MXU とより高いクロック レートで最大 257 TFLOPS を実現。
  • 2 倍の転移および置換の帯域幅。
  • Common Memory(Cmem)のロードストア メモリ アクセス モデル。
  • MXU 重み付け読み込みの高速化と 8 ビットモードのサポートにより、バッチサイズを小さくし、推論のレイテンシを改善しました。

チップ間相互接続:

チップあたり 6 個の相互接続リンクにより、ネットワークダイヤグラムがより小さいネットワーク トポロジを実現。

その他:

  • ホスト用の x16 PCIE gen3 インターフェース(直接接続)。
  • セキュリティ モデルの改善。
  • エネルギー効率の向上。

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 v4 の構成

TPU v4 Pod は、再構成可能な高速リンクで相互接続された 4,096 のチップで構成されます。TPU v4 の柔軟なネットワーキングでは、さまざまな方法で同じサイズのスライスにチップを接続できます。v4 TPU の構成には AcceleratorConfig を使用します。 AcceleratorConfig を使用すると、トポロジ フィールドを通じてチップの観点からスライスサイズを指定できます。

TPU v4 トポロジは、TPU チップの数、形状、相互接続を記述する 3 タプルを使用して指定されます。次の図は、一般的な TPU v4 トポロジを示しています。

イメージ

より大きなスライスを、1 つまたは複数の 4x4x4 の「キューブ」チップから構築できます。

特定の数のチップの TPU スライスは、さまざまな方法で構成できます。たとえば、AcceleratorType が v4-1024 の TPU スライスを 4x4x32、4x8x16、8x8x8 として構成できます。v4-4096 の TPU では、柔軟性がさらに向上します(4x4x128、4x8x64、4x16x32、8x16x16、8x16x16_twisted)。詳細については、TPU タイプとトポロジをご覧ください。

TPU v2 と v3 の構成

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

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

シングル TPU ボード

シングルボード TPU 構成は、4 個の TPU チップ(8 個の Tensor コア)を備えたスタンドアロン ボードで、他の 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 ノード アーキテクチャは、gRPC を介して TPU ホストと通信するユーザー VM で構成されています。このアーキテクチャを使用する場合、TPU ホストに直接アクセスすることはできません。これにより、トレーニング エラーや TPU エラーのデバッグが難しくなります。

イメージ

TPU VM

TPU VM アーキテクチャでは、ユーザー VM が不要になり、TPU デバイスに物理的に接続されている 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 TPU VM の詳細については、Cloud TPU VM ユーザーガイドをご覧ください。

TPU ノード

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

イメージ

PyTorch

TPU VM

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

画像

PyTorch と Cloud TPU の詳細については、PyTorch/XLA ユーザーガイドをご覧ください。

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 に似た形で動作します。

画像

次のステップ