El TPU v2 es la segunda generación de la Unidad de Procesamiento Tensorial, un circuito integrado de aplicación específica (ASIC) personalizado desarrollado por Google para acelerar cargas de trabajo de aprendizaje automático en redes neuronales. Anunciado en mayo de 2017, sucedió al TPU original e introdujo mejoras arquitectónicas significativas, destacando especialmente la capacidad de realizar cálculos en coma flotante y la inclusión del formato bfloat16, inventado por Google Brain. Esto hizo que el TPU v2 fuera adecuado tanto para el entrenamiento como para la inferencia de modelos de aprendizaje automático, a diferencia de su predecesor, que se limitaba a operaciones con enteros. El TPU v2 se puso a disposición de terceros a través del servicio Cloud TPU de Google Cloud, y se ha utilizado ampliamente en los propios sistemas de producción de Google, incluidos modelos basados en transformadores y grandes modelos de lenguaje.
El TPU v2 fue diseñado para abordar las limitaciones de ancho de banda de memoria de la primera generación de TPU. Al incorporar 16 GB de memoria de alto ancho de banda (HBM), el diseño de segunda generación aumentó el ancho de banda de memoria a 600 GB/s y ofreció 45 teraFLOPS de rendimiento por chip. Estos chips se organizaron luego en módulos de cuatro chips, logrando un rendimiento combinado de 180 teraFLOPS. Además, 64 de estos módulos se ensamblaron en pods de 256 chips, proporcionando 11,5 petaFLOPS de rendimiento total. Esta arquitectura escalable permitió a Google desplegar sistemas TPU v2 para tareas de aprendizaje automático a gran escala, incluidas aquellas que involucran redes neuronales y aprendizaje profundo.
Desarrollo e Historia
El desarrollo del TPU v2 fue liderado por Amir Salek, quien fundó y dirigió el grupo de Silicio Personalizado de Google, y Norman P. Jouppi actuó como líder técnico y arquitecto principal del programa general de TPU. El TPU original se desplegó en producción en solo 15 meses, y el v2 se basó en esa base. Según Jonathan Ross, uno de los ingenieros originales del TPU y más tarde fundador de Groq, el diseño del TPU, que utiliza una arquitectura de matriz sistólica, fue seleccionado entre tres propuestas competidoras de aceleradores de IA en Google. El TPU v2 fue el primer chip de entrenamiento de aprendizaje profundo en producción de la industria, marcando un hito en el silicio personalizado para IA. Broadcom co-desarrolló el TPU v2, traduciendo la arquitectura de Google en silicio fabricable y gestionando la fabricación a través de fundiciones como TSMC.
Especificaciones Técnicas
El TPU v2 es un motor de multiplicación de matrices que admite aritmética tanto de enteros como de coma flotante. Su innovación clave fue la introducción del formato bfloat16, una representación de coma flotante de 16 bits que proporciona un rango dinámico más amplio que los formatos tradicionales de 16 bits, al tiempo que requiere menos memoria y ancho de banda. Esto lo hizo particularmente efectivo para entrenar redes neuronales profundas. Cada chip TPU v2 contiene una matriz sistólica de multiplicadores y acumuladores, junto con 16 GB de HBM que proporcionan 600 GB/s de ancho de banda. Los chips están interconectados en módulos de cuatro chips, y estos módulos se combinan en pods más grandes. La configuración de pod de 256 chips ofrece 11,5 petaFLOPS, permitiendo un cálculo paralelo masivo para tareas como entrenar transformadores y otros grandes modelos de lenguaje.
Aplicaciones y Uso
Google utilizó sistemas TPU v2 en una variedad de aplicaciones de producción. Se emplearon en la serie de partidas de Go entre AlphaGo y Lee Sedol, así como en el sistema AlphaZero, que aprendió a jugar ajedrez, shogi y Go solo a partir de las reglas del juego. Los TPU también se usaron para el procesamiento de texto de Google Street View, permitiendo extraer todo el texto de la base de datos de Street View en menos de cinco días. En Google Photos, un solo TPU podía procesar más de 100 millones de fotos al día. Además, los TPU v2 se usaron en RankBrain, el sistema de Google para proporcionar resultados de búsqueda. A través del servicio Cloud TPU en Google Cloud, terceros podían acceder a los TPU v2 para sus propias cargas de trabajo de aprendizaje automático, con soporte para marcos como TensorFlow, JAX y PyTorch.
Comparación con GPUs
En comparación con las GPUs, los TPU están diseñados para cómputo de alto volumen y baja precisión con más operaciones de entrada/salida por julio. Carecen de hardware para rasterización o mapeo de texturas, centrándose en cambio en operaciones de matrices. Los TPU son adecuados para redes neuronales convolucionales (CNN), mientras que las GPUs tienen beneficios para algunas redes neuronales totalmente conectadas, y las CPUs pueden tener ventajas para redes neuronales recurrentes (RNN). En flujos de trabajo que involucran redes neuronales basadas en transformadores para grandes modelos de lenguaje, los TPU se usan a menudo para inferencia, mientras que las GPUs se usan frecuentemente para entrenamiento. Esta división del trabajo refleja las diferentes fortalezas de cada tipo de procesador, con los TPU sobresaliendo en los patrones específicos de cómputo que se encuentran en muchos modelos modernos de IA.
Legado e Impacto
El TPU v2 estableció una plantilla para las generaciones posteriores del hardware TPU de Google, incluido el TPU v3, anunciado el 8 de mayo de 2018, con el doble de rendimiento por chip y pods con cuatro veces más chips. La introducción de bfloat16 tuvo un impacto duradero en el ecosistema de hardware de IA, influyendo en otros diseños de aceleradores. A partir de 2025, Google Cloud genera ingresos significativos por productos a partir de la venta de sistemas TPU, y Google ha estado en conversaciones con varias "neoclouds" y empresas como Meta sobre el despliegue de TPU en sus centros de datos. El TPU v2 sigue siendo una plataforma fundamental en la historia del hardware de inteligencia artificial, demostrando el valor del silicio personalizado para el aprendizaje automático.