- La Google TPU v1 es un ASIC desarrollado en 15 meses desde finales de 2013 para soportar el costo y la escala de la inferencia en servicios de deep learning, con el objetivo de lograr una relación costo-rendimiento 10 veces mejor que una GPU
- La clave fue implementar el arreglo sistólico de H.T. Kung y Charles E. Leiserson de 1978 en una estructura MAC de 256×256, reduciendo los viajes de ida y vuelta a memoria de los resultados intermedios durante la multiplicación de matrices
- La TPU v1 se comunica con el host por PCIe y almacena los pesos en DDR3-2133, organizando el flujo de inferencia con unas pocas instrucciones como
Read_Host_Memory, Read_Weights, Matrix_Multiply / Convolve, Activate y Write_Host_Memory
- Usa multiplicación entera de 8-bit×8-bit y cuantización para evitar el costo en área de silicio del cómputo en punto flotante, y la ejecución es controlada por un User Space Driver y un Kernel Driver integrados con TensorFlow
- Los puntos de comparación de la época eran la CPU Intel Haswell y la GPU Nvidia K80; la TPU v1 era aproximadamente 15 a 30 veces más rápida en inferencia y tenía una mejora relativa de rendimiento incremental por Watt de 25 a 29 veces frente a la GPU, pero no era un dispositivo para entrenamiento
El punto de partida de la TPU v1 para bajar el costo de la inferencia
- El objetivo del proyecto TPU v1 era crear rápidamente un ASIC para inferencia que redujera el enorme costo de hardware requerido por los servicios basados en deep learning
- La meta era una ventaja de 10 veces en costo-rendimiento en inferencia frente a una GPU
- Se requerían al mismo tiempo desarrollo rápido, alto rendimiento, despliegue a gran escala, soporte inmediato para nuevas cargas de trabajo y eficiencia de costos
- El nombre TPU viene de que es un dispositivo para acelerar operaciones de tensores
- En la práctica, los cálculos centrales que maneja el hardware TPU v1 son principalmente operaciones de vectores y matrices
- Las capas ocultas y de salida de una red neuronal pueden expresarse como el resultado de aplicar una función de activación al producto entre un vector de entrada y una matriz de pesos
- Cuando entran varios datos de entrada, esto toma la forma de aplicar la función de activación a cada elemento del resultado de una multiplicación de matrices
Procesamiento de multiplicación de matrices con arreglos sistólicos
- La TPU v1 usa el concepto de sistema sistólico presentado en el artículo de 1978 de H.T. Kung y Charles E. Leiserson Systolic Arrays (for VLSI)
- Es una estructura donde múltiples procesadores calculan y transfieren datos con un ritmo regular
- Cada procesador realiza un cálculo corto mientras mueve continuamente los datos hacia adentro y hacia afuera
- En una multiplicación simple de matrices 2×2, si los valores de entrada se introducen desde arriba y desde la izquierda en el orden correcto, el resultado aparece de forma natural en un arreglo MAC 2×2
- Cada MAC realiza multiplicación y acumulación
- Las sumas parciales se almacenan dentro del arreglo, y el resultado final aparece en forma de diagonal en movimiento
- El ejemplo 2×2 requiere 4 etapas, pero en la práctica, tan pronto como el MAC superior izquierdo queda libre, puede iniciarse la siguiente multiplicación de matrices, así que es posible una nueva multiplicación cada 2 ciclos
- La idea central es que, si los datos se suministran al arreglo sistólico en el orden correcto, el propio flujo de valores y resultados crea el orden de cálculo necesario
- No hace falta guardar los resultados intermedios en la memoria principal y luego volver a traerlos
- Gracias a la estructura de la unidad de multiplicación de matrices y al orden de entrada, los resultados intermedios pueden usarse automáticamente en el momento necesario
Configuración del sistema TPU v1
- La TPU v1 se comunica con la computadora host mediante el bus serial de alta velocidad PCIe y accede directamente a su propia DRAM DDR3
- Los componentes principales son los siguientes
- DDR3 DRAM / Weight FIFO
- Los pesos se almacenan en chips de RAM DDR3 conectados mediante una interfaz DDR3-2133
- Se cargan previamente desde la memoria del host por PCIe y luego se mueven al Weight FIFO para que los use la Matrix Multiply Unit
- Matrix Multiply Unit
- Es un arreglo sistólico compuesto por 256×256 MAC
- Recibe 256 pesos desde arriba y 256 entradas de datos desde la izquierda
- Accumulators
- Almacenan los resultados que salen de la unidad matricial sistólica por la parte inferior
- Activation
- Es la etapa que aplica la función de activación de la red neuronal
- Unified Buffer / Systolic Data Setup
- Almacena el resultado tras aplicar la función de activación y lo prepara para volver a alimentar la entrada de la Matrix Multiply Unit para el cálculo de la siguiente capa
Formato de cálculo y conjunto de instrucciones
- La Matrix Multiply Unit de la TPU v1 realiza multiplicación entera de 8-bit×8-bit
- Usa cuantización para evitar el cómputo en punto flotante, que requeriría un área de silicio mayor
- El conjunto de instrucciones es un diseño CISC con alrededor de 20 instrucciones
- Las instrucciones no se traen desde memoria; la computadora host las envía por PCIe
- La mayor parte del flujo de inferencia se compone de 5 instrucciones principales
Read_Host_Memory
- Lee las entradas desde la memoria del host hacia el Unified Buffer por PCIe
Read_Weights
- Lee los pesos desde la memoria de pesos hacia el Weight FIFO
Matrix_Multiply / Convolve
- Envía las entradas del Unified Buffer a los Accumulators mientras realiza multiplicación de matrices o convolución
- Multiplica una entrada B×256 por una entrada de pesos constantes 256×256 para producir una salida B×256, y requiere B ciclos de pipeline
Activate
- Aplica a la entrada de los Accumulators funciones no lineales de neuronas artificiales como ReLU o Sigmoid y envía el resultado al Unified Buffer
Write_Host_Memory
- Escribe el resultado del Unified Buffer en la memoria del host por PCIe
- Este flujo puede verse aproximadamente así
Read_Host_Memory
Read_Weights
Loop_Start
Matrix_Multiply
Activate
Loop_End
Write_Host_Memory
- La unidad matricial usa ejecución sistólica para ahorrar energía al reducir las lecturas y escrituras del Unified Buffer
- Los datos entran por la izquierda y los pesos se cargan desde arriba
- Las operaciones MAC de 256 elementos atraviesan la matriz como un frente de onda diagonal
TensorFlow y la pila de drivers
- Para usar el hardware TPU v1 en servicios reales, se necesitaba una pila de software que lo soportara
- Como Google desarrollaba y usaba TensorFlow, un paso clave fue crear drivers para que TensorFlow funcionara con la TPU v1
- La pila de software de TPU tenía que ser compatible con las pilas para CPU y GPU
- Las aplicaciones debían poder migrarse rápidamente a TPU
- La parte de la aplicación que se ejecuta en TPU normalmente se escribe en TensorFlow y se compila a una API que puede ejecutarse en GPU o TPU
- Igual que en GPU, la pila de TPU también se divide en User Space Driver y Kernel Driver
- El Kernel Driver se mantiene liviano, maneja solo memoria e interrupciones y apunta a estabilidad de largo plazo
- El User Space Driver cambia con frecuencia y se encarga de la configuración y control de ejecución de la TPU, del reformateo de datos para ajustarlos a la secuencia de la TPU, de convertir llamadas de API en instrucciones TPU y de generar binarios de aplicación
Proceso de 28 nm y distribución del die
- La TPU v1 se fabricó con el relativamente maduro proceso de 28 nm de TSMC
- En ese momento, los chips Intel Haswell CPU y Nvidia K80 GPU usados en los centros de datos de Google estaban fabricados con procesos más avanzados
- Según Google, el área del die de la TPU v1 era menos de la mitad del área de esos dies
- La ISA simple redujo el overhead de silicio necesario para la decodificación y tareas relacionadas
- El área de
control ocupa solo 2% del die
- La Matrix Multiply Unit ocupa 24% y el Unified Buffer 29%
Comparación de rendimiento y límites claros
- La TPU v1 es un dispositivo para inferencia pensado para usar de forma más eficiente modelos ya entrenados en servicios reales a escala de Google
- No fue diseñada para mejorar la velocidad ni la eficiencia del entrenamiento
- La inferencia y el entrenamiento plantean desafíos distintos en el desarrollo de hardware especializado
- En 2013, los principales puntos de comparación eran la CPU Intel Haswell y la GPU Nvidia K80
- La TPU v1 tiene 25 veces más MAC que la GPU K80
- La TPU v1 tiene 3.5 veces más memoria on-chip que la GPU K80
- La TPU v1 es aproximadamente 15 a 30 veces más rápida en inferencia que la GPU K80 y la CPU Haswell
- Su mejora relativa de rendimiento incremental por Watt frente a la GPU es de 25 a 29 veces
- Gracias a su arquitectura personalizada, la TPU v1 logró un mayor rendimiento de inferencia y menor consumo de energía que las CPU y GPU de la época
- Como fue un diseño de primera generación enfocado en un único objetivo —inferencias rápidas y eficiencia energética—, quedó la limitación de no haber sido diseñada para entrenamiento
1 comentarios
Opiniones en Hacker News
El CEO de Groq, Jonathan Ross, contó recientemente en una entrevista de podcast la historia de cuando estaban creando el TPU inicial en Google: originalmente era un FPGA que hizo durante su 20% time, después de sentarse cerca de un equipo que tenía problemas con la velocidad de inferencia.
Después de hacer algo que funcionaba, Jeff Dean hizo los cálculos y decidieron ir por un ASIC.
Creo que, si fuera hoy, Google debería escindir el equipo de TPU como una empresa aparte. Es el único competidor realmente creíble frente a Nvidia, y su soporte de software también está en el nivel inmediatamente después de Nvidia.
https://open.spotify.com/episode/0V9kRgNS7Ds6zh3GjdXUAQ?si=q...
Pero cada una tiene sus debilidades. La capacidad de fundición está ajustada, pero Nvidia puede sacrificar el mercado de GPU de consumo si puede vender chips de IA más caros. Si un competidor empezó a apostar fuerte desde hace unos años, o si una empresa con mucha capacidad de producción, como Intel, cambia sus prioridades, esa ventaja desaparece.
Es cómodo que el software propietario sea el estándar de la industria, pero su importancia real depende mucho del caso de uso. El diseño de hardware para TPU parece inherentemente mucho más simple que el de una GPU: no necesita ray tracing, texture samplers ni rasterización; en su mayoría basta con muchas multiplicaciones de matrices y memoria.
Las relaciones con los clientes sirven para seguir participando en la conversación, pero en un mercado que busca cualquier ventaja, el proveedor de hardware con más FLOPS por dólar conseguirá suficientes clientes para llenar su capacidad de producción. Por eso creo que, en unos años, la competencia se volverá realidad bastante rápido.
Dicho eso, el TPU también es un recurso de cómputo relativamente escaso dentro de Google, y es muy probable que incluso les cueste cubrir la demanda interna.
Casi nunca digo que respeto tanto a una empresa, pero a Groq realmente la respeto.
Google inventó el TPU y Google Research incluso publicó el paper de los LLM, así que no entiendo por qué NVDA y las startups de IA se llevaron casi el 100% del valor.
En el caso de Google, la analogía sería que si alguien proponía ofrecer ampliamente algo como ChatGPT, eso podía canibalizar los anuncios pagados del buscador y los ingresos publicitarios de sitios que la gente ya no tendría que visitar. Así que quizá decidieron introducirlo con cuidado, solo cuando fuera necesario por competencia y de una forma menos disruptiva.
En realidad seguramente no es tan simple, pero si esa fuera la razón sería bastante gracioso.
https://www.linkedin.com/posts/eolver_googles-defense-agains...
Trabajo en Google; si hace tiempo que no miran los TPU, les recomiendo revisar v5. Ahora soporta PyTorch/JAX, así que es mucho más fácil de usar que cuando era solo para TensorFlow.
Este artículo conectó muy bien varias piezas que estaban dispersas de forma abstracta y mostró cómo fluyen realmente dentro del silicio.
Me gustó especialmente ver cómo instrucciones CISC simples se corresponden casi directamente con las etapas de inferencia de un LLM.
Tal vez sea una pregunta tonta que revela mi ignorancia, pero en el lado de consumo sigo escuchando que los chips M1 a M4 son buenos para algunas tareas de IA.
Hoy lo más importante para mí son herramientas como Photoshop y Resolve, y he visto que en los nuevos chips propios de Apple corren mucho más rápido que en mi máquina vieja.
Quizá esto no se traduzca bien a lo que pueden hacer este chip o un H100, pero me da curiosidad saber si se traduce en alguna medida. Claro que Apple no vende sus chips propios por separado, así que para hacerlo práctico tendría que lanzar algún producto tipo servidor externo cargado de GPU y chips de IA.
Los chips M* usan memoria unificada y, en particular, los Pro/Max/Ultra tienen un ancho de banda de memoria muy alto incluso comparados con GPU como la 1080. El ancho de banda de memoria del M1 Ultra está más o menos entre una 2080 y una 3090.
Con tamaños de batch pequeños, especialmente batch 1 como en la mayoría del trabajo local, la inferencia queda limitada por el ancho de banda de memoria más que por la capacidad de cómputo. Por eso se dice que los chips M* son buenos para machine learning.
Pero el H100 se usa principalmente para entrenamiento con tamaños de batch enormes, y para entrenar modelos grandes se necesitan muchas interconexiones. A esa escala, la intensidad aritmética es muy alta, así que aunque se pudieran conectar chips M* en red, no serían muy competitivos. Simplemente eligieron otro punto en la curva de Pareto de potencia/eficiencia frente a chips de alto consumo como el H100.
Lo que Google realmente tendría que hacer es entrar en el terreno de 2 nm EUV y bajar de los 2 nm.
Si llega a tener algo así, ya sea litografía electrónica o la tecnología que ASML imprime en los chips, se vuelve un actor realmente peligroso. Creo que necesitaría un proyecto moonshot hardcore al estilo Google X.
O quizá sí tenga unos 500 millones de dólares para comprar una máquina. Si el TPU es realmente tan bueno, podría ser un buen negocio integrarse verticalmente con tecnología propia y hasta con su propia fab.
Las fabs de vanguardia actuales son rentables gracias a décadas de talento e ingeniería acumulados, fabricando dispositivos de cómputo de propósito general para diversos clientes y usos. Además, los clientes impulsan innovación independiente en áreas clave, como las mejoras de rendimiento de HDI chip-on-chip de Micron, y el fabric de comunicación entre dies y el diseño de sustratos multichip de Xilinx.
El TPU jamás podría generar el volumen necesario, ni atraer clientes que creen economías de escala rentables. Google también tendría que ofrecer un precio atractivo frente a sus competidores.
Si hubiera un caso de negocio lo suficientemente convincente, las fabs existentes estarían encantadas de asignarle capacidad. El TPU todavía no es, ni de cerca, tan convincente.
Escuché una presentación de Jim Keller de TensTorrent donde explicaba otro enfoque para crear núcleos de IA. Consiste en usar 5 núcleos RISC-V: uno para cargar datos, otro para subir datos y los demás dedicados a operaciones matriciales.
También mencionó el TPU de Google, y dijo que programarlo se sentía como lidiar con VLIW, y que había unas 500 personas trabajando en el compilador.
En el texto original dice que “TPU v1 es un diseño CISC con unas 20 instrucciones”, y me da risa cómo CISC/RISC parece haber pasado de ser una observación aguda a un programa de investigación, luego una tecnología revolucionaria, después una palabra de moda de marketing y, finalmente, algo completamente carente de significado.
Supongo que así es el ciclo de vida de un término.
Así que aunque el TPU tenga pocas instrucciones, si cada una es bastante compleja, podría ser CISC. Dicho eso, la última vez que cursé arquitectura de computadoras fue hace 15 años en posgrado, así que mi memoria es borrosa. Además, pasé la mayor parte de ese semestre trabajando en cosas relacionadas con Itanium que ahora ya no sirven para nada.
La demanda de capacidad de foundries parece enorme, así que me pregunto cómo Microsoft o Google logran ponerse al frente de la fila cuando diseñan sus propios chips y necesitan producirlos.
¿Son lo bastante simples como para fabricarse en fabs “viejas y con menos demanda”? Tengo entendido que Apple y Nvidia ya tienen reservada mucha capacidad de foundry.
https://en.wikipedia.org/wiki/Tensor_Processing_Unit#Product...
También tienen una presencia e inversión considerables en áreas como HBM, y SemiAnalysis tiene buenos artículos al respecto.
Me da curiosidad cómo cambiará el hardware si los LLM realmente despegan con cuantización -1, 0, 1.