A primeira Tensor Processing Unit (TPU) do Google: arquitetura

• A primeira TPU (Tensor Processing Unit) v1 do Google teve como objetivo desenvolver um ASIC (Application Specific Integrated Circuit) que oferecesse uma vantagem de custo-desempenho de 10x, considerando as oportunidades de novos serviços com deep learning e a escala e o custo do hardware necessário.
• A TPU foi projetada para acelerar operações com tensores, e tensor aqui se refere a operações que mapeiam relações de arrays multidimensionais entre vetores, escalares e outros tensores.
• A multiplicação de matrizes é a operação central das redes neurais, usada para calcular camadas ocultas e de saída ao multiplicar o vetor de entrada pela matriz de pesos e, em seguida, aplicar a função de ativação.

A primeira unidade de processamento de tensores do Google: origens

• O projeto da TPU v1 começou no fim de 2013 e foi desenvolvido em 15 meses, o que o torna um tema interessante para quem se interessa por aceleradores de redes neurais.
• O contexto e os objetivos do desenvolvimento da TPU v1 foram abordados no post anterior, e neste post o foco é examinar com mais detalhes a estrutura e o desempenho da TPU v1.

Estrutura e desempenho

• O design da TPU v1 é baseado na arquitetura apresentada no artigo de 1978 de H.T Kung e Charles E. Leiserson, "Systolic Arrays (for VLSI)".
• Um systolic array é composto por uma rede de processadores que calcula dados periodicamente e os transmite pelo sistema, realizando multiplicação de matrizes de forma eficiente.
• O systolic array da TPU v1 é composto por unidades de multiplicação/acumulação 256x256 e gera automaticamente os resultados necessários sem precisar salvar e buscar resultados intermediários na memória principal.

Unidades de multiplicação/acumulação (MACs)

• A TPU v1 realiza multiplicação inteira de 8 bits x 8 bits e utiliza quantização, que exige menos área de die do que cálculos em ponto flutuante.

Conjunto de instruções

• A TPU v1 usa um conjunto de aproximadamente 20 instruções complexas (CISC), que são transmitidas do computador host por meio da interface PCIe.
• As principais instruções incluem leitura da memória do host, leitura de pesos, multiplicação/composição de matrizes, ativação e escrita na memória do host.

Software

• O hardware da TPU v1 oferece suporte a uma stack de software como o Tensorflow, permitindo portar rapidamente para a TPU aplicações desenvolvidas em CPU e GPU.

Fabricação e die

• A TPU v1 foi fabricada usando o processo de 28nm da TSMC, e sua área de die era menos da metade da dos chips Intel Haswell CPU e Nvidia K80 GPU que o Google usava em seus data centers.

Desempenho

• A TPU v1 foi projetada para tornar a inferência mais eficiente e oferece velocidade de inferência e eficiência energética cerca de 15 a 30 vezes melhores que a GPU Nvidia K80 e a CPU Haswell.

Lições

• A arquitetura personalizada da TPU v1 permite desempenho e uso de energia muito superiores aos de CPUs e GPUs modernas.
• A TPU v1 foi projetada com o objetivo de tornar a inferência rápida e eficiente em consumo de energia, e não para treinamento.

Opinião do GN⁺

• O desenvolvimento da TPU v1 marca um ponto de inflexão importante na migração do Google para hardware de propósito específico a fim de otimizar cargas de trabalho de inferência em data centers de grande escala.
• A arquitetura systolic array otimiza o fluxo de dados e minimiza o acesso à memória, o que contribui para ganhos de eficiência energética e desempenho.
• O sucesso da TPU v1 levou posteriormente ao desenvolvimento de várias versões da TPU e teve papel importante na evolução do hardware de inferência de IA.
• Outros produtos que oferecem funcionalidades semelhantes à TPU incluem GPUs da Nvidia com Tensor Core, o Nervana NNP da Intel e o Inferentia da Amazon.
• Ao adotar a tecnologia TPU, é preciso considerar a compatibilidade com a infraestrutura existente, a stack de software e o nível de otimização para aplicações específicas. Os benefícios de escolher a TPU incluem maior velocidade nas tarefas de inferência e eficiência energética, mas pode ser necessária otimização para modelos ou tipos de dados específicos.

[GN⁺]

Comentários no Hacker News

• O CEO da Groq, Jonathon Ross, falou em uma entrevista de podcast sobre o processo de criação do protótipo do TPU no Google. No início, ele criou um FPGA em seu tempo de 20% para resolver um problema de velocidade de inferência de uma equipe vizinha. Depois que Jeff Dean fez os cálculos matemáticos, decidiram mudar para um ASIC. Há quem defenda que o Google deveria separar a equipe de TPU em uma empresa independente. O TPU é o único concorrente realmente confiável da NVidia, e o suporte de software é o segundo melhor, atrás apenas da NVidia.
• Levanta-se a questão de como, mesmo tendo inventado o TPU e publicado os principais artigos sobre LLM pelo Google Research, o Google acabou vendo a NVidia e startups de IA ficarem com a maior parte do valor de mercado.
• Um funcionário do Google recomenda conferir o TPU v5. Ele oferece suporte a PyTorch/JAX, o que o torna muito mais fácil de usar do que quando só era possível usar TensorFlow.
• Um artigo explica muito bem como as partes abstratas se conectam por meio do silício. É interessante ver como instruções CISC são mapeadas para a etapa de inferência de LLM.
• Há a opinião de que o Google deveria entrar na tecnologia EUV de 2 nm e avançar para abaixo de 2 nm. Se conseguisse garantir a tecnologia de litografia eletrônica oferecida pela ASML, poderia obter resultados muito poderosos. Isso poderia ser um projeto ousado do Google X e, se o TPU for realmente tão bom, uma boa oportunidade de negócio com fábrica própria e tecnologia própria.
• Pergunta sobre qual criatura marinha seria a origem do nome atual do TPU.
• Citação do OP (autor do post original): "O TPU v1 usava um design CISC com cerca de 20 instruções." Em resposta, há um comentário sarcástico dizendo que CISC/RISC passou por vários estágios, de programa de pesquisa a termo de marketing, e agora virou apenas um som sem significado.
• Um comentário diz não entender por que se usa o termo "tensor" se, na prática, o hardware do TPU só realiza operações de vetor e matriz.
• [comentário excluído]
• [comentário excluído]