Tipos e topologias de TPU

Configurações de TPU

Ao criar uma TPU, é necessário especificar uma configuração dela. É possível especificar uma configuração de TPU em termos de TensorCores ou chips de TPU. A especificação de uma configuração de TPU em termos de chips só é necessária em cenários avançados em que você esteja treinando um modelo grande e ajustando o desempenho dele.

Configuração baseada no TensorCore

É possível usar uma configuração baseada no TensorCore com as TPUs v2 e v4. Descreva as configurações baseadas no TensorCore com uma versão de TPU e o número de TensorCores que você quer usar. Por exemplo, v4-128 especifica uma configuração com TPUs v4 e 128 TensorCores. Especifique uma configuração baseada no TensorCore com a sinalização --accelerator-type do comando "gcloud compute tpus tpu-vm create":

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

Configuração de TPU baseada em chip

É possível usar uma configuração baseada em chip de TPU com a TPU v4 e versões posteriores. Você descreve configurações baseadas em chip com uma versão de TPU e uma topologia (uma disposição física de chips de TPU). Especifique uma versão de TPU com uma string como "v4". Especifique uma topologia com três tuplas que descreva como os chips são organizados em três dimensões. Especifique uma configuração baseada em chip com as sinalizações --version e --topology do comando gcloud compute tpus tpu-vm create. Por exemplo, o comando a seguir especifica a TPU v4 como a versão da TPU e 2x2x4 (16 chips, cada um com dois TensorCores) como a topologia: 

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=2x2x4
    --version=tpu-vm-tf-2.12.0

Informações adicionais sobre as diferentes versões e topologias do Cloud TPU podem ser encontradas em Arquitetura do sistema.

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

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

Variantes de topologia

A topologia padrão associada a uma determinada contagem de ícones é a mais semelhante a um cubo. Consulte Formas de topologia. 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 de dados ou particionamento de espaço de uma simulação). Essas cargas de trabalho terão melhor desempenho se a topologia for correspondida ao paralelismo usado. Por exemplo, colocar o paralelismo de modelo de quatro vias na dimensão X e de 256 vias nas dimensões Y e Z corresponde a uma topologia 4x16x16.

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

Topologias pequenas da v4

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

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

Topologias grandes da v4

As fatias da TPU v4 estão disponíveis em incrementos de 64 chips, com formas que são múltiplas de 4 em todas as três dimensões. As dimensões também precisam estar em ordem crescente. Vários exemplos são mostrados na tabela a seguir. Algumas dessas topologias são topologias "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 comumente usada.

Nome (com base na contagem do TensorCore) Número de chips 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-4096 2048 8x16x16
… … …

API Topology

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

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

image

Compatibilidade com tipos de TPU

É possível alterar o tipo de TPU para outro que tenha o mesmo número de tensores ou chips (por exemplo, v3-128 e v4-128) e executar o 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, será necessário executar ajustes e otimizações significativos. Para mais informações, consulte Treinamento em pods de TPU.

A seguir