Arquitetura do sistema
As Unidades de Processamento de Tensor (TPUs) são aceleradores de ML desenvolvidos pelo Google. A Cloud TPU as disponibiliza como um recurso escalonável do Google Cloud. É possível executar cargas de trabalho de machine learning em Cloud TPUs usando frameworks de machine learning, como TensorFlow, Pytorch e JAX (links em inglês).
Cada TensorCore em um chip de TPU consiste em uma ou mais unidades de multiplicação de matriz (MXUs, na sigla em inglês), uma unidade vetorial e uma unidade escalar.
Uma MXU é composta de multiplicação/acumuladores 128 x 128 em uma matriz sistólica. As MXUs fornecem a maior parte da potência de computação em um TensorCore. Cada MXU é capaz de realizar 16 mil operações de multiplicação e acumulação em cada ciclo. Todas as multiplicações recebem entradas bfloat16, mas todas as acumulações são realizadas no formato de número FP32.
A unidade vetorial é usada para computação geral, como ativações e softmax. A unidade escalar é usada para controlar o fluxo, calcular o endereço de memória e outras operações de manutenção.
Versões de TPU
O layout exato de uma TPU depende da versão de TPU usada. Detalhes de arquitetura e características de desempenho da TPU v2 e v3 estão disponíveis em Um supercomputador específico para o domínio para treinar redes neurais profundas.
TPU v4
Cada chip de TPU v4 contém dois TensorCores. Cada TensorCore tem quatro MXUs, uma unidade vetorial e uma unidade escalar. A tabela a seguir mostra as principais especificações de um pod de TPU v4.
| Principais especificações | Valores do Pod v4 |
|---|---|
| Computação máxima por chip | 275 teraflops (bf16 ou int8) |
| HBM2 e largura de banda | 32 GiB, 1200 GBps |
| Potência mínima/média/máxima medida | 90/170/192 O |
| Tamanho do Pod de TPU | 4.096 chips |
| Topologia de interconexão | Toro 3D |
| Computação de pico 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 secção por pod | 24 TB/s |
O diagrama a seguir ilustra um chip TPU v4.
TPU v3
Cada chip de TPU v3 contém dois TensorCores. Cada TensorCore tem dois MXUs, uma unidade vetorial e uma unidade escalar. A tabela a seguir mostra as principais especificações e os respectivos valores para um pod de TPU v3.
| Principais especificações | Valores do Pod v3 |
|---|---|
| Computação máxima por chip | 123 teraflops (bf16) |
| HBM2 e largura de banda | 32 GiB, 900 GBps |
| Potência mínima/média/máxima medida | 123/220/262 O |
| Tamanho do Pod de TPU | 1.024 chips |
| Topologia de interconexão | Toro 2D |
| Computação de pico por pod | 126 petaflops (bf16) |
| Largura de banda de redução total por pod | 340 TB/s |
| Largura de banda de secção por pod | 6,4 TB/s |
O diagrama a seguir ilustra um chip TPU v3.
TPU v2
A menor configuração de TPU v2 contém quatro chips de TPU e 16 GiB de HBM. Cada chip de TPU contém dois TensorCores. Cada TensorCore tem uma MXU, uma unidade vetorial e uma unidade escalar. O diagrama a seguir ilustra um chip TPU v2.
O Cloud TPU oferece as seguintes configurações de TPU:
- Um único dispositivo de TPU
- Um pod de TPU: um grupo de dispositivos de TPU conectados por interconexões de alta velocidade
- Uma fração da TPU: uma subdivisão de um pod de TPU
Benefícios do desempenho da TPU v4 em relação à v3
Sistema de memória:
O acesso à memória não uniforme (NUMA, na sigla em inglês) é uma arquitetura de memória para máquinas com várias CPUs. Cada CPU tem acesso direto a um bloco de memória de alta velocidade. Uma CPU e sua memória são chamadas de nós NUMA. Os nós NUMA estão conectados a nós NUMA diretamente adjacentes. 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 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 memória. Para mais informações sobre NUMA, consulte Acesso não uniforme à memória na Wikipédia.
O ambiente de execução da TPU v4 não reconhece nativamente NUMA. No entanto, você pode aproveitar os benefícios da localidade NUMA vinculando seu script de treinamento ao nó 0 do NUMA.
Para ativar a vinculação de nó NUMA:
Instale a ferramenta de linha de comando numactl.
$ $ sudo apt-get update $ $ sudo apt-get install numactl
Use
numactl --cpunodebind=0ao iniciar o script de treinamento. Isso vincula seu código de script ao NUMA nó 0.$ $ numactl --cpunodebind=0 python3 your-training-script
Ative a vinculação de nós NUMA se:
- Se a carga de trabalho tiver uma forte dependência de cargas de trabalho de CPU (por exemplo, classificação de imagens, cargas de trabalho de recomendação), independentemente do framework.
- Se você estiver usando uma versão de 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 de TPU v4 têm um espaço de memória unificado de HBM de 32 GiB em todo o chip, possibilitando uma melhor coordenação entre os dois TensorCores no chip.
- Melhor desempenho de HBM usando padrões e velocidades de memória mais recentes.
- Perfil de desempenho de DMA aprimorado com suporte integrado para etapas de alto desempenho em granularidades de 512 B.
TensorCores:
- Duas vezes o número de MXUs e uma taxa de clock mais alta, entregando no máximo 275 TFLOPS.
- Largura de banda de transposição e permutação de 2x.
- Modelo de acesso à memória de armazenamento de carga para memória comum (Cmem).
- Largura de banda de carregamento do 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 ativar topologias de rede que têm diâmetros de rede menores.
Outro:
- Interface PCI x16 gen3 para host (conexão direta).
- Modelo de segurança aprimorado.
- Melhoria na eficiência energética.
Benefícios de desempenho da TPU v3 em relação à v2
O aumento de FLOPS por TensorCore e da capacidade de memória nas configurações da TPU v3 pode melhorar o desempenho dos seus modelos das seguintes maneiras:
As configurações de TPU v3 oferecem benefícios significativos de desempenho por TensorCore para modelos vinculados a computação. Os modelos vinculados à memória nas configurações da TPU v2 podem não ter essa mesma melhoria de desempenho se também estiverem vinculados à memória nas configurações da TPU v3.
Nos casos em que os dados não cabem na memória nas configurações da TPU v2, a TPU v3 pode fornecer melhor desempenho e redução de recomputação de valores intermediários (rematerialização).
As configurações da TPU v3 executam novos modelos com tamanhos de lotes que não cabem nas configurações da TPU v2. Por exemplo, a TPU v3 pode permitir ResNets mais profundos e imagens maiores com RetinaNet.
Os modelos que são quase limitados à entrada ("alimentação") na TPU v2, porque as etapas de treinamento estão aguardando pela entrada, também podem ficar limitados à entrada no Cloud TPU v3. O guia de desempenho do pipeline pode ajudar a resolver problemas de alimentação.
Configurações da TPU v4
Um pod da TPU v4 é composto de 4.096 chips interconectados com links reconfiguráveis de alta velocidade. A rede flexível da TPU v4 permite que você conecte os chips em uma fatia do mesmo tamanho de várias maneiras. Ao criar uma fatia de TPU, especifique a versão e o número de recursos de TPU necessários. Ao criar uma fatia da TPU v4, é possível especificar o tipo e o tamanho dela de duas maneiras:
Como usar o
AcceleratorTypeIsso é recomendado quando você não especifica nenhuma topologia.
Como usar o
AcceleratorConfigRecomendamos o uso de
AcceleratorConfigquando você quiser personalizar a topologia física. Geralmente, isso é necessário para o ajuste do desempenho com fatias maiores que 256 chips.
Para configurar TPUs v4 usando AcceleratorType, use a sinalização --accelerator-type ao criar a fatia 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 fatia v4 com 32 TensorCores, use --accelerator-type=v4-32.
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 da TPU na fatia. Por exemplo, para criar uma fatia v4 com 16 chips, use --version=v4 e --topology=2x2x4.
A topologia da TPU v4 é especificada com três tuplas, que descreve o número, a forma e as interconexões entre os chips de TPU. As ilustrações a seguir mostram algumas topologias comuns da TPU v4.
Fatias maiores podem ser criadas a partir de um ou mais "cubos" 4 x 4 x 4 de chips.
Fatias de TPU de um determinado número de chips podem ser configuradas de maneiras diferentes. Por exemplo, uma fatia de TPU com chips 512 (1024 TensorCores) pode ser configurada usando as seguintes topologias: 4x4x32, 4x8x16 ou 8x8x8.
Uma fatia da TPU com 2048 chips (4096 TensorCores) oferece ainda mais opções de topologia: 4x4x128, 4x8x64, 4x16x32 e 8x16x16. Para mais informações, consulte Tipos e topologias de TPU.
Configurações da TPU v2 e v3
As TPUs estão disponíveis nas seguintes configurações:
- Uma única placa de TPU
- Uma fração do pod de TPU
Única placa de TPU
Uma configuração de TPU de placa única é uma placa independente com quatro chips de TPU (oito TensorCores) sem conexões de rede com outras placas de TPU. As TPUs de placa única não fazem parte de uma configuração de pod de TPU e não ocupam uma parte de um pod de TPU.
Fatias de Pods do TPU
Em uma fração do Pod de TPU, os chips do TPU são conectados usando uma interconexão de alta velocidade. Cada chip de TPU se comunica diretamente com os outros chips no dispositivo. O software da TPU gerencia automaticamente a distribuição de dados para cada TensorCore em uma fração de pod. As fatias do Pod da TPU v3 estão disponíveis com TensorCores de 32, 128, 512, 1024 ou 2048. As fatias do Pod da TPU v2 estão disponíveis com TensorCores de 32, 128, 256 ou 512.
Arquiteturas de VM do Cloud TPU
As TPUs são projetadas para executar operações de matriz rapidamente. Cada placa de TPU é conectada a uma máquina host com base em CPU para realizar operações que não podem ser executadas na TPU. As máquinas host são responsáveis por carregar dados do Cloud Storage, pré-processar e enviar dados para a TPU.
Em um pod de TPU, há um host de TPU para cada placa de TPU.
A interação com o host e a placa da TPU depende da arquitetura da VM da TPU usada: nós ou VMs da TPU.
Nós da TPU
A arquitetura do nó da TPU consiste em uma VM de usuário que se comunica com o host da TPU por meio do gRPC. Ao usar essa arquitetura, não é possível acessar diretamente o host da TPU, o que dificulta a depuração de erros de treinamento e TPU.
VMs de TPU
A arquitetura de VM da TPU permite que você se conecte via SSH à VM fisicamente conectada ao dispositivo da TPU. Você tem acesso raiz à VM, então pode executar um código arbitrário. É possível acessar registros e mensagens de erro do compilador e do ambiente de execução.
Frameworks compatíveis
Frameworks como JAX, PyTorch e TensorFlow acessam TPUs usando uma biblioteca compartilhada chamada libtpu, presente em todas as VMs de TPU. libtpu inclui o compilador XLA, o software de ambiente de execução da TPU e o driver da TPU.
TensorFlow
VMs de TPU
Com VMs de TPU, em vez do código Python ser executado em uma VM de usuário, ele pode ser executado diretamente no host da TPU.
Para mais informações sobre como executar o TensorFlow na Cloud TPU, consulte o guia de início rápido do TensorFlow.
Nós da TPU
A arquitetura do sistema do nó do Cloud TPU foi criada originalmente para o TensorFlow. Os hosts da TPU ficam inacessíveis ao usuário e executam uma cópia headless do servidor do TensorFlow. Eles não executam o Python ou qualquer código de usuário não representado como um gráfico do TensorFlow. O código do usuário é executado em uma VM separada, que se comunica com os hosts da TPU na rede gRPC.
PyTorch
VMs de TPU
Com as VMs da TPU, o código do PyTorch é executado diretamente nos hosts da TPU.
Para mais informações sobre o PyTorch e o Cloud TPU, consulte o guia do usuário do PyTorch/XLA.
Nós da TPU
O PyTorch é executado na arquitetura de nó do Cloud TPU usando uma biblioteca chamada XRT. O XRT envia gráficos XLA e instruções de ambiente de execução pelo gRPC para serem executados nos servidores do TensorFlow. Uma VM de usuário é necessária para cada host de TPU.
JAX
VMs de TPU
Com as VMs da TPU, seu código JAX é executado diretamente nos hosts da TPU.
Para mais informações sobre como executar o JAX no Cloud TPU, consulte o Guia de início rápido do JAX.
Nós da TPU
O JAX em nós do Cloud TPU é executado de maneira semelhante ao PyTorch, já que uma VM de usuário separada é necessária para cada VM de host.
