Imagens de VM da TPU

Ao criar recursos de TPU, você transmite o --version ou o --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 seu código em TPUs. Este documento fornece orientações sobre como selecionar a imagem de VM de TPU apropriada ao criar e Cloud TPUs.

PyTorch e JAX

Use as seguintes imagens de 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 de início rápido do PyTorch/XLA e do JAX para ver as instruções de instalação.

TensorFlow

Há imagens de VM de TPU específicas para cada versão do TensorFlow. O as seguintes versões do TensorFlow são compatíveis com 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 da TPU. com base na API do dispositivo (PJRT ou executor de stream) que você está usando.

Treinamento em v5p e v5e

A TPU v5e e v5p oferecem suporte ao TensorFlow 2.15.0 e versões mais recentes. Você especifica a VM da TPU imagem usando o formato: tpu-vm-tf-x.y.z-{pod}-pjrt, em que x é o principal a versão do TensorFlow, y é a secundária e z é a Versão de patch do TensorFlow. Adicione pod depois da versão do TensorFlow se você estiver usando uma TPU com vários hosts. Por exemplo, se você usa o TensorFlow 2.16.0 em uma TPU com vários hosts, use a imagem de VM de TPU tpu-vm-tf-2.16.0-pod-pjrt. Para outras versões do TensorFlow, substitua 2.16.0 pela versão principal e pelo patch. diferentes do TensorFlow que você está usando. Se você estiver usando uma TPU de host único, omitir pod.

Veiculação na v5e

Há imagens Docker de veiculação que contêm todos os requisitos de software necessários para com 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 de treinamento nas versões v5p e v5e (em inglês). Se você estiver usando TensorFlow 2.10.0 ou anterior, use uma imagem de VM de TPU específica para a 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, use a imagem de VM da TPU que corresponde à versão o TensorFlow que você está usando. Por exemplo, se você usa o TensorFlow 2.14.1, use a imagem de TPU tpu-vm-tf-2.14.1. Para outras versões do TensorFlow, substitua 2.14.1 pela versão do TensorFlow estão usando. Se você estiver usando um TPU com vários hosts, anexe o pod ao final da TPU imagem, 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. Por exemplo, se você estiver usando o TensorFlow 2.16.1 com a API PJRT, use a imagem de TPU tpu-vm-tf-2.16.1-pjrt. Se você estiver usando a API do executor de stream com a mesma versão do TensorFlow, use a TPU tpu-vm-tf-2.16.1-se imagem. As versões do TensorFlow anteriores à 2.15.0 são compatíveis apenas com o executor de stream.

Suporte ao PJRT do TensorFlow

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

Accelerator Recurso Suporte à PJRT Suporte ao executor de streaming
TPU v2 a v4 Computação densa (sem API de incorporação de TPU) Sim Sim
TPU v2 a v4 API Dense Compute + API Embedding de TPU Não Sim
TPU v2 a v4 tf.summary/tf.print com posicionamento de dispositivo flexível 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 nas 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 de TPU do TensorFlow e as Versões do libtpu:

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