TPU runtimes

This document provides guidance on selecting the appropriate runtime version for your Cloud TPU, depending on the framework and TPU architecture (VM or Node) you use.

TPU VM runtime versions

This section describes the TPU software versions you should use for a TPU with the TPU VM architecture. For the TPU Node architecture, see TPU Node software versions.

TPU software versions are available for TensorFlow, PyTorch, and JAX frameworks.

TensorFlow

Use the TPU software version that matches the version of TensorFlow with which your model was written.

Beginning with TensorFlow 2.15.0, you must also specify either the stream executor (SE) runtime or the PJRT runtime. For example, if you are using TensorFlow 2.16.1 with the PJRT runtime, use the tpu-vm-tf-2.16.1-pjrt TPU software version. Versions prior to TensorFlow 2.15.0 only support stream executor. For more information about PJRT, see TensorFlow PJRT support.

The current supported TensorFlow TPU VM software versions are:

  • 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

For more information on TensorFlow patch versions, see Supported TensorFlow patch versions.

TensorFlow PJRT support

Beginning with TensorFlow 2.15.0, you can use the PJRT interface for TensorFlow on TPU. PJRT features automatic device memory defragmentation and simplifies the integration of hardware with frameworks. For more information about PJRT, see PJRT: Simplifying ML Hardware and Framework Integration on the Google Open Source Blog.

Not all features of TPU v2, v3, and v4 have been migrated to the PJRT runtime. The following table describes which features are supported on PJRT or steam executor.

Accelerator Feature Supported on PJRT Supported on stream executor
TPU v2-v4 Dense compute (no TPU embedding API) Yes Yes
TPU v2-v4 Dense compute API + TPU embedding API No Yes
TPU v2-v4 tf.summary/tf.print with soft device placement No Yes
TPU v5e Dense compute (no TPU embedding API) Yes No
TPU v5e TPU embedding API N/A - TPU v5e doesn't support TPU embedding API N/A
TPU v5p Dense compute (no TPU embedding API) Yes No
TPU v5p TPU embedding API Yes No

TPU v4 with TensorFlow versions 2.10.0 and earlier

If you are training a model on TPU v4 with TensorFlow, TensorFlow versions 2.10.0 and earlier use v4-specific versions shown in the following table. If the TensorFlow version you're using is not shown in the table, follow the guidance in the TensorFlow section.

TensorFlow version TPU software version
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 versions

TPU VMs are created with TensorFlow and the corresponding Libtpu library preinstalled. If you are creating your own VM image, specify the following TensorFlow TPU software versions and corresponding libtpu versions:

TensorFlow version libtpu.so version
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 and JAX

Use the following common base images for PyTorch and JAX, then install the framework version you need with pip.

  • tpu-ubuntu2204-base (default)
  • v2-alpha-tpuv5 (TPU v5p)

Refer to the quickstart documents for PyTorch/XLA and JAX for installation instructions.

TPU Node runtime versions

This section describes the TPU software versions you should use for a TPU with the TPU Node architecture. For the TPU VM architecture, see TPU VM software versions.

TPU software versions are available for TensorFlow, PyTorch, and JAX frameworks.

TensorFlow

Use the TPU software version that matches the version of TensorFlow with which your model was written. For example, if you are using TensorFlow 2.12.0, use the 2.12.0 TPU software version. The TensorFlow specific TPU software versions are:

  • 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

For more information on TensorFlow patch versions, see Supported TensorFlow patch versions.

When you create a TPU Node, the latest version of TensorFlow is preinstalled on the TPU Node.

What's next