DruckFin

Transcripción de Google: Co-diseño de TPUs, sistemas distribuidos y modelos de IA de frontera

Fireside Chat de Google Dev Labs con Jeff Dean y Bill Jia sobre el escalado de sistemas de IA, 20 de septiembre de 2026

Introducción y los primeros días de Google

Moderador: Hola a todos, bienvenidos. Muchas gracias por acompañarnos y por todo su entusiasmo. Este es nuestro séptimo Dev Labs, y cada vez que lo hacemos se vuelve más grande y mejor. Y este es, por mucho, el mejor hasta ahora. Hubo muchísimo interés, así que gracias por sus comentarios. La dinámica será la siguiente: vamos a conversar sobre un par de preguntas, hablar sobre la evolución de las TPUs. Él es Jeff, quien no necesita presentación alguna. Fue el empleado número 30 en Google. Así que empecemos con esa historia, y también tenemos aquí a Bill, quien hará algunas preguntas, pero a quien también queremos pedirle que nos comparta algunos detalles sobre la ejecución de la infraestructura de IA, la visión y cómo se está ejecutando ahora y en los próximos años, si les parece bien. Muy bien, manos a la obra. Jeff, cuéntanos un poco más. Fuiste el empleado número 30 en Google.

Jeff Dean: Sí, así es.

Moderador: 1999. La mayoría de los presentes aquí ni habían nacido, o tal vez sí, o estaban en pañales. ¿Cómo era aquello?

Jeff Dean: No, fue muy divertido. O sea, creo que cuando entré, estábamos todos en una oficina diminuta en el centro de Palo Alto, metidos justo arriba de lo que hoy es una tienda de T-Mobile, más o menos a esa escala de oficina. Y estábamos intentando construir un producto de búsqueda de muy alta calidad. Yo había tenido cierta exposición a eso en un trabajo anterior sobre recuperación de información, en particular en la web, utilizando la estructura de grafos de la web para complementar la información existente en las páginas y demás. Así que Google parecía un lugar natural para intentar eso en un entorno menos académico y más enfocado en personas que usaran mi trabajo directamente, lo cual era súper divertido. Había una energía increíble y duplicábamos nuestro tráfico como un 7% cada semana o algo así. Y si calculas 1.07 elevado a la 52, sabes que te estás derritiendo cada año e intentando no caerte los martes al mediodía. Así que fue muy divertido. Y a medida que la compañía ha seguido creciendo, hemos hecho más y más cosas, lo cual siempre es emocionante de ver.

Moderador: ¿Y qué pasó en 2001, si avanzamos rápido en el tiempo?

Jeff Dean: Para evitar derriternos, una de las cosas que hacíamos era reescribir constantemente todo el sistema de búsqueda, indexación y atención de consultas para que fuera más eficiente, usara diferentes estructuras de datos y representaciones más compactas de las cosas. Y llegó un punto en 2001 en el que nos dimos cuenta de que la forma en que estábamos escalando tanto el tamaño del índice como la capacidad que teníamos para servir el índice consistía en crear más y más particiones del índice —a las que llamábamos shards— para tener un índice más grande. Y luego, para cada uno de esos shards, creabas más y más copias para proporcionar la capacidad de atender más consultas simultáneamente. Y al final llegamos al punto en que, en lugar de 7 shards y 10 copias cada uno, teníamos como 60 shards y 20 copias de cada uno en cada centro de datos. Y finalmente hicimos cuentas y nos dimos cuenta de que podíamos poner, en lugar de tener el índice en disco, todo el índice de la web en la memoria de las 1.200 máquinas que de otro modo serían 20 copias de las 60 particiones. Y eso fue muy divertido. Aceleramos todo nuestro sistema de búsqueda unas 4 veces en 3 días, lo cual fue bastante divertido, sí.

La génesis de la TPU: del reconocimiento de voz a los ASIC personalizados

Moderador: Y luego, en 2013, hiciste algunos cálculos rápidos en una servilleta.

Jeff Dean: Sí, exacto. Empecé a trabajar en la infraestructura de entrenamiento de deep learning y en el entrenamiento de modelos de deep learning porque parecía la abstracción correcta. De hecho, había tenido contacto con las redes neuronales como estudiante universitario en 1990, e hice mi tesis de licenciatura sobre el entrenamiento paralelo de redes neuronales porque simplemente sentía que, vaya, esto es lo único que necesitamos. Solo necesitamos conseguir 32 procesadores modestos en lugar de uno, y entonces podremos entrenar modelos increíbles. Resulta que necesitábamos un millón de veces más potencia de cálculo, no 32. Pero eventualmente, a través de las mejoras de la ley de Moore y demás, empezamos a contar con eso a partir de 2008 o 2009. Y luego, en 2011, iniciamos un proyecto para intentar realmente escalar y hacer entrenamiento distribuido de modelos para voz, visión y lenguaje utilizando —en ese momento teníamos CPUs en nuestros centros de datos, muchísimas CPUs—. Así que usamos unos 16.000 núcleos de CPU para entrenar un modelo de visión grande, y obtuvimos mejoras asombrosas. Y lo estábamos haciendo también para modelos de voz, y los modelos representaban mejoras increíbles en calidad. Así que obtuvimos el equivalente a los últimos 20 años de investigación en voz en términos de mejoras en la tasa de error por palabra en los dos meses que tardamos en construir un modelo acústico profundo de voz. Y fue como: "¡Oh, por Dios, esto va a ser maravilloso!". Pero, ¿cómo lo implementamos para dar servicio?

Y en ese momento hice algunos cálculos donde dije: "Muy bien, supongamos que 100 millones de personas empiezan a hablarle a su teléfono durante unos minutos al día, ¿cómo atenderemos eso?". Y resulta que habríamos necesitado duplicar el número de ordenadores que tenía Google solo para implementar el mejor modelo de reconocimiento de voz en una parte diminuta y oscura de Google. Y eso parecía un poco excesivo, por no mencionar que era impracticable. Así que decidimos que el hardware especializado era el camino a seguir. Y esa es un poco la génesis de la familia TPU: queríamos construir aceleradores especializados que básicamente fueran muy, muy buenos en álgebra lineal de baja precisión y en nada más. Y así es como puedes servir los modelos de voz, los modelos de visión y demás de forma mucho más eficiente. Al final resultó ser entre 30 y 80 veces mejor rendimiento por vatio que las CPU y GPU contemporáneas de la época, en 2015, cuando tuvimos el chip disponible. Y esa fue la primera iteración.

Bill Jia: Sí, la TPU v1, que estaba enfocada principalmente en la inferencia, ¿verdad? Y las generaciones posteriores se han enfocado sobre todo en el entrenamiento y también en la inferencia. Y más recientemente, incluso hemos empezado a separar esas líneas de nuevo porque el entrenamiento y la inferencia son un poco diferentes.

Moderador: Y cuando publicaron el artículo en 2017, fue aceptado. Es el artículo más citado en los 50 años de historia de...

Jeff Dean: Sí, de ISCA, exacto. El artículo sobre la TPU v1 tiene muchos autores; probablemente sean como 35 autores o algo así. Pero está muy bien, considerando que es relativamente reciente en los 50 años de historia de la arquitectura de computadoras.

Evolución de la arquitectura de las TPU, pilas de software y conmutación óptica

Bill Jia: Es realmente fascinante que Google fuera considerada una empresa de software de primer nivel, pero que luego empezara a impulsar su propia hoja de ruta de hardware. Es verdaderamente increíble. Y fabricar hardware es una cosa, ¿claro?, porque para hacer que la TPU sea fantástica se requiere mucho diseño de placas, de conjuntos de chips y de la propia TPU. Pero también necesitamos que el compilador y los frameworks funcionen a la perfección. Al principio, ustedes impulsaron XLA, que es nuestro compilador para TPU, y también impulsaron TensorFlow. Si pudieras contarnos cómo fueron los inicios al impulsar esta pila de software para la TPU.

Jeff Dean: Sí. O sea, creo que lo que buscas como desarrollador o investigador de machine learning es simplemente poder expresar tu idea de alto nivel y lograr que se ejecute mágicamente en un sistema a gran escala sin tener que pensar demasiado en todos los aspectos de rendimiento que podrías necesitar gestionar bajo el capó para conseguir un verdadero buen rendimiento. Y eso significa contar con un buen compilador, interconexiones de altísima calidad y frameworks que expresen la abstracción de: ¿me estoy comunicando con 4 chips en mi máquina local o me estoy comunicando con 10.000 chips en una configuración distribuida masiva? Te gustaría no tener que pensar demasiado en eso ni hacer cosas distintas para esos dos escenarios.

Moderador: Háblanos de las iteraciones de las diferentes TPU y hacia dónde crees que se dirige todo esto.

Jeff Dean: Sí. Después de la TPU v1, que era básicamente una tarjeta PCIe que cabía en una ranura —y de hecho compramos un montón de ellas sin saber muy bien cómo íbamos a usarlas; arrinconé al director financiero de entonces y le dije: "Tenemos que comprar muchas de estas porque las vamos a usar; no estamos seguros para qué, ¿pero por favor podemos comprar muchas?" Y así lo hicimos—. Luego, la TPU v2 fue prácticamente el primer sistema que diseñamos pensando en el entrenamiento, no solo con un chip individual, sino con todo un sistema similar a una supercomputadora compuesto por muchos chips interconectados con un toroide en 2D. A lo largo de las generaciones, hemos introducido la refrigeración líquida, a partir de la TPU v3. Así que siempre es emocionante cuando tienes tuberías llegando a la superficie de los chips en tu sistema; las fugas nunca son algo bueno.

Y después vino la TPU v4: a medida que escalabas el tamaño del pod —las TPU v2 y v3 tenían un cableado fijo entre 256 chips y luego 1.024 chips—, al empezar a manejar sistemas cada vez más grandes, comienzas a experimentar fallas en chips, placas o sistemas individuales. Necesitas tener una mayor tolerancia a fallas porque sigues queriendo aprovechar toda la topología de un pod mucho más grande, incluso si algo se rompe. Y por eso la TPU v4 introdujo este concepto de red óptica reconfigurable entre racks de máquinas.

Moderador: Así que uno puede ensamblar algo parecido a bloques de Lego desde cualquier rincón del piso del centro de datos, ¿verdad?

Jeff Dean: Con pequeños conmutadores curiosos que tienen espejos micromecánicos ajustables para dar la impresión de que estos ocho racks están justo uno al lado del otro, aunque estén dispersos por varios pasillos del centro de datos. Y eso ha sido de una utilidad tremenda. La TPU v5, v6 y v7 introducen formatos de cálculo de precisión cada vez menor, lo cual ha sido sumamente útil. Obtienes un rendimiento mucho mayor gracias a ello, aunque siempre resulta extraño comparar los FLOPS cuando algunos FLOPS son de 16 bits, otros de 8 bits y otros de 4 bits. Si logras aprovecharlos en tus algoritmos de machine learning, los formatos de menor precisión son una forma fantástica de conseguir todavía más rendimiento.

Pathways: orquestación de la irregularidad y entrenamiento distribuido masivo

Moderador: Y entonces la primera TPU fue en 2015; en 2017 se publicó el artículo. Ya es toda una leyenda. Y luego, en 2018, llegó Pathways. Hablemos un poco de Pathways y de por qué es tan importante.

Jeff Dean: Sí. Empezamos a construir Pathways fundamentalmente para gestionar el entrenamiento de modelos mucho más irregulares y dispersos, y también para que sirviera como infraestructura de sistemas subyacente que nos permitiera contar con un único modelo de programación, un único proceso capaz de impulsar un sistema de entrenamiento distribuido a gran escala. Como programador, es mucho más agradable si tu proceso único en Python simplemente parece tener 8.000 dispositivos conectados y puedes decir: "Genial, solo quiero hacer un AllReduce en los 8.000 dispositivos", y el sistema lo hace realidad sin que tengas que hacer nada especial.

El software subyacente, tanto en el compilador XLA como de forma un poco más abstracta en un nivel superior, hace que Pathways orchestre el movimiento de datos entre una gran cantidad de chips diferentes. Así que tienes chips dentro del mismo pod de TPU, y si necesitas comunicarte ahí, utilizará los enlaces de alta velocidad ICI que se encuentran en el pod de TPU. Pero si necesitas algo que abarque múltiples pods, o comunicarte desde este chip en este pod hacia este otro chip en otro pod, Pathways se encargará de orquestar esa transferencia utilizando la red del centro de datos, o incluso una red de área amplia (WAN) si estás ensamblando un trabajo de entrenamiento combinando pods en Oklahoma, Texas e Iowa, por ejemplo. Y eso te otorga esa bonita abstracción de: simplemente tengo una enorme cantidad de capacidad de cómputo, y voy a permitir que el sistema determine la mejor manera de utilizarla.

Co-diseño de los modelos Gemini con el hardware y la infraestructura

Bill Jia: Sí, es increíble. Así que, Jeff, desde muy al principio estuviste muy involucrado en el diseño de redes para la TPU, en el diseño de la propia TPU, luego en el compilador y el framework, y por supuesto en Pathways, que se encarga de orquestar todo el movimiento de datos, la reconfiguración del tráfico de red, absolutamente todo. Toda esa es la pila de infraestructura de IA. Ahora estás muy profundamente involucrado en el diseño de los modelos Gemini, los modelos Gemini actuales y también los futuros modelos Gemini. De este modo, los modelos Gemini se pueden co-diseñar junto con toda la pila de infraestructura. Desde tu perspectiva, ¿qué tipo de beneficios ya hemos obtenido y qué ideas de futuro tienes para que podamos seguir empujando los límites tecnológicos de los modelos?

Jeff Dean: Creo que todo el conjunto de generaciones de TPU se ha beneficiado enormemente del hecho de que nosotros mismos somos grandes usuarios del cómputo de machine learning, y tenemos a muchísimas personas intentando superar los límites de nuevas ideas de investigación que podrían poner a prueba el hardware de maneras que las cargas de trabajo existentes no hacen. Así que al tener todo bajo el mismo techo de Alphabet, podemos fomentar una gran interacción entre los diseñadores de las futuras generaciones de TPU, los ingenieros de software de los compiladores y la infraestructura, y los investigadores de machine learning que pueden decir: "Oigan, este enfoque parece funcionar a pequeña escala; creemos que será muy crucial para futuros ciclos de entrenamiento a gran escala dentro de un año o dos, y queremos asegurarnos de que nuestro hardware lo soporte".

Porque como arquitecto de computadoras, si haces esto de forma aislada, intentando adivinar hacia dónde se dirige el vertiginoso campo de la IA en el horizonte de dos a seis años en el que el chip específico que estás diseñando hoy debe seguir siendo relevante, resulta sumamente difícil. Cuantas más perspectivas puedas obtener sobre qué funcionará —y cuantas más iteraciones realces— mejor. No se trata solo de pensar que necesitamos hacer algo; a menudo ocurre que los ingenieros de hardware dicen: "Bueno, eso es difícil, pero podríamos hacer esto otro, ¿sería útil?". Y se genera un tira y afloja en el que dices: "Sí, eso nos permitiría hacer algo muy cercano a esto, o quizás incluso mejor". Y eso es fundamental para lograr un co-diseño eficaz.

Bill Jia: Ayer estuve en el escenario dando la apertura. Estuve hablando de que, a medida que empezamos a escalar los modelos Gemini —recuerdo, recuerdo perfectamente cómo cada mes nos sentábamos en las mismas reuniones para hablar de confiabilidad. Así que yo estoy del lado de la infraestructura y Jeff representa a Google DeepMind—. Decíamos: "Oigan, los modelos Gemini son cada vez más grandes, así que utilizamos montones de TPU para entrenarlos". Al final, llegamos a utilizar 100.000 TPU para entrenar los modelos. Esa es la métrica principal que vigilamos juntos, la cual se denomina *goodput* (rendimiento útil).

Al principio, nuestro *goodput* no era bueno. Recuerdo lo que decíamos juntos: el *badput* dominaba sobre el *goodput*. Al comienzo era como del 60%. A medida que escalamos más y más TPU para realizar el preentrenamiento de ese modelo, la situación se vuelve muy compleja. Necesitamos mayor escala, más datos, más paralelismo. El *goodput* era del 50%, 60%, 70% al principio, pero simplemente no era suficiente, ¿verdad? Pero para abreviar la historia, hoy utilizamos una infraestructura de entrenamiento a escala masiva, entrenamos modelos mucho más complejos y el *goodput* puede alcanzar hasta el 95% e incluso el 98%.

Jeff Dean: Sí, y creo que eso es en realidad la combinación de muchas cosas: mejores prácticas operativas, mejores pruebas de control de calidad en los chips cuando se despliegan, soluciones de software que lidian mejor con las fallas y permiten seguir avanzando aunque una parte del sistema falle. Todos estos factores juntos importan muchísimo y marcan una gran diferencia.

Bill Jia: Sí, creo que esa es otra ventaja de ser dueños del modelo, de los datos, del hardware completo, de las redes, de toda la pila de software y del equipo de operaciones.

Jeff Dean: Puedes decir: "Oye, esto es sumamente importante", porque si piensas en una bandeja individual de TPU y crees que no forma parte de un sistema más grande, dices: "Bueno, no pasa nada, podemos repararla en una semana cuando pasemos por aquí". Pero si es parte de un pod que está en uso y tiene bandejas dañadas, quieres ir a repararlas de inmediato, porque el radio de impacto de lo que necesitas arreglar es en realidad mucho mayor de lo que parece.

Moderador: Entonces, cuando hablamos de 100.000 chips, la confiabilidad ya no es un problema de parche de software. ¿Dirías que se trata más de diseño de sistemas? ¿Está más relacionado con la física?

Jeff Dean: Creo que la confiabilidad es verdaderamente una propiedad de todo el sistema, y hay muchos aspectos diferentes que deseas volver robustos. Una cosa que puedes hacer es construir componentes robustos a partir de piezas poco confiables. Esto viene desde los mismos inicios de Google. Comprábamos ordenadores de consumo baratos ("El Cheapo") para atender nuestro tráfico de búsqueda, y construíamos sistemas de software robustos sobre ellos que nos permitían gestionar fallas en máquinas individuales y aun así ofrecer la funcionalidad que el sistema debía proporcionar. Quizás sea un sistema de archivos distribuido, por lo que replicas los datos en múltiples máquinas y siempre tienes algunas disponibles aunque algunas réplicas de un fragmento de datos fallen.

Y creo que se puede hacer exactamente lo mismo en sistemas de entrenamiento de IA a gran escala. Si tienes una configuración predeterminada de 20 pods, pero uno de ellos está caído, puedes seguir avanzando con los 19 pods restantes mientras reparas el vigésimo.

Coordinación operativa y mitigación de la corrupción silenciosa de datos

Moderador: Y Bill, tú tocaste un poco este tema ayer en tu discurso principal sobre OCS y Jupiter. ¿Quieres hablarnos un poco de eso y de cómo influye en la confiabilidad?

Bill Jia: Sí, creo que ha sido un largo camino. Lo primero es examinar las interrupciones por cada 10.000 chips, por cada componente individual: cuántas interrupciones diarias tenemos. Tenemos que minimizar eso, porque si las interrupciones son demasiadas a medida que escalamos el clúster de entrenamiento, con un exceso de interrupciones, la situación no es muy favorable, ¿verdad? Así que una de las cosas que hicimos fue mejorar la confiabilidad del hardware. Y no solo eso, sino que ayer mencionaba que incluso antes de que comience el preentrenamiento, el software escanea toda la flota, cada uno de los componentes: ¿cuál es tu signo vital? Si el signo vital presenta algún problema, lo reparamos o lo excluimos antes de que el entrenamiento siquiera empiece. Esa es la primera parte.

La segunda parte: el entrenamiento da comienzo. Pero una vez iniciado, incluso si el hardware muestra signos vitales muy sólidos, el entrenamiento dura semanas y aún pueden surgir problemas. ¿Cómo podemos identificar qué componentes y servidores específicos tienen problemas para luego corregirlos, excluirlos, reemplazarlos y solucionarlos? Esa es la segunda parte.

Y ahora también, como mencionó Jeff, hay una enorme coordinación operativa. Si el centro de datos está realizando labores de mantenimiento eléctrico en toda una fila del centro de datos, y esa misma fila está ejecutando el entrenamiento en ese momento, por supuesto que no es una buena idea. Así que nos coordinamos estrechamente con las configuraciones del centro de datos, los ingenieros de confiabilidad de sitios (SRE) y todos los investigadores e ingenieros de machine learning. Requiere muchísima coordinación; tenemos que estar sincronizados.

Y creo que Jeff también lidera una gran parte del diseño de los modelos Gemini y del diseño de software. A veces las cosas fallan, pero el entrenamiento combina el paralelismo de datos y el paralelismo de modelos de manera simultánea. Si eso ocurre dentro de una misma réplica de datos, tal vez la réplica adicional simplemente continúe su curso, ¿no? Promedias los pesos automáticamente para no tener que preocuparte por la caída de esa única réplica. Así que se trata de un conjunto muy completo de mejoras. ¿Algo que quieras agregar, Jeff?

Jeff Dean: No. O sea, creo que pueden suceder todo tipo de imprevistos y es muy difícil predecir todas las formas posibles en que algo podría fallar. Por lo tanto, construir sistemas robustos capaces de detectar fallas —incluso cuando no estás necesariamente seguro de qué las causó, pero puedes detectarlas— es clave. Ocasionalmente te encuentras con chips que, tal vez cuando la temperatura sube, empiezan a mostrar problemas de confiabilidad, como sumar 2 más 2 y obtener 5 o algo así.

Bill Jia: Ese es el peor de todos, un error de corrupción silenciosa de datos. Es horrible, horrendo.

Jeff Dean: Sí, y eso puede ocurrir en el propio chip, en un enlace de red inestable o en muchos otros lugares. Estamos aplicando muchos de los mismos principios que utilizábamos en los primeros días de Google, donde, como comprábamos PCs de consumo baratas, estas no solo carecían de memoria ECC, sino que tampoco tenían paridad en su memoria. Por lo tanto, si utilizas una gran cantidad de computadoras para hacer algo y ninguna tiene paridad, vas a obtener inversiones aleatorias de bits en las cosas. Muchos de los cálculos que realizábamos requerían ser tolerantes a inversiones de bits. Una de las formas de lograrlo era: muy bien, si estoy procesando mil millones de páginas web y descarto una de ellas, probablemente no sea el fin del mundo. Así que podías simplemente ignorarlo mediante la comprobación de sumas (*checksumming*) en el software por encima de la capa de hardware para mantener la robustez, incluso si una máquina en particular o el enlace de red no eran confiables.

Estrategia de código abierto: JAX, StableHLO y PyTorch en TPU

Moderador: Ahora bien, Google cree firmemente en el código abierto, así que voy a cambiar un poco de tema para hablar de open source. Sé que tenemos a muchos académicos en la audiencia y también a muchos de nuestros socios. Hablemos un poco sobre la estrategia de infraestructura de IA de Google en lo que respecta a garantizar que no se convierta en un recurso cerrado bajo llave. Todo esto de lo que hablamos queremos asegurarnos de compartirlo con el mundo. Hablemos de JAX, StableHLO y de esa evolución del código abierto en su conjunto, y de cómo resulta tan relevante hoy en día para asegurar que esto no se convierta en un coto privado.

Jeff Dean: Sí, creo que para nosotros es de suma importancia interactuar con el ecosistema en general. Abrir el código para que la gente pueda examinarlo, modificarlo, ayudarnos a mejorarlo y convertirlo en un proyecto colaborativo que abarque a múltiples organizaciones, en lugar de ser algo que nos guardamos para nosotros, es sumamente importante. Esa es la razón por la cual liberamos TensorFlow, por la cual hicimos de código abierto a JAX y la representación StableHLO para el compilador XLA. Y seguiremos abriendo más cosas en el futuro. Hemos sido grandes contribuyentes al código abierto en general, al núcleo de Linux y demás, durante muchos años. Por lo general, somos una de las organizaciones más grandes que contribuyen a los esfuerzos colectivos de open source porque creemos en ello y consideramos que todo el ecosistema se beneficia cuando todos trabajan juntos.

Moderador: Es la transición perfecta para que hables de PyTorch y de tu experiencia en Meta.

Bill Jia: Sí, estuve hablando con Jeff justo antes de esta charla. Jeff ha liderado gran parte de la estrategia de código abierto en infraestructura de IA. Como mencionó Jeff, en el pasado Google abrió el código de TensorFlow y JAX; Kubernetes también fue creado y liberado por Google, al igual que Android y muchas otras cosas fantásticas. Ahora, a medida que impulsamos la TPU hacia la comunidad y enfatizamos su uso en Google Cloud, estamos redoblando nuestra apuesta por esta estrategia de código abierto.

De hecho, no solo liberamos el núcleo de JAX, sino que también construimos muchas bibliotecas de nivel superior, como la forma de hacer aprendizaje por refuerzo —lo llamamos TuneX, que en realidad está construido sobre el núcleo de JAX—. También lo hemos hecho de código abierto. Cómo realizar la gestión de puntos de control (*checkpointing*), cómo gestionar la inferencia: construimos muchas bibliotecas y frameworks de nivel superior sobre JAX y los estamos abriendo al público porque queremos que la gente los utilice y aproveche esta estrategia de código abierto para usar la TPU directamente. Ese es un aspecto.

Y por supuesto, queremos encontrarnos allí donde está el cliente, ¿verdad? Así que también queremos abrazar cualquier producto maduro que ya exista en la comunidad de código abierto. PyTorch es muy maduro y se utiliza ampliamente en la comunidad. ¿Saben qué? Adoptemos PyTorch en la TPU. De hecho, tenemos un proyecto llamado Torch-TPU. Actualmente se encuentra en versión de vista previa privada ("private preview"). Queremos lanzar la versión de vista previa pública el próximo trimestre, en apenas un par de meses, y luego, en el cuarto trimestre, queremos ofrecer una versión pública en GitHub para que todo el mundo pueda usarla. Y como mencioné ayer, si usamos Torch-TPU, se trata de cambiar apenas unas pocas líneas de código sencillas: el dispositivo backend pasa a ser la TPU y, con suerte, todo el entrenamiento y la inferencia se ejecutan en la TPU. De hecho, admitimos tanto el modo imperativo (*eager mode*) como el modo compilado. Eso a nivel de framework.

También nos relacionamos —estoy seguro de que en la audiencia hay muchos usuarios de vLLM y SGLang—. Hablamos también con los desarrolladores de vLLM y SGLang para asegurarnos de que este tipo de frameworks de inferencia de código abierto de nivel superior también puedan utilizarse en la TPU. Así que hay muchísimo trabajo en marcha: abrimos el código de nuestros propios desarrollos y también adoptamos productos de código abierto maduros de la comunidad.

Moderador: Por cierto, tenemos muchos estudiantes de doctorado de primer y segundo año en la audiencia también. ¿Cuántos de ustedes son estudiantes de PhD? ¿En qué cosas deberían estar enfocándose? Algunos también cursan carreras universitarias de grado, para que lo sepan.

Jeff Dean: Muy bien. Siempre es emocionante estar en esa etapa de tu carrera, porque creo que puedes encontrar cosas que realmente disfrutas y consideras importantes, y descubrir la manera de impulsar algún área para generar un impacto en el mundo. Creo que es una etapa sumamente divertida.

Moderador: ¿Hay algún cuello de botella físico en el que estés pensando desde la perspectiva de la infraestructura?

Jeff Dean: Hay muchísimos. Soy bastante optimista respecto al hardware cada vez más especializado, porque creo que esa es la manera de conseguir sistemas verdaderamente mucho más eficientes. Y ahora tenemos cargas de trabajo donde un puñado de ellas representarán una parte enorme del cómputo en todo el mundo, ¿cierto? Y si lo piensas, eso clama a gritos por especialización. El problema con la especialización es que si lo que deseas hacer cambia en el futuro, aquello que esculpiste con tanto cariño en hardware durante tal vez dos años puede dejar de ser tan relevante.

Por lo tanto, para lograr que la especialización funcione de verdad, necesitas automatizar mucho más el proceso de diseño de hardware. En la actualidad, la forma en que se diseña el hardware consiste en reunir a un gran equipo de personas. Algunas toman la especificación de alto nivel y construyen RTL de bajo nivel a partir de ella. Luego, debido a que se ha traducido manualmente, terminas necesitando otro equipo de personas cuya tarea es verificar que el primer equipo hizo lo correcto. Y después tienes a otro grupo que diseña físicamente el chip en detalle. Ante esto, parece evidente que si logras crear bucles más automatizados que puedan ser explorados mediante aprendizaje por refuerzo u otras técnicas evolutivas, y consigues que esos bucles se ejecuten con la suficiente rapidez —algo para lo que las herramientas EDA actuales normalmente no están diseñadas—, entonces tendrás la oportunidad de contar con un bucle de exploración del proceso de diseño mucho más automatizado, y quizás comprimir drásticamente el ciclo de diseño.

Si pudieras, por ejemplo, diseñar un nuevo chip con 10 personas en 3 meses en lugar de 150 personas y 2 años, verías mucho más hardware especializado en el mundo. Y tendrías que apostar mucho menos por el futuro: la incógnita sobre qué tipo de cómputo querrás realizar de 2 a 6 años en el futuro se convertiría más bien en un horizonte de 3 a 6 meses a 4 años, y esa es una apuesta mucho más fácil de hacer.

Bill Jia: Si retrocedemos uno o dos años, en aquella época gran parte de la industria, incluido Google junto con otros laboratorios punteros, centraba su enfoque principal en cómo hacer que el modelo fuera extraordinario. Había un enorme enfoque en lograr que el preentrenamiento y el postratrenamiento funcionaran bien, por lo que gran parte de la estrategia de hardware se centraba en el lado del entrenamiento. Pero ahora, a medida que los modelos grandes van madurando, una gran cantidad de tráfico comienza a dirigirse hacia el mundo de los agentes y hacia el lado de la inferencia. Así que, Jeff, si pudieras compartir alguna reflexión desde el punto de vista del diseño de hardware: ¿cómo podemos diseñar hardware para enfocarnos y priorizar el tráfico de inferencia?

Jeff Dean: La inferencia y el entrenamiento son algo diferentes. Para la inferencia, básicamente tienes un modelo y solo quieres atender una gran cantidad de solicitudes para él. Por lo tanto, buscas mover la menor cantidad posible de información que no cambia. Lo que no cambia son los pesos del modelo, y lo que sí cambia son las solicitudes, el caché de claves y valores (KV cache), etc. Así que necesitas diseñar un sistema que sea verdaderamente eficiente minimizando ese movimiento de datos.

Y también creo que, a medida que ves no solo un aviso (*prompt*) y luego una respuesta, sino algo mucho más independiente —un agente hace algunas cosas, decide invocar una herramienta, la herramienta se ejecuta, los resultados de la herramienta regresan, integras todo eso en el contexto del modelo y dejas que decida qué hacer a continuación—, nos daremos cuenta de que todas nuestras herramientas son demasiado lentas, porque fueron diseñadas para iteraciones a la velocidad humana. Si estás compilando código como parte de tu herramienta, te frustrará mucho que tu compilador sea lento. Si haces que tu hardware de inferencia sea súper rápido, podrá generar código quizás incluso más rápido de lo que puedes compilarlo, y ciertamente más rápido de lo que puedes ejecutarlo.

Así que una de las cosas que hemos estado haciendo es acelerar algunas de nuestras herramientas internas para optimizarlas. Resulta que en realidad se puede traducir de un lenguaje de programación a otro de manera bastante eficaz, porque dispones de una especificación completa y totalmente detallada de lo que deseas. Tienes el programa entero escrito en un lenguaje interpretado como Python, y solo quieres el programa equivalente exacto en Go, Rust, C++ o el que sea. Y un agente puede hacer eso bastante bien, lo cual difiere bastante del tipo habitual de interacción de codificación que tenemos con un agente, que suele ser del tipo: "Por favor, hazme un servidor web", obligándolo a rellenar todo tipo de detalles con supuestos que quizás no eran los que deseabas. Pero con una herramienta totalmente especificada, el agente puede hacer un trabajo excelente: puede ejecutar todas las pruebas unitarias, traducir todas las pruebas unitarias, ejecutarlas en el nuevo sistema y verificar el comportamiento de forma simultánea para asegurarse de que se comporte exactamente igual.

Bill Jia: De hecho, tenemos un ejemplo en vivo internamente en Google al que llamamos Proyecto Tern. Debido a que muchos modelos se construyeron sobre TensorFlow, pero estamos migrando a JAX, traducimos todos los modelos de TensorFlow y los migramos a JAX, realizando todas las pruebas unitarias y de código de manera automatizada.

Moderador: Oigan, hubo un artículo determinado en NeurIPS que fue rechazado y que ustedes habían presentado, el cual señalaba que tenía muy poco impacto. Era un artículo sobre destilación.

Jeff Dean: Ah, sí. Es como una historia de motivación; se supone que debe ser una inspiración. Resulta que la destilación es importante. Este es un artículo que mis colegas Geoff Hinton, Oriol Vinyals y yo presentamos sobre cómo tomar un modelo de una forma y usar la destilación para transferirlo a un modelo estudiante.

Bill Jia: ¿En qué año fue?

Jeff Dean: Fue en 2015 o 2016. Originalmente estábamos pensando en ello en el contexto de entrenar un gran conjunto de diferentes tipos de modelos especializados para visión. Ese es uno de los conjuntos de experimentos que incluimos en el artículo: tienes 20.000 clases de visión, pero entrenas uno especial para animales, otro para automóviles y otro para lo que sea, y luego destilas todo eso en un solo modelo que sea bueno en todas esas tareas. Y fue rechazado, pero no pasa nada; lo subimos a arXiv. La gente lo leyó de todos modos.

Moderador: No se desanimen.

Jeff Dean: Exactamente.

Preguntas y respuestas: Hardware genérico frente a supercomputadoras especializadas

Moderador: Muy bien, nos quedan unos pocos minutos. Vamos a tomar un par de preguntas. Hay algunas personas caminando con micrófonos. Allá vamos, una por aquí. Tal vez podamos empezar en esta zona. Si quieres ponerte de pie, presentarte y darnos una breve descripción de tus antecedentes en una sola línea.

Miembro de la audiencia (Ying): Mi nombre es Ying. Estuve en Google Brain durante una década. Qué bueno verte de nuevo. Mi pregunta para Bill y Jeff es la siguiente: hablamos sobre software genérico y cómputo genérico desde los primeros días de Google; así fue como Google escaló a principios de los 2000. Hoy en día, la supercomputadora de IA se vuelve cada vez más especializada, pareciéndose más al marco de supercomputación original. Entonces, ¿cómo ven ustedes los dos patrones de diseño diferentes para los sistemas de hardware y machine learning? ¿Creen que en el futuro deberíamos orientarnos más hacia hardware genérico o deberíamos seguir construyendo máquinas supercomputacionales? Gracias.

Jeff Dean: La razón por la cual pudimos escalar con hardware genérico para la búsqueda es que la búsqueda es un problema muy agradable donde, si lo divides, prácticamente no hay comunicación entre las máquinas y tienes mucho trabajo que realizar en una sola máquina. Así que no necesitas nada exótico en términos de interconexión. De hecho, teníamos Ethernet de 100 megabits en las máquinas y compartíamos un enlace ascendente de 1 gigabit por rack entre 40 máquinas, por lo que en realidad teníamos una sobresuscripción de 4 a 1. Pero eso funcionaba bien porque envías consultas como "restaurantes en Palo Alto" a cada máquina, y lo que obtienes de regreso es un pequeño fragmento que dice que estos son los 10 resultados. En cambio, el entrenamiento realmente empuja los límites al extremo: necesitas una conectividad masiva.

Te encantaría poder entrenar en un solo chip —esa sería la opción ideal—, pero eso toma demasiado tiempo o simplemente no cabe. Y por lo tanto, terminas teniendo que particionar el problema en muchos chips. La forma en que lo divides, sin importar cómo lo recortes —paralelismo de modelos o paralelismo de datos—, por lo general genera una cantidad considerable de comunicación. Y esa es la razón por la cual existen estas interconexiones más exóticas entre las máquinas. Quieres el máximo rendimiento por chip para no tener que distribuirlo entre tantos chips, pero aun así tienes que repartirlo entre muchísimos. Por esa razón estamos utilizando refrigeración líquida, para lograr que cada chip sea lo más potente posible en ese escenario.

Creo que la inferencia se parecerá más a esto: puedes tener hardware especializado, pero no es tan exótico, particularmente en el caso de modelos pequeños. Sin embargo, a medida que empiezas a manejar modelos más grandes, surgen necesidades de comunicación que empiezan a parecer bastante exóticas en comparación con, digamos, el hardware y las redes genéricas de Ethernet, aunque definitivamente sigue siendo algo mucho más convencional y estándar en comparación con el entrenamiento. ¿Tiene sentido?

Preguntas y respuestas: Contribuciones de código abierto, paridad con CUDA y agentes de IA para desarrolladores

Moderador: Creo que tenemos otra pregunta por aquí.

Miembro de la audiencia (Andra): Jeff, mi nombre es Andra, de Uber. Tengo una pregunta relacionada con el ecosistema de JAX y OpenXLA, especialmente dado este ecosistema de código abierto que la gente ha construido. Veo que esta es prácticamente una era de infraestructura de IA similar a lo que fueron Android e iOS. ¿Qué opinas sobre cómo los jóvenes ingenieros pueden contribuir a las funciones de biblioteca de OpenXLA? Porque para la mayoría de los usuarios, son usuarios intensivos de CUDA o construyen bibliotecas y funciones basadas en bibliotecas de CUDA. En el caso de OpenXLA, ¿en qué consideras que los jóvenes ingenieros pueden aportar?

Bill Jia: En primer lugar, OpenXLA y JAX, y por encima de JAX, toda la pila de JAX ha sido abierta en su totalidad. Ese es un bloque de trabajo. Otro bloque de trabajo es Torch-TPU. Actualmente, de forma interna, estamos colaborando con Meta; lo estamos desarrollando, pero eventualmente lo subiremos a GitHub y permitiremos que todos los investigadores e ingenieros de la comunidad contribuyan. Esos son nuestros dos esfuerzos paralelos. Cualquier joven ingeniero de la comunidad que desee contribuir al repositorio de código abierto es más que bienvenido a trabajar con Google. Podemos debatir cómo co-desarrollar este repositorio de código abierto.

Además, creo que a la larga muchos usuarios utilizarán la pila de JAX o la pila de PyTorch, ya sea en TPU o en GPU. Hoy en día, si las personas en la comunidad experimentan algún problema al utilizarlas, muchos llaman a Google o llaman a Nvidia porque cuentan con ingenieros experimentados, o publican sus preguntas en foros de discusión de la comunidad con la esperanza de que alguien comprenda la duda y tenga una experiencia similar para responderla. Pero, en mi opinión, todo eso es muy lento. En donde la comunidad puede ayudar a contribuir dentro de la estrategia de código abierto es: imaginen que pudiéramos contar con un agente de JAX y un agente de PyTorch que se ejecuten en GPU o TPU, da igual. Si todos en la comunidad aportan fuentes de datos para ayudar a entrenar a ese agente y volverlo verdaderamente potente, imaginen si tuviera un agente colega especializado en TPU. Cuando estoy ejecutando un entrenamiento, postratrenamiento o inferencia, no tengo que acudir a Google ni a un foro de discusión; utilizo directamente este agente para que me ayude. Si el agente es capaz de responder a una gran parte de las dudas, sería fantástico. Pero esto requiere que toda la comunidad contribuya, porque necesitamos una gran cantidad de puntos de datos y experiencias.

Jeff Dean: Tal vez añadiré una observación general sobre las contribuciones de código abierto. Hay muchas maneras en que las personas pueden contribuir. Es muy recomendable interactuar con algunos de los propietarios de los diferentes repositorios para intentar identificar: "Estoy pensando en hacer esto, ¿sería útil?" o "Tienen alguna idea sobre en qué debería trabajar?". Antes de limitarse a arrojar 5.000 líneas de código a alguien y decir: "Ey, aquí está", creo que es una excelente idea conseguir cierto respaldo o guía acerca de lo que realmente resulta útil. Si tienes muchas ganas de contribuir pero no tienes necesariamente un tema específico en mente, a menudo existe una larga lista de tareas pendientes que la gente tiene identificadas y que serían muy útiles.

Miembro de la audiencia (Andra): Gracias por la respuesta. Una de las dudas con las que me he encontrado es que muchas funciones o *kernels* de módulos todavía dependen en gran medida de CUDA. ¿Ven una tendencia hacia la implementación de alternativas equivalentes utilizando JAX, OpenXLA o Torch-TPU de modo que alcancemos algún tipo de compatibilidad generalizada en comparación con las funciones de CUDA? ¿Sería esa una buena dirección?

Bill Jia: Sí, creo que eventualmente llegaremos a eso. Internamente en Google, cuando analizamos la pila de JAX interna y también la pila de PyTorch en la TPU, revisamos todas las funciones de CUDA y también las nuestras. Nos aseguramos de estar al menos a la par, si no es que superándolas. Por lo tanto, creo que esa es una necesidad que debemos resolver por completo. Y en esa capa, de hecho, queremos abrir el código del SDK de TPU de bajo nivel.

Jeff Dean: También considero que la gente quiere pensar en un nivel de abstracción mucho más alto. Pensar en términos de expresiones de JAX o PyTorch es lo que debes buscar, y no necesariamente: ¿cómo puedo paralelizar esto con algún código de *kernel*, ya sea un lenguaje de *kernel* de TPU o de GPU, para obtener el máximo rendimiento? Con suerte, el compilador y el sistema subyacente te proporcionarán eso, permitiéndote pensar en abstracciones elegantes como la multiplicación de matrices.

Preguntas y respuestas: Escalado de redes, toroides en 3D y entrenamiento sincrónico

Moderador: Creo que tenemos tiempo para una última pregunta. Justo una última pregunta, adelante.

Miembro de la audiencia (John): Hola, soy John, de la Universidad de Nueva York (NYU). Llevo muchos años trabajando en redes. Me impresionó mucho la charla de Bill Jia y su colega respecto a la escalabilidad aquí. Sé que estas TPU están conectadas en un toroide en 3D. En comparación con la Ultra Ethernet Consortium y otras iniciativas que analizan interconexiones en árbol gordo (*fat-tree*), desde el punto de vista de la escalabilidad, cuando escalas TPU o GPU hasta decenas de miles o incluso un millón, ¿este toroide en 3D sigue siendo capaz de gestionar la comunicación de todo a todo (*all-to-all*) a esa gran escala?

Jeff Dean: Nuestros pods más grandes para la TPU v5e o la TPU v4, si mal no recuerdo, tienen 9.600 y pico de chips. Esa es la escala a la que operamos con toroides en 3D. Más allá de eso, utilizamos Pathways como una abstracción de software por encima de muchos de esos pods conectados mediante toroides en 3D, y eso emplea la red del centro de datos —ya sea la estructura (*fabric*) que tengas en tu centro de datos, o incluso una configuración de entrenamiento multi-área metropolitana donde quizás tengas 5 pods en este edificio en Oklahoma y 8 más en este edificio en Iowa, con un enlace WAN de alta velocidad entre ellos—. Eso nos ha funcionado bastante bien para escalar.

A veces te interesa mapear el cálculo que estás realizando de modo que tengas dimensiones de paralelismo de datos y de paralelismo de modelos. Por lo general, es preferible que el aspecto de paralelismo de modelos se mantenga dentro de un único pod o una división (*slice*) de un pod, y luego tener réplicas de paralelismo de datos a través de esos pods o divisiones de pods. Eso nos ha dado muy buenos resultados. Lo bueno de las redes basadas en toroides es que son sumamente fáciles de conectar a nivel local, a diferencia de requerir un cableado sumamente complejo en el piso del centro de datos.

Moderador: A excepción de los divertidos conmutadores ópticos con espejos.

Jeff Dean: ¿Algo que agregar, Bill?

Bill Jia: Sí, es muy cierto. Contamos con ICI dentro del rack y dentro del cubo. Luego utilizamos el OCS para escalar hasta 9.600. Y avanzando más allá, disponemos de nuestra estructura de red de centro de datos, orquestada por el software Pathways. Para rebasar la barrera de los 100.000 chips, empleamos redes de centros de datos en la nube para conectar todo entre sí. Este tipo de escalabilidad nos está funcionando de maravilla.

Jeff Dean: Y diría que, incluso a esa escala, somos capaces de realizar un entrenamiento totalmente sincrónico, lo cual resulta excelente desde la perspectiva de la interpretabilidad y la reproducibilidad en machine learning. En algún momento el entrenamiento asincrónico volverá a cobrar relevancia, pero hasta ahora hemos podido llevar el entrenamiento sincrónico bastante lejos, y ya veremos qué sucede más adelante.

Moderador: Esa es una excelente manera de terminar. Se nos ha acabado el tiempo. Muchísimas gracias, Jeff. Gracias, Bill, por su tiempo. Tenemos algunos artículos promocionales (*swag*) muy geniales afuera, junto al patio, que incluyen algunos de los famosos memes de Jeff Dean impresos en un bolígrafo y una camiseta de Dev Labs. Muchas gracias a todos.