Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.

Arquitetura do sistema

As Unidades de Processamento de Tensor (TPUs) são aceleradores de ML desenvolvidos pelo Google. A Cloud TPU disponibiliza as TPUs 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 matrizes (MXU, na sigla em inglês) é composta de multiplicação/acumuladores de 128 x 128 em uma matriz sistólica. As MXUs fornecem a maior parte da capacidade de computação em um TensorCore (núcleo 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 uma TPU depende da versão da TPU que você usa. 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 treinamento de redes neurais profundas.

TPU v4

Cada chip de v4 TPU 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 e os valores deles para um pod de TPU v4.

Principais especificações Valores de pod v4
Pico de computação por chip 275 teraflops (bf16 ou int8)
Capacidade e largura de banda HBM2 32 GiB, 1.200 GB/s
Potência mínima/média/máxima medida 90/170/192 W
Tamanho do pod da TPU 4.096 ícones
Topologia de interconexão Torus 3D
Pico de computação por pod 1.1 exaflops (bf16 ou int8)
Largura de banda com redução total por pod 1,1 PB/s
Largura de banda bisseccional por pod 24 TB/s

O diagrama a seguir ilustra um chip TPU v4.

image

TPU v3

Cada chip de TPU v3 contém dois TensorCores. Cada TensorCore tem duas MXUs, uma unidade vetorial e uma unidade escalar. A tabela a seguir mostra as principais especificações e os valores deles para um pod de TPU v3.

Principais especificações Valores de pod v3
Pico de computação por chip 123 teraflops (bf16)
Capacidade e largura de banda HBM2 32 GiB, 900 GB/s
Potência mínima/média/máxima medida 123/220/262 W
Tamanho do pod da TPU 1.024 ícones
Topologia de interconexão Torus 2D
Pico de computação por pod 126 petaflops (bf16)
Largura de banda com redução total por pod 340 TB/s
Largura de banda bisseccional por pod 6,4 TB/s

O diagrama a seguir ilustra um chip TPU v3.

image

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 ícone da TPU v2.

image

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:

  • Os chips de TPU v4 têm um espaço de memória unificado de HBM de 32 GiB 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.
  • Perfil de desempenho de DMA aprimorado com suporte nativo para desempenho com alto desempenho, com granularidade de 512 bilhões.

TensorCores:

  • Duas vezes o 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 2x.
  • 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 suporte ao modo de 8 bits para permitir tamanhos de lote mais baixos e melhoria na latência de inferência.

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 " PCIe gen3 gen3" para hospedar (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 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 de 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 fração do mesmo tamanho de várias maneiras. Configure TPUs v4 usando AcceleratorConfig. O AcceleratorConfig permite especificar o tamanho da fatia em termos de ícones no campo topology.

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

image

Fatias maiores podem ser criadas com um ou mais ícones de 4x4 x 4 e cubos.

Fatias de TPU de um determinado número de ícones podem ser configuradas de várias maneiras. Por exemplo, uma fração de TPU com o AcceleratorType da v4-1024 agora pode ser configurada como 4x4x32, 4x8x16 ou 8x8x8.

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 (8 TensorCores) sem conexão 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 processa automaticamente a distribuição de dados para cada núcleo de TPU em um pod ou uma 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.

image

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 de 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 a depuração de erros de treinamento e TPU.

image

VMs de TPU

A arquitetura da VM da TPU elimina a necessidade de uma VM do usuário, e é possível estabelecer uma conexão SSH diretamente com a VM que está conectada fisicamente ao dispositivo da TPU. Você tem acesso raiz à VM, então executa um código arbitrário. Você pode acessar mensagens de erro e registros de depuração do compilador e do ambiente de execução.

image

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.

image

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.

image

Para mais informações sobre o TensorFlow e a VM da TPU do Cloud TPU, consulte o Guia do usuário da VM do 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.

image

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 a 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, 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.

image

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.

image

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.

image

A seguir