TPU v4

このドキュメントでは、Cloud TPU v4 のアーキテクチャとサポートされている構成について説明します。

システムアーキテクチャ

各 TPU v4 チップには、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 TFLOPS（bf16 または int8）
Pod あたりの all-reduce 帯域幅	1.1 PB/s
Pod あたりの二分割帯域幅	24 TB/s

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

TPU v4 チップの図

3D メッシュと 3 次元トーラス

v4 TPU は 3 次元の最近傍チップに直接接続され、ネットワーク接続の 3D メッシュを形成します。接続は、トポロジ AxBxC が 2A=B=C または 2A=2B=C（各次元は 4 の倍数）のスライス上で 3 次元トーラスとして構成されます。たとえば、4x4x8、4x8x8、12x12x24 です。一般に、3 次元トーラス構成のパフォーマンスは 3D メッシュ構成よりも優れています。詳細については、ツイストトリトポロジをご覧ください。

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

このセクションでは、TPU v4 のパフォーマンス上の利点について説明します。

メモリシステム

不均一メモリアクセス（NUMA）は、複数の CPU を搭載したマシンのコンピュータメモリアーキテクチャです。各 CPU は、高速メモリのブロックに直接アクセスできます。CPU とそのメモリは NUMA ノードと呼ばれます。NUMA ノードは、互いに直接隣接している別の NUMA ノードに接続されます。ある NUMA ノードの CPU が別の NUMA ノードのメモリにアクセスすることは可能ですが、こうしたアクセスは NUMA ノード内のメモリにアクセスするよりも時間がかかります。

マルチ CPU マシン上で稼働するソフトウェアは、CPU が必要とするデータをその NUMA ノード内に配置し、メモリのスループットを向上させます。NUMA の詳細については、Wikipedia の不均一メモリアクセスをご覧ください。

トレーニングスクリプトを NUMA ノード 0 にバインドすることで、NUMA-locality の利点を活用できます。

NUMA ノードバインディングを有効にするには:

numactl コマンドラインツールをインストールします。
```
 $ sudo apt-get update
 $ sudo apt-get install numactl
```
スクリプトコードを 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 のパフォーマンスプロファイルを向上。

TensorCores

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

チップ間相互接続

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

その他

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

構成

TPU v4 Pod は、再構成可能な高速リンクで相互接続された 4,096 のチップで構成されます。TPU v4 の柔軟なネットワーキングでは、さまざまな方法で同じサイズの Pod スライスにチップを接続できます。TPU Pod スライスを作成する場合は、TPU のバージョンと必要な TPU リソースの数を指定します。TPU v4 Pod スライスを作成する際は、AcceleratorType または AccleratorConfig のいずれかでタイプとサイズを指定できます。

`AcceleratorType` の使用

トポロジを指定していない場合は、AcceleratorType を使用します。AcceleratorType を使用して v4 TPU を構成するには、TPU Pod スライスの作成時に --accelerator-type フラグを使用します。--accelerator-type を、TPU のバージョンと使用する TensorCore の数を含む文字列に設定します。たとえば、32 個の TensorCore を使用して v4 Pod スライスを作成するには、--accelerator-type=v4-32 を使用します。

次のコマンドは、--accelerator-type フラグを使用して 512 個の TensorCore を持つ v4 TPU Pod スライスを作成します。

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

TPU バージョン（v4）の後の数字は、TensorCore の数を指定しています。v4 TPU には 2 つの TensorCore があるため、TPU チップの数は 512/2 = 256 になります。

`AcceleratorConfig` の使用

TPU スライスの物理トポロジをカスタマイズする場合は、AcceleratorConfig を使用します。これは、通常、256 チップを超える Pod スライスのパフォーマンスチューニングに必要です。

AcceleratorConfig を使用して v4 TPU を構成するには、--version フラグと --topology フラグを使用します。--version を、使用する TPU バージョンに設定し、--topology を Pod スライスの TPU チップの物理配置に設定します。

3 タプル AxBxC を使用して TPU トポロジを指定します。ここで、A<=B<=C で、A、B、C はすべて 4 以下または 4 の整数倍です。値 A、B、C は、3 つのディメンションのチップ数です。たとえば、16 チップの v4 Pod スライスを作成するには、--version=v4 と --topology=2x2x4 を設定します。

次のコマンドは、128 個の TPU チップを 4x4x8 のアレイに配置した v4 TPU Pod スライスを作成します。

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

2A=B=C または 2A=2B=C のトポロジでも、すべての通信用に最適化されたトポロジバリアントがあります（例: 4×4×8、8×8x16、12×12×24）。これらをツイストトリトポロジといいます。

次の図は、一般的な TPU v4 トポロジを示しています。

一般的な TPU v4 構成のトポロジ

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

ツイストトリトポロジ

一部の v4 3 次元トーラススライスの形状には、ツイストトーラストポロジとして知られるオプションを使用できます。たとえば、2 つの v4 キューブは 4x4x8 スライスまたは 4x4x8_ツイストとして配置できます。ツイストトポロジでは、非常に大きな二分帯域幅が提供されます。二分帯域幅が増加すると、グローバル通信パターンを使用するワークロードに役立ちます。ツイストトポロジでは、ほとんどのモデルでパフォーマンスが向上しますが、最もメリットがあるのは大規模な TPU 埋め込みワークロードです。

データ並列処理を唯一の並列処理戦略として使用するワークロードの場合、ツイストトポロジのほうがパフォーマンスがわずかに向上する可能性があります。LLM の場合、ツイストトポロジを使用したパフォーマンスは、並列処理の種類（DP、MP など）によって異なります。モデルに最適なパフォーマンスが得られるかどうか判断するために、ツイストトポロジの有無にかかわらず LLM をトレーニングすることをおすすめします。FSDP MaxText モデルの一部の実験では、ツイストトポロジを使用して MFU が 1～2 回改善されました。

ツイストトポロジの主なメリットは、非対称トーラストポロジ（4×4×8 など）を密接に関連する対称トポロジに変換することです。対称トポロジには次のような多くの利点があります。

ロードバランシングの改善
二分帯域幅の向上
パケットルートの短縮

これらのメリットは、最終的に多くのグローバル通信パターンのパフォーマンス向上につながります。

TPU ソフトウェアは、各ディメンションのサイズが最小のディメンションのサイズ（例: 4x4x8、4×8×8、12x12x24）に等しいか、2 倍であるスライスのツイストトリをサポートします。

たとえば、スライス内に（X、Y）座標でラベル付けされた TPU がある次の 4×2 トーラストポロジについて考えてみます。

このトポロジグラフのエッジは、わかりやすくするために無向エッジとして示されています。実際には、各エッジは TPU 間の双方向接続です。図に示すように、このグリッドの一方の側と反対側の間にあるエッジをラップアラウンドエッジと呼びます。

図形描画

このトポロジをツイストすると、完全に対称な 4×2 ツイストトーラストポロジになります。

図形描画

この図と前の図で異なる点は、Y ラップアラウンドエッジです。同じ X 座標の別の TPU に接続するのではなく、座標 X+2 mod 4 の TPU に接続するようシフトされています。

同じ考え方を、さまざまなディメンションサイズとさまざまな数のディメンションに一般化できます。結果として得られるネットワークは、各ディメンションが最小ディメンションのサイズと等しいか、2 倍である限り、対称になります。

Cloud TPU の作成時にツイストトリ構成を指定する方法の詳細については、AcceleratorConfig の使用をご覧ください。

次の表は、サポートされているツイストトポロジと、ツイストされていないトポロジを使用した場合の二分帯域幅の理論的な増加を示しています。

チップトポロジ	二分帯域幅の理論上の増加とツイストされていないトポロジ
4×4×8_ツイスト	最大 70%
8x8x16_ツイスト
12×12×24_ツイスト
4×8×8_ツイスト	～40%
8×16×16_ツイスト

TPU v4 トポロジバリアント

同じ数のチップを含む一部のトポロジは、さまざまな方法で配置できます。たとえば、512 チップ（1,024 Tensor コア）の TPU Pod スライスは、4x4x32、4x8x16、8x8x8 のトポロジを使用して構成できます。2,048 チップ（4,096 TensorCores）の TPU Pod スライスは、さらに多くの 4x4x128、4x8x64、4x16x32、8x16x16 のトポロジオプションを提供します。

特定のチップ数に関連付けられるデフォルトのトポロジは、キューブに最も近いものです。この形状は、データの並列 ML トレーニングに最適です。他のトポロジ（モデルとデータの並列処理、シミュレーションの空間パーティショニングなど）は、複数の種類の並列処理のワークロードで役立ちます。これらのワークロードは、トポロジが使用された並列処理に一致する場合に最も効果的です。たとえば、X 次元に 4 方向モデル並列処理、Y ディメンションと Z ディメンションに 256 方向データ並列処理を配置すると、4x16x16 トポロジと一致します。

並列処理のディメンションが複数あるモデルは、並列処理ディメンションが TPU トポロジディメンションにマッピングされていると、最も効果的です。これらは通常、データとモデル並列大規模言語モデル（LLM）です。たとえば、トポロジ 8x16x16 の TPU v4 Pod スライスで、TPU トポロジディメンションは 8、16、16 です。8 方向または 16 方向のモデル並列処理（物理 TPU トポロジディメンションの 1 つにマッピング）を使用したほうがパフォーマンスが向上します。4 方向モデル並列処理は、TPU トポロジのディメンションと整合していないため、このトポロジには最適ではありませんが、同じ数のチップで 4x16x32 のトポロジを使用すると最適です。

TPU v4 構成は、64 チップより小さいトポロジ（小規模トポロジ）と 64 チップを超えるトポロジ（大規模トポロジ）の 2 つのグループで構成されています。

小規模な v4 トポロジ

Cloud TPU は、次の 64 チップよりも小さい TPU v4 Pod スライス（4x4x4 キューブ）をサポートしています。これらの小規模な v4 トポロジは、TensorCore ベースの名前（v4-32 など）またはトポロジ（2x2x4 など）を使用して作成できます。

名前（TensorCore のカウントに基づく）	チップ数	トポロジ
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

大規模な v4 トポロジ

TPU v4 Pod スライスは 64 チップ単位で利用でき、3 つのディメンションすべてで 4 の倍数の形状になっています。ディメンションは昇順にする必要があります。次の表に、いくつかの例を示します。これらのトポロジの一部は、--type フラグと --topology フラグを使用してのみ作成できる「カスタム」トポロジです。これは、チップを配置する方法は複数あるためです。

名前（TensorCore のカウントに基づく）	チップ数	トポロジ
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
カスタムトポロジ: `--type` フラグと `--topology` フラグを使用する必要があります。	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
カスタムトポロジ: `--type` フラグと `--topology` フラグを使用する必要があります。	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…