TPU v4

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

arquitetura do sistema.

Cada chip de 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.

Principais especificações	Valores do 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)
Reduzir totalmente a largura de banda por pod	1,1 PB/s
Largura de banda de bissecção por pod	24 TB/s

O diagrama a seguir ilustra um chip de TPU v4.

Malha 3D e toro 3D

As TPUs v4 têm uma conexão direta com os chips vizinhos mais próximos em 3 o que resulta em uma malha 3D de conexões de rede. As conexões pode ser configurado como um toro 3D em frações onde 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 um 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 alta velocidade memória. A CPU e a memória dela são chamadas de nó NUMA. NUMA nós estão conectados a outros NUMA nós diretamente adjacentes. 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.

Um software executado em uma máquina com várias CPUs pode colocar os dados necessários para uma CPU seu nó NUMA, aumentando a capacidade de processamento da 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 para o nó NUMA.

Para ativar a vinculação de nós do 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ó 0 do NUMA.

    $ 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.
• Melhor desempenho do HBM usando os padrões e as velocidades de memória mais recentes.
• Perfil de desempenho de DMA aprimorado com suporte integrado para alto desempenho com granularidades de 512B.

TensorCores

• O dobro do número de MXUs e uma taxa de clock maior, oferecendo um máximo de 275 TFLOPS.
• 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 menor aos diâmetros de rede.

Outro

• Interface x16 do PCIE gen3 para o 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 chaves reconfiguráveis links de alta velocidade. A rede flexível da TPU v4 permite conectar os chips 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. Quando você criar uma fração do Pod da TPU v4, será possível especificar o tipo e o tamanho dela de duas maneiras: AcceleratorType e AccleratorConfig.

Observação: AcceleratorType e AccleratorConfig são mutuamente exclusivos. é possível use apenas 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 o Fração do Pod de TPU. Defina --accelerator-type como uma string que contenha a TPU. versão e o número de TensorCores que você quer usar. Por exemplo, para criar um 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.17.0-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 o pod frações maiores que 256 ícones.

Para configurar TPUs v4 usando AcceleratorConfig, use --version e o --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.

Você especifica uma topologia de TPU usando uma tupla 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 um pod v4 com 16 ícones, você definiria --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

Topologias em que 2A=B=C ou 2A=2B=C também têm variantes de topologia otimizada para comunicação "tudo para todos", por exemplo, 4 × 4 × 8, 8 × 8 × 16 e 12 × 12 × 24. Elas são conhecidas como topologias de torri torcida.

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

Frações maiores de pods podem ser criadas a partir de um ou mais "cubos" 4x4x4 de chips.

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 no formato de um bloco fatia ou 4x4x8_twisted. As topologias torcidas oferecem bissecção significativamente maior largura de banda larga. O aumento da largura de banda com bissecção é útil para cargas de trabalho que usam e os padrões de comunicação. As topologias tortas podem melhorar o desempenho da maioria com grandes cargas de trabalho de incorporação de TPU que mais se beneficiam.

Para cargas de trabalho que usam o paralelismo de dados como a única estratégia de paralelismo, distorcidas pode ter um desempenho um pouco melhor. Para LLMs, o desempenho a topologia pode variar de acordo com o tipo de paralelismo (DP, MP etc.). Melhor a prática é treinar seu LLM com e sem uma topologia torcida para determinar o que oferece o melhor desempenho para seu 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 do Torus (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
• Rotas de pacotes mais curtas

Esses benefícios acabam se traduzindo em um melhor desempenho para muitos e os padrões de comunicação.

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. Em prática, cada borda é uma conexão bidirecional entre TPUs. Nos referimos ao as bordas entre um lado dessa grade e o lado oposto como bordas circundantes, como mostrado no diagrama.

Ao girar essa topologia, temos um toro 4x2 trançado completamente simétrico topologia:

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 o mesmo X coordenadas, elas foram alteradas para se conectar à TPU com coordenadas X+2 mod 4.

A mesma ideia se generaliza para diferentes tamanhos de dimensão e diferentes números de dimensões. 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 de como para especificar uma configuração de tori torcida ao criar um Cloud TPU.

A tabela a seguir mostra as topologias torcidas suportadas e um modelo aumento da largura de banda em bissecção com eles em comparação com topologias sem torção.

Topologia torcida	Aumento teórico na bissecção largura de banda em comparação com um toro não trançado
4×4×8_torcido	Aprox. 70%
8x8x16_trançado
12×12×24_torcido
4×8×8_torcido	Aprox. 40%
8×16×16_torcido

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

Variantes de topologia da TPU v4

Algumas topologias com o mesmo número de ícones podem ser organizadas de várias formas. 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. Esse formato provavelmente é a melhor escolha para treinamento de ML com paralelismo de dados. Outra opção as topologias de rede podem ser úteis para cargas de trabalho com vários tipos de paralelismo exemplo, paralelismo de modelo e dados ou particionamento de espaço de uma simulação). O desempenho dessas cargas de trabalho é melhor quando a topologia corresponde ao paralelismo usado. Por exemplo, colocar o paralelismo de modelos de 4 vias na dimensão X e nos dados de 256 vias o paralelismo 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. Elas geralmente são e modelar modelos de linguagem grande (LLMs) paralelos de modelos. 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. É melhor desempenho para usar o paralelismo de modelos de 8 ou 16 vias (associado a um dos e dimensões de topologia de TPU física). Um paralelismo de modelo de 4 vias não seria o ideal com essa topologia, já que ela não está alinhada com nenhuma das dimensões de topologia de TPU, mas seria ideal com uma topologia 4x16x32 com o mesmo número de chips.

As configurações da TPU v4 consistem em dois grupos, aqueles com topologias menores que 64 chips (topologias pequenas) e aqueles com topologias maiores que 64 chips (topologias grandes).

Topologias pequenas 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. Você pode criar essas pequenas topologias v4 usando Nome baseado no TensorCore (por exemplo, v4-32) ou topologia (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

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.

Nome (com base na contagem do TensorCore)	Número de ícones	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 as sinalizações --type e --topology.	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…