Configurações de TPU

Configurações de TPU v5p

Um pod TPU v5p é composto por 8.960 chips interconectados com links reconfiguráveis de alta velocidade. A rede flexível da TPU v5p permite conectar os chips em uma fração do mesmo tamanho de várias maneiras. Ao criar uma fração de TPU usando o comando gcloud compute tpus tpu-vm create, especifique o tipo e a forma usando os parâmetros AcceleratorType ou AcceleratorConfig.

A tabela a seguir mostra as formas de fatia única mais comuns compatíveis com a v5p, além da maioria (mas não todas) das formas completas de cubo maiores que 1 cubo. O formato máximo do v5p é 16x16x24 (6.144 ícones, 96 cubos).

Formato da fatia Tamanho da VM No de núcleos # de ícones No de máquinas No de cubos Oferece suporte ao Twisted?
2x2x1 Host completo 8 4 1 N/A N/A
2x2x2 Host completo 16 8 2 N/A N/A
2x4x4 Host completo 64 32 8 N/A N/A
4x4x4 Host completo 128 64 16 1 N/A
4x4x8 Host completo 256 128 32 2 Sim
4x8x8 Host completo 512 256 64 4 Sim
8x8x8 Host completo 1024 512 128 8 N/A
8x8x16 Host completo 2048 1024 256 16 Sim
8x16x16 Host completo 4096 2048 512 32 Sim
16x16x16 Host completo 8192 4096 1024 64 N/A
16x16x24 Host completo 12.288 6144 1.536 96 N/A

O treinamento de fração única é compatível com até 6.144 ícones. Ele é extensível para 18432 chips usando Multislice. Para mais detalhes, consulte a Visão geral de vários pedaços do Cloud TPU.

Como usar o parâmetro AcceleratorType

Ao alocar recursos de TPU, use o argumento --accelerator-type para especificar o número de TensorCores em uma fração. --accelerator-type é uma string formatada "v$VERSION_NUMBERp-$CORES_COUNT". Por exemplo, v5p-32 especifica uma fração de TPU v5p com 32 TensorCores (16 chips).

Para provisionar TPUs para um job de treinamento v5p, use um dos seguintes tipos de acelerador na solicitação de criação da CLI ou da API TPU:

  • v5p-8
  • v5p-16
  • v5p-32
  • v5p-64
  • v5p-128 (um cubo/rack completo)
  • v5p-256 (2 cubos)
  • v5p-512
  • v5p-1024, v5p-12288

Como usar o parâmetro AcceleratorConfig

Nas versões v5p e posteriores do Cloud TPU, AcceleratorConfig é usado da mesma forma que no Cloud TPU v4. A diferença é que, em vez de especificar o tipo de TPU como --type=v4, é preciso defini-lo como a versão do TPU que você está usando (por exemplo, --type=v5p para a versão v5p).

Resiliência de ICI do Cloud TPU

A resiliência de ICI ajuda a melhorar a tolerância a falhas de links ópticos e chaves de circuito óptico (OCS, na sigla em inglês) que conectam TPUs entre cubos. As conexões ICI dentro de um cubo usam elos de cobre que não são afetados. A resiliência da ICI permite que as conexões da ICI sejam roteadas em torno das falhas de OCS e ICI ópticas. Como resultado, ele melhora a disponibilidade de programação de fatias de TPU, com a compensação da degradação temporária no desempenho do ICI.

Semelhante ao Cloud TPU v4, a resiliência de ICI é ativada por padrão para frações v5p que têm um cubo ou maior:

  • v5p-128 ao especificar o tipo de acelerador.
  • 4x4x4 ao especificar a configuração do acelerador

Propriedades da VM, do host e da fração

Propriedade Valor em uma TPU
Número de ícones v5p 4
No de vCPUs 208 (apenas metade pode ser usada ao usar a vinculação NUMA para evitar penalidade de desempenho entre NUMA)
RAM (GB) 448 (apenas metade pode ser usada ao usar a vinculação NUMA para evitar penalidade de desempenho entre NUMA)
No de nós NUMA 2
Capacidade de processamento de NIC (Gbps) 200

Relação entre o número de TensorCores, chips, hosts/VMs e cubos em um pod:

Núcleos Ícones Hosts/VMs Cubos
Organizador 8 4 1
Cube (também conhecido como rack) 128 64 16 1
Maior fração com suporte 12.288 6144 1.536 96
Pod completo v5p 17920 8960 2240 140

Configurações da TPU v5e

O Cloud TPU v5e é um produto combinado de treinamento e inferência (veiculação). Para diferenciar um ambiente de treinamento e de inferência, use as sinalizações AcceleratorType ou AcceleratorConfig com a API TPU ou a sinalização --machine-type ao criar um pool de nós do GKE.

Os jobs de treinamento são otimizados para capacidade e disponibilidade, enquanto os jobs de exibição são otimizados para latência. Assim, um job de treinamento em TPUs provisionadas para exibição pode ter menor disponibilidade e, da mesma forma, um job de exibição executado em TPUs provisionadas para treinamento pode ter maior latência.

Use AcceleratorType para especificar o número de TensorCores que você quer usar. Especifique o AcceleratorType ao criar uma TPU usando a CLI gcloud ou o console do Google Cloud. O valor especificado para AcceleratorType é uma string com o formato: v$VERSION_NUMBER-$CHIP_COUNT.

Também é possível usar AcceleratorConfig para especificar o número de TensorCores que você quer usar. No entanto, como não há variantes de topologia 2D personalizadas para a TPU v5e, não há diferença entre usar AcceleratorConfig e AcceleratorType.

Para configurar uma TPU v5e usando AcceleratorConfig, use as sinalizações --version e --topology. Defina --version como a versão da TPU que você quer usar e --topology como a disposição física dos chips de TPU na fração. O valor especificado para AcceleratorConfig é uma string com o formato AxB, em que A e B são as contagens de ícones em cada direção.

Os seguintes formatos de fatias 2D são compatíveis com a v5e:

topologia Número de chips do TPU Número de hosts
1x1 1 1/8
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Cada VM de TPU em uma fração de TPU v5e contém 1, 4 ou 8 chips. Em partes menores e com quatro chips, todos os chips de TPU compartilham o mesmo nó de acesso não uniforme à memória (NUMA, na sigla em inglês).

Para VMs de TPU de 8 chips v5e, a comunicação entre CPU e TPU será mais eficiente nas partições NUMA. Por exemplo, na figura a seguir, a comunicação com CPU0-Chip0 será mais rápida que a comunicação com CPU0-Chip4.

Comunicação do nó NUMA

Tipos do Cloud TPU v5e para veiculação

A veiculação de host único é suportada para até oito ícones v5e. As seguintes configurações são compatíveis: frações 1x1, 2x2 e 2x4. Cada fatia tem 1, 4 e 8 ícones, respectivamente.

Configurações da TPU v5e compatíveis com a veiculação: 1 x 1, 2 x 2 e 2 x 4.

Para provisionar TPUs para um job de exibição, use um dos seguintes tipos de acelerador na solicitação de criação de TPU da CLI ou API:

AcceleratorType (API TPU) Tipo de máquina (API GKE)
v5litepod-1 ct5lp-hightpu-1t
v5litepod-4 ct5lp-hightpu-4t
v5litepod-8 ct5lp-hightpu-8t

A disponibilização em mais de oito chips v5e, também chamada de exibição de vários hosts, é aceita usando Sax. Para mais informações, consulte Veiculação de modelos de linguagem grandes.

Tipos do Cloud TPU v5e para treinamento

O treinamento oferece suporte para até 256 chips.

Para provisionar TPUs para um job de treinamento v5e, use um dos seguintes tipos de acelerador na solicitação de criação de TPU de CLI ou API:

AcceleratorType (API TPU) Tipo de máquina (API GKE) topologia
v5litepod-16 ct5lp-hightpu-4t 4x4
v5litepod-32 ct5lp-hightpu-4t 4x8
v5litepod-64 ct5lp-hightpu-4t 8x8
v5litepod-128 ct5lp-hightpu-4t 8x16
v5litepod-256 ct5lp-hightpu-4t 16x16

Comparação dos tipos de VM da TPU v5e:

Tipo de VM n2d-48-24-v5lite-tpu n2d-192-112-v5lite-tpu n2d-384-224-v5lite-tpu
Número de ícones v5e 1 4 8
No de vCPUs 24 112 224
RAM (GB) 48 192 384
No de nós NUMA 1 1 2
Aplicável a v5litepod-1 v5litepod-4 v5litepod-8
Interrupção Alta Média Baixa

Para liberar espaço para cargas de trabalho que exigem mais chips, os programadores podem forçar a interrupção de VMs com menos chips. Portanto, as VMs de 8 chips provavelmente vão forçar a interrupção das VMs de 1 e 4 chips.

Configurações da TPU v4

Um Pod TPU v4 é composto por 4096 chips interconectados com links reconfiguráveis de alta velocidade. A rede flexível da TPU v4 permite conectar os chips em uma fração de pod do mesmo tamanho de várias maneiras. Ao criar uma fração do Pod de TPU, especifique a versão da TPU e o número de recursos de TPU necessários. Ao criar uma fração do Pod de TPU v4, é possível especificar o tipo e o tamanho de duas maneiras: AcceleratorType e AcceleratorConfig.

Como usar o AcceleratorType

Use AcceleratorType quando não estiver especificando uma topologia. Para configurar TPUs v4 usando AcceleratorType, use a sinalização --accelerator-type ao criar sua fração de Pod de TPU. Defina --accelerator-type como uma string que contenha a versão da TPU e o número de TensorCores que você quer usar. Por exemplo, para criar uma fração do pod v4 com 32 TensorCores, você usaria --accelerator-type=v4-32.

O comando a seguir cria uma fração do Pod de TPU v4 com 512 TensorCores usando a sinalização --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-512
    --version=tpu-vm-tf-2.16.1-pod-pjrt

O número após a versão da TPU (v4) especifica o número de TensorCores. Como há dois TensorCores em uma TPU v4, o número de chips de TPU seria 512/2 = 256.

Como usar o AcceleratorConfig

Use AcceleratorConfig quando quiser personalizar a topologia física da fração de TPU. Isso geralmente é necessário para o ajuste de desempenho com frações de pod maiores que 256 chips.

Para configurar TPUs v4 usando AcceleratorConfig, use as sinalizações --version e --topology. Defina --version como a versão da TPU que você quer usar e --topology como a disposição física dos chips de TPU na fração do pod.

Especifique uma topologia de TPU usando uma de três tuplas, AxBxC, em que A<=B<=C e A, B, C são todos <= 4 ou são todos múltiplos inteiros de 4. Os valores A, B e C são as contagens de ícones em cada uma das três dimensões. Por exemplo, para criar uma fatia de pod v4 com 16 chips, defina --version=v4 e --topology=2x2x4.

O comando a seguir cria uma fração do Pod de TPU v4 com 128 chips de TPU organizados em uma matriz de 4x4x8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=4x4x8
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Topologias em que 2A=B=C ou 2A=2B=C também têm variantes de topologia otimizadas para comunicação total, por exemplo, 4×4×8, 8×8×16 e 12×12×24. Elas são conhecidas como topologias tori trançadas.

As ilustrações a seguir mostram algumas topologias comuns da TPU v4.

imagem

Fatias de pod maiores podem ser criadas com um ou mais "cubos" de chips 4x4x4.

Topologias torcidas de Tori

Algumas formas de frações de toro 3D v4 têm a opção de usar o que é conhecido como topologia de toro trançado. Por exemplo, dois cubos v4 podem ser organizados como uma fatia 4x4x8 ou 4x4x8_torcida. As topologias torcidas oferecem uma largura de banda com bissecção significativamente maior. O aumento da largura de banda de bissecção é útil para cargas de trabalho que usam padrões de comunicação globais. As topologias trançadas podem melhorar o desempenho da maioria dos modelos, sendo que grandes cargas de trabalho de incorporação de TPU são os mais beneficiadas.

Para cargas de trabalho que usam o paralelismo de dados como a única estratégia de paralelismo, as topologias torcidas podem ter um desempenho um pouco melhor. Para LLMs, o desempenho com uma topologia torcida pode variar dependendo do tipo de paralelismo (DP, MP etc.). A prática recomendada é treinar seu LLM com e sem uma topologia torcida para determinar qual oferece o melhor desempenho para o modelo. Alguns experimentos no modelo FSDP MaxText (link em inglês) tiveram de 1 a 2 melhorias de MFU usando uma topologia torcida.

O principal benefício das topologias torcidas é que elas transformam uma topologia de torus assimétrica (por exemplo, 4×4×8) em uma topologia simétrica intimamente relacionada. A topologia simétrica tem muitos benefícios:

  • Balanceamento de carga aprimorado
  • Largura de banda com bissecção maior
  • Rotas de pacotes mais curtas

Esses benefícios se traduzem em melhor desempenho para muitos padrões de comunicação globais.

O software da TPU é compatível com tori trançado em frações em que o tamanho de cada dimensão é igual ou duas vezes o tamanho da menor dimensão. Por exemplo, 4x4x8, 4×8×8 ou 12x12x24.

Por exemplo, considere esta topologia de toro 4×2 com TPUs rotuladas com as coordenadas (X,Y) na fração:

As bordas neste gráfico de topologia são mostradas como arestas não direcionadas para maior clareza. Na prática, cada borda é uma conexão bidirecional entre as TPUs. Chamamos as bordas entre um lado dessa grade e o oposto como bordas circulares, conforme observado no diagrama.

desenho

Ao torcer essa topologia, acabamos com uma topologia de torus trançado 4×2 completamente simétrica:

desenho

Tudo o que mudou entre esse diagrama e o anterior são as bordas em torno em Y. Em vez de se conectarem a outra TPU com a mesma coordenada X, elas foram deslocadas para se conectar à TPU com a coordenada X+2 mod 4.

A mesma ideia se aplica a diferentes tamanhos de dimensão e números de dimensões diferentes. A rede resultante é simétrica, desde que cada dimensão seja igual ou duas vezes o tamanho da menor dimensão.

A tabela a seguir mostra as topologias trançadas com suporte e um aumento teórico na largura de banda de bissecção delas em comparação com as topologias não trançadas.

Topologia torcida Aumento teórico da largura de banda com bissecção
em comparação com um toro não trançado
4×4×8_torcido Aprox. 70%
8x8x16_twisted
12×12×24_torcido
4×8×8_torrado Aprox. 40%
8×16×16_trançado

Variantes da topologia de TPU v4

Algumas topologias que contêm o mesmo número de ícones podem ser organizadas de maneiras diferentes. Por exemplo, uma fração do Pod de TPU com 512 chips (1024 TensorCores) pode ser configurada usando as seguintes topologias: 4x4x32, 4x8x16 ou 8x8x8. Uma fração de Pod de TPU com 2048 chips (4096 TensorCores) oferece ainda mais opções de topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16. Uma fração do Pod de TPU com 2.048 chips (4096 TensorCores) oferece ainda mais opções de topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16.

A topologia padrão associada a determinada contagem de ícones é a mais semelhante a um cubo (consulte Topologia v4). Esse formato é provavelmente a melhor escolha para treinamento de ML com paralelismo de dados. Outras topologias podem ser úteis para cargas de trabalho com vários tipos de paralelismo. Por exemplo, paralelismo de modelo e dados ou o particionamento espacial de uma simulação. Essas cargas de trabalho funcionam melhor quando a topologia corresponde ao paralelismo usado. Por exemplo, colocar o paralelismo de modelos de 4 vias na dimensão X e o paralelismo de dados de 256 vias nas dimensões Y e Z corresponde a uma topologia 4x16x16.

Os modelos com várias dimensões de paralelismo têm melhor desempenho com as dimensões de paralelismo associadas às dimensões de topologia de TPU. Eles geralmente são modelos de linguagem grandes (LLMs) paralelos de dados + modelo. Por exemplo, para uma fração de pod da TPU v4 com topologia 8x16x16, as dimensões da topologia da TPU são 8, 16 e 16. É mais eficiente usar paralelismo de modelos de 8 ou 16 direções (mapeado para uma das dimensões de topologia física da TPU). Um paralelismo de modelo de quatro direções não seria ideal para essa topologia, já que não está alinhado com nenhuma das dimensões de topologia da TPU, mas seria ideal com uma topologia 4x16x32 no mesmo número de chips.

As configurações da TPU v4 consistem em dois grupos: os com topologias menores que 64 chips (topologias pequenas) e aqueles com topologias maiores que 64 chips (topologias grandes).

Topologias pequenas da v4

O Cloud TPU é compatível com as seguintes frações do Pod TPU v4 menores que 64 chips, um cubo 4x4x4. É possível criar essas pequenas topologias da v4 usando o nome baseado no TensorCore (por exemplo, v4-32) ou a topologia (por exemplo, 2x2x4):

Nome (com base na contagem do TensorCore) Número de fichas topologia
v4-8 4 2x2x1
v4-16 8 2x2x2
v4-32 16 2x2x4
v4-64 32 2x4x4

Topologias grandes da v4

As frações do pod da TPU v4 estão disponíveis em incrementos de 64 ícones, com formas que são múltiplos de 4 em todas as três dimensões. As dimensões também precisam estar em ordem crescente. Confira alguns exemplos na tabela a seguir. Algumas dessas topologias são "personalizadas" que só podem ser criadas usando as flags --type e --topology, porque há mais de uma maneira de organizar os ícones.

O comando a seguir cria uma fração do Pod de TPU v4 com 512 chips de TPU organizados em uma matriz de 8x8x8:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=8x8x8
    --version=tpu-vm-tf-2.16.1-pod-pjrt

É possível criar uma fração do Pod de TPU v4 com o mesmo número de TensorCores usando --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-1024
    --version=tpu-vm-tf-2.16.1-pod-pjrt
Nome (com base na contagem do TensorCore) Número de fichas topologia
v4-128 64 4x4x4
v4-256 128 4x4x8
v4-512 256 4x8x8
N/A: é necessário usar as sinalizações --type e --topology. 256 4x4x16
v4-1024 512 8x8x8
v4-1536 768 8x8x12
v4-2048 1024 8x8x16
N/A: é necessário usar sinalizações --type e --topology. 1024 4x16x16
v4-4096 2048 8x16x16
… … …

Configurações da TPU v3

Um Pod TPU v3 é composto por 1024 chips interconectados por links de alta velocidade. Para criar um dispositivo TPU v3 ou uma fração de pod, use a sinalização --accelerator-type para o comando gcloud compute tpus tpu-vm. Para especificar o tipo de acelerador, você especifica a versão da TPU e o número de núcleos de TPU. Para uma única TPU v3, use --accelerator-type=v3-8. Para uma fração do Pod v3 com 128 TensorCores, use --accelerator-type=v3-128.

O comando a seguir mostra como criar uma fração do Pod de TPU v3 com 128 TensorCores:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v3-128
    --version=tpu-vm-tf-2.16.1-pjrt

A tabela a seguir lista os tipos de TPU v3 compatíveis:

Versão da TPU Fim da compatibilidade
v3-8 (data final ainda não definida)
v3-32 (data final ainda não definida)
v3-128 (data final ainda não definida)
v3-256 (data final ainda não definida)
v3-512 (data final ainda não definida)
v3-1024 (data final ainda não definida)
v3-2048 (data final ainda não definida)

Para mais informações sobre como gerenciar TPUs, consulte Gerenciar TPUs. Para mais informações sobre as diferentes versões do Cloud TPU, consulte Arquitetura do sistema.

Configurações da TPU v2

Um Pod TPU v2 é composto por 512 chips interconectados com links reconfiguráveis de alta velocidade. Para criar uma fração do Pod de TPU v2, use a sinalização --accelerator-type para o comando gcloud compute tpus tpu-vm. Para especificar o tipo de acelerador, especifique a versão da TPU e o número de núcleos de TPU. Para uma única TPU v2, use --accelerator-type=v2-8. Para uma fração de pod v2 com 128 TensorCores, use --accelerator-type=v2-128.

O comando a seguir mostra como criar uma fração do Pod de TPU v2 com 128 TensorCores:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v2-128
    --version=tpu-vm-tf-2.16.1-pjrt

Para mais informações sobre como gerenciar TPUs, consulte Gerenciar TPUs. Para mais informações sobre as diferentes versões do Cloud TPU, consulte Arquitetura do sistema.

A tabela a seguir lista os tipos de TPU v2 compatíveis

Versão da TPU Fim da compatibilidade
v2-8 (data final ainda não definida)
v2-32 (data final ainda não definida)
v2-128 (data final ainda não definida)
v2-256 (data final ainda não definida)
v2-512 (data final ainda não definida)

Compatibilidade do tipo de TPU

É possível alterar o tipo de TPU para outro que tenha o mesmo número de TensorCores ou chips (por exemplo, v3-128 e v4-128) e executar seu script de treinamento sem alterações no código. No entanto, se você mudar para um tipo de TPU com um número maior ou menor de TensorCores, será necessário realizar ajustes e otimização significativos. Para mais informações, consulte Treinamento em pods de TPU.

Versões de software 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.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 Não Sim
TPU v2 a 4 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 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 Não
TPU v5p API TPU embedding Sim Não

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.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

Use a versão do software da TPU que corresponde à versão do PyTorch com que seu modelo foi gravado. Por exemplo, se você estiver usando o PyTorch 1.13 e a TPU v2 ou v3, use a versão de software do TPU tpu-vm-pt-1.13. Se você estiver usando a TPU v4, utilize a versão tpu-vm-v4-pt-1.13 do software da TPU. A mesma versão do software da TPU é usada para pods de TPU (por exemplo,v2-32, v3-128, v4-32). As versões atuais do software de TPU compatíveis são:

TPU v2/v3:

  • tpu-vm-pt-2.0 (pytorch-2.0)
  • tpu-vm-pt-1.13 (pytorch-1.13)
  • tpu-vm-pt-1.12 (pytorch-1.12)
  • tpu-vm-pt-1.11 (pytorch-1.11)
  • tpu-vm-pt-1.10 (pytorch-1.10)
  • v2-alpha (pytorch-1.8.1)

TPU v4:

  • tpu-vm-v4-pt-2.0 (pytorch-2.0)
  • tpu-vm-v4-pt-1.13 (pytorch-1.13)

TPU v5 (v5e e v5p):

  • v2-alpha-tpuv5 (pytorch-2.0)

Quando você cria uma VM de TPU, a versão mais recente do PyTorch é pré-instalada nela. A versão correta do libtpu.so é instalada automaticamente quando você instala o PyTorch.

Para alterar a versão atual do software PyTorch, consulte Como alterar a versão do PyTorch.

JAX

Você precisa instalar manualmente o JAX na sua VM de TPU. Não existe uma versão de software de TPU específica do JAX para a TPU v2 e v3. Para versões posteriores da TPU, use as seguintes versões de software:

  • TPU v4: tpu-vm-v4-base
  • TPU v5e: v2-alpha-tpuv5
  • TPU v5p: v2-alpha-tpuv5

A versão correta do libtpu.so é automaticamente instalada quando você instala o JAX.

Versões de software 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.

PyTorch

Use a versão do software de TPU que corresponde à versão do PyTorch com que o modelo foi gravado. Por exemplo, se estiver usando o PyTorch 1.9, utilize a versão de software pytorch-1.9.

As versões do software de TPU específicas para PyTorch são:

  • pytorch-2.0
  • pytorch-1,13
  • pytorch-1,12
  • pytorch-1,11
  • pytorch-1,10
  • pytorch-1.9
  • pytorch-1.8
  • pytorch-1.7

  • pytorch-1.6

  • pytorch-nightly

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

JAX

Como o JAX precisa ser instalado manualmente na VM de TPU, não há uma versão de software de TPU específica do JAX pré-instalada. É possível usar qualquer uma das versões de software listadas para o TensorFlow.

A seguir