TPU ランタイム

このドキュメントでは、使用するフレームワークと TPU アーキテクチャ（VM またはノード）に応じて、Cloud TPU に適したランタイム バージョンを選択するためのガイダンスを示します。

TPU VM ランタイム バージョン

このセクションでは、TPU VM アーキテクチャで TPU に対して使用する TPU ソフトウェア バージョンについて説明します。 TPU ノードのアーキテクチャについては、TPU ノードのソフトウェア バージョンをご覧ください。

TPU ソフトウェア バージョンは、TensorFlow、PyTorch、JAX の各フレームワークで使用できます。

TensorFlow

モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。

TensorFlow 2.15.0 以降では、ストリーム エグゼキュータ（SE）ランタイムまたは PJRT ランタイムも指定する必要があります。たとえば、PJRT ランタイムで TensorFlow 2.16.1 を使用している場合は、tpu-vm-tf-2.16.1-pjrt TPU ソフトウェア バージョンを使用します。TensorFlow 2.15.0 より前のバージョンは、ストリーム エグゼキュータのみをサポートしています。PJRT の詳細については、TensorFlow PJRT のサポートをご覧ください。

現在サポートされている TensorFlow TPU VM ソフトウェア バージョンは次のとおりです。

  • tpu-vm-tf-2.16.1-pjrt
  • tpu-vm-tf-2.16.1-se
  • tpu-vm-tf-2.15.1-pjrt
  • tpu-vm-tf-2.15.1-se
  • tpu-vm-tf-2.15.0-pjrt
  • tpu-vm-tf-2.15.0-se
  • tpu-vm-tf-2.14.1
  • tpu-vm-tf-2.14.0
  • tpu-vm-tf-2.13.1
  • 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

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

TensorFlow PJRT のサポート

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

TPU v2、v3、v4 のすべての機能が PJRT ランタイムに移行されたわけではありません。次の表に、PJRT またはストリーム エグゼキュータでサポートされている機能を示します。

アクセラレータ 特徴 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 v5e では TPU 埋め込み API がサポートされません なし
TPU v5p 高密度なコンピューティング（TPU 埋め込み API なし） ○ ×
TPU v5p TPU 埋め込み API ○ ×

TensorFlow バージョン 2.10.0 以前を使用した TPU v4

TensorFlow で TPU v4 を使用してモデルをトレーニングする場合、TensorFlow バージョン 2.10.0 以前では、次の表に示す v4 固有のバージョンを使用します。使用している TensorFlow のバージョンが表に表示されていない場合は、TensorFlow セクションのガイダンスに従ってください。

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

Libtpu のバージョン

TPU VM は、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.* 1.2.0
2.7.3 1.1.2

PyTorch と JAX

PyTorch と JAX 用の次の共通のベースイメージを使用してから、pip を使用して必要なフレームワーク バージョンをインストールします。

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

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

TPU ノード ランタイム バージョン

このセクションでは、TPU ノード アーキテクチャで TPU に対して使用する TPU ソフトウェア バージョンについて説明します。TPU VM アーキテクチャについては、TPU VM ソフトウェア バージョンをご覧ください。

TPU ソフトウェア バージョンは、TensorFlow、PyTorch、JAX の各フレームワークで使用できます。

TensorFlow

モデルの書き込まれた TensorFlow のバージョンと一致する TPU ソフトウェア バージョンを使用します。たとえば、TensorFlow 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 ノードにプリインストールされます。

次のステップ