Versões do software de TPU

Este documento oferece orientações sobre como selecionar uma versão adequada de software de TPU ao criar Cloud TPUs.

Ao criar recursos de TPU, especifique 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 de 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 CLI do Google Cloud, especifique a versão do software de TPU com o parâmetro --version ou --runtime-version. Se você estiver usando o console do Google Cloud , 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 PyTorch e JAX e instale o framework que você quer usar.

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

Para saber como instalar e começar a usar o PyTorch ou o JAX, consulte Executar um cálculo em uma VM do Cloud TPU usando o PyTorch e Executar um cálculo em uma VM do Cloud TPU usando o JAX.

TensorFlow

As versões do software de TPU para 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 principais, secundárias 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, como -pjrt se você estiver usando a API PJRT.

Consulte as seções a seguir para saber como especificar uma versão do software de TPU.

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

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 disponíveis.

TPU v6e, v5p e v5e

As TPUs v6e, v5e e v5p aceitam o TensorFlow 2.15.0 e versões mais recentes. Você especifica a versão do software de TPU usando o formato tpu-vm-tf-x.y.z-{pod}-pjrt, em que x é a versão principal, y é a versão secundária e z é a versão de patch do TensorFlow. Adicione pod depois da versão do TensorFlow ao usar 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 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 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 para a v4:

Versão do TensorFlow Versão do software de 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 em uso. 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 "pod" ao final da versão do software de TPU, por exemplo, tpu-vm-tf-2.14.1-pod.

A partir do TensorFlow 2.15.0, também é necessário especificar uma API de 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 Stream Executor com a mesma versão do TensorFlow, use a versão do software de TPU tpu-vm-tf-2.16.1-se. As versões do TensorFlow anteriores à 2.15.0 só aceitam o Stream Executor.

Disponibilidade do PJRT no TensorFlow

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

Acelerador Recurso Disponibilidade para o PJRT Disponibilidade para o Stream Executor
TPU v2 a v4 Computação densa (sem API de embedding de TPU) Sim Sim
TPU v2 a v4 API de computação densa + API de embedding de TPU Não Sim
TPU v2 a v4 tf.summary/tf.print com posicionamento flexível de dispositivo Não Sim
TPU v5e Computação densa (sem API de embedding de TPU) Sim Não
TPU v5e API de embedding de TPU N/D Não
TPU v5p Computação densa (sem API de embedding de TPU) Sim Não
TPU v5p API de embedding de TPU Sim Não

A seguir