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 de nuvem escalonável do GCP. É possível executar cargas de trabalho de machine learning em Cloud TPUs usando frameworks de machine learning, como TensorFlow, Pytorch e JAX.

A unidade de multiplicação de matriz (MXU, na sigla em inglês) é composta de multiplicação/acumuladores 128 x 128 em uma matriz sistólica. As MXUs fornecem a maior parte da capacidade de computação em um TensorCore (TPU Core). Cada MXU é capaz de executar 16 mil operações de multiplicação seguida de acumulação em cada ciclo usando o formato de número bfloat16 (em inglês).

A VPU é usada para computação geral, como ativações, softmax etc. 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 depende da versão da TPU usada. Detalhes de arquitetura e características de desempenho da TPU v2 e v3 estão disponíveis em Um supercomputador específico do domínio para treinar redes neurais profundas.

TPU v4

A menor configuração da TPU v4 contém quatro chips da TPU e 32 GiB de HBM e 128 MiB de memória comum compartilhada. Cada chip de TPU contém dois TensorCores. Cada TensorCore tem quatro MXUs, uma unidade vetorial e uma unidade escalar. O diagrama a seguir ilustra um ícone da TPU v4.

imagem

TPU v3

A menor configuração da TPU v3 contém quatro chips da TPU e 32 GiB de HBM. Cada chip de TPU contém dois TensorCores. Cada TensorCore tem duas MXUs, uma unidade vetorial e uma unidade escalar. O diagrama a seguir ilustra um chip TPU v3.

imagem

TPU v2

A menor configuração da TPU v2 contém quatro chips da TPU e 16 GiB de HBM. Cada chip de TPU contém dois TensorCores. Cada TensorCore tem um MXU, uma unidade vetorial e uma unidade escalar. O diagrama a seguir ilustra um chip TPU v2.

imagem

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 de desempenho da TPU v4 em relação à v3

Sistema de memória:

  • Os chips de TPU v4 têm um espaço de memória HBM de 32 GiB unificado em todo o chip, permitindo uma melhor coordenação entre os dois TensorCores no chip.
  • Melhor desempenho de HBM usando padrões e velocidades de memória mais recentes.
  • Melhoramos o perfil de desempenho de DMA com compatibilidade nativa com passos de alto desempenho em 512 bilhões de granularidades.

TensorCores:

  • O dobro do número de MXUs e uma taxa de relógio mais alta, entregando no máximo 275 TFLOPS.
  • Largura de banda de transposição e permutação duas vezes maior.
  • 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 compatibilidade com o 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 com diâmetros de rede menores.

Outro:

  • Interface x16 PCIe gen3 para hospedar (conectar diretamente).
  • Modelo de segurança aprimorado.
  • Maior eficiência energética.

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

O aumento de FLOPS por núcleo e capacidade de memória nas configurações da TPU v3 melhora o desempenho dos seus modelos das seguintes maneiras:

  • As configurações da TPU v3 fornecem benefícios de desempenho significativos por núcleo para modelos vinculados à computação. Os modelos limitados à memória nas configurações da TPU v2 talvez não alcancem essa mesma melhoria de desempenho se eles também estiverem limitados à 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 melhorar o desempenho e reduzir a necessidade de recalcular 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 você a resolver problemas de alimentação.

Configurações da TPU v4

Um pod da TPU v4 é composto por 4.096 chips interconectados com links de alta velocidade reconfiguráveis. Com a rede flexível da TPU v4, é possível conectar os chips em uma fatia do mesmo tamanho de várias maneiras. Configure TPUs v4 usando a AcceleratorConfig. Com o AcceleratorConfig, é possível especificar o tamanho da fração em termos de chips por meio do campo de topologia.

A topologia de TPU v4 é especificada usando uma tupla de três tuplas que descreve o número, o formato e as interconexões entre os chips de TPU. As ilustrações a seguir mostram algumas topologias comuns da TPU v4.

imagem

Fatias maiores podem ser criadas com um ou mais cubos 4x4 x 4.

Fatias de TPU de um determinado número de chips podem ser configuradas de diferentes maneiras. Por exemplo, uma fração de TPU com o AcceleratorType v4-1024 agora pode ser configurada como: 4x4x32, 4x8x16, 8x8x8. Uma TPU de v4-4096 oferece ainda mais flexibilidade: 4x4x128, 4x8x64, 4x16x32, 8x16x16, 8x16x16_torcido. 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:

  1. Uma única placa de TPU
  2. Um pod de TPU
  3. 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. TPUs de placa única não fazem parte de uma configuração do pod de TPU e não ocupam uma parte dele.

Pods e frações de TPU

Em uma fração de pod de TPU ou em um pod de TPU, os chips de TPU são conectados por meio de uma interconexão de alta velocidade, e cada chip de TPU se comunica diretamente com os outros chips no dispositivo de TPU. O software da TPU administra automaticamente a distribuição de dados para cada núcleo de TPU em um pod ou fatia. As frações de pod estão disponíveis com 32, 128, 512, 1.024 ou 2.048 núcleos.

Arquiteturas de VM do Cloud TPU

As TPUs só podem executar operações de matriz. Portanto, cada placa de TPU é conectada a uma máquina host baseada em CPU para executar 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 dados e enviar dados à TPU.

imagem

Em um pod de TPU, há um host de TPU para cada placa de TPU.

imagem

A forma como você interage com o host de TPU (e a placa de TPU) depende da arquitetura de VM da TPU usada: nós ou VMs.

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. Isso pode dificultar o treinamento de erros de depuração e da TPU.

imagem

VMs de TPU

A arquitetura da VM de TPU elimina a necessidade de usar a VM do usuário, e é possível executar o SSH diretamente na VM conectada fisicamente ao dispositivo de TPU. Você tem acesso raiz à VM para poder executar código arbitrário. Você pode acessar os registros de depuração do compilador e do ambiente de execução e as mensagens de erro.

imagem

Frameworks compatíveis

Frameworks como JAX, PyTorch e TensorFlow acessam TPUs por uma biblioteca compartilhada chamada libtpu, presente em todas as VMs de TPU. Essa biblioteca inclui o compilador XLA usado para compilar programas de TPU, o software de TPU usado para executar programas compilados e o driver de TPU usado pelo ambiente de execução para acesso de baixo nível à TPU.

imagem

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.

imagem

Para mais informações sobre o TensorFlow e o Cloud TPU, consulte Como executar modelos do TensorFlow no Cloud TPU.

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 de usuário separada que se comunica com os hosts de TPU pela rede gRPC.

imagem

PyTorch

VMs de TPU

Com as VMs de TPU, seu código PyTorch é executado diretamente nos hosts da TPU.

images

Para mais informações sobre o PyTorch e o Cloud TPU, consulte Como executar modelos PyTorch no Cloud TPU.

Nós da TPU

O PyTorch é executado na arquitetura de nó do Cloud TPU usando uma biblioteca chamada XRT, que permite enviar gráficos XLA e instruções de ambiente de execução por gRPC do TensorFlow e executá-los nos servidores do TensorFlow. Uma VM de usuário é necessária para cada host de TPU.

imagem

JAX

VMs de TPU

Com as VMs de TPU, não é necessário ter VMs de usuário porque é possível executar o código diretamente nos hosts de TPU.

imagem

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.

imagem

A seguir