TPU VM イメージ

TPU リソースを作成するときに、TPU VM イメージを指定する --version パラメータまたは --runtime-version パラメータを渡します。TPU VM イメージには、オペレーティング システム（Ubuntu）と、TPU でコードを実行するために必要なその他のソフトウェアが含まれています。このドキュメントでは、Cloud TPU の作成時に適切な TPU VM イメージの選択に関するガイダンスを示します。

PyTorch と JAX

PyTorch と JAX に次の一般的な TPU VM ベースイメージを使用し、使用するフレームワークをインストールします。

  • tpu-ubuntu2204-base（デフォルト）
  • v2-alpha-tpuv5（TPU v5p）

インストール手順については、PyTorch/XLA と JAX のクイックスタート ドキュメントをご覧ください。

TensorFlow

TensorFlow の各バージョンに固有の TPU VM イメージがあります。Cloud TPU では、次の TensorFlow バージョンがサポートされています。

  • 2.16.1
  • 2.15.1
  • 2.15.0
  • 2.14.1
  • 2.14.0
  • 2.13.1
  • 2.13.0
  • 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.3
  • 2.8.0
  • 2.7.4
  • 2.7.3

TensorFlow パッチ バージョンの詳細については、サポートされている TensorFlow パッチ バージョンをご覧ください。

TensorFlow バージョン 2.15.0 以降には、使用しているデバイス API（PJRT またはストリーム エグゼキュータ）に基づく TPU VM イメージのバリアントがあります。

v5p と v5e でのトレーニング

TPU v5e と v5p は TensorFlow 2.15.0 以降をサポートします。TPU VM イメージは、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 VM イメージを使用します。TensorFlow の他のバージョンの場合は、2.16.0 を、使用している TensorFlow のメジャー バージョンとパッチ バージョンに置き換えます。単一ホストの TPU を使用している場合は、pod を省略します。

v5e で提供する

TensorFlow、PyTorch、JAX で処理を行うために必要なすべてのソフトウェア要件を含む Docker イメージが用意されています。詳細については、Cloud TPU v5e 推論の概要をご覧ください。

TPU v4

TPU v4 と TensorFlow 2.15.0 以降を使用している場合は、v5p と v5e でのトレーニングの手順に従ってください。TensorFlow 2.10.0 以前を使用している場合は、v4 固有の TPU VM イメージを使用します。

TensorFlow のバージョン TPU VM イメージのバージョン
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 および v3

TPU v2 または v3 を使用している場合は、使用している TensorFlow のバージョンと一致する TPU VM イメージを使用します。たとえば、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 イメージ tpu-vm-tf-2.16.1-pjrt を使用します。同じバージョンの TensorFlow でストリーム エグゼキュータ API を使用している場合は、tpu-vm-tf-2.16.1-se TPU イメージを使用します。バージョン 2.15.0 より前の TensorFlow は、ストリーム エグゼキュータのみをサポートしています。

TensorFlow PJRT のサポート

TensorFlow 2.15.0 以降では、TPU で TensorFlow に PJRT インターフェースを使用できます。PJRT は自動デバイス メモリ デフラグメンテーション機能を備えており、ハードウェアとフレームワークのインテグレーションを簡素化します。PJRT の詳細については、PJRT: ML ハードウェアとフレームワークの統合の簡素化をご覧ください。

アクセラレータ 特徴 PJRT のサポート ストリーム エグゼキュータのサポート
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 なし ×
TPU v5p 高密度なコンピューティング（TPU 埋め込み API なし） ○ ×
TPU v5p TPU 埋め込み API ○ ×

Libtpu のバージョン

TPU VM の TensorFlow イメージには、特定の TensorFlow バージョンと、対応する libtpu ライブラリが含まれています。独自の VM イメージを作成する場合は、次の TensorFlow TPU ソフトウェア バージョンと対応する Libtpu バージョンを使用します。

TensorFlow のバージョン libtpu.so のバージョン
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

次のステップ