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