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
sistema operacional, Docker e outros softwares necessários para executar seu código em TPUs.
Neste documento, fornecemos orientações sobre como selecionar a imagem de VM de TPU apropriada ao
a criação de 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-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 têm suporte para o Cloud TPU:
- 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 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 do 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 TPU v2 ou v3, use a imagem de VM da TPU que corresponda à versão do
o 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 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 como
parte do nome da imagem. 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 o
Imagem da 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 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 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 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 para VMs 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.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
- Saiba mais sobre a arquitetura de TPU na página Sistemas Arquitetura.
- Consulte Quando usar TPUs para saber mais sobre os tipos de modelos adequados para o Cloud TPU.