Imagens de VM da TPU

Ao criar recursos de TPU, você transmite o parâmetro --version ou --runtime-version, que especifica uma imagem de VM da TPU. As imagens de VM da TPU contêm o sistema operacional (Ubuntu) e, opcionalmente, outro software necessário para executar o código em TPUs. Neste documento, fornecemos orientações sobre como selecionar a imagem de VM de TPU apropriada ao criar Cloud TPUs.

PyTorch e JAX

Use as seguintes imagens base de VM de TPU comuns para PyTorch e JAX e instale o framework que você quer usar.

  • tpu-ubuntu2204-base (padrão)
  • v2-alpha-tpuv5 (TPU v5p)

Consulte os documentos do guia de início rápido sobre PyTorch/XLA e JAX para ver as instruções de instalação.

TensorFlow

Há imagens de VM de TPU específicas para cada versão do TensorFlow. As seguintes versões do TensorFlow são compatíveis com as Cloud TPUs:

  • 2.16.1
  • 2.15.1
  • 2.15.0
  • 2.14.1
  • 2.14.0
  • 2.13.1
  • 2.13.0
  • 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.3
  • 2.8.0
  • 2.7.4
  • 2.7.3

Para mais informações sobre as versões de patch do TensorFlow, consulte Versões de patch do TensorFlow compatíveis.

Para as versões 2.15.0 e mais recentes do TensorFlow, há variantes de imagem de VM de TPU com base na API do dispositivo (PJRT ou executor de stream) que você está usando.

Como treinar nas versões v5p e v5e

As TPUs v5e e v5p são compatíveis com o TensorFlow 2.15.0 e versões mais recentes. Especifique a imagem da VM da TPU usando o formato: 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 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 imagem de VM da TPU tpu-vm-tf-2.16.0-pod-pjrt. 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 única TPU de host, omita pod.

Veiculação na v5e

Há imagens Docker exibidas que contêm todos os requisitos de software necessários para disponibilização com o TensorFlow, o PyTorch e o JAX. Para mais informações, consulte Introdução à inferência do Cloud TPU v5e.

TPU v4

Se você estiver usando a TPU v4 e o TensorFlow 2.15.0 ou uma versão mais recente, siga as instruções para treinar nas versões v5p e v5e. Se você estiver usando o TensorFlow 2.10.0 ou anterior, use uma imagem de VM da TPU específica da v4:

Versão do TensorFlow Versão da imagem da VM 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, utilize a imagem de VM da TPU que corresponda à versão do TensorFlow em uso. Por exemplo, se você estiver usando o TensorFlow 2.14.1, use a imagem de TPU tpu-vm-tf-2.14.1. Substitua 2.14.1 pela versão do TensorFlow que você está usando para outras versões. Se você estiver usando um pod de anexação de TPU com vários hosts ao final da imagem 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. Por exemplo, se você estiver usando o TensorFlow 2.16.1 com a API PJRT, use a imagem da TPU tpu-vm-tf-2.16.1-pjrt. Se você estiver usando a API executor de stream com a mesma versão do TensorFlow, use a imagem de TPU tpu-vm-tf-2.16.1-se. As versões do TensorFlow anteriores à 2.15.0 oferecem suporte apenas ao executor de stream.

Compatibilidade com PJRT do TensorFlow

A partir do TensorFlow 2.15.0, é possível usar a interface PJRT para TensorFlow na TPU. O PJRT oferece desfragmentação automática da memória do dispositivo e simplifica a integração do hardware com frameworks. Saiba mais sobre o PJRT em PJRT: Como simplificar o hardware e a integração de framework de ML (em inglês).

Accelerator Engenharia de Suporte a PJRT Suporte ao 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 embedding de TPU No Sim
TPU v2 a v4 tf.summary/tf.print com posicionamento de dispositivo de software No Sim
TPU v5e Computação densa (sem API de incorporação de TPU) Sim No
TPU v5e API TPU embedding N/A No
TPU v5p Computação densa (sem API de incorporação de TPU) Sim No
TPU v5p API TPU embedding Sim No

Versões do Libtpu

As imagens do TensorFlow das VMs de TPU contêm uma versão específica do TensorFlow e a biblioteca libtpu correspondente. Se você estiver criando sua própria imagem de VM, use as seguintes versões de software da TPU do TensorFlow e as versões do libtpu correspondentes:

Versão do TensorFlow Versão do libtpu.so
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

A seguir