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

システム アーキテクチャ

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

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

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

TPU のバージョン

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

TPU v4

各 v4 TPU チップには、2 つの TensorCore が含まれています。各 TensorCore には、4 つの MXU、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。次の表に、v4 TPU Pod の主な仕様と値を示します。

主な仕様 v4 Pod の値
チップあたりのピーク コンピューティング 275 TFLOPS(bf16 または int8)
HBM2 の容量と帯域幅 32 GiB、1200 GBps
最小/平均/最大電力として測定 90/170/192 W
TPU Pod のサイズ 4096 チップ
相互接続のトポロジ 3D トルー
Pod あたりのピーク コンピューティング 1.1 exaflops(bf16 または int8)
ポッドごとのすべて削減帯域幅 1.1 PB/s
ポッドあたりの全二分帯域幅 24 TB/s

次の図は、TPU v4 チップを示しています。

イメージ

TPU v3

各 v3 TPU チップには、2 つの TensorCore が含まれています。各 TensorCore には、2 つの MXU、1 つのベクトル ユニット、1 つのスカラー ユニットがあります。次の表に、v3 TPU Pod の主な仕様と値を示します。

主な仕様 v3 Pod の値
チップあたりのピーク コンピューティング 123 TFLOPS(bf16)
HBM2 の容量と帯域幅 32 GiB、900 GBps
最小/平均/最大電力として測定 123/220/262 W
TPU Pod のサイズ 1024 チップ
相互接続のトポロジ 2D トルー
Pod あたりのピーク コンピューティング 126 PFLOPS(BF16)
ポッドごとのすべて削減帯域幅 340 TB/s
ポッドあたりの全二分帯域幅 6.4 TB/s

次の図は、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 とより高いクロック レートで最大 275 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 を使用すると、topology フィールドを通じてチップの観点からスライスサイズを指定できます。

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

イメージ

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

特定の数のチップの TPU スライスは、さまざまな方法で構成できます。たとえば、AcceleratorType が v4-1024 の TPU スライスを 4x4x32、4x8x16、8x8x8 として構成できます。

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 アーキテクチャを使用すると、TPU デバイスに物理的に接続されている VM に SSH 接続できます。VM への root アクセス権があるため、任意のコードを実行できます。コンパイラとランタイム デバッグログ、エラー メッセージにアクセスできます。

イメージ

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

JAX、PyTorch、TensorFlow などのフレームワークでは、すべての TPU VM で提供される libtpu という共有ライブラリを介して TPU にアクセスします。libtpu には、XLA コンパイラ、TPU ランタイム ソフトウェア、TPU ドライバが含まれます。

イメージ

TensorFlow

TPU VM

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

イメージ

TensorFlow と Cloud 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 ノード アーキテクチャで実行されます。XRT は、TensorFlow サーバーで実行するために、gRPC を介して XLA グラフとランタイム命令を送信します。各 TPU ホストには、ユーザー VM が必要です。

イメージ

JAX

TPU VM

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

イメージ

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

TPU ノード

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

画像

次のステップ