Ha habido mucha discusión tras la presentación del kernel MLA Jalapeño en HotChips y los comentarios posteriores de SemiAnalysis. Como equipo de hardware de OpenAI, apenas rozamos esta pequeña joya: el hecho de que la IA está escribiendo nuestros kernels y que, cuando lo hace, realmente no necesitamos entender línea por línea lo que hace el kernel. Dejamos de lado de manera evidente: ¿cómo es posible algo así? ¿Cuál es la forma correcta de pensar en esto, en comparación con un método más tradicional de generación de código? ¿Es el kernel optimizado tan sólido como el no optimizado?
Por mis antecedentes, he trabajado en compiladores para aceleradores durante más de una década. Inicié XLA, que es una excelente infraestructura de compilador con un equipo y esfuerzo estelar entre empresas. En los últimos 2+ años en OpenAI, he estado tratando de reconceptualizar cómo deberían funcionar los compiladores en la era de la IA. Las nuevas formulaciones de compiladores aprovecharán las fortalezas existentes, pero es imposible negar que hay una herramienta nueva y poderosa que aprovechar en el kit de herramientas.
Este será un viaje un tanto largo, pero espero aclarar cómo se está utilizando la IA para la automatización de la mejora de programas informáticos; es decir, la compilación optimizadora. Creo que, mediante la IA, podríamos experimentar algo que consideramos como "compiladores 2.0". La IA está fundamentalmente menos limitada en lo que puede proponer, y lo que propone es el resultado del entrenamiento y el contexto del modelo, lo que me lleva a clasificarla como un "optimizador estocástico"; esto puede presentar desafíos, pero, como veremos, también es una fuente de grandes fortalezas...
Una gran cantidad de investigación académica y aplicación industrial ya se dirige en esta dirección, y está descubriendo rápidamente el potencial de la participación de la IA en el ámbito de los compiladores optimizadores, pero estamos en un punto en el que merece una explicación general.
Antecedentes
Los compiladores toman programas y generan versiones traducidas o mejoradas de esos programas.
Los programas, tanto en el lado de entrada como en el de salida, tienen semántica que nos dice qué significan los programas, qué podrían hacer y cómo razonar sobre esas cosas que podrían hacer.
Quienes trabajamos en compiladores los consideramos como funciones puras: toman una estructura de datos y generan una estructura de datos que debería tener una semántica correspondiente.
A veces nuestros compiladores se centran en "reducir" o "traducir". Por ejemplo, pueden tomar C y generar ensamblador x86-64, lo que a menudo consideramos de "bajo nivel". Pero a menudo hacen más que solo traducir como una subparte de ese proceso...
Nuestros compiladores, en la práctica, se centran en "optimizar". Pueden tomar una estructura de datos que representa el programa – en nuestra jerga una "Representación Intermedia" (IR) – y tratan de producir una versión mejorada de ese programa. A veces "mejor" significa que requiere menos ciclos para ejecutarse, a veces significa que tendrá menos código innecesario, a veces significa especializarse para cosas que podemos demostrar que "deben ser ciertas" sobre el programa (evaluación parcial).
Ahora, brevemente, considera que los LLM se crearon originalmente para traducir texto humano de un idioma a otro. Claramente, la traducción está en su área de especialización. Y podemos ver, mediante el uso de LLM en tareas cotidianas, que también pueden escribir nuevas soluciones y mejorar las existentes. Muchos de nosotros, programadores, también tenemos experiencia pidiendo a un LLM que "optimice este fragmento de código" y, notablemente, pueden hacerlo. (Sin embargo, necesitamos saber que optimizaron el código correctamente, ¡a lo que llegaremos!) Esto es simplemente para resaltar que los LLM tienen las capacidades que buscamos en un compilador optimizador.
Optimización y Optimalidad
Los compiladores optimizadores, como era de esperar, intentan aumentar la optimalidad del programa en el que están trabajando, según algún objetivo (generalmente el tiempo de ejecución). Eso es tan difícil de hacer en el caso general, para un programa arbitrario, que existe un teorema llamado el teorema de pleno empleo para ingenieros de compiladores. (Solo me enteré de esto después de elegir ser ingeniero de compiladores, ¡pero igual me dio consuelo!)
Los "superoptimizadores" son un subcampo increíble de los compiladores optimizadores. Imagina que hay un programa dado, y podemos decir lo que hace mediante la semántica. ¿Cuál es el programa más óptimo que tiene esa misma semántica? Eso es lo que los superoptimizadores intentan abordar, y es efectivamente un problema de búsqueda...
Imagina que estoy tratando de encontrar el programa más corto que tuviera esa misma semántica, y tuviera una forma de preguntar si un programa candidato tenía la misma semántica. Hipotéticamente, podría enumerar cada programa en orden objetivo y elegir el más pequeño que tuviera la misma semántica.
Sin embargo, enumerar cada programa en orden objetivo parece bastante intratable. Uno de mis artículos académicos favoritos, realizado en 2013 titulado "STOKE" (Superoptimización Estocástica), preguntó: "bueno, ¿y si simplemente modificamos programas al azar una y otra vez, eventualmente observamos el mejor programa?" Propusieron que mediante un paseo aleatorio (y con nuestro amigo OG del aprendizaje automático, Monte Carlo de Cadena de Markov / Metropolis-Hastings), eventualmente verías ese programa óptimo.
La modificación Monte Carlo suele ser tonta (eligues una modificación al azar), pero también rápida. Los LLM son muy inteligentes (muchos tokens de razonamiento), pero comparativamente lentos.
¿Qué pasaría si, en lugar de la modificación Monte Carlo tonta/rápida, tuviéramos LLM que descubrieran las direcciones en las que llevar los programas? Tendríamos un optimizador estocástico que fuera muy inteligente, guiando nuestro programa a través del espacio de programas optimizados.
Intuiciones para la Optimización
Retrocedamos un paso. Piensa en la persona que conoces que mejor personifica "optimiza al máximo los fragmentos de código". Para abreviar, llamémosle "Ollie el optimizador". Ollie probablemente tiene una intuición innata sobre qué tipos de modificaciones de código podrían dar frutos. Ollie probablemente prueba algunas cosas para ver si funcionan, y si no funcionan, las revierte y prueba otra cosa. Pero tienen cierta intuición sobre qué tipos de cosas son posibles y cómo podrían superar al compilador.
Estas intuiciones que tiene Ollie a menudo van más allá de lo que hacen los compiladores. Aunque los compiladores optimizadores modernos son bastante impresionantes en sus resultados, se basan en reglas y heurísticas bastante simples. En jerga técnica, se basan en la idea de una transformación de flujo de datos local que se ejecuta hasta un punto fijo. También ordenamos las fases de las consideraciones; es decir, construimos tuberías de compilador para considerar A y luego B, pero no el problema compuesto AB. Los planificadores y asignadores de registros son un ejemplo notorio de esto; se han intentado muchos doctorados en el planificador-asignador de registros compuesto (para obtener los beneficios de colapsar el orden de fases), pero ha sido un desafío hacerlos funcionar en la práctica.
Por eso la experiencia de Ollie es valiosa. A menudo, Ollie sabe cómo equilibrar varios problemas NP-completos con heurísticas que son hechas a medida para la situación. Por lo tanto, hay una mayor conciencia y sensibilidad contextual hecha a medida. Ollie también puede emplear técnicas que los compiladores optimizadores pueden no aplicar de manera rentable, especialmente en combinación, cosas como extraer o crear ABI personalizados o transformaciones para habilitar la vectorización, o la otra gran cantidad de cosas que nos hacen refunfuñar "ojalá el compilador tuviera una forma de hacer esto..."
Ahora considera que la IA, a través de las instalaciones de razonamiento que tenga, puede actuar como un mini Ollie. Puede que no tenga la intuición equivalente en términos de lo que dará resultado, pero tiene una idea de lo que puede ser rentable, y puede hacer muchas, muchas intentos.
Con este enfoque, a diferencia del artículo de STOKE, no podemos garantizar que a medida que el tiempo tiende a infinito podamos ver el programa óptimo, pero debido a que la IA tiene instalaciones de razonamiento "más humanas", puede lograr un progreso significativo similar al humano por unidad de tiempo.
Vinculándolo: Kernel MLA
Permítanme comenzar diciendo: yo no sé qué código de bajo nivel generó la IA para nuestro kernel MLA Jalapeño, pero sí sé cómo escribir el numpy para MLA.
En el compilador XLA en el que trabajé anteriormente, fusionábamos esas operaciones de numpy en grupos, y luego usábamos un metaprograma llamado "emisor" para reducirlas a bucles, instrucciones y primitivas de bajo nivel.
Cuando el compilador/programa emisor de XLA hacía eso, no necesitaba preocuparme por qué ensamblador salía por detrás. Para nuestro optimizador estocástico, la IA conceptualmente toma el lugar del metaprograma emisor: tanto reduce como optimiza, y podemos pedirle que optimice cada vez más hacia el roofline.
Espero que esto aclare dónde encaja la IA y cómo es análoga a un componente en un sistema de compilador optimizador existente. También es útil pensar: qué capa consideramos "código ensamblador" ahora está subiendo. Cuando escribes C++ normal y lo compilas con -O3 (el nivel de optimización típico más alto), no esperas entender el ensamblador que sale, incluso si entendías el C++ que habías escrito. Estamos haciendo algo análogo aquí, pero con una especificación de entrada de nivel superior y más matemática.
Ahora, una pregunta clave es cómo verificamos que el programa que obtenemos de la IA es efectivamente equivalente a la descripción de alto nivel / numpy. Ese mecanismo de verificación establece la solidez del proceso de optimización estocástica de la IA. Espero que una futura publicación de blog entre en más detalle sobre esto, pero por ahora, baste decir que probar la equivalencia semántica es posible y lo hacemos. Los programas de aceleradores son particularmente adecuados para contratos sólidos y completos que podemos verificar "son exactamente lo que hace el programa optimizado por IA", ya que son bastante matemáticos y orientados al flujo de datos en su contexto amplio.
Ten en cuenta que muchas técnicas relevantes en esta área fueron pioneras por esfuerzos en el subcampo de la síntesis de programas. Mientras que los compiladores optimizadores dicen: "aquí hay un programa con semántica, ¡mejóralo pero con semántica equivalente!", la síntesis de programas dice: "creo que existe un programa con esta semántica, por favor intenta encontrar el mejor que puedas". La síntesis de programas es un problema más difícil que la compilación optimizadora, pero también está fundamentalmente menos limitada. Es efectivamente lo que hacen humanos como Ollie cuando superan al compilador optimizador, y es algo que la IA ahora puede ayudarnos a automatizar. La IA puede obtener "inspiración" del programa original, pero no necesita solo realizar transformaciones locales menores en él. Los compiladores optimizadores clásicos no verán "oh, escribiste un bubble sort" y, al comprender el contrato, lo cambiarán a un quick-sort, pero tanto Ollie como la IA pueden hacer eso. Esto es lo que nos sitúa más en el régimen de síntesis de programas con optimización estocástica que en el régimen clásico de compiladores optimizadores.
Todo esto se une en el hecho de que puedes comenzar con algo que "no está muy lejos del numpy", esperar 48 horas y tener un kernel optimizado con la misma semántica, como mostramos en nuestra charla de HotChips:
Como también señala la diapositiva, en nuestra máquina a menudo podemos observar a la IA superar el rendimiento de nuestros expertos humanos incluso en los kernels que considerábamos bastante bien ajustados. A menudo, todavía queda un porcentaje alcanzable decente simplemente debido a las muchas variedades de combinaciones/permutaciones que pueden necesitar ser exploradas. Estas suelen ser intratablemente tediosas para un ingeniero de rendimiento humano.
Resumen y Conclusión
Un compilador, al final del día, es solo una función. Le damos nuestro programa a esa función, y obtenemos una versión mejorada de nuestro programa. Esperamos que el programa que obtenemos y el programa que ingresamos tengan la misma semántica.
Los compiladores optimizadores tradicionales mejoran los programas mediante reglas de flujo de datos y heurísticas. Estos son completamente comprensibles en su procedencia, pero también pueden ser más limitados en los movimientos que pueden realizar.
En contraste, la IA, como optimizador estocástico, solo tiene que "pensar intensamente" y generar algo. Sus movimientos no están tan fundamentalmente limitados, lo que los hace más análogos a nuestro experto humano optimizador. Necesitamos formas de verificar que los programas que genera sean sólidos e implementen la misma semántica que ingresamos, y tenemos esas formas implementadas. Y este tipo de optimización por IA es particularmente adecuada para operaciones matemáticas que tienen contratos muy sólidos. Los contratos evitan la necesidad de entender lo que hace el kernel línea por línea.
¡Así es como obtuvimos el kernel MLA generado por IA!
