Versões do software do TPU

Este documento fornece orientações sobre a seleção da versão de software da TPU adequada quando cria TPUs na nuvem.

Quando cria recursos de TPU, especifica a versão do software, também denominada versão de tempo de execução, que se refere ao ambiente de software pré-instalado na VM da TPU. Isto inclui o sistema operativo Ubuntu, o Docker e outro software necessário para executar o seu código em TPUs.

Se estiver a usar a CLI Google Cloud, especifique a versão do software da TPU através do parâmetro --version ou --runtime-version. Se estiver a usar a Google Cloud consola, selecione uma versão do software de TPU na lista Versão do software de TPU.

PyTorch e JAX

Use as seguintes versões de software de TPU comuns para o PyTorch e o JAX e, em seguida, instale a framework que quer usar.

Versão da TPU Versão do software do TPU
Trillium (v6e) v2-alpha-tpuv6e
v5p v2-alpha-tpuv5
v5e v2-alpha-tpuv5-lite
v4 e mais antigas tpu-ubuntu2204-base

Para mais informações sobre a instalação e o início da utilização do PyTorch ou do JAX, consulte os artigos Execute um cálculo numa VM do Cloud TPU com o PyTorch e Execute um cálculo numa VM do Cloud TPU com o JAX.

TensorFlow

As versões de software da TPU para o TensorFlow seguem uma convenção de nomenclatura específica:

tpu-vm-tf-x.y.z[-{pod}][-{device_api}]

  • x.y.z: representa as versões principal, secundária e de patch do TensorFlow.
  • -pod (opcional): indica que está a usar uma fatia de TPU com vários anfitriões.
  • -{device_api} (opcional): especifica a API do dispositivo, por exemplo, -pjrt (se estiver a usar a API PJRT).

Consulte as secções seguintes para ver mais informações sobre como especificar uma versão de software da TPU.

Existem versões de software de TPU específicas para cada versão do TensorFlow. A tabela seguinte mostra as versões do TensorFlow suportadas e as versões do libtpu associadas:

Versão do TensorFlow Versão de libtpu.so
2.18.0 1.12.0
2.17.1 1.11.1
2.17.0 1.11.0
2.16.2 1.10.1
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.0 1.2.0
2.7.3 1.1.2

Para mais informações sobre as versões de patch do TensorFlow, consulte o artigo Versões de patch do TensorFlow suportadas.

TPU v6e, v5p e v5e

As TPUs v6e, v5e e v5p suportam o TensorFlow 2.15.0 e mais recente. Especifica a versão do software da TPU através do formulário: tpu-vm-tf-x.y.z-{pod}-pjrt, em que x é a versão principal do TensorFlow, y é a versão secundária e z é a versão de patch do TensorFlow. Adicione pod após a versão do TensorFlow se estiver a usar uma TPU multi-host. Por exemplo, se estiver a usar o TensorFlow 2.16.0 numa TPU multi-anfitrião, use a versão do software de TPU tpu-vm-tf-2.16.0-pod-pjrt. Para outras versões do TensorFlow, substitua 2.16.0 pelas versões principal e de patch do TensorFlow que está a usar. Se estiver a usar uma TPU de anfitrião único, omita pod.

TPU v4

Se estiver a usar a TPU v4 e o TensorFlow 2.10.1 ou mais recente, siga as instruções para a TPU v2 e v3. Se estiver a usar o TensorFlow 2.10.0 ou anterior, use uma versão de software de TPU específica da v4:

Versão do TensorFlow Versão do software do TPU
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

TPUs v2 e v3

Se estiver a usar a TPU v2 ou v3, use a versão do software de TPU que corresponda à versão do TensorFlow que está a usar. Por exemplo, se estiver a usar o TensorFlow 2.14.1, use a