Tesla quiere llevar el silicio de aprendizaje automático al Dojo

Para saciar la sed de modelos de inteligencia artificial y aprendizaje automático cada vez más grandes, Tesla ha revelado una gran cantidad de detalles en Hot Chips 34 sobre su arquitectura de supercomputación totalmente personalizada llamada Dojo.

El sistema es esencialmente una supercomputadora componible masiva, aunque a diferencia de lo que vemos en el Top 500, está construido a partir de una arquitectura completamente personalizada que abarca la computación, las redes y el silicio de entrada/salida (E/S) hasta la arquitectura del conjunto de instrucciones (ISA), el suministro de energía, el empaquetado y la refrigeración. Todo se hizo con el propósito expreso de ejecutar algoritmos de entrenamiento de aprendizaje automático específicos y personalizados a escala.

“El procesamiento de datos del mundo real solo es factible a través de técnicas de aprendizaje automático, ya sea procesamiento de lenguaje natural, conducción en calles hechas para la visión humana o robótica que interactúa con el entorno cotidiano”, dijo Ganesh Venkataramanan, director senior de ingeniería de hardware en Tesla. durante su discurso de apertura.

Sin embargo, argumentó que los métodos tradicionales para escalar las cargas de trabajo distribuidas no lograron acelerar al ritmo necesario para mantenerse al día con las demandas del aprendizaje automático. En efecto, la Ley de Moore no es suficiente y tampoco los sistemas están disponibles para el entrenamiento de IA/ML a escala, es decir, alguna combinación de CPU/GPU o, en circunstancias más raras, mediante el uso de aceleradores de IA especializados. 

“Tradicionalmente construimos chips, los colocamos en paquetes, los paquetes van en PCB, que van a los sistemas. Los sistemas van a los bastidores”, dijo Venkataramanan. El problema es que cada vez que los datos se mueven del chip al paquete y fuera del paquete, se incurre en una penalización de latencia y ancho de banda.

Un sándwich de centro de datos

Entonces, para sortear las limitaciones, Venkataramanan y su equipo comenzaron de cero.

“Desde mi entrevista con Elon, me preguntó qué se puede hacer que sea diferente de las CPU y las GPU para la IA. Siento que todo el equipo todavía está respondiendo esa pregunta”.

Esto condujo al desarrollo del mosaico de entrenamiento Dojo, un clúster de computación autónomo que ocupa medio pie cúbico con capacidad para 556 TFLOPS de rendimiento FP32 en un paquete refrigerado por líquido de 15kW.

Cada mosaico está equipado con 11 GB de SRAM y está conectado a través de un tejido de 9 TB/s utilizando un protocolo de transporte personalizado en toda la pila.

“Este mosaico de capacitación representa cantidades incomparables de integración de la computadora a la memoria, a la entrega de energía, a la comunicación, sin necesidad de interruptores adicionales”, dijo Venkataramanan.

En el corazón del mosaico de entrenamiento se encuentra el D1 de Tesla, un dado de transistores de 50 mil millones, basado en el proceso de 7nm de TSMC. Tesla dice que cada D1 es capaz de 22 TFLOPS de rendimiento FP32 a un TDP de 400W. Sin embargo, Tesla señala que el chip es capaz de ejecutar una amplia gama de cálculos de coma flotante, incluidos algunos personalizados.

“Si comparas los transistores por milímetro cuadrado, esta es probablemente la vanguardia de todo lo que existe”, dijo Venkataramanan.

Luego, Tesla tomó 25 D1, los clasificó en troqueles buenos conocidos y luego los empaquetó utilizando la tecnología de sistema en oblea de TSMC para "lograr una gran cantidad de integración informática con una latencia muy baja y un ancho de banda muy alto", dijo.

Sin embargo, el diseño de sistema en oblea y la arquitectura apilada verticalmente presentaron desafíos en lo que respecta a la entrega de energía.

Según Venkataramanan, la mayoría de los aceleradores de hoy colocan la energía directamente junto al silicio. Y aunque está probado, este enfoque significa que una gran parte del acelerador debe dedicarse a esos componentes, lo que lo hace poco práctico para Dojo, explicó. En cambio, Tesla diseñó sus chips para entregar energía directamente a través de la parte inferior de la matriz. 

Cómo aplicar todos los conceptos

“Podríamos construir un centro de datos completo o un edificio completo a partir de este mosaico de capacitación, pero el mosaico de capacitación es solo la parte de cómputo. También necesitamos alimentarlo”, dijo Venkataramanan.

Para esto, Tesla también desarrolló el Procesador de interfaz Dojo (DIP), que funciona como un puente entre la CPU anfitriona y los procesadores de entrenamiento. El DIP también sirve como fuente de memoria compartida de gran ancho de banda (HBM) y como NIC de alta velocidad de 400 Gbit/seg.

Cada DIP cuenta con 32 GB de HBM y hasta cinco de estas tarjetas se pueden conectar a un mosaico de entrenamiento a 900 GB/s para un total de 4.5 TB/s al host para un total de 160 GB de HBM por mosaico.

Los pares de configuración V1 de Tesla de estos mosaicos, o 150 troqueles D1, en matriz admitieron cuatro CPU de host, cada una equipada con cinco tarjetas DIP para lograr un exaflop declarado de rendimiento BF16 o CFP8.

En conjunto, Venkataramanan dice que la arquitectura, detallada en profundidad aquí by La próxima plataforma – permite a Tesla superar las limitaciones asociadas con los aceleradores tradicionales de Nvidia y AMD.

“Cómo funcionan los aceleradores tradicionales, por lo general intentas colocar un modelo completo en cada acelerador. Replíquelo y luego haga fluir los datos a través de cada uno de ellos”, dijo. “¿Qué pasa si tenemos modelos cada vez más grandes? Estos aceleradores pueden fracasar porque se quedan sin memoria”.

Este no es un problema nuevo, anotó. El conmutador NV de Nvidia, por ejemplo, permite agrupar la memoria en grandes bancos de GPU. Sin embargo, Venkataramanan argumenta que esto no solo agrega complejidad, sino que introduce latencia y compromete el ancho de banda.

“Pensamos en esto desde el principio. Nuestros mosaicos de computación y cada uno de los troqueles se hicieron para adaptarse a modelos grandes”, dijo Venkataramanan.

Software

Una arquitectura informática tan especializada exige una pila de software especializada. Sin embargo, Venkataramanan y su equipo reconocieron que la programabilidad haría o arruinaría el Dojo.

“La facilidad de programación para las contrapartes de software es primordial cuando diseñamos estos sistemas”, dijo. “Los investigadores no esperarán a que su gente de software escriba un kernel escrito a mano para adaptarse a un nuevo algoritmo que queremos ejecutar”.

Para hacer esto, Tesla abandonó la idea de usar núcleos y diseñó la arquitectura de Dojo en torno a los compiladores.

“Lo que hicimos fue usar PiTorch. Creamos una capa intermedia, que nos ayuda a paralelizar para escalar el hardware debajo de ella. Debajo de todo hay código compilado”, dijo. “Esta es la única forma de crear pilas de software que se adapten a todas esas cargas de trabajo futuras”.

A pesar del énfasis en la flexibilidad del software, Venkataramanan señala que la plataforma, que actualmente se ejecuta en sus laboratorios, está limitada al uso de Tesla por el momento.

“Nos enfocamos primero en nuestros clientes internos”, dijo. “Elon ha hecho público que, con el tiempo, pondremos esto a disposición de los investigadores, pero no tenemos un marco de tiempo para eso. ®