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
- Saiba mais sobre a arquitetura de TPU na página Arquitetura do sistema.
- Consulte Quando usar TPUs para conhecer os tipos de modelos adequados para o Cloud TPU.