Ambientes de execução da TPU

Neste documento, fornecemos orientações sobre como selecionar a versão apropriada do ambiente de execução para o Cloud TPU, de acordo com o framework e a arquitetura da TPU (VM ou nó) usada.

Versões de ambiente de execução da VM da TPU

Nesta seção, descrevemos as versões do software de TPU que precisam ser usadas em uma TPU com a arquitetura de VM da TPU. Para saber mais sobre a arquitetura do nó de TPU, consulte Versões de software do nó de TPU.

As versões do software da TPU estão disponíveis para os frameworks do TensorFlow, PyTorch e JAX.

TensorFlow

Use a versão de software do TPU que corresponde à versão do TensorFlow com que seu modelo foi gravado.

A partir do TensorFlow 2.15.0, também é necessário especificar o ambiente de execução do executor de stream (SE, na sigla em inglês) ou PJRT. Por exemplo, se você estiver usando o TensorFlow 2.16.1 com o ambiente de execução PJRT, use a versão tpu-vm-tf-2.16.1-pjrt do software de TPU. As versões anteriores ao TensorFlow 2.15.0 oferecem suporte apenas ao executor de stream. Para mais informações sobre PJRT, consulte o suporte a PJRT do TensorFlow.

As versões de software atuais da VM de TPU do TensorFlow compatíveis são:

  • tpu-vm-tf-2.16.1-pjrt
  • tpu-vm-tf-2.16.1-se
  • tpu-vm-tf-2.15.1-pjrt
  • tpu-vm-tf-2.15.1-se
  • tpu-vm-tf-2.15.0-pjrt
  • tpu-vm-tf-2.15.0-se
  • tpu-vm-tf-2.14.1
  • tpu-vm-tf-2.14.0
  • tpu-vm-tf-2.13.1
  • tpu-vm-tf-2.13.0
  • tpu-vm-tf-2.12.1
  • tpu-vm-tf-2.12.0
  • tpu-vm-tf-2.11.1
  • tpu-vm-tf-2.11.0
  • tpu-vm-tf-2.10.1
  • tpu-vm-tf-2.10.0
  • tpu-vm-tf-2.9.3
  • tpu-vm-tf-2.9.1
  • tpu-vm-tf-2.8.4
  • tpu-vm-tf-2.8.3
  • tpu-vm-tf-2.8.0
  • tpu-vm-tf-2.7.4
  • tpu-vm-tf-2.7.3

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

Compatibilidade com 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 da memória do dispositivo e simplifica a integração do hardware com frameworks. Para mais informações sobre o PJRT, consulte PJRT: Simplificando o hardware de ML e a integração de framework (em inglês) no blog do Google Open Source.

Nem todos os recursos da TPU v2, v3 e v4 foram migrados para o ambiente de execução PJRT. A tabela a seguir descreve quais recursos têm suporte no PJRT ou no executor do stream.

Accelerator Engenharia de Compatível com PJRT Compatível com o executor de streaming
TPU v2 a 4 Computação densa (sem API de incorporação de TPU) Sim Sim
TPU v2 a 4 API Dense compute + API de embedding de TPU No Sim
TPU v2 a 4 tf.summary/tf.print com posicionamento de dispositivo flexível No Sim
TPU v5e Computação densa (sem API de incorporação de TPU) Sim No
TPU v5e API TPU embedding N/A: o TPU v5e não é compatível com a API de incorporação de TPU N/A
TPU v5p Computação densa (sem API de incorporação de TPU) Sim No
TPU v5p API TPU embedding Sim No

TPU v4 com o TensorFlow versão 2.10.0 e anteriores

Se você estiver treinando um modelo na TPU v4 com o TensorFlow, as versões 2.10.0 e anteriores do TensorFlow usarão as versões específicas de v4 mostradas na tabela a seguir. Se a versão do TensorFlow que você está usando não aparecer na tabela, siga as orientações na seção do TensorFlow.

Versão do TensorFlow Versão de software da TPU
2.10.0 tpu-vm-tf-2.10.0-v4 e tpu-vm-tf-2.10.0-pod-v4
2.9.3 tpu-vm-tf-2.9.3-v4 e tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4 e tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4 e tpu-vm-tf-2.9.1-pod-v4

Versões do Libtpu

As VMs de TPU são criadas com o TensorFlow e a biblioteca Libtpu correspondente pré-instalada. Se você estiver criando sua própria imagem de VM, especifique 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* 1.2.0
2.7.3 1.1.2

PyTorch e JAX

Use as seguintes imagens de base comuns para PyTorch e JAX e instale a versão do framework necessária com pip.

  • 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 instruções de instalação.

Versões de ambiente de execução do nó da TPU

Nesta seção, descrevemos as versões do software da TPU que precisam ser usadas para uma TPU com a arquitetura de nó de TPU. Para a arquitetura de VM da TPU, consulte Versões do software da VM da TPU.

As versões do software da TPU estão disponíveis para os frameworks do TensorFlow, PyTorch e JAX.

TensorFlow

Use a versão de software do TPU que corresponde à versão do TensorFlow com que seu modelo foi gravado. Por exemplo, se você estiver usando o TensorFlow 2.12.0, use a versão 2.12.0 do software de TPU. As versões específicas do software de TPU do TensorFlow são:

  • 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.2
  • 2.7.3

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

Quando você cria um nó da TPU, a versão mais recente do TensorFlow é pré-instalada nele.

A seguir