システム アーキテクチャ
Tensor Processing Unit（TPU）は、Google が設計した ML アクセラレータです。Cloud TPU では、TPU をスケーラブルな Google Cloud クラウド リソースとして利用できます。TensorFlow、Pytorch、JAX などの機械学習フレームワークを使用して、Cloud TPU で機械学習ワークロードを実行できます。
TPU チップの各 TensorCore は、1 つ以上のマトリックス乗算ユニット（MXU）、ベクトル ユニット、スカラー ユニットから構成されます。
MXU は、シストリック アレイ内の 128 x 128 の乗算アキュムレータで構成されています。 MXU は、TensorCore の大部分の計算能力を提供します。 各 MXU により、サイクルごとに 16,000 の乗累算演算を処理できます。 すべての乗算は bfloat16 の入力を取りますが、すべての累積は FP32 数値形式で実行されます。
ベクトル ユニットは、活性化やソフトマックスなどの一般的な計算に使用されます。スカラー ユニットは、制御フロー、メモリアドレスの計算、その他のメンテナンス オペレーションに使用されます。
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 のパフォーマンス上の利点
メモリシステム:
不均一メモリアクセス（NUMA）は、複数の CPU を搭載したマシンのコンピュータ メモリ アーキテクチャです。各 CPU は、高速メモリのブロックに直接アクセスできます。CPU とそのメモリは NUMA ノードと呼ばれます。NUMA ノードは、互いに直接隣接する NUMA ノードに接続されます。ある NUMA ノードの CPU は別の NUMA ノードのメモリにアクセスできますが、このアクセスは NUMA ノードのメモリにアクセスするよりも時間がかかります。
マルチ CPU マシン上で稼働するソフトウェアは、CPU が必要とするデータをその NUMA ノード内に配置し、メモリのスループットを向上させます。NUMA の詳細については、Wikipedia の 不均一メモリアクセスをご覧ください。
TPU v4 ランタイムは、NUMA-aware 専用ではありません。ただし、トレーニング スクリプトを NUMA ノード 0 にバインドすることで、NUMA-locality の利点を活用できます。
NUMA ノード バインディングを有効にするには:
numactl コマンドライン ツールをインストールします。
$ $ sudo apt-get update $ $ sudo apt-get install numactl
トレーニング スクリプトの起動時に
numactl --cpunodebind=0を使用します。これにより、スクリプト コードが NUMA ノード 0 にバインドされます。$ $ numactl --cpunodebind=0 python3 your-training-script
次の場合は NUMA ノード バインディングを有効にします。
- フレームワークに関係なく、ワークロードが CPU ワークロードに大きく依存している場合（例: 画像分類、レコメンデーション ワークロード）。
- -pod 接尾辞のない TPU ランタイム バージョン（例: tpu-vm-tf-2.10.0-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 構成では、TensorCore あたりの FLOPS とメモリ容量が増えているため、以下の方法でモデルのパフォーマンスを改善できます。
計算依存型モデルの場合、TPU v3 構成を使用する TensorCore ごとのパフォーマンスが大幅に向上します。メモリ依存型モデルの場合、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 の柔軟なネットワーキングでは、さまざまな方法で同じサイズのスライスにチップを接続できます。TPU スライスを作成するときには、TPU のバージョンと必要な TPU リソースの数を指定します。TPU v4 スライスを作成する場合、次の 2 つの方法のいずれかでタイプとサイズを指定できます。
AcceleratorTypeの使用この方法は、トポロジを指定していない場合におすすめします。
使用中:
AcceleratorConfig物理トポロジをカスタマイズする場合は、
AcceleratorConfigを使用することをおすすめします。これは通常、スライス数が 256 を超えるパフォーマンス チューニングに必要です。
AcceleratorType を使用して v4 TPU を構成するには、TPU スライスの作成時に --accelerator-type フラグを使用します。--accelerator-type を、TPU バージョンと使用する TensorCore の数を含む文字列に設定します。たとえば、32 個の TensorCore を使用して v4 スライスを作成するには、--accelerator-type=v4-32 を使用します。
AcceleratorConfig を使用して v4 TPU を構成するには、--version フラグと --topology フラグを使用します。--version を、使用する TPU バージョンに設定し、--topology をスライス内の TPU チップの物理配置に設定します。たとえば、16 チップの v4 スライスを作成するには、--version=v4 と --topology=2x2x4 を使用します。
TPU v4 トポロジは、TPU チップの数、形状、相互接続を記述する 3 タプルを使用して指定されます。次の図は、一般的な TPU v4 トポロジを示しています。
より大きなスライスを、1 つまたは複数の 4x4x4 の「キューブ」チップから構築できます。
特定の数のチップの TPU スライスは、さまざまな方法で構成できます。たとえば、512（1024 TensorCores）チップの TPU スライスは、4x4x32、4x8x16、8x8x8 のトポロジを使用して構成できます。
2,048 チップ（4,096 TensorCores）の TPU スライスは、さらに多くの 4x4x128、4x8x64、4x16x32、8x16x16 のトポロジ オプションを提供します。詳細については、TPU タイプとトポロジをご覧ください。
TPU v2 と v3 の構成
TPU は、次の構成で使用できます。
- シングル TPU ボード
- TPU Pod スライス
シングル TPU ボード
シングルボード TPU 構成は、4 個の TPU チップ（8 個の Tensor コア）を備えたスタンドアロン ボードで、他の TPU ボードへのネットワーク接続はありません。シングルボード TPU は TPU Pod の構成の一部ではなく、TPU Pod の一部分を占有することもありません。
TPU Pod スライス
TPU Pod スライスでは、TPU チップが高速相互接続を使用して接続されます。 各 TPU チップは、TPU デバイス上の他のチップと直接通信します。TPU ソフトウェアは、Pod スライスの各 TensorCore へのデータの分散を自動的に処理します。TPU v3 Pod スライスは、32、128、512、1024、2048 の TensorCore で使用できます。TPU v2 Pod スライスは、32、128、256、512 の TensorCore で使用できます。
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 ホストで直接実行できます。
Cloud TPU で TensorFlow を実行する方法の詳細については、TensorFlow のクイックスタートをご覧ください。
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 に似た形で動作します。
