TPU ソフトウェア バージョン
このドキュメントでは、Cloud TPU を作成するときに適切な TPU ソフトウェア バージョンを選択する方法について説明します。
TPU リソースの作成時に、ソフトウェア バージョン（ランタイム バージョンとも呼ばれます）を指定します。これは、TPU VM にプリインストールされているソフトウェア環境を指します。これには、Ubuntu オペレーティング システム、Docker、TPU でコードを実行するために必要なその他のソフトウェアが含まれます。
Google Cloud CLI を使用している場合は、--version または --runtime-version パラメータを使用して TPU ソフトウェア バージョンを指定します。 Google Cloud コンソールを使用している場合は、[TPU ソフトウェア バージョン] リストから TPU ソフトウェア バージョンを選択します。
PyTorch と JAX
PyTorch と JAX で一般的な次の TPU ソフトウェア バージョンを使用し、使用するフレームワークをインストールします。
| TPU バージョン | TPU ソフトウェア バージョン |
|---|---|
| Trillium（v6e） | v2-alpha-tpuv6e |
| v5p | v2-alpha-tpuv5 |
| v5e | v2-alpha-tpuv5-lite |
| v4 以前 | tpu-ubuntu2204-base |
PyTorch または JAX のインストールと使用開始の詳細については、PyTorch を使用して Cloud TPU VM で計算を実行すると JAX を使用して Cloud TPU VM で計算を実行するをご覧ください。
TensorFlow
TensorFlow の TPU ソフトウェア バージョンは、特定の命名規則に従います。
tpu-vm-tf-x.y.z[-{pod}][-{device_api}]
x.y.z: TensorFlow のメジャー バージョン、マイナー バージョン、パッチ バージョンを表します。-pod（省略可）: マルチホスト TPU スライスを使用していることを示します。-{device_api}（省略可）: デバイス API を指定します。例: PJRT API を使用している場合は-pjrt。
TPU ソフトウェア バージョンを指定する方法の詳細については、以降のセクションをご覧ください。
TensorFlow のバージョンごとに固有の TPU ソフトウェア バージョンがあります。次の表に、サポートされている TensorFlow バージョンと関連する libtpu バージョンを示します。
| TensorFlow のバージョン | libtpu.so のバージョン |
|---|---|
| 2.18.0 | 1.12.0 |
| 2.17.1 | 1.11.1 |
| 2.17.0 | 1.11.0 |
| 2.16.2 | 1.10.1 |
| 2.16.1 | 1.10.1 |
| 2.15.1 | 1.9.0 |
| 2.15.0 | 1.9.0 |
| 2.14.1 | 1.8.1 |
| 2.14.0 | 1.8.0 |
| 2.13.1 | 1.7.1 |
| 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.0 | 1.2.0 |
| 2.7.3 | 1.1.2 |
TensorFlow のパッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。
TPU v6e、v5p、v5e
TPU v6e、v5e、v5p は、TensorFlow 2.15.0 以降をサポートしています。TPU ソフトウェア バージョンは、tpu-vm-tf-x.y.z-{pod}-pjrt の形式で指定します。ここで、x は TensorFlow のメジャー バージョン、y はマイナー バージョン、z は TensorFlow のパッチ バージョンです。マルチホスト TPU を使用している場合は、TensorFlow バージョンの後に pod を追加します。たとえば、マルチホスト TPU で TensorFlow 2.16.0 を使用している場合は、tpu-vm-tf-2.16.0-pod-pjrt TPU ソフトウェア バージョンを使用します。TensorFlow の他のバージョンについては、2.16.0 を使用している TensorFlow のメジャー バージョンとパッチ バージョンに置き換えます。単一ホストの TPU を使用している場合は、pod を省略します。
TPU v4
TPU v4 と TensorFlow 2.10.1 以降を使用している場合は、TPU v2 と v3 の手順に沿って操作します。TensorFlow 2.10.0 以前を使用している場合は、v4 固有の TPU ソフトウェア バージョンを使用します。
| TensorFlow のバージョン | TPU ソフトウェア バージョン |
|---|---|
| 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 |
TPU v2 と TPU v3
TPU v2 または v3 を使用している場合は、使用している TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TensorFlow 2.14.1 を使用している場合は、tpu-vm-tf-2.14.1 TPU ソフトウェア バージョンを使用します。TensorFlow の他のバージョンについては、2.14.1 を使用している TensorFlow バージョンに置き換えます。マルチホスト TPU を使用している場合は、TPU ソフトウェア バージョンの末尾に pod を追加します（例: tpu-vm-tf-2.14.1-pod）。
TensorFlow 2.15.0 以降では、ソフトウェア バージョン名の一部としてデバイス API も指定する必要があります。たとえば、PJRT API で TensorFlow 2.16.1 を使用している場合は、tpu-vm-tf-2.16.1-pjrt TPU ソフトウェア バージョンを使用します。同じバージョンの TensorFlow で stream executor API を使用している場合は、tpu-vm-tf-2.16.1-se TPU ソフトウェア バージョンを使用します。2.15.0 より前の TensorFlow バージョンは、stream executor のみをサポートしています。
TensorFlow PJRT のサポート
TensorFlow 2.15.0 以降では、TPU 上の TensorFlow に PJRT インターフェースを使用できます。PJRT は自動デバイスメモリ デフラグメンテーション機能を備えており、ハードウェアとフレームワークのインテグレーションを簡素化します。PJRT の詳細については、PJRT: ML ハードウェアとフレームワークのインテグレーションの簡素化をご覧ください。
| アクセラレータ | 機能 | PJRT のサポート | stream executor のサポート |
|---|---|---|---|
| TPU v2 - v4 | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | ○ |
| TPU v2 - v4 | 高密度なコンピューティング API + TPU 埋め込み API | × | ○ |
| TPU v2 - v4 | ソフトデバイス配置を使用した tf.summary/tf.print | × | ○ |
| TPU v5e | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | × |
| TPU v5e | TPU 埋め込み API | N/A | × |
| TPU v5p | 高密度なコンピューティング（TPU 埋め込み API なし） | ○ | × |
| TPU v5p | TPU 埋め込み API | ○ | × |
次のステップ
- TPU アーキテクチャを参照して、TPU アーキテクチャについて確認する。
- TPU の用途を参照して、Cloud TPU に最適なモデルの種類について確認する。