Arquitetura da TPU
As Tensor Processing Units (TPUs) são circuitos integrados específicos da aplicação (ASICs) concebidos pela Google para acelerar as cargas de trabalho de aprendizagem automática. A Cloud TPU é um Google Cloud serviço que disponibiliza TPUs como um recurso escalável.
As TPUs foram concebidas para realizar operações de matriz rapidamente, o que as torna ideais para cargas de trabalho de aprendizagem automática. Pode executar cargas de trabalho de aprendizagem automática em TPUs usando frameworks como o Pytorch e o JAX.
Como funcionam as TPUs?
Para compreender como funcionam as TPUs, é útil saber como outros aceleradores resolvem os desafios computacionais da preparação de modelos de AA.
Como funciona uma CPU
Uma CPU é um processador de uso geral baseado na arquitetura de von Neumann. Isto significa que uma CPU funciona com software e memória da seguinte forma:
A maior vantagem das CPUs é a sua flexibilidade. Pode carregar qualquer tipo de software numa CPU para muitos tipos diferentes de aplicações. Por exemplo, pode usar uma CPU para processamento de texto num PC, controlar motores de foguetes, executar transações bancárias ou classificar imagens com uma rede neural.
Uma CPU carrega valores da memória, faz um cálculo sobre os valores e armazena o resultado novamente na memória para cada cálculo. O acesso à memória é lento em comparação com a velocidade de cálculo e pode limitar o débito total das CPUs. Isto é frequentemente denominado gargalo de von Neumann.
Como funciona uma GPU
Para obter um débito mais elevado, as GPUs contêm milhares de unidades lógicas aritméticas (ULAs) num único processador. Normalmente, uma GPU moderna contém entre 2500 e 5000 ALUs. O grande número de processadores significa que pode executar milhares de multiplicações e adições em simultâneo.
Esta arquitetura de GPU funciona bem em aplicações com paralelismo massivo, como operações de matriz numa rede neural. Na verdade, numa carga de trabalho de preparação típica para aprendizagem profunda, uma GPU pode oferecer um débito uma ordem de magnitude superior ao de uma CPU.
No entanto, a GPU continua a ser um processador de uso geral que tem de suportar muitas aplicações e software diferentes. Por conseguinte, as GPUs têm o mesmo problema que os CPUs. Para cada cálculo nos milhares de UALs, uma GPU tem de aceder a registos ou memória partilhada para ler operandos e armazenar os resultados dos cálculos intermédios.
Como funciona uma TPU
A Google concebeu as TPUs na nuvem como um processador de matriz especializado para cargas de trabalho de redes neurais. As TPUs não podem executar processadores de texto, controlar motores de foguetes nem executar transações bancárias, mas podem processar operações de matrizes massivas usadas em redes neurais a velocidades rápidas.
A tarefa principal 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 multiplicadores-acumuladores que estão diretamente ligados entre si para formar uma grande matriz física. Isto é denominado arquitetura de matriz sistólica. As Cloud TPUs v3 contêm duas matrizes sistólicas de 128 x 128 ALUs num único processador.
O anfitrião da TPU transmite dados para uma fila de entrada. A TPU carrega dados da fila de entrada e armazena-os na memória HBM. Quando a computação estiver concluída, a TPU carrega os resultados na fila de saída. Em seguida, o anfitrião da TPU lê os resultados da fila de saída e armazena-os na memória do anfitrião.
Para realizar as operações de matriz, a TPU carrega os parâmetros da memória HBM para a unidade de multiplicação de matrizes (MXU).
Em seguida, a TPU carrega dados da memória HBM. À medida que cada multiplicação é executada, o resultado é transmitido ao multiplicador-acumulador seguinte. O resultado é a soma de todos os resultados da multiplicação entre os dados e os parâmetros. Não é necessário acesso à memória durante o processo de multiplicação de matrizes.
Como resultado, as TPUs podem alcançar um elevado débito computacional nos cálculos da rede neural.
Arquitetura do sistema de TPU
As secções seguintes descrevem os conceitos-chave de um sistema de TPU. Para mais informações sobre termos comuns de aprendizagem automática, consulte o Glossário de aprendizagem automática.
Se está a usar o Cloud TPU pela primeira vez, consulte a página inicial da documentação da TPU.
Chip TPU
Um chip de TPU contém um ou mais TensorCores. O número de TensorCores depende da versão do chip TPU. Cada TensorCore consiste numa ou mais unidades de multiplicação de matrizes (MXUs), numa unidade vetorial e numa unidade escalar. Para mais informações sobre os TensorCores, consulte o artigo Um supercomputador específico do domínio para a preparação de redes neurais profundas.
Uma MXU é composta por 256 x 256 (TPU v6e) ou 128 x 128 (versões de TPU anteriores à v6e) multiplicadores-acumuladores numa matriz sistólica. As MXUs fornecem a maior parte da capacidade de computação num TensorCore. Cada MXU é capaz de realizar 16 mil operações de multiplicação e acumulaçã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 computação geral, como ativações e softmax. A unidade escalar é usada para o fluxo de controlo, o cálculo de endereços de memória e outras operações de manutenção.
pod TPU
Um Pod de TPUs é um conjunto contíguo de TPUs agrupadas numa rede especializada. O número de chips de TPU num TPU Pod depende da versão da TPU.
Dividir
Uma fatia é um conjunto de chips localizados no mesmo TPU Pod ligados por interconexões de alta velocidade entre chips (ICI). As fatias são descritas em termos de chips ou TensorCores, consoante a versão da TPU.
A forma do chip e a topologia do chip também se referem às formas das fatias.
Vários segmentos vs. um único segmento
O Multislice é um grupo de fatias que expande a conetividade da TPU para além das ligações de interconexão entre chips (ICI) e tira partido da rede do centro de dados (DCN) para transmitir dados para além de uma fatia. Os dados em cada fatia continuam a ser transmitidos pela ICI. Com esta conetividade híbrida, o Multislice permite o paralelismo entre fatias e permite-lhe usar um número maior de núcleos da TPU para uma única tarefa do que o que uma única fatia pode acomodar.
As TPUs podem ser usadas para executar uma tarefa numa única fatia ou em várias fatias. Consulte a introdução ao Multislice para ver mais detalhes.
Cubo da TPU
Uma topologia de 4x4x4 chips de TPU interligados. Isto só se aplica a topologias 3D (a partir da TPU v4).
SparseCore
Os SparseCores são processadores de fluxo de dados que aceleram os modelos através de operações esparsas. Um exemplo de utilização principal é a aceleração dos modelos de recomendações, que dependem fortemente de incorporações. A v5p inclui quatro SparseCores por chip e a v6e inclui dois SparseCores por chip. Para uma explicação detalhada sobre como os SparseCores podem ser usados, consulte o artigo Uma análise detalhada do SparseCore para modelos de incorporação grandes (LEM). Pode controlar a forma como o compilador XLA usa SparseCores através de flags XLA. Para mais informações, consulte: Flags XLA da TPU.
Resiliência da ICI do Cloud TPU
A resiliência da ICI ajuda a melhorar a tolerância a falhas das ligações óticas e dos comutadores de circuitos óticos (OCS) que ligam as TPUs entre cubos. (As ligações ICI num cubo usam ligações de cobre que não são afetadas). A resiliência da ICI permite que as ligações ICI sejam encaminhadas em torno de falhas da OCS e da ICI ótica. Como resultado, melhora a disponibilidade de agendamento de fatias de TPU, com a contrapartida de uma degradação temporária no desempenho da ICI.
Para o Cloud TPU v4 e v5p, a resiliência de ICI está ativada por predefinição para fatias que sejam um cubo ou maiores, por exemplo:
- v5p-128 quando especifica o tipo de acelerador
- 4x4x4 quando especificar a configuração do acelerador
Versões da TPU
A arquitetura exata de um chip da TPU depende da versão da TPU que usa. Cada versão da TPU também suporta diferentes tamanhos e configurações de fatias. Para mais informações sobre a arquitetura do sistema e as configurações suportadas, consulte as seguintes páginas:
Arquitetura de nuvem da TPU
Google Cloud Disponibiliza TPUs como recursos de computação através de VMs da TPU. Pode usar diretamente VMs de TPUs para as suas cargas de trabalho ou usá-las através do Google Kubernetes Engine ou do Vertex AI. As secções seguintes descrevem os componentes principais da arquitetura de nuvem da TPU.
Arquitetura de VM da TPU
A arquitetura de VM da TPU permite-lhe estabelecer ligação diretamente à VM fisicamente ligada ao dispositivo TPU através do SSH. Uma VM da TPU, também conhecida como um trabalhador, é uma máquina virtual que executa o Linux e tem acesso às TPUs subjacentes. Tem acesso root à VM, pelo que pode executar código arbitrário. Pode aceder aos registos de depuração e às mensagens de erro do compilador e do tempo de execução.
Anfitrião único, vários anfitriões e subanfitrião
Um anfitrião de TPU é uma VM executada num computador físico ligado ao hardware de TPU. As cargas de trabalho da TPU podem usar um ou mais anfitriões.
Uma carga de trabalho de anfitrião único está limitada a uma VM da TPU. Uma carga de trabalho com vários anfitriões distribui a preparação por várias VMs da TPU. Uma carga de trabalho de sub-host não usa todos os chips numa VM da TPU.
Arquitetura de nós da TPU (descontinuada)
A arquitetura de nós da TPU consiste numa VM do utilizador que comunica com o anfitrião da TPU através do gRPC. Quando usa esta arquitetura, não pode aceder diretamente ao anfitrião da TPU, o que dificulta a depuração de erros de treino e da TPU.
Transição da arquitetura de nós da TPU para a arquitetura de VMs da TPU
Se tiver TPUs que usam a arquitetura de nós da TPU, siga estes passos para as identificar, eliminar e voltar a aprovisionar como VMs da TPU.
Aceda à página de TPUs:
Localize a TPU e a respetiva arquitetura no cabeçalho Arquitetura. Se a arquitetura for "VM da TPU", não tem de fazer nada. Se a arquitetura for "Nó da TPU", tem de eliminar e voltar a aprovisionar a TPU.
Elimine e volte a aprovisionar a TPU.
Consulte o artigo Gerir TPUs para ver instruções sobre como eliminar e reaprovisionar TPUs.