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
- Learn more about TPU architecture in the System Architecture page.
- See When to use TPUs to learn about the types of models that are well suited to Cloud TPU.