Versões do software de TPU

Este documento fornece orientações sobre como selecionar a versão do software da TPU adequada ao criar Cloud TPUs.

Ao criar recursos de TPU, você especifica a versão do software, também chamada de versão do ambiente de execução, que se refere ao ambiente de software pré-instalado na VM da TPU. Isso inclui o sistema operacional Ubuntu, o Docker e outros softwares necessários para executar o código em TPUs.

Se você estiver usando a Google Cloud CLI, especifique a versão do software TPU usando o parâmetro --version ou --runtime-version. Se você estiver usando o console Google Cloud , selecione uma versão do software da TPU na lista Versão do software da TPU.

PyTorch e JAX

Use as versões comuns de software de TPU para PyTorch e JAX abaixo e instale o framework que você quer usar.

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

Para mais informações sobre a instalação e como começar a usar o PyTorch ou o JAX, consulte Executar um cálculo em uma VM da Cloud TPU usando o PyTorch e Executar um cálculo em uma VM da Cloud TPU usando o JAX.

TensorFlow

As versões do software de 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 você está usando uma fração de TPU de vários hosts.
  • -{device_api} (opcional): especifica a API do dispositivo, por exemplo, -pjrt (se você estiver usando a API PJRT).

Consulte as seções a seguir para mais informações sobre como especificar uma versão de software da TPU.

Há versões do software de TPU específicas para cada versão do TensorFlow. A tabela a seguir mostra as versões do TensorFlow e as versões do libtpu associadas:

Versão do TensorFlow Versão do 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 Versões de patch do TensorFlow compatíveis.

TPU v6e, v5p e v5e

A TPU v6e, v5e e v5p são compatíveis com o TensorFlow 2.15.0 e versões mais recentes. Especifique a versão do software da TPU usando o 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 do patch do TensorFlow. Adicione pod após a versão do TensorFlow se você estiver usando uma TPU de vários hosts. Por exemplo, se você estiver usando o TensorFlow 2.16.0 em uma TPU de vários hosts, use a versão do software tpu-vm-tf-2.16.0-pod-pjrt da TPU. Para outras versões do TensorFlow, substitua 2.16.0 pelas versões principais e de patch do TensorFlow que você está usando. Se você estiver usando uma TPU de host único, omita pod.

TPU v4

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

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

TPU v2 e v3

Se você estiver usando a TPU v2 ou v3, use a versão do software de TPU que corresponde à versão do TensorFlow que você está usando. Por exemplo, se você estiver usando o TensorFlow 2.14.1, use a versão do software de TPU tpu-vm-tf-2.14.1. Para outras versões do TensorFlow, substitua 2.14.1 pela versão do TensorFlow que você está usando. Se você estiver usando uma TPU de vários hosts, adicione o pod ao final da versão do software da TPU, por exemplo, tpu-vm-tf-2.14.1-pod.

A partir do TensorFlow 2.15.0, também é necessário especificar uma API do dispositivo como parte do nome da versão do software. Por exemplo, se você estiver usando o TensorFlow 2.16.1 com a API PJRT, use a versão do software de TPU tpu-vm-tf-2.16.1-pjrt. Se você estiver usando a API executor de fluxo com a mesma versão do TensorFlow, use a versão do software TPU tpu-vm-tf-2.16.1-se. As versões do TensorFlow anteriores à 2.15.0 só oferecem suporte ao executor de fluxo.

Suporte do TensorFlow para PJRT

A partir do TensorFlow 2.15.0, é possível usar a interface PJRT para TensorFlow em TPU. O PJRT tem desfragmentação automática da memória do dispositivo e simplifica a integração de hardware com frameworks. Para mais informações sobre o PJRT, consulte PJRT: simplificando a integração de hardware e framework de ML.

Accelerator Recurso Suporte a PJRT Suporte a executor de stream
TPU v2 a v4 Computação densa (sem API de incorporação de TPU) Sim Sim
TPU v2 a v4 API Dense Compute + API de incorporação de TPU Não Sim
TPU v2 a v4 tf.summary/tf.print com posicionamento flexível do dispositivo Não Sim
TPU v5e Computação densa (sem API de incorporação de TPU) Sim Não
TPU v5e API de incorporação de TPU N/A Não
TPU v5p Computação densa (sem API de incorporação de TPU) Sim Não
TPU v5p API de incorporação de TPU Sim Não

A seguir