TPU v4

Neste documento, descrevemos a arquitetura e as configurações compatíveis do Cloud TPU v4.

arquitetura do sistema.

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

Especificações da chave	Valores de pod v4
Pico de computação por chip	275 teraflops (bf16 ou int8)
Capacidade e largura de banda do HBM2	32 GiB, 1.200 GBps
Potência mínima/média/máxima medida	90/170/192 W
Tamanho do pod de TPU	4096 ícones
Topologia de interconexão	Malha 3D
Pico de computação por pod	1.1 exaflops (bf16 ou int8)
Largura de banda de redução total por pod	1,1 PB/s
Largura de banda de bissecção por pod	24 TB/s

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 frações 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 torus 3D será melhor do que uma configuração de malha 3D. Para mais informações, consulte Topologias de tori trançadas.

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

Nesta seção, descrevemos os benefícios de desempenho da TPU v4

Sistema de memória

O acesso não uniforme à memória (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. Uma CPU e a memória dela são chamadas de nó NUMA. Os nós NUMA estão conectados a outros nós diretamente adjacentes um ao outro. Uma CPU de um nó NUMA pode acessar a memória em outro nó NUMA, mas esse acesso é mais lento do que o acesso à memória dentro de um nó NUMA.

O software em execução em uma máquina com várias CPUs pode colocar os dados necessários para uma CPU no nó NUMA, aumentando a capacidade de processamento da memória. Para saber mais sobre NUMA, consulte Acesso não uniforme à memória na Wikipédia.

Para aproveitar os benefícios da localidade NUMA, vincule seu script de treinamento ao nó 0 do NUMA.

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 seu código de 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 tem uma forte dependência das cargas de trabalho da CPU (por exemplo, cargas de trabalho de classificação de imagem ou recomendação), independentemente do framework.
• Se você estiver usando uma versão de ambiente de execução do TPU sem um sufixo -pod (por exemplo, tpu-vm-tf-2.10.0-v4).

Outras diferenças entre sistemas 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 no chip.
• Desempenho do HBM aprimorado usando os padrões e velocidades mais recentes de memória.
• Perfil de desempenho de DMA aprimorado com suporte integrado para avanço de alto desempenho em granularidades de 512 bilhões.

TensorCores

• O dobro de MXUs e uma frequência de relógio maior oferecendo 275 TFLOPS no máximo.
• 2x largura de banda de transposição e permutação.
• Modelo de acesso à memória de armazenamento de carga para memória comum (Cmem).
• Largura de banda de carregamento de peso de MXU mais rápida e suporte ao modo de 8 bits para permitir tamanhos de lote menores 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 de rede menores.

Outro

• Interface PCIE gen3 x16 para host (conexão direta).
• Modelo de segurança aprimorado.
• Maior eficiência energética.

Configurações

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

Observação :AcceleratorType e AccleratorConfig são mutuamente exclusivos. Só é possível usar um ao criar uma TPU.

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.

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

Topologias de tori trançadas

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.

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

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.

Consulte como usar o AcceleratorConfig para ver detalhes sobre como especificar uma configuração de tori torcido ao criar um Cloud TPU.

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

Observação: a topologia torcida não será usada automaticamente, a menos que o sufixo _twisted seja especificado com a sinalização --topology. Por exemplo, para especificar uma topologia torcida 4x4x8, especifique a sinalização topology como 4x4x8_twisted.

Variantes de topologia da 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.

A topologia padrão associada a determinada contagem de ícones é mais semelhante a um cubo. 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 dados e modelos de linguagem grandes (LLMs) paralelos. 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 os com topologias maiores que 64 (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 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.

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