Arquitetura do sistema

Neste documento, descrevemos a arquitetura dos componentes de hardware e software do sistema da Cloud TPU.

As Unidades de Processamento de Tensor (TPUs, na sigla em inglês) são circuitos integrados de aplicativos específicos (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Eles são usados para acelerar as cargas de trabalho de machine learning (ML). As TPUs foram projetadas do zero com o benefício da ampla experiência e liderança do Google em machine learning.

Uma única placa de TPU contém quatro chips TPU. Cada chip contém dois núcleos de TPU. Os recursos disponíveis em um núcleo da TPU variam de acordo com a versão. Cada núcleo de TPU tem unidades escalares, vetoriais e multiplicadores de matriz (MXUs, na sigla em inglês). 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 de multiplicação em cada ciclo com precisão bfloat16 reduzida. bfloat16 é uma representação de ponto flutuante de 16 bits que oferece melhor treinamento e precisão do modelo que a representação de meia precisão IEEE. Cada um dos núcleos em uma placa de TPU pode executar cálculos de usuários de maneira independente. As interconexões de alta largura de banda permitem que os chips se comuniquem diretamente entre si. Veja informações técnicas mais detalhadas em Um supercomputador específico para domínios para treinamento de redes neurais profundas.

As TPUs foram projetadas para escalonar horizontalmente em um pod de TPU. Um pod de TPU é um supercomputador que pode ter até 2.048 núcleos de TPU, permitindo distribuir a carga de processamento por várias placas de TPU. Em uma configuração de pod da TPU, as interfaces de rede dedicadas de alta velocidade conectam vários dispositivos TPU para fornecer um número maior de núcleos de TPU e um pool maior de memória da TPU para seu cargas de trabalho de machine learning.

Cada placa TPU é conectada a uma máquina host com base em CPU de alto desempenho para, por exemplo, carregar e pré-processar dados para alimentar as TPUs.

Cloud TPU

O Cloud TPU é um serviço que concede acesso às TPUs pelo Google Cloud Platform (GCP). É possível usar o Cloud TPU para executar cargas de trabalho de machine learning no hardware do acelerador de TPU do Google. A Cloud TPU foi criada para ter o máximo de desempenho e flexibilidade, ajudando pesquisadores, desenvolvedores e empresas a treinar cargas de trabalho de ML.

Use a API Cloud TPU para automatizar o gerenciamento de TPUs para suas TPUs. Dessa maneira, é fácil escalonar verticalmente verticalmente em clusters de computação grandes, executar suas cargas de trabalho e reduzi-las quando as cargas de trabalho estiverem concluídas. A compatibilidade de hardware integrada aos chips resulta em um escalonamento de desempenho efetivamente linear em uma ampla variedade de cargas de trabalho de aprendizado profundo.

Configurações de TPU

O Cloud TPU fornece acesso a diferentes configurações de TPU:

  • Um pod de TPU
  • Um fatia de TPU
  • Uma placa de TPU

Pods de TPU

Em um pod de TPU, todos os chips de TPU em um pod de TPU são conectados diretamente entre si por uma interconexão de alta velocidade que ignora os atrasos de comunicação pelos hosts da CPU. Os chips são conectados em um toro 2-D, com cada chip se comunicando diretamente com quatro vizinhos. Essa arquitetura proporciona um desempenho especialmente excelente para padrões comuns de comunicação em cargas de trabalho de ML, como redução total. Os detalhes de arquitetura e as características de desempenho da TPU v2 e v3 são mostrados na tabela a seguir, com mais informações publicadas em Um supercomputador específico para domínios para treinamento de redes neurais profundas (em inglês).

Recurso TPU v2 TPU v3
Links de rede x Gbits/s/Chip 4 x 496 4 x 656
Pod com terabits/largura total de largura de banda de bisection 15.9 42

Fatias da TPU

Um Slice de TPU é uma parte de um pod de TPU. Se você não precisar dos recursos de um pod inteiro, poderá usar uma parte deles. Há várias configurações de Slices diferentes disponíveis para mais informações. Consulte a seção de preços do Cloud TPU no Cloud TPU.

Placa única 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. Leia a página de tipos de TPU para ver quais configurações de TPU de placa única estão disponíveis.

Versões de TPU

A versão da TPU define a arquitetura de cada núcleo da TPU, a quantidade de memória de alta largura de banda (HBM, na sigla em inglês) de cada núcleo da TPU. As interconexões entre os núcleos em cada placa TPU e as interfaces de rede disponíveis para }comunicado entre dispositivos. As versões disponíveis da TPU são v2 e v3.

TPU v2

Uma placa TPU v2 contém quatro chips TPU com dois núcleos. Há 8 GiB de HBM para cada núcleo de TPU e uma MXU por núcleo. Um pod TPU v2 tem até 512 núcleos de TPU e 4 TiB de memória.

imagem

TPU v3.

Uma placa TPU v3 contém quatro chips TPU com dois núcleos. Há 16 GiB de HBM para cada núcleo de TPU e dois MXUs para cada núcleo. Um pod TPU v3 tem até 2.048 núcleos TPU e 32 TiB de memória.

imagem

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 vinculados à memória nas configurações da TPU v2 podem não alcançar 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 melhorar o desempenho e reduzir o 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 permite ResNets mais profundos e imagens maiores com o RetinaNet.

Os modelos que estão quase limitados pela entrada ("alimentação") na TPU v2 porque as etapas de treinamento estão aguardando a entrada também podem ser limitadas pela Cloud TPU v3. O guia de desempenho do pipeline pode ajudar você a resolver problemas de alimentação.

Arquiteturas de VM do Cloud TPU

Cada placa de TPU está fisicamente conectada a uma máquina host (host de 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 da TPU (e a placa da TPU) depende da arquitetura de VM da TPU usada: nós da TPU ou VMs da TPU.

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 pelo gRPC. não há acesso direto ao host da TPU.

imagem

VMs de TPU

Ao usar VMs de TPU, você faz SSH diretamente para uma VM do Google Compute Engine em execução no host da TPU. Você recebe acesso raiz à máquina para executar qualquer código que quiser. Você tem acesso a registros de depuração e mensagens de erro diretamente pelo compilador e ambiente de execução da TPU. As VMs de TPU são compatíveis com novos casos de uso que não são possíveis com os nós de TPU. Por exemplo, é possível executar operações personalizadas no pipeline de entrada e usar o armazenamento local.

imagem

Como não há VM do usuário, não é necessário ter uma rede, uma nuvem privada virtual ou um firewall entre o código e a VM da TPU, o que melhora o desempenho do pipeline de entrada. Além disso, as VMs de TPU são mais baratas porque você não precisa pagar pelas VMs de usuário.

Frameworks compatíveis

Frameworks como JAX, PyTorch e TensorFlow acessam TPUs por meio de uma biblioteca compartilhada chamada libtpu, que está presente em todas as VMs de TPU. Essa biblioteca inclui o compilador XLA usado para compilar programas TPU, o ambiente de execução do TPU usado para executar programas compilados e o driver 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 em execução 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 Cloud TPU Node foi criada originalmente para o TensorFlow e para o modelo de programação distribuído. Os hosts da TPU ficam inacessíveis ao usuário e executam apenas uma cópia headless de um 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 de TPU pela rede.

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 tempo de execução por meio de conexões 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