Tipos e topologias de TPU

Informações gerais

Ao criar uma configuração de TPU, é preciso especificar um tipo e tamanho de TPU ou um tipo de TPU e uma topologia (uma organização lógica de ícones). Há duas maneiras de fazer isso em uma operação gcloud create:

• Crie uma única string accelerator-type que especifica o tipo (v2, v3 ou v4) e o número de TensorCores que estão sendo solicitados.
• Especifique o parâmetro type da TPU com a sinalização topology.

Consulte Tipos de TPU para ver exemplos de como criar TPUs usando a string accelerator-type ou as sinalizações type e topology.

Para mais informações sobre as diferentes versões e topologias do Cloud TPU, consulte o documento Arquitetura do sistema.

Tipos de TPU

Use a sinalização "accelerator-type" do comando gcloud para especificar uma configuração de TPU. Uma configuração de TPU é composta por uma versão de TPU e pelo número de TensorCores. Por exemplo, o comando gcloud a seguir cria uma TPU v3 com 128 TensorCores:

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

Também é possível especificar a configuração de TPU como o número de chips de TPU e a disposição lógica deles (a topologia). Para mais informações, consulte Configurações de TPU baseadas em chip.

Configurações de TPU baseadas no TensorCore

Para todos os tipos de TPU, a versão é seguida pelo número de TensorCores (por exemplo, 8, 32, 128). Por exemplo, --accelerator-type=v2-8 especifica uma TPU v2 com 8 TensorCores, e v3-1024 especifica uma TPU v3 com 1024 TensorCores (uma fração de um pod v3).

Configurações de TPU baseadas em chips

As configurações da v4 são especificadas com uma versão de TPU e uma topologia. A topologia é especificada com três tuplas, que descrevem como os ícones são organizados. Por exemplo, a configuração do 32-TensorCore pode ser representada em ícones como 2x2x4 (16 chips, cada um com dois TensorCores). A tupla de três tuplas especifica apenas a topologia, e não a configuração, que requer uma versão de TPU.

A representação de uma topologia v4 é AxBxC, em que A<=B<=C e A,B,C são todos <=4 ou vários múltiplos de 4. Os valores A, B e C são as contagens de ícones das três dimensões. Topologias em que 2A=B=C ou 2A=2B=C também têm variantes otimizadas para comunicação geral, por exemplo, 4×4×8, 8×8×16 e 12×12×24.

Para mais informações sobre as topologias da v4, consulte Configurações da TPU v4.

Como criar topologias baseadas em chip (v4)

Com a v4, você vê dispositivos em vez de TensorCores, um dispositivo por chip. Para especificar topologias baseadas em topologia de chip, a v4 introduziu formatos de topologia tridimensionais (por exemplo, 4x4x4) com uma sinalização --topology para a operação de criação de TPU gcloud. As tabelas mostradas em Variantes de topologia incluem exemplos de outras topologias baseadas em topologia de chip.

Importante: é recomendável especificar uma topologia baseada em topologia de chip para personalizar a topologia física. Geralmente, isso é necessário para ajustar o desempenho com fatias maiores que 256 chips. Entre em contato com o representante da conta do Google Cloud para ativar esse recurso.

Observação :as topologias baseadas em TensorCore e em topologia de chip são mutuamente exclusivas. Só é possível especificar uma ao criar uma TPU.

É possível especificar um tipo e topologia de v4 usando as sinalizações type e topology para a operação de criação gcloud. Por exemplo, o comando gcloud a seguir cria uma TPU v4 com 64 chips em uma topologia de cubo 4x4.

  $ gcloud alpha compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=4x4x4
    --version=tpu-vm-tf-2.12.0

O valor aceito para --type é v4.

Variantes de topologia

A topologia padrão associada a um determinado TensorCore ou contagem de ícones é a mais semelhante a um cubo. Consulte Formas de topologia. Esse formato provavelmente será a melhor escolha para o 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 dados e modelos ou particionamento de espaço de uma simulação). Essas cargas de trabalho vão ter um desempenho melhor se houver correspondência entre o formato da fatia e o paralelismo usado. Por exemplo, colocar o paralelismo de modelo de quatro direções no eixo 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 paralelas mapeadas para dimensões de toro. Normalmente, são modelos de dados grandes (LLMs, na sigla em inglês) paralelos a dados e modelos. Por exemplo, para uma fração de pod da TPU v4 com topologia de 8x16x16, as dimensões do toro são 8, 16 e 16. É mais eficiente usar o paralelismo de modelo de 8 ou 16 vias (mapeado para uma das dimensões físicas do tronco). Um paralelismo de modelo de quatro direções não seria o ideal com essa topologia, já que não está alinhado a nenhuma dimensão de torus, mas seria ideal com uma topologia de 4x16x32 no mesmo número de chips.

Topologias v4 pequenas

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

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

Observação: "links próximos" são links sem sentido. Nenhuma das topologias pequenas inclui links de ajuste.

Topologias grandes da v4

As frações da TPU v4 estão disponíveis em incrementos de 64 chips, com formas que sejam múltiplos de 4 em todas as três dimensões. As dimensões também precisam estar em ordem crescente. A tabela a seguir mostra vários exemplos. Algumas dessas topologias são "personalizadas" que só podem ser iniciadas usando a API de topologia porque têm o mesmo número de chips que uma topologia nomeada mais usada.

Nome (com base na contagem do TensorCore)	Número de ícones	topologia
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Consulte API Topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
Consulte API Topology	1024	4x16x16
v4-4.096	2048	8x16x16
…	…	…

Observação :todas as frações grandes de pod v4, com pelo menos 64 chips, têm links de união em todas as três dimensões, proporcionando maior largura de banda para operações de comunicação coletiva.

API Topology

Para criar frações Pod do Cloud TPU com topologia personalizada, a API gcloud TPU pode ser usada da seguinte maneira:

     $ gcloud alpha compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --subnetwork=tpusubnet \
        --type=v4 \
        --topology=4x4x16 \
        --version=runtime-version

Formas de topologia

A imagem a seguir mostra três possíveis formas de topologia e as configurações de sinalização topology associadas.

Compatibilidade com tipos de TPU

É possível alterar o tipo de TPU para outro tipo de TPU que tenha o mesmo número de tensores 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 chips ou TensorCores, precisará realizar ajustes e otimizações significativos. Para mais informações, consulte Treinamento em pods de TPU.

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.