El Tensor Processing Unit (TPU) es un circuito integrado de aplicación específica diseñado por Google específicamente para acelerar cargas de trabajo de aprendizaje automático, en particular las multiplicaciones de matrices que son el núcleo del entrenamiento de redes neuronales y la inferencia. A diferencia de una GPU de propósito general, una TPU está diseñada específicamente para operaciones tensoriales y no está pensada para renderizar gráficos.
Historia
Google comenzó a implementar TPU internamente en 2015 para acelerar servicios como la clasificación de búsqueda, el reconocimiento de texto de Street View y la evaluación de movimientos de AlphaGo durante sus partidas, y anunció el chip públicamente en su conferencia I/O de 2016. La primera generación de TPU se centró en la inferencia de modelos ya entrenados; la segunda generación, lanzada en 2017, añadió soporte para entrenamiento, y Google puso las TPU a disposición de clientes externos a través de Google Cloud. Las generaciones sucesivas, TPU v3 (2018), v4 (2021), v5 (2023) y Trillium (2024), aumentaron cada una el rendimiento y la memoria, y a mediados de la década de 2020 las TPU formaban grandes conjuntos interconectados de miles de chips utilizados para entrenar los propios modelos fundacionales de Google, incluida su familia Gemini, además de alquilarse a clientes externos a través de Google Cloud.
Diseño y comparación con las GPU
Las TPU se basan en una arquitectura de matriz sistólica, una cuadrícula de elementos de procesamiento que pasan datos a elementos vecinos en un ritmo fijo, lo que es muy eficiente para las multiplicaciones de matrices repetidas utilizadas en el aprendizaje profundo, pero mucho menos flexible que una GPU para la computación paralela de propósito general. Esta especialización otorga a las TPU una ventaja de eficiencia por vatio y por dólar para el entrenamiento y la inferencia a gran escala de redes neuronales, a costa de ser menos adaptables a cargas de trabajo no tensoriales. Mientras que las GPU de Nvidia funcionan con la plataforma de software ampliamente adoptada CUDA, las TPU se programan principalmente a través de TensorFlow y, más tarde, PyTorch y JAX, la biblioteca de cálculo numérico propia de Google, lo que les da un ecosistema de software más pequeño pero en crecimiento.
Importancia
Las TPU representan uno de los ejemplos más claros de un gran laboratorio de IA que integra verticalmente su propio diseño de chips en lugar de depender únicamente de silicio comercial de proveedores como Nvidia, una estrategia que luego replicaron otras empresas que diseñan aceleradores de IA personalizados, como Trainium de Amazon y MTIA de Meta, a medida que la demanda de capacidad de entrenamiento e inferencia de modelos de lenguaje grandes crecía durante el auge de la IA de la década de 2020. Dado que las TPU están disponibles principalmente a través de Google Cloud y no se venden como hardware independiente, han mantenido una cuota menor del mercado general de aceleradores de IA en comparación con las GPU, pero sustentan la infraestructura interna de IA de Google, incluidos los entrenamientos de sus modelos más grandes, y Google las cita como una razón clave para poder competir con rivales que dependen más del suministro externo de GPU.