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 de TPU. As imagens de VM do TPU contêm o sistema operacional Ubuntu, o Docker e outros softwares necessários para executar o código em TPUs. Este documento fornece orientações sobre como selecionar a imagem de VM da TPU adequada 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-tpuv6e (TPU v6e)
  • v2-alpha-tpuv5 (TPU v5p)
  • v2-alpha-tpuv5-lite (TPU v5e)

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

TensorFlow

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

  • 2.18.0
  • 2.17.0
  • 2.16.2
  • 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 TPU com base na API do dispositivo (PJRT ou executor de stream) que você está usando.

Treinamento nas versões 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 imagem da VM 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 de 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 imagem da VM 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 TPU de host único, omita pod.

Veiculação na v6e e v5e

Há imagens do Docker que contêm todos os requisitos de software necessários para servir com o TensorFlow, PyTorch e 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 mais recente, siga as instruções para treinamento na v6e, v5p e v5e. Se você estiver usando o TensorFlow 2.10.0 ou anterior, use uma imagem de VM de TPU específica da v4:

Versão do TensorFlow Versão da imagem da VM 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 imagem da VM TPU que corresponde à versão do TensorFlow que você está usando. Por exemplo, se você estiver usando o TensorFlow 2.14.1, use a imagem 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 um pod de adição de TPU de vários hosts, adicione ao final da imagem do 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 imagem. Por exemplo, se você estiver usando o TensorFlow 2.16.1 com a API PJRT, use a imagem 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 imagem TPU tpu-vm-tf-2.16.1-se. As versões do TensorFlow anteriores à 2.15.0 têm suporte apenas para o 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

Versões do Libtpu

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

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

A seguir