Arquitetura do sistema

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

As unidades de processamento de tensor (TPUs, na sigla em inglês) são circuitos integrados de aplicação específica (ASICs, na sigla em inglês) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning (ML). As TPUs foram desenvolvidas desde o início aproveitando a extensa experiência e liderança do Google em machine learning.

Um único painel de TPU contém quatro chips de TPU. Cada chip contém dois núcleos de TPU. Os recursos disponíveis em um núcleo de TPU variam de acordo com a versão. Cada núcleo de TPU tem unidades escalares escalares, vetoriais e de multiplicação de matrizes (MXUs). Os MXUs fornecem a maior parte da potência de computação em um chip de TPU. Cada MXU é capaz de realizar 16 mil operações de multiplicação seguida de acumulação em cada ciclo com precisão de bfloat16. bfloat16 é uma representação de ponto flutuante de 16 bits que proporciona melhor treinamento e acurácia do modelo do 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. Informações técnicas mais detalhadas podem ser encontradas em Um supercomputador específico do domínio para treinamento de redes neurais profundas.

As TPUs foram projetadas para serem escalonadas para um pod de TPU. Um pod de TPU é um supercomputador que pode ter até 2.048 núcleos de TPU, o que permite distribuir a carga de processamento por várias placas de TPU. Em uma configuração de Pod de TPU, 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 de TPU para o cargas de trabalho de machine learning.

Cada placa de TPU está conectada a uma máquina host de alto desempenho baseada em CPU para itens como carregar e pré-processar dados para alimentar as TPUs.

Cloud TPU

A Cloud TPU é um serviço que fornece acesso a TPUs por meio do Google Cloud Platform (GCP). É possível usar o Cloud TPU para executar cargas de trabalho de machine learning no hardware acelerador de TPUs do Google. O Cloud TPU foi desenvolvido para proporcionar 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 TPU para as TPUs. Como resultado, é fácil expandir para clusters de computação massivos, executar as cargas de trabalho e dimensionar esses clusters quando as cargas de trabalho estiverem completas. O suporte de hardware incorporado nos chips resulta em escalonamento de desempenho efetivamente linear em uma ampla variedade de cargas de trabalho de aprendizado profundo.

Configurações de TPU

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

  • Um pod de TPU
  • Um slice de TPU
  • Uma placa da TPU

Pods de TPU

Em um Pod de TPU, todos os chips de TPU em um Pod de TPU são conectados diretamente uns aos outros em uma interconexão de alta velocidade que pula os atrasos de comunicação por meio dos hosts de CPU. Os chips são conectados em uma névoa de 2D, com cada chip se comunicando diretamente com quatro vizinhos. Essa arquitetura gera um desempenho especialmente alto para padrões de comunicação comuns em cargas de trabalho de ML, como o tamanho todo. Os detalhes da arquitetura e as características de desempenho do TPU v2 e v3 são mostrados na tabela a seguir, com mais informações publicadas em Um supercomputador específico do domínio para redes neurais profundas.

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

Slices do 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 de um pod. Há várias configurações de Slices disponíveis para mais informações, consulte a seção Preços do Cloud TPU no Cloud TPU.

Quadro de TPU único

Uma configuração de TPU de placa única é uma placa independente com quatro chips (oito núcleos da TPU) 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. 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 de TPU define a arquitetura para cada núcleo de TPU, a quantidade de memória de alta largura de banda (HBM, na sigla em inglês) para cada núcleo de TPU, as interconexões entre os núcleos em cada placa de TPU e as interfaces de rede disponíveis para {101 }comunicação entre dispositivos. As versões de TPU disponíveis são: v1 e v2.

TPU v2

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

imagem

TPU v3

Uma placa da TPU v3 contém quatro chips da TPU, cada um com dois núcleos. Há 16 GiB de HBM para cada núcleo de TPU e dois MXUs para cada núcleo. Um pod do TPU v3 tem até 2.048 núcleos de 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 ligados à memória nas configurações da TPU v2 podem não alcançar essa mesma melhoria de desempenho se eles também estiverem ligados à 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 o RetinaNet.

Os modelos que foram quase limitados pela entrada ("alimentação") na TPU v2 (porque os passos de treinamento estão aguardando pela entrada) podem ficar limitados pela entrada no Cloud TPU v3. O guia de desempenho do pipeline pode ajudar a resolver problemas de alimentação.

Arquitetura de VMs do Cloud TPU

Cada placa de TPU está fisicamente conectada a uma máquina host (host da TPU).

imagem

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

imagem

A maneira como você interage com o host de TPU (e a placa de TPU) depende da arquitetura de VM da TPU que você está usando: nós da TPU ou VMs de TPU.

Nós da TPU

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

imagem

VMs da TPU

Ao usar VMs de TPU, você se conecta SSH diretamente a uma VM do Google Compute Engine em execução no host da TPU. Você tem acesso raiz à máquina para executar qualquer código que quiser. Você tem acesso a registros de depuração e mensagens de erro diretamente do compilador e do 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 armazenamento local.

imagem

Como não há uma VM do usuário, não é necessário ter uma rede, uma nuvem privada virtual ou um firewall entre seu código e a VM de 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 por 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 cada VM da TPU. Essa biblioteca inclui o compilador XLA usado para compilar programas de TPU, o ambiente de execução da TPU usado para executar programas compilados e o driver da TPU usado pelo ambiente de execução para acesso de baixo nível à TPU.

imagem

TensorFlow

VMs da TPU

Com VMs de TPU, em vez de seu 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 do Cloud TPU Node foi criada originalmente para o TensorFlow e o modelo de programação distribuído. Os hosts de TPU ficam inacessíveis para o usuário e executam apenas uma cópia sem cabeçalho de um servidor do TensorFlow. Elas não executam 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 na rede.

imagem

PyTorch

VMs da TPU

Com VMs de TPU, o 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 do 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 do XLA e instruções de tempo de execução sobre conexões gRPC do TensorFlow e executá-las nos servidores do TensorFlow. É necessário ter uma VM de usuário para cada host de TPU.

imagem

JAX

VMs da TPU

Com as VMs de TPU, não é necessário ter VMs de usuário, porque você executa seu 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

Os JAX em nós do Cloud TPU são executados de maneira semelhante ao PyTorch, em que uma VM de usuário separada é necessária para cada VM do host.

imagem

A seguir