Arquitetura do sistema

As Unidades de Processamento de Tensor (TPUs) são aceleradores de ML desenvolvidos pelo Google. O Cloud TPU disponibiliza as TPUs como um recurso escalonável de nuvem do GPC. É possível executar cargas de trabalho de machine learning em Cloud TPUs usando frameworks de machine learning, como TensorFlow, Pytorch e JAX.

Um único dispositivo de TPU contém quatro chips, cada um com dois núcleos de TPU. Um núcleo de TPU contém uma ou mais unidades de multiplicação de matriz (MXU, na sigla em inglês), uma unidade de processamento vetorial (VPU, na sigla em inglês) e uma unidade escalar.

A MXU é composta de multiplicadores/acumuladores de 128 x 128 em uma matriz sistólica. As MXUs fornecem a maior parte da capacidade de computação em um chip de TPU. 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 de um dispositivo de TPU depende da versão da TPU usada. Os detalhes de arquitetura e as características de desempenho da TPU v2 e v3 estão disponíveis em Um supercomputador específico para domínios para treinamento de redes neurais profundas.

TPU v2

Uma placa TPU v2 contém quatro chips TPU e 16 GiB de HBM. Cada chip de TPU contém dois núcleos. Cada núcleo tem uma unidade MXU, uma unidade vetorial e uma unidade escalar.

imagem

TPU v3

Uma placa TPU v3 contém quatro chips de TPU e 16 GiB de HBM. Cada chip de TPU contém dois núcleos. Cada núcleo tem uma unidade MXU, uma unidade vetorial e uma unidade escalar.

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 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 de TPU

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 autônoma com quatro chips de TPU (oito núcleos de TPU) 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 ambiente de execução da TPU gerencia automaticamente a distribuição de dados para cada núcleo de TPU em um pod ou fração. 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

Os nós da TPU são a experiência original da TPU. Elas requerem uma VM de usuário extra que se comunique com o host da TPU por gRPC. Não há acesso direto ao host da TPU.

imagem

VMs de TPU

Ao usar VMs de TPU, você efetua SSH diretamente com a VM do Google Compute Engine que está fisicamente conectada ao dispositivo de TPU. Você tem acesso raiz à VM para executar código arbitrário. É possível acessar registros de depuração do compilador e do ambiente de execução e 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 ambiente de execução 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 remota separada que se comunica com os hosts da 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