TPU v4

Este documento descreve 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 escalar. A tabela a seguir mostra as principais especificações de um Pod de TPU v4.

Principais especificações	Valores do pod v4
Pico de computação por chip	275 teraflops (bf16 ou int8)
Capacidade e largura de banda da HBM2	32 GiB, 1200 GBps
Potência mínima/média/máxima medida	90/170/192 W
Tamanho do Pod de TPU	4.096 chips
Topologia de interconexão	Malha 3D
Tempo de computação máximo 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 bisection 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 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 toro retorcido.

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. Uma 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ó NUMA pode acessar a memória em outro nó NUMA, mas esse acesso é mais lento do que acessar a memória 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 de memória não uniforme na Wikipédia.

É possível 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 tiver uma dependência pesada de cargas de trabalho da CPU (por exemplo, classificação de imagens, cargas de trabalho de recomendação), independentemente do framework.
• Se você estiver usando uma versão do ambiente de execução da TPU sem um sufixo -pod (por exemplo, tpu-vm-tf-2.10.0-v4).

Outras diferenças do sistema de memória:

• Os chips 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.
• 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 carga-armazenamento 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 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âmetro menor.

Outro

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

Configurações

Um pod de TPU v4 é composto por 4.096 chips interconectados com links de alta velocidade reconfiguráveis. 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 um podslice de TPU, especifique a versão da TPU 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.

Observação: AcceleratorType e AccleratorConfig são mutuamente exclusivos. Você só pode usar um deles 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 flag --accelerator-type ao criar a fração do pod de TPU. Defina --accelerator-type como uma string que contém 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 flag --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.18.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 da 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 da TPU que você quer usar e --topology como a disposição física dos chips da TPU na fração do pod.

Você especifica uma topologia de TPU usando uma terna de três 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 fatia de pod de TPU v4 com 128 chips de TPU organizados em uma matriz 4x4x8:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.18.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 de torção

Algumas formas de fatia de toro 3D da v4 têm a opção de usar o que é conhecido como uma topologia de toro retorcido. Por exemplo, dois cubos v4 podem ser organizados como um slice 4x4x8 ou 4x4x8_twisted. As topologias torcidas oferecem uma largura de banda de bipartição significativamente maior. O aumento da largura de banda de bipartição é ú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 usando uma torção pode variar dependendo do 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 no modelo MaxText do FSDP tiveram uma a duas melhorias de 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
• Largura de banda bisseccional maior
• Trajetos de pacotes mais curtos

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

O software da TPU oferece suporte a torções torcidas em fatias em que o tamanho de cada dimensão é igual ou o dobro do tamanho da menor. Por exemplo, 4x4x8, 4×8×8 ou 12x12x24.

Como exemplo, considere esta topologia de anel 4×2 com TPUs marcadas com as coordenadas (X,Y) na fração:

Para maior clareza, as arestas neste gráfico de topologia são mostradas como arestas 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:

A única coisa que mudou entre este diagrama e o anterior são as bordas de contorno em Y. 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.

Topologia de cabos torcidos	Aumento teórico na largura de banda de bipartição em comparação com um toro não torcido
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

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

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, um slice de pod de TPU com 512 chips (1.024 TensorCores) pode ser configurado usando as seguintes topologias: 4x4x32, 4x8x16 ou 8x8x8. Uma fatia de pod de TPU com 2.048 chips (4.096 TensorCores) oferece ainda mais opções de topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16.

A topologia padrão associada a uma determinada contagem de chips é a mais 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.

Os modelos com várias dimensões de paralelismo têm melhor desempenho quando as dimensões de paralelismo são mapeadas para as dimensões de topologia do TPU. Geralmente, são dados e modelos de linguagem grandes paralelos (LLMs). Por exemplo, para uma fatia de pod de TPU v4 com topologia 8x16x16, as dimensões da topologia de 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 oferece suporte às seguintes fatias de pod do TPU v4 menores que 64 chips, um cubo 4x4x4. É possível criar essas pequenas topologias v4 usando o nome baseado em TensorCore (por exemplo, v4-32) ou a topologia (por exemplo, 2x2x4):

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

Topologias v4 grandes

As fatias de pods da 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 precisam estar em ordem crescente. Confira alguns exemplos na tabela a seguir. Algumas dessas topologias são "personalizadas" e só podem ser criadas usando as flags --type e --topology, porque há mais de uma maneira de organizar os chips.

Nome (com base na contagem de TensorCore)	Número de chips	Topologia
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Topologia personalizada: é necessário usar as flags --type e --topology.	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
Topologia personalizada: é necessário usar as flags --type e --topology.	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…