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 da TPU contêm o sistema operacional (Ubuntu) e, opcionalmente, outro software necessário para executar o código nas 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 de base de VM de TPU comuns para PyTorch e JAX. Em seguida, 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. 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 baseadas 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. 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
depois da 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 com 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 do Docker 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 mais recente, siga as instruções de treinamento na v5p e v5e. Se você estiver usando o 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 da VM da TPU que corresponde à 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
. Para outras versões do TensorFlow, substitua 2.14.1
pela versão do TensorFlow que você está usando. Se você estiver usando um TPU com vários hosts, anexe o pod 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 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 imagem da TPU tpu-vm-tf-2.16.1-se
. 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 o 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 saber mais, consulte PJRT: como simplificar a integração de hardware e framework de ML (em inglês).
Accelerator | Engenharia de | 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 | No | Sim |
TPU v2 a v4 | 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 de incorporação de TPU | N/A | No |
TPU v5p | Computação densa (sem API de incorporação de TPU) | Sim | No |
TPU v5p | API de incorporação de TPU | 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 de TPU do TensorFlow e versões de 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
- 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.