TPU の構成
TPU v4 の構成
TPU v4 Pod は、再構成可能な高速リンクで相互接続された 4,096 のチップで構成されます。TPU v4 の柔軟なネットワーキングでは、さまざまな方法で同じサイズの Pod スライスにチップを接続できます。TPU Pod スライスを作成する場合は、TPU のバージョンと必要な TPU リソースの数を指定します。TPU v4 Pod スライスを作成する際は、次の 2 つの方法のいずれかでタイプとサイズを指定できます。
AcceleratorTypeの使用トポロジを指定していない場合、この方法をおすすめします。
AcceleratorConfigの使用物理トポロジをカスタマイズする場合は、
AcceleratorConfigを使用することをおすすめします。これは、通常、Pod スライス数が 256 チップ以上のパフォーマンス チューニングに必要です。
AcceleratorType を使用して v4 TPU を構成するには、TPU Pod スライスの作成時に --accelerator-type フラグを使用します。--accelerator-type を、TPU のバージョンと使用する TensorCore の数を含む文字列に設定します。たとえば、32 個の TensorCore を使用して v4 Pod スライスを作成するには、--accelerator-type=v4-32 を使用します。
AcceleratorConfig を使用して v4 TPU を構成するには、--version フラグと --topology フラグを使用します。--version を、使用する TPU バージョンに設定し、--topology を Pod スライスの TPU チップの物理配置に設定します。たとえば、16 チップの v4 Pod スライスを作成するには、--version=v4 と --topology=2x2x4 を使用します。
3 タプル AxBxC を使用して TPU トポロジを指定します。ここで、A<=B<=C で、A、B、C はすべて 4 以下または 4 の整数倍です。値 A、B、C は、3 つのディメンションのチップ数です。2A=B=C または 2A=2B=C のトポロジでも、すべての通信用に最適化されたトポロジ バリアントがあります（例: 4×4×8、8×8x16、12×12×24）。これらを「ツイスト トリ」トポロジと呼びます。ツイスト トリ構成は、それをサポートできるトポロジに対して指定できます。サポートされているツイスト トポロジのリストについては、バリアント構成をご覧ください。
次のコマンドは、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.13.0
--accelerator-type を使用して、同じ数のチップ/TensorCore を持つ v4 TPU Pod スライスを作成できます。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-512
--version=tpu-vm-tf-2.13.0
TPU バージョン（v4）の後の数字は、TensorCore の数を指定しています。v4 TPU には 2 つの TensorCore があるため、TPU チップの数は 512/2 = 256 になります。
次の図は、一般的な TPU v4 トポロジを示しています。
より大きな Pod スライスを、1 つまたは複数の 4x4x4 の「キューブ」チップから構築できます。
トポロジのバリアント
同じ数のチップを含む一部のトポロジは、さまざまな方法で配置できます。たとえば、512 チップ（1,024 Tensor コア）の TPU Pod スライスは、4x4x32、4x8x16、8x8x8 のトポロジを使用して構成できます。2,048 チップ（4,096 TensorCores）の TPU Pod スライスは、さらに多くの 4x4x128、4x8x64、4x16x32、8x16x16 のトポロジ オプションを提供します。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 フラグを使用してのみ作成できる「カスタム」トポロジです。これは、チップを配置する方法は複数あるためです。
次のコマンドは、512 個の TPU チップを 8x8x8 のアレイに配置した v4 TPU Pod スライスを作成します。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--type=v4
--topology=8x8x8
--version=tpu-vm-tf-2.13.0
--accelerator-type を使用して、同じ数のチップ/TensorCore を持つ v4 TPU Pod スライスを作成できます。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v4-1024
--version=tpu-vm-tf-2.13.0
| 名前（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 |
| … | … | … |
TPU v3 の構成
TPU v3 Pod は、高速リンクで相互接続された 1,024 個のチップで構成されます。TPU v3 デバイスまたは Pod スライスを作成するには、gcloud compute tpus tpu-vm コマンドに --accelerator-type フラグを使用します。アクセラレータ タイプを指定するには、TPU のバージョンと TPU コアの数を指定します。単一の v3 TPU の場合は、--accelerator-type=v3-8 を使用します。128 個の TensorCore を使用する v3 Pod スライスの場合は、--accelerator-type=v3-128 を使用します。
次のコマンドは、128 個の TensorCore を使用して v3 TPU Pod スライスを作成する方法を示しています。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v3-128
--version=tpu-vm-tf-2.13.0
次の表に、サポートされている v3 TPU タイプを示します。
| TPU バージョン | サポートの終了 |
|---|---|
| v3-8 | （終了日未定） |
| v3-32 | （終了日未定） |
| v3-128 | （終了日未定） |
| v3-256 | （終了日未定） |
| v3-512 | （終了日未定） |
| v3-1024 | （終了日未定） |
| v3-2048 | （終了日未定） |
TPU の管理の詳細については、TPU の管理をご覧ください。Cloud TPU のさまざまなバージョンの詳細については、システム アーキテクチャをご覧ください。
TPU v2 の構成
TPU v2 Pod は、再構成可能な高速リンクで相互接続された 512 のチップで構成されます。TPU v2 Pod スライスを作成するには、gcloud compute tpus tpu-vm コマンドに --accelerator-type フラグを使用します。アクセラレータ タイプを指定するには、TPU のバージョンと TPU コアの数を指定します。単一の v2 TPU の場合は、--accelerator-type=v2-8 を使用します。128 個の TensorCore を使用する v2 Pod スライスの場合は、--accelerator-type=v2-128 を使用します。
次のコマンドは、128 個の TensorCore を使用して v2 TPU Pod スライスを作成する方法を示しています。
$ gcloud compute tpus tpu-vm create tpu-name
--zone=zone
--accelerator-type=v2-128
--version=tpu-vm-tf-2.13.0
TPU の管理の詳細については、TPU の管理をご覧ください。Cloud TPU のさまざまなバージョンの詳細については、システム アーキテクチャをご覧ください。
次の表に、サポートされている v2 TPU タイプを示します。
| TPU バージョン | サポートの終了 |
|---|---|
| v2-8 | （終了日未定） |
| v2-32 | （終了日未定） |
| v2-128 | （終了日未定） |
| v2-256 | （終了日未定） |
| v2-512 | （終了日未定） |
TPU タイプの互換性
TPU タイプを、同じ TensorCore 数を持つ別の TPU タイプ（たとえば、v3-128 や v4-128）に変更し、コードを変更せずにトレーニング スクリプトを実行できます。ただし、チップ数または TensorCore 数の増減を実施した TPU タイプに変更する場合は、かなりの調整と最適化を行う必要があります。詳細については、TPU Pod でのトレーニングをご覧ください。
TPU ソフトウェア バージョン
使用する TPU ソフトウェアのバージョンは、TPU アーキテクチャ（TPU VM または TPU ノード）、および使用している ML フレームワーク（TensorFlow、PyTorch、JAX）によって異なります。
TPU VM
TPU ソフトウェア バージョンは、TensorFlow、PyTorch、JAX のフレームワークで使用できます。
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TensorFlow 2.13.0 を使用している場合は、tpu-vm-tf-2.13.0 TPU ソフトウェア バージョンを使用します。TPU で現在サポートされている TensorFlow TPU VM ソフトウェア バージョンは次のとおりです。
- tpu-vm-tf-2.13.0
- tpu-vm-tf-2.12.1
- tpu-vm-tf-2.12.0
- tpu-vm-tf-2.11.1
- tpu-vm-tf-2.11.0
- tpu-vm-tf-2.10.1
- tpu-vm-tf-2.10.0
- tpu-vm-tf-2.9.3
- tpu-vm-tf-2.9.1
- tpu-vm-tf-2.8.4
- tpu-vm-tf-2.8.3
- tpu-vm-tf-2.8.0
- tpu-vm-tf-2.7.4
- tpu-vm-tf-2.7.3
Pod スライスを使用している場合は、使用する TPU ソフトウェア バージョンに -pod を追加します。例: tpu-vm-tf-2.13.0-pod
TPU VM と TPU v4
TensorFlow を使用して TPU v4 でモデルをトレーニングする場合は、次の表に示す v4 バージョンのいずれかを使用します。
| TensorFlow のバージョン | TPU ソフトウェア バージョン |
|---|---|
| 2.13.0 | tpu-vm-tf-2.13.0、tpu-vm-tf-2.13.0-pod |
| 2.12.1 | tpu-vm-tf-2.12.1, tpu-vm-tf-2.12.1-pod |
| 2.12.0 | tpu-vm-tf-2.12.0、tpu-vm-tf-2.12.0-pod |
| 2.11.1 | tpu-vm-tf-2.11.1, tpu-vm-tf-2.11.1-pod |
| 2.11.0 | tpu-vm-tf-2.11.0、tpu-vm-tf-2.11.0-pod |
| 2.10.0 | tpu-vm-tf-2.10.0-v4、tpu-vm-tf-2.10.0-pod-v4 |
| 2.10.1 | tpu-vm-tf-2.10.1, tpu-vm-tf-2.10.1-pod |
| 2.9.3 | tpu-vm-tf-2.9.3-v4、tpu-vm-tf-2.9.3-pod-v4 |
| 2.9.2 | tpu-vm-tf-2.9.2-v4、tpu-vm-tf-2.9.2-pod-v4 |
| 2.9.1 | tpu-vm-tf-2.9.1-v4、tpu-vm-tf-2.9.1-pod-v4 |
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU VM は、TensorFlow を使用して作成され、対応する Libtpu ライブラリがプリインストールされます。独自の VM イメージを作成する場合は、次の TensorFlow TPU ソフトウェア バージョンと対応する Libtpu バージョンを指定します。
| TensorFlow のバージョン | libtpu.so のバージョン |
|---|---|
| 2.13.0 | 1.7.0 |
| 2.12.1 | 1.6.1 |
| 2.12.0 | 1.6.0 |
| 2.11.1 | 1.5.1 |
| 2.11.0 | 1.5.0 |
| 2.10.1 | 1.4.1 |
| 2.10.0 | 1.4.0 |
| 2.9.3 | 1.3.2 |
| 2.9.1 | 1.3.0 |
| 2.8.3 | 1.2.3 |
| 2.8.* | 1.2.0 |
| 2.7.3 | 1.1.2 |
PyTorch
モデルの書き込まれた PyTorch のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、PyTorch 1.13 と TPU v2 または v3 を使用している場合は、tpu-vm-pt-1.13 TPU ソフトウェア バージョンを使用します。TPU v4 を使用している場合は、tpu-vm-v4-pt-1.13 TPU ソフトウェア バージョンを使用します。TPU Pod に同じ TPU ソフトウェア バージョン（例: v2-32、v3-128、v4-32）が使用されます。現在サポートされている TPU ソフトウェア バージョンは次のとおりです。
TPU v2/v3:
- tpu-vm-pt-2.0（pytorch-2.0）
- tpu-vm-pt-1.13（pytorch-1.13）
- tpu-vm-pt-1.12（pytorch-1.12）
- tpu-vm-pt-1.11（pytorch-1.11）
- tpu-vm-pt-1.10（pytorch-1.10）
- v2-alpha（pytorch-1.8.1）
TPU v4:
- tpu-vm-v4-pt-2.0（pytorch-2.0）
- tpu-vm-v4-pt-1.13 (pytorch-1.13)
TPU VM を作成すると、TPU VM に最新バージョンの PyTorch がプリインストールされます。PyTorch をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
現在の PyTorch ソフトウェア バージョンを変更するには、PyTorch バージョンの変更をご覧ください。
JAX
JAX 固有の TPU ソフトウェア バージョンがないため、TPU VM に JAX を手動でインストールする必要があります。TPU バージョンではすべて、tpu-ubuntu2204-base が使用されます。JAX をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
TPU ノード
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TF 2.12.0 を使用している場合は、2.12.0 TPU ソフトウェア バージョンを使用します。TensorFlow 固有の TPU ソフトウェア バージョンは次のとおりです。
- 2.12.1
- 2.12.0
- 2.11.1
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.2
- 2.7.3
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU ノードを作成すると、最新バージョンの TensorFlow が TPU ノードにプリインストールされます。
PyTorch
モデルの書き込まれた PyTorch のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、PyTorch 1.9 を使用している場合は、pytorch-1.9 ソフトウェア バージョンを使用します。
PyTorch 固有の TPU ソフトウェア バージョンは次のとおりです。
- pytorch-2.0
- pytorch-1.13
- pytorch-1.12
- pytorch-1.11
- pytorch-1.10
- pytorch-1.9
- pytorch-1.8
- pytorch-1.7
pytorch-1.6
pytorch-nightly
TPU ノードを作成すると、TPU ノードに最新バージョンの PyTorch がプリインストールされます。
JAX
TPU VM に JAX を手動でインストールする必要があるため、JAX 固有の TPU ソフトウェア バージョンはプリインストールされていません。TensorFlow に含まれるソフトウェア バージョンはすべて使用できます。
次のステップ
- TPU アーキテクチャの詳細については、システム アーキテクチャ ページをご覧ください。
- TPU の用途を参照して、Cloud TPU に最適なモデルの種類について確認する。