Supported TPU versions

To determine the most effective TPU hardware version for training your model, see the TPU hardware versions documentation on the System Architecture page. See the supported models page for a list of reference models supported by Cloud TPU.

The following table shows the supported TPU versions.

TPU version Support started Support ends
v2-8 February 12, 2018 (End date not yet set)
v2-32 May 7, 2019 (End date not yet set)
v2-128 May 7, 2019 (End date not yet set)
v2-256 May 7, 2019 (End date not yet set)
v2-512 May 7, 2019 (End date not yet set)
v3-8 October 10, 2018 (End date not yet set)
v3-32 May 7, 2019 (End date not yet set)
v3-128 May 7, 2019 (End date not yet set)
v3-256 May 7, 2019 (End date not yet set)
v3-512 May 7, 2019 (End date not yet set)
v3-1024 May 7, 2019 (End date not yet set)
v3-2048 May 7, 2019 (End date not yet set)

TPU software versions

The version of TPU software you should use depends upon the TPU architecture and the ML framework you are using.

TPU Node

TensorFlow

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

  • 1.15.3
  • 1.15.4
  • 1.15.5
  • 2.1
  • 2.1.1
  • 2.1.2
  • 2.1.3
  • 2.1.4
  • 2.2
  • 2.2.1
  • 2.2.2
  • 2.2.3
  • 2.3
  • 2.3.1
  • 2.3.2
  • 2.3.3
  • 2.3.4
  • 2.4.0
  • 2.4.1
  • 2.4.2
  • 2.4.3
  • 2.5.0
  • 2.5.1
  • 2.6.0
  • nightly

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

The PyTorch specific TPU software versions are:

  • pytorch-1.6
  • pytorch-1.7
  • pytorch-1.8
  • pytorch-1.9
  • pytorch-nightly

JAX

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

TPU VM

TensorFlow

Use the TPU software version that matches the version of TensorFlow with which your model was written. For example, if you are using TF 2.6.0, use the tpu-vm-tf-2.6.0 TPU software version. If you are using a TPU Pod, use tpu-vm-tf-2.6.0-pod.

PyTorch

The v2-alpha TPU version has PyTorch 1.8 preinstalled. If you want to use PyTorch 1.9 or later, see Changing PyTorch version.

JAX

You must manually install JAX on your TPU VM, so there is no JAX-specific TPU software version. Use the v2-alpha TPU version.