TPU VM イメージ

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

PyTorch と JAX

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

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

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

TensorFlow

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

  • 2.17.0
  • 2.16.2
  • 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.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

次のステップ