サポートされている TPU 構成
モデルのトレーニングに最適な TPU 構成を決定するには、システム アーキテクチャのドキュメントの TPU ハードウェア バージョンのセクションをご覧ください。Cloud TPU でサポートされている参照モデルの一覧については、サポートされている参照モデルページをご覧ください。
このページでは、TPU タイプと TPU ランタイム バージョンのリストを示します。TensorFlow フレームワークでプログラミングするとき、使用するランタイムは TPU VM アーキテクチャと TPU ノード アーキテクチャのどちらを使用するかによって異なります。2 つの VM アーキテクチャの違いについては、システム アーキテクチャのページをご覧ください。
次の表に、サポートされている TPU ハードウェア タイプを示します。
TPU v4 の構成
TPU v4 構成は、64 チップより小さいトポロジ(小規模トポロジ)と 64 チップを超えるトポロジ(大規模トポロジ)の 2 つのグループで構成されています。各トポロジでサポートされている構成については、次のセクションで説明します。
小規模な v4 トポロジ
Cloud TPU は、次の 64 チップよりも小さい TPU v4 スライス(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 スライスは 64 チップ単位で利用でき、3 つのディメンションすべてで 4 の倍数の形状になっています。ディメンションも昇順でなければなりません。次の表にいくつかの例を示します。これらのトポロジの一部は、トポロジ API を使用してのみ起動できる「カスタム」トポロジです。これは、一般的に使用される名前付きトポロジと同じチップ数を持つためです。
名前(TensorCore のカウントに基づく) | チップ数 | トポロジ |
v4-128 | 64 | 4x4x4 |
v4-256 | 128 | 4x4x8 |
v4-512 | 256 | 4x8x8 |
Topology API をご覧ください | 256 | 4x4x16 |
v4-1024 | 512 | 8x8x8 |
v4-1536 | 768 | 8x8x12 |
v4-2048 | 1024 | 8x8x16 |
Topology API をご覧ください | 1024 | 4x16x16 |
v4-4096 | 2048 | 8x16x16 |
… | … | … |
トポロジ API
カスタム トポロジを使用して Cloud TPU Pod スライスを作成するには、gcloud
TPU API を次のように使用します。
$ gcloud alpha compute tpus tpu-vm create tpu-name \
--zone=us-central2-b \
--subnetwork=tpusubnet \
--type=v4 \
--topology=4x4x16 \
--version=runtime-version
TPU v2 と v3 の構成
次の表に、サポートされている TPU v2 と v3 のタイプを示します。
TPU バージョン | サポートの終了 |
---|---|
v2-8 | (終了日未定) |
v2-32 | (終了日未定) |
v2-128 | (終了日未定) |
v2-256 | (終了日未定) |
v2-512 | (終了日未定) |
v3-8 | (終了日未定) |
v3-32 | (終了日未定) |
v3-128 | (終了日未定) |
v3-256 | (終了日未定) |
v3-512 | (終了日未定) |
v3-1024 | (終了日未定) |
v3-2048 | (終了日未定) |
TPU ソフトウェア バージョン
使用する TPU ソフトウェアのバージョンは、TPU アーキテクチャ、TPU VM または TPU ノード、使用している ML フレームワーク(TensorFlow、PyTorch、JAX)によって異なります。
TPU VM
TPU VM を作成すると、TPU VM に最新バージョンの TensorFlow がプリインストールされます。
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TensorFlow 2.11.0 を使用している場合は、tpu-vm-tf-2.11.0
TPU ソフトウェア バージョンを使用します。TPU Pod を使用している場合は、tpu-vm-tf-2.11.0-pod
を使用します。TPU で現在サポートされている TensorFlow TPU VM ソフトウェア バージョンは次のとおりです。
- 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
- tpu-vm-tf-2.6.5
TPU VM Pod を指定するには、使用する TPU ソフトウェア バージョンに -pod
を追加します(例: tpu-vm-tf-2.11.0-pod
)。
TPU v4 を搭載した TPU VM
TensorFlow を使用して TPU VM v4 でモデルをトレーニングする場合は、以下に示す v4
のバージョンのいずれかを使用します。
2.10.0 - tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4 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 2.10.0 - tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4
TensorFlow 2.10.1 以降は、イメージは 1 つのみであるため、-v4
をバージョン番号で指定する必要はありません。
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU VM は、TensorFlow を使用して作成され、対応する Libtpu ライブラリがプリインストールされます。独自の VM イメージを作成する場合は、次の TensorFlow TPU ソフトウェア バージョンと対応する Libtpu バージョンを指定します。
TensorFlow のバージョン | libtpu.so のバージョン |
---|---|
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 を使用している場合、v2 と v3 の場合は tpu-vm-pt-1.13
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-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-1.13 (pytorch-1.13)
TPU VM を作成すると、TPU VM に最新バージョンの PyTorch がプリインストールされます。PyTorch をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
現在の PyTorch ソフトウェア バージョンを変更するには、PyTorch バージョンの変更をご覧ください。
JAX
JAX 固有の TPU ソフトウェア バージョンがないため、TPU VM に JAX を手動でインストールする必要があります。v2 と v3 の構成では、tpu-vm-base
TPU ソフトウェア バージョンを使用します。v4 構成の場合は、tpu-vm-v4-base
を使用します。JAX をインストールすると、libtpu.so の正しいバージョンが自動的にインストールされます。
TPU ノード
TensorFlow
モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TF 2.11.0 を使用している場合は、2.11.0
TPU ソフトウェア バージョンを使用します。TensorFlow 固有の TPU ソフトウェア バージョンは次のとおりです。
- 2.11.0
- 2.10.1
- 2.10.0
- 2.9.3
- 2.9.1
- 2.8.4
- 2.8.2
- 2.7.3
- 2.6.5
- 2.5.3
- 2.4.3
- 2.3.4
- 2.2.3
- 2.1.4
TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU ノードを作成すると、最新バージョンの TensorFlow が TPU ノードにプリインストールされます。
PyTorch
モデルの書き込まれた PyTorch のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、PyTorch 1.9 を使用している場合は、pytorch-1.9
ソフトウェア バージョンを使用します。
PyTorch 固有の TPU ソフトウェア バージョンは次のとおりです。
- 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 に含まれるソフトウェア バージョンはすべて使用できます。