Arquitetura da TPU
As unidades de processamento de tensor (TPUs) são circuitos integrados de aplicativo específico (ASICs) desenvolvidos especialmente pelo Google. Elas são usadas para acelerar as cargas de trabalho de machine learning. O Cloud TPU é um serviço do Google Cloud que disponibiliza TPUs como um recurso escalonável.
As TPUs foram projetadas para realizar operações de matriz com rapidez, o que as torna ideais para cargas de trabalho de machine learning. É possível executar cargas de trabalho de machine learning em TPUs usando frameworks como Pytorch e JAX.
Como as TPUs funcionam?
Para entender como as TPUs funcionam, vale a pena entender como outros aceleradores lidam com os desafios computacionais do treinamento de modelos de ML.
Como funciona uma CPU
Ela é um processador de uso geral baseado na arquitetura de von Neumann. Isso significa que uma CPU trabalha com software e memória, desta forma:
O maior benefício das CPUs é a flexibilidade. É possível carregar qualquer tipo de software em uma CPU para vários tipos diferentes de aplicativos. Por exemplo, é possível usar uma CPU para processar textos em um PC, controlar mecanismos de foguetes, executar transações bancárias ou classificar imagens com uma rede neural.
Uma CPU carrega valores da memória, realiza um cálculo com base nos valores e armazena o resultado de cada cálculo na memória. O acesso à memória é lento quando comparado à velocidade do cálculo, e pode limitar a capacidade de processamento total das CPUs. Isso geralmente é chamado de gargalo de von Neumann.
Como funciona uma GPU
Para aumentar a capacidade de processamento, as GPUs contêm milhares de unidades lógicas aritméticas (ALUs) em um único processador. Em geral, uma GPU moderna contém entre 2.500 e 5.000 ALUs. Esse grande número de processadores significa que é possível executar milhares de multiplicações e adições simultaneamente.
Essa arquitetura de GPU funciona bem em aplicativos com enorme paralelismo, como as operações de matriz em uma rede neural. Na verdade, em uma carga de trabalho de treinamento típica para aprendizado profundo, uma GPU fornece uma ordem de grandeza maior do que uma CPU.
No entanto, a GPU ainda é um processador de uso geral que precisa aceitar vários aplicativos e softwares diferentes. Portanto, as GPUs e as CPUs têm o mesmo problema. Para cada cálculo nas milhares de ALUs, uma GPU precisa acessar registros ou a memória compartilhada para ler operandos e armazenar os resultados de cálculos intermediários.
Como funciona uma TPU
O Google projetou os Cloud TPUs como um processador de matriz especializado em cargas de trabalho de redes neurais. As TPUs não podem executar processadores de texto, controlar mecanismos de foguetes ou executar transações bancárias, mas podem processar grandes operações de matriz usadas em redes neurais em velocidades muito rápidas.
A principal tarefa das TPUs é o processamento de matrizes, que é uma combinação de operações de multiplicação e acumulação. As TPUs contêm milhares de acumuladores de multiplicação que estão diretamente conectados entre si para formar uma matriz física grande. Isso é chamado de arquitetura de matriz sistólica. O Cloud TPU v3 contém duas matrizes sistólicas de ALUs de 128x128 em um único processador.
O host da TPU transmite os dados para uma fila de entrada. A TPU carrega os dados da fila de entrada e os armazena na memória HBM. Quando o cálculo é concluído, a TPU carrega os resultados na fila de saída. Em seguida, o host da TPU lê os resultados da fila de saída e os armazena na memória do host.
Para executar as operações da matriz, a TPU carrega os parâmetros da memória HBM na unidade de multiplicação de matrizes (MXU).
Em seguida, a TPU carrega os dados da memória HBM. Conforme cada multiplicação é executada, o resultado é transmitido para o próximo acumulador de multiplicação. A saída é a soma de todos os resultados de multiplicação entre os dados e os parâmetros. Não é necessário acessar a memória durante o processo de multiplicação de matrizes.
Como resultado, as TPUs podem alcançar uma alta capacidade de processamento computacional em cálculos de redes neurais.
Arquitetura do sistema de TPU
As seções a seguir descrevem os principais conceitos de um sistema de TPU. Para mais informações sobre termos comuns de machine learning, consulte o Glossário de machine learning.
Se você não tem experiência com o Cloud TPU, confira a página inicial da documentação sobre TPUs.
Chip de TPU
Um chip de TPU contém um ou mais TensorCores. O número de TensorCores depende da versão do chip de TPU. Cada TensorCore consiste em uma ou mais unidades de multiplicação de matrizes (MXUs), uma unidade vetorial e uma escalar. Para mais informações sobre TensorCores, consulte Um supercomputador específico de domínio voltado ao treinamento de redes neurais profundas.
Uma MXU é composta de acumuladores de multiplicação de 256x256 (TPU v6e) ou de 128x128 (versões de TPU anteriores à v6e) em uma matriz sistólica. As MXUs fornecem a maior parte da capacidade de computação em um TensorCore. Cada MXU é capaz de realizar 16 mil operações de acumulação de multiplicação por ciclo. Todas as multiplicações usam entradas bfloat16, mas todas as acumulações são realizadas no formato de número FP32.
A unidade vetorial é usada para cálculos gerais, como ativações e softmax. A unidade escalar é usada para fluxo de controle, cálculo de endereços de memória e outras operações de manutenção.
Pod de TPU
Um Pod de TPU é um conjunto contíguo de TPUs agrupadas em uma rede especializada. O número de chips de TPU em um Pod de TPU depende da versão da TPU.
Fração
Uma fração é um conjunto de chips no mesmo Pod de TPU conectados por interconexões entre chips (ICI) de alta velocidade. As frações são descritas como chips ou TensorCores, dependendo da versão da TPU.
Formato de chip e topologia de chip também se referem a formatos de fração.
Várias frações x fração única
Um ambiente de várias frações é um grupo de frações que estende a conectividade da TPU além das conexões de interconexão entre chips (ICI) e utiliza a rede do data center (DCN) para transmitir dados além de uma fração. Os dados em cada fração ainda são transmitidos pela ICI. Usando essa conectividade híbrida, o ambiente de várias frações permite o paralelismo entre as frações e possibilita o uso de um número maior de núcleos de TPU para um único job do que é possível com uma única fração.
As TPUs podem ser usadas para executar um job em uma única fração ou em várias. Consulte a introdução ao ambiente de várias frações para mais detalhes.
Cubo de TPU
Uma topologia 4x4x4 de chips de TPU interconectados. Isso só é aplicável a topologias 3D (a partir da TPU v4).
SparseCore
Os SparseCores são processadores de fluxo de dados que aceleram os modelos usando operações esparsas. Um caso de uso principal é a aceleração de modelos de recomendação, que dependem muito de embeddings. A v5p inclui quatro SparseCores por chip, e a v6e inclui dois SparseCores por chip. Para saber como usar SparseCores, consulte Uma análise detalhada sobre SparseCores para modelos de embedding grandes (LEM). Para controlar como o compilador XLA usa SparseCores, use as flags do XLA. Para mais informações, consulte Flags do XLA para TPUs.
Resiliência da ICI do Cloud TPU
A resiliência da ICI ajuda a melhorar a tolerância a falhas de links ópticos e de interruptores de circuito óptico (OCS) que conectam as TPUs entre os cubos. As conexões da ICI em um cubo usam links de cobre que não são afetados. A resiliência da ICI permite que as conexões da ICI sejam roteadas em torno de falhas do OCS e de falhas ópticas da ICI. Isso melhora a disponibilidade de programação das frações de TPU, com o efeito negativo de uma degradação temporária no desempenho da ICI.
Para o Cloud TPU v4 e v5p, a resiliência da ICI é ativada por padrão para frações de um cubo ou maiores, por exemplo:
- v5p-128 ao especificar o tipo de acelerador
- 4x4x4 ao especificar a configuração do acelerador
Versões da TPU
A arquitetura exata de um chip de TPU depende da versão usada. Cada versão de TPU também aceita diferentes tamanhos e configurações de fração. Para mais informações sobre a arquitetura do sistema e as configurações disponíveis, consulte estas páginas:
Arquitetura de nuvem da TPU
OGoogle Cloud disponibiliza TPUs como recursos de computação por meio de VMs de TPU. É possível usar diretamente as VMs de TPU para cargas de trabalho ou usá-las pelo Google Kubernetes Engine ou pela Vertex AI. As seções a seguir descrevem os principais componentes da arquitetura de nuvem da TPU.
Arquitetura da VM de TPU
A arquitetura da VM de TPU permite que você se conecte por SSH diretamente à VM fisicamente conectada ao dispositivo de TPU. Uma VM de TPU, também conhecida como worker, é uma máquina virtual que executa o Linux e tem acesso às TPUs. Você tem acesso raiz à VM para executar códigos arbitrários. É possível acessar registros de depuração do compilador e do ambiente de execução, além de mensagens de erro.
Host único, vários hosts e sub-host
Um host de TPU é uma VM executada em um computador físico conectado ao hardware de TPU. As cargas de trabalho de TPU podem usar um ou mais hosts.
Uma carga de trabalho de host único é limitada a uma VM de TPU. Uma carga de trabalho de vários hosts distribui o treinamento entre várias VMs de TPU. Uma carga de trabalho de sub-host não usa todos os chips em uma VM de TPU.
Arquitetura de nó de TPU (descontinuada)
A arquitetura de nós de TPU consiste em uma VM do usuário que se comunica com o host da TPU por gRPC. Ao usar essa arquitetura, não é possível acessar diretamente o host da TPU, o que dificulta a depuração de erros de treinamento e da TPU.
Como migrar da arquitetura de nó de TPU para a de VM de TPU
Se você tiver TPUs que usam a arquitetura de nós de TPU, siga estas etapas para a identificação, a exclusão e o novo provisionamento das VMs de TPU.
Acesse a página TPUs:
Localize a TPU e a arquitetura dela abaixo do cabeçalho Arquitetura. Se a arquitetura for VM de TPU, não será necessário fazer nada. Se a arquitetura for Nós de TPU, exclua a TPU e faça o novo provisionamento.
Exclua e provisione novamente a TPU.
Consulte Como gerenciar TPUs para saber como excluir e provisionar novamente suas TPUs.