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.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.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

Use the TPU software version that matches the version of PyTorch with which your model was written. For example, if you are using PyTorch 1.13 and TPU v2 or v3, use the tpu-vm-pt-1.13 TPU software version. If you are using TPU v4 use the tpu-vm-v4-pt-1.13 TPU software version. The same TPU software version is used for TPU Pods (for example,v2-32, v3-128, v4-32). The current supported TPU software versions are:

TPU v2/v3:

  • tpu-vm-pt-2.0 (pytorch-2.0)
  • tpu-vm-pt-1.13 (pytorch-1.13)
  • tpu-vm-pt-1.12 (pytorch-1.12)
  • tpu-vm-pt-1.11 (pytorch-1.11)
  • tpu-vm-pt-1.10 (pytorch-1.10)
  • v2-alpha (pytorch-1.8.1)

TPU v4:

  • tpu-vm-v4-pt-2.0 (pytorch-2.0)
  • tpu-vm-v4-pt-1.13 (pytorch-1.13)

TPU v5 (v5e and v5p):

  • v2-alpha-tpuv5 (pytorch-2.0)

When you create a TPU VM, the latest version of PyTorch is preinstalled on the TPU VM. The correct version of libtpu.so is automatically installed when you install PyTorch.

To change the current PyTorch software version, see Changing PyTorch version.

JAX

You must manually install JAX on your TPU VM. There is no JAX-specific TPU software (runtime) version for TPU v2 and v3. For later TPU versions, use the following software versions:

  • TPU v4: tpu-vm-v4-base
  • TPU v5e: v2-alpha-tpuv5
  • TPU v5p: v2-alpha-tpuv5

The correct version of libtpu.so is automatically installed when you install JAX.

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.

PyTorch

Use the TPU software version that matches the version of PyTorch with which your model was written. For example, if you are using PyTorch 1.9, use the pytorch-1.9 software version.

The PyTorch specific TPU software versions are:

  • pytorch-2.0
  • pytorch-1.13
  • pytorch-1.12
  • pytorch-1.11
  • pytorch-1.10
  • pytorch-1.9
  • pytorch-1.8
  • pytorch-1.7
  • pytorch-1.6

  • pytorch-nightly

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

JAX

You must manually install JAX on your TPU VM, so there is no pre-installed JAX-specific TPU software version. You can use any of the software versions listed for TensorFlow.

What's next