TPU v4

Este documento descreve a arquitetura e as configurações compatíveis do Cloud TPU v4.

arquitetura do sistema.

Cada chip TPU v4 contém dois TensorCores. Cada TensorCore tem quatro unidades de multiplicação de matrizes (MXUs), uma uma unidade de vetor e uma unidade escalar. A tabela a seguir mostra as principais especificações para um Pod de TPU v4.

O diagrama a seguir ilustra um chip TPU v4.

Malha 3D e toro 3D

As TPUs v4 têm uma conexão direta com os chips vizinhos mais próximos em três dimensões, resultando em uma malha 3D de conexões de rede. As conexões podem ser configuradas como um toro 3D em fatias em que a topologia, AxBxC, é 2A=B=C ou 2A=2B=C, em que cada dimensão é um múltiplo de 4. Por exemplo, 4x4x8, 4x8x8 ou 12x12x24. Em geral, o desempenho de uma configuração de toro 3D é melhor do que uma configuração de malha 3D. Para mais informações, consulte Topologias de tori trançado.

Benefícios de desempenho da TPU v4 em relação à v3

Esta seção descreve os benefícios de desempenho da TPU v4

Sistema de memória

O acesso à memória não uniforme (NUMA, na sigla em inglês) é uma arquitetura de memória de computador para máquinas com várias CPUs. Cada CPU tem acesso direto a um bloco de memória de alta velocidade. A CPU e a memória dela são chamadas de nó NUMA. Os nós NUMA são conectados a outros nós NUMA que são adjacentes uns aos outros. Uma CPU de um nó do NUMA pode acessar a memória em outro nó do NUMA, mas esse acesso é mais lento em um nó NUMA.

O software executado em uma máquina com várias CPUs pode colocar os dados necessários por uma CPU no nó NUMA, aumentando a taxa de transferência de memória. Para mais informações sobre NUMA, consulte Acesso à memória não uniforme na Wikipédia.

Você pode aproveitar os benefícios da localidade NUMA vinculando seu script de treinamento ao nó NUMA 0.

Para ativar a vinculação de nós NUMA:

1. Instale a ferramenta de linha de comando numactl.

    $ sudo apt-get update
    $ sudo apt-get install numactl

2. Vincule o código do script ao nó NUMA 0.

    $ numactl --cpunodebind=0 python3 your-training-script

Ative a vinculação de nós NUMA se:

• Se a carga de trabalho depende muito das cargas de trabalho da CPU (por exemplo, cargas de trabalho de classificação, cargas de trabalho de recomendação) independentemente do framework.
• Se você estiver usando uma versão do ambiente de execução da TPU sem o sufixo -pod (por exemplo, tpu-vm-tf-2.10.0-v4).

Outras diferenças do sistema de memória:

• Os chips de TPU v4 têm um espaço de memória HBM unificado de 32 GiB em todo o chip, permitindo uma melhor coordenação entre os dois TensorCores do chip.
• Melhoria no desempenho da HBM usando os padrões e velocidades de memória mais recentes.
• Melhoria no perfil de desempenho de DMA com suporte integrado para escalonamento de alto desempenho com granularidades de 512B.

TensorCores

• O dobro do número de MXUs e uma taxa de clock mais alta, oferecendo 275 TFLOPS de pico.
• Largura de banda de transposição e permutação 2x.
• Modelo de acesso à memória de armazenamento de carga para memória comum (Cmem).
• Largura de banda de carregamento de peso MXU mais rápida e suporte ao modo de 8 bits para permitir tamanhos de lote e latência de inferência aprimorada.

Interconexão entre chips

Seis links de interconexão por chip para permitir topologias de rede com diâmetros menores.

Outro

• Interface x16 PCIE de 3ª geração para host (conexão direta).
• Modelo de segurança aprimorado.
• Melhoria na eficiência energética.

Configurações

Um pod TPU v4 é composto por 4.096 chips interconectados com camadas reconfiguráveis links 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 e o número de recursos de TPU necessários. Ao criar uma fatia de pod da TPU v4, é possível especificar o tipo e o tamanho dela de duas maneiras: AcceleratorType e AccleratorConfig.

Como usar o AcceleratorType

Use AcceleratorType quando não estiver especificando uma topologia. Para configurar TPUs v4 usando AcceleratorType, use a flag --accelerator-type ao criar a fração do pod de TPU. Defina --accelerator-type como uma string que contenha a versão do TPU e o número de TensorCores que você quer usar. Por exemplo, para criar uma fração de pod v4 com 32 TensorCores, use --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.17.0-pod-pjrt

O número após a versão da TPU (v4) especifica o número de TensorCores. Há dois TensorCores em uma TPU v4, então 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 fatias de pod maiores que 256 chips.

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

Especifique uma topologia de TPU usando uma terna de 3 elementos, 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 chips em cada uma das três dimensões. Por exemplo, para criar um segmento 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.17.0-pod-pjrt

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

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

As fatias de pods maiores podem ser criadas a partir de um ou mais "cubos" de chips 4x4x4.

Topologias tori torcidas

Algumas formas de fatias de toro 3D v4 têm a opção de usar o que é conhecido como topologia de torus trançado. Por exemplo, dois cubos v4 podem ser organizados como um slice 4x4x8 ou 4x4x8_twisted. As topologias torcidas oferecem bissecção significativamente maior largura de banda larga. A largura de banda de bipartição aumentada é útil para cargas de trabalho que usam padrões de comunicação globais. As topologias torcidas podem melhorar o desempenho da maioria dos modelos, com grandes cargas de trabalho de incorporação de TPU se beneficiando mais.

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 a topologia pode variar de acordo com o tipo de paralelismo (DP, MP etc.). A prática recomendada é treinar o LLM com e sem uma topologia torcida para determinar qual oferece o melhor desempenho para o modelo. Alguns experimentos O modelo FSDP MaxText tem e observaram de 1 a 2 melhorias na MFU usando uma topologia torcida.

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

• Balanceamento de carga aprimorado
• Maior largura de banda com bissecção
• Trajetos de pacotes mais curtos

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

O software de TPU oferece suporte a tori torcido em fatias em que o tamanho de cada é igual ou duas vezes o tamanho da menor dimensão. Para por exemplo, 4x4x8, 4×8×8 ou 12x12x24.

Como exemplo, considere esta topologia de torus 4×2 com TPUs rotuladas com seus (X,Y) na fração:

Para maior clareza, as bordas neste gráfico de topologia são mostradas como bordas não direcionadas. Na prática, cada aresta é uma conexão bidirecional entre TPUs. Chamamos as bordas entre um lado da grade e o lado oposto de bordas de contorno, conforme indicado no diagrama.

Ao torcer essa topologia, chegamos a uma topologia de anel torcido 4×2 completamente simétrica:

Tudo o que mudou entre o diagrama e o anterior foi o Y ao redor das bordas. Em vez de se conectar 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 e números de dimensões. A rede resultante é simétrica, desde que cada dimensão seja igual ou duas vezes o tamanho da menor.

Consulte como usar o AcceleratorConfig para saber como especificar uma configuração de tori retorcido ao criar uma Cloud TPU.

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

Variantes de topologia da TPU v4

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

A topologia padrão associada a uma determinada contagem de ícones é aquela com maior semelhante a um cubo. Essa forma é provavelmente a melhor escolha para o treinamento de ML paralelo 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 espacial de uma simulação). Essas cargas de trabalho têm melhor desempenho se a topologia corresponder ao paralelismo usado. Por exemplo, colocar o paralelismo de modelo 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.

Modelos com várias dimensões de paralelismo têm melhor desempenho dimensões de paralelismo mapeadas para as dimensões de topologia de TPU. Geralmente, são dados e modelos de modelos de linguagem grandes (LLMs) paralelos. Por exemplo, para uma TPU v4 Fração do pod 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 da topologia de TPU física). Um paralelismo de modelo de 4 vias seria subótimo com essa topologia, já que não está alinhado com nenhuma das dimensões da 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: aqueles com topologias menores que 64 chips (pequenas) e aqueles com topologias maiores que 64 chips (grandes).

Pequenas topologias v4

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

Topologias v4 grandes

Frações do Pod de TPU v4 estão disponíveis em incrementos de 64 chips, com formas múltiplos de 4 nas três dimensões. As dimensões devem estar em em ordem crescente. Confira alguns exemplos na tabela a seguir. Alguns essas topologias são "personalizadas" que só podem ser criadas usando o as sinalizações --type e --topology, porque há mais de uma maneira de organizar. os chips.