Versões do software do TPU
Este documento fornece orientações sobre a seleção da versão de software da TPU adequada quando cria TPUs na nuvem.
Quando cria recursos de TPU, especifica a versão do software, também denominada versão de tempo de execução, que se refere ao ambiente de software pré-instalado na VM da TPU. Isto inclui o sistema operativo Ubuntu, o Docker e outro software necessário para executar o seu código em TPUs.
Se estiver a usar a CLI Google Cloud, especifique a versão do software da TPU através do parâmetro --version ou --runtime-version. Se estiver a usar a
Google Cloud consola, selecione uma versão do software de TPU na lista Versão do software de TPU.
PyTorch e JAX
Use as seguintes versões de software de TPU comuns para o PyTorch e o JAX e, em seguida, instale a framework que quer usar.
| Versão da TPU | Versão do software da TPU |
|---|---|
| Trillium (v6e) | v2-alpha-tpuv6e |
| v5p | v2-alpha-tpuv5 |
| v5e | v2-alpha-tpuv5-lite |
| v4 e mais antigas | tpu-ubuntu2204-base |
Para mais informações sobre a instalação e o início da utilização do PyTorch ou do JAX, consulte os artigos Executar um cálculo numa VM do TPU do Google Cloud com o PyTorch e Executar um cálculo numa VM do TPU do Google Cloud com o JAX.
TensorFlow
As versões de software da TPU para o TensorFlow seguem uma convenção de nomenclatura específica:
tpu-vm-tf-x.y.z[-{pod}][-{device_api}]
x.y.z: representa as versões principal, secundária e de patch do TensorFlow.-pod(opcional): indica que está a usar uma fatia de TPU com vários anfitriões.-{device_api}(opcional): especifica a API do dispositivo, por exemplo,-pjrt(se estiver a usar a API PJRT).
Consulte as secções seguintes para ver mais informações sobre como especificar uma versão de software da TPU.
Existem versões de software de TPU específicas para cada versão do TensorFlow. A tabela seguinte mostra as versões do TensorFlow suportadas e as versões libtpu associadas:
| Versão do TensorFlow | Versão de libtpu.so |
|---|---|
| 2.18.0 | 1.12.0 |
| 2.17.1 | 1.11.1 |
| 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 |
Para mais informações sobre as versões de patch do TensorFlow, consulte o artigo Versões de patch do TensorFlow suportadas.
TPU v6e, v5p e v5e
As TPUs v6e, v5e e v5p suportam o TensorFlow 2.15.0 e mais recente. Especifica a versão do software da TPU através do formulário: 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 após a versão do TensorFlow se estiver a usar uma TPU multi-host. Por exemplo, se estiver a usar o TensorFlow 2.16.0 numa TPU multi-anfitrião, use a versão do software de 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 está a usar. Se estiver a usar uma TPU de anfitrião único, omita pod.
TPU v4
Se estiver a usar a TPU v4 e o TensorFlow 2.10.1 ou mais recente, siga as instruções para a TPU v2 e v3. Se estiver a usar o TensorFlow 2.10.0 ou anterior, use uma versão de software de TPU específica da v4:
| Versão do TensorFlow | Versão do software 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 |
TPUs v2 e v3
Se estiver a usar a TPU v2 ou v3, use a versão do software da TPU que corresponda à versão do TensorFlow que está a usar. Por exemplo, se estiver a usar o TensorFlow 2.14.1, use a tpu-vm-tf-2.14.1 versão do software de TPU. Para
outras versões do TensorFlow, substitua 2.14.1 pela
versão do TensorFlow que está a usar. Se estiver a usar uma TPU multi-host, acrescente pod ao final da versão do software da TPU, por exemplo, tpu-vm-tf-2.14.1-pod.
A partir do TensorFlow 2.15.0, também tem de especificar uma API de dispositivo como parte do nome da versão do software. Por exemplo, se estiver a usar o TensorFlow 2.16.1 com a API PJRT, use a versão de software da TPU tpu-vm-tf-2.16.1-pjrt. Se estiver a usar a API stream executor com a mesma versão do TensorFlow, use a versão do software da TPU.tpu-vm-tf-2.16.1-se As versões do TensorFlow anteriores à 2.15.0 só suportam o stream
executor.
Compatibilidade com PJRT do TensorFlow
A partir do TensorFlow 2.15.0, pode usar a interface PJRT para o TensorFlow na TPU. O PJRT inclui a desfragmentação automática da memória do dispositivo e simplifica a integração do hardware com frameworks. Para mais informações acerca da PJRT, consulte o artigo PJRT: simplificar a integração de hardware e frameworks de ML.
| Accelerator | Funcionalidade | Compatibilidade com PJRT | Compatibilidade com o executor de streams |
|---|---|---|---|
| TPU v2 a v4 | Computação densa (sem API TPU Embedding) | Sim | Sim |
| TPU v2 a v4 | API Dense Compute + API TPU Embedding | Não | Sim |
| TPU v2 a v4 | tf.summary/tf.print com posicionamento flexível de dispositivos | Não | Sim |
| TPU v5e | Computação densa (sem API TPU Embedding) | Sim | Não |
| TPU v5e | API TPU Embedding | N/A | Não |
| TPU v5p | Computação densa (sem API TPU Embedding) | Sim | Não |
| TPU v5p | API TPU Embedding | Sim | Não |
O que se segue?
- Consulte o artigo Arquitetura da TPU para saber mais sobre a arquitetura da TPU.
- Consulte Quando usar TPUs para saber mais sobre os tipos de modelos adequados para a Cloud TPU.