Crear un programador de clúster para priorizar la investigación de alto impacto mientras se mantiene una ocupación completa
En el equipo de Infraestructura de IA de Ai2, somos responsables de proporcionar la capacidad de cálculo de GPU del instituto, específicamente dirigida a trabajos de entrenamiento distribuidos de gran tamaño. Consideramos esta tarea como una pirámide de cuatro métricas que se construyen sobre sí mismas.
La base es disponibilidad: cuántas veces el hardware está en buen estado y listo para trabajar. Por encima de esto está ocupación: la fracción del tiempo disponible asignado a una carga de trabajo específica. Luego viene impacto: cuántas veces las cargas de trabajo más valiosas son elegidas para recibir recursos. El punto clave de la pirámide es utilización: la fracción de la capacidad de la GPU utilizada durante el ciclo de vida de una carga de trabajo.
Este artículo trata sobre cómo mejorar el impacto de nuestras decisiones de programación. Recientemente, reemplazamos un programador basado en prioridades por un sistema que incluye presupuestos de tiempo para GPU, asignación jerárquica de uso equitativo y un contrato de segmentación de tiempo. Como resultado, cambiamos el debate sobre cuánto tiempo de GPU merece cada proyecto de investigación a un proceso transparente de presupuestación administrativa.
Sobrecarga
En Ai2, gestionamos miles de GPUs NVIDIA H100, B200 y B300, organizadas en clústeres cuya cantidad varía entre 88 y 1024 unidades. Estos clústeres se construyen para el entrenamiento distribuido a gran escala de modelos de IA, y sirven a un grupo de aproximadamente 150 investigadores internos cuyo trabajo abarca una variedad de dominios de IA, incluyendo el flujo completo de entrenamiento de LLM y VLM, la simulación del aprendizaje refuerzo en robótica (RL), y el proceso posterior al entrenamiento para casos de uso científicos agenticos.
Como muchos laboratorios, tenemos una demanda de tiempo en GPU que supera con creces la oferta. Basándose en las cargas de trabajo presentadas, en cualquier momento tenemos solicitudes pendientes que requieren de 2-3 veces más GPUs de las que están disponibles. Una forma de entenderlo es que cada hora de GPU disponible en nuestro clúster tiene 2-3 cargas de trabajo de investigación diferentes compitiendo por ella.
Históricamente, utilizábamos un programador de prioridades y permitíamos que las cargas de trabajo se excluyeran de la preemptabilidad. Cada equipo tenía un límite en las GPU concurrentes que podían ser utilizadas por las cargas de trabajo protegidas contra la preemptación. Las cargas de trabajo preemptables podían superar ese límite en las GPUs inactivas. Esta estrategia generaba patologías predecibles. Por ejemplo, observamos casos de “squatting” de GPU, donde los usuarios almacenaban cargas de trabajo sin operaciones que podían conectarse cuando surgiera la necesidad. Esto ocurrió porque los investigadores descubrieron que no podían lanzar cargas de trabajo de depuración con una latencia lo suficientemente baja para resolver problemas en tiempo real. También observamos una inflación de prioridades, donde, al final, el 100% de las cargas de trabajo programadas utilizaba un nivel de prioridad ALTO. Esto significaba que los niveles de prioridad más bajos se quedaban completamente sin tiempo en la GPU. Dado que la preemptibilidad era opcional, también descubrimos que nuestros ingenieros en turno pasaban la mayor parte del tiempo de respuesta de sus tickets negociando el cierre organizado de las cargas de trabajo no preemptibles que se ejecutaban en hosts con problemas conocidos de mantenimiento.
La tragedia del común
Cuando estos problemas surgieron, tardamos en identificar sus causas raíz. Nuestros intentos iniciales para asegurar que la tarea más importante recibiera tiempo en GPU se centraron en un control más estricto de cómo se establecen las prioridades, y, en última instancia, en solucionar el programador basado en prioridades asignando explícitamente monopolios de GPU a los proyectos importantes. Aunque al principio no lo reconocimos, habíamos creado un laboratorio perfecto para observar la “tragedia del común”. Las personas competían por un recurso escaso y compartido, y al intentar maximizar los resultados individuales, se obtenía un resultado global no óptimo y se abusaba del recurso subyacente.
Estamos lejos de ser los primeros en observar este tipo de interacción. La asignación de recursos es un campo de investigación fascinante que combina el desarrollo de algoritmos, la economía y la gestión de sistemas. Un problema central es que los usuarios a menudo conocen mejor el valor de sus propios trabajos que la organización, pero pueden tener incentivos para ocultar ese valor o retener recursos incluso cuando eso perjudica el rendimiento total. Por ejemplo, en su artículo de 2011 titulado Dominant Resource Fairness, Ghodsi y colaboradores relatan una anécdota en la que una empresa de búsqueda proporcionaba máquinas especializadas para los trabajos solo si sus usuarios podían garantizar un alto nivel de utilización. Pronto descubrieron que “los usuarios escribían bucles infinitos en su código para aumentar artificialmente los niveles de utilización”. Los cambios en el hardware persisten, pero los problemas fundamentales que complican la asignación de recursos siguen presentes.
>Presupuestos no programas
La solución clásica para una tragedia comunal es privatizar el recurso compartido; los propietarios reciben incentivos para maximizar el valor de su propiedad. Cuando asignamos a los equipos monopolios sobre conjuntos de GPUs, ya estábamos haciendo una versión de esto, pero era demasiado rudimentario. Esto causó que las GPUs permanecieran inactivas debido a la estacionalidad de la investigación. Los equipos están listos para realizar experimentos y entrenamientos en diferentes momentos, por lo que asignar un monopolio aseguraría que haya tiempos en los que no haya trabajos listos para ejecutarse, dejando otro equipo esperando por capacidad.
Tratamos manualmente de resolver un problema de bolso, intentando adaptar las necesidades de investigación que cambian dinámicamente a un horario estático. Queríamos el incentivo de propiedad, pero también queríamos mantener una ocupación total de los GPUs.
Decidimos iterar en el modelo de propiedad. En lugar de emitir GPUs para los equipos, optamos por asignar una parte del tiempo de GPU. Predecir la demanda en el futuro requeriría conocer los resultados de experimentos científicos innovadores, por lo que no se puede pronosticar con precisión. Sin embargo, la priorización de los esfuerzos de investigación es una cuestión de estrategia, y puede ser más fácilmente debatida y decidida de antemano. En lugar de intentar resolver el rompecabezas del programación, permitimos que la dirección piense como inversores. Antes de que existan las cargas de trabajo, decida cómo financiar cada esfuerzo de investigación con el tiempo de GPU, basándose en su juicio sobre el posible impacto. El planificador podría luego utilizar esa información al priorizar las cargas de trabajo que llegan.
Con esto en mente, desarrollamos un sistema jerárquico en el que los gerentes podían asignar de manera proporcional el tiempo del GPU a los proyectos y investigadores para los cuales eran responsables. Como ilustra el diagrama a continuación, esto convierte la estrategia del programa directamente en una participación garantizada de tiempo del GPU. El Proyecto A1 sabe que tiene un 35% de derecho sobre la capacidad total, independientemente de cuántos otros proyectos estén en espera en otro lugar.
Los valores entre corchetes representan la capacidad total del clúster asignada a un proyecto de hoja.
En este sistema, cada solicitud de tiempo en GPU debe ser financiada con un presupuesto; de lo contrario, no está protegida contra la preempresión. En el sistema antiguo, la alta prioridad no implicaba costos, y la no preempresabilidad permitía a los equipos llenar su límite simultáneo de GPU indefinidamente, por lo que todos los utilizaban. Ahora, nada es gratis, así que cualquier truco para obtener tiempo en GPU proviene de la asignación del usuario beneficiado. Una carga de trabajo no utilizada está gastando el presupuesto del equipo en algo inútil. Nuestra estrategia es hacer que el programador de juegos sea más costoso que participar honestamente en el debate con un presupuesto mayor. Estamos constantemente iterando este proceso de revisión de presupuesto, pero los requisitos clave son que haya oportunidades frecuentes para que los investigadores defiendan el tiempo que necesitan, y las decisiones se toman por parte de los gerentes que tienen el mayor contexto sobre los compromisos en cuestión. Esto significa que las decisiones de asignación dentro de un proyecto de investigación son tomadas por el investigador líder, dentro de un programa de investigación por el investigador principal, y entre los programas por el gerente líder del programa o por el CEO.
Fair-share
Junto con esta herramienta de presupuestación de tiempo para GPU, construimos un planificador de reparto justo jerárquico para gestionar la ocupación real de las asignaciones a lo largo del árbol del programa. El algoritmo no es nuevo: el reparto justo jerárquico en un marco temporal forma parte de una línea que se remonta al Hadoop Fair Scheduler en 2009, y el mismo enfoque se utiliza actualmente en el Fair Tree de SLURM y el Fair Scheduler de YARN. Lo nuevo para nosotros son las entradas: el árbol refleja la estructura del programa de investigación, y los pesos son presupuestos establecidos por los gerentes, en lugar de cuotas estáticas.
El planificador registra la ocupación en un período de retroceso deslizante (por defecto, 7 días) y ordena las cargas de trabajo según las asignaciones subutilizadas frente a las asignaciones sobreutilizadas. De esta manera, durante un rango de tiempo de una semana, podemos esperar que cada grupo reciba el tiempo asignado a la GPU siempre y cuando envíe activamente cargas de trabajo con demanda suficiente.
“El nuevo programador de tareas hace que parezca que tenemos un 30% adicional de capacidad de procesamiento. En el antiguo programador, si había momentos en los que no necesitábamos todo el límite de tiempo asignado, esa capacidad se perdía básicamente. Ahora, con el nuevo programador, si eso ocurre, podemos posteriormente superar el límite de asignación y, aún así, ver que nuestras tareas se programan rápidamente sin interrupciones, esencialmente permitiéndonos recuperar esa capacidad de procesamiento. Nuestras cargas de trabajo son a menudo volátiles, por lo que esto nos ha devuelto una cantidad significativa de capacidad de procesamiento.” — Chris Clark
El planificador distingue dos tipos de ocupación. Occupación asignada es el tiempo durante el cual una carga de trabajo se carga en un presupuesto. Esto se basa en las asignaciones del propietario de la carga de trabajo, lo que afecta el cálculo del presupuesto de participación justa, y estas cargas de trabajo están protegidas contra la preemisión durante su ventana de tiempo mínimo de ejecución. Occupación no asignada no se carga en ningún presupuesto, no está protegida desde el principio y puede ser preemptada por cualquier solicitud asignada. Esto permite que las GPU estén completamente ocupadas incluso cuando las asignaciones no coinciden adecuadamente con la demanda, y evita que los equipos renuncien a los ciclos gratuitos de GPU.
El contrato de programación
Una característica adicional del entrenamiento distribuido que dificulta la asignación justa de recursos es que las cargas de trabajo pueden ejecutarse durante mucho tiempo. Las tareas de entrenamiento se ejecutan regularmente durante horas, días e incluso semanas. Una vez programada, una carga de trabajo puede permanecer en las GPU asignadas durante una semana o más, sin dar oportunidad a otros de obtener el tiempo presupuestado. Esta es la propiedad del sistema que hace posible el uso prolongado de las GPU. También fue lo que obligó a los ingenieros en turno de guardia a negociar con los propietarios de trabajos de larga duración para resolver los problemas de mantenimiento en curso.
Para abordar estos problemas, introdujimos un “contrato de programación”. A cambio del acceso al clúster, una carga de trabajo debe declarar su tiempo de ejecución mínimo, o la cantidad mínima de ocupación necesaria para lograr un progreso significativo. Durante este período, la carga de trabajo está protegida contra la preempción. Esto proporciona al investigador una garantía de progreso, mientras que da al programador el derecho de reequilibrar la carga de trabajo una vez que se acumula ese progreso, volviendo a enqueue automáticamente las cargas de trabajo recuperables. Otra opción es que el usuario establezca el tiempo de ejecución mínimo en cero, lo que indica que el tiempo del GPU debe estar no asignado. Estas cargas de trabajo siempre están sujetas a preembarco, pero también están libres en el sentido de que no se cobran desde ningún presupuesto.
El ciclo de vida del trabajo se sigue este patrón:
- La carga de trabajo se presenta con un tiempo de ejecución mínimo y indica si es o no reanudable.
- La carga de trabajo se programa según el algoritmo de reparto justo, ponderado por una proporción entre la ocupación real y el tiempo asignado en la ventana de retroceso.
- La carga de trabajo se ejecuta durante su tiempo de ejecución mínimo, lo cual se carga en sus asignaciones.
- La carga de trabajo puede continuar ejecutándose mientras las asignaciones asociadas siguen dándole prioridad sobre otras. Este tiempo también se carga en sus asignaciones.
- Puede ser interrumpido y reenviado, lo que vuelve al paso 2.
- La carga de trabajo se completa, liberando su demanda sobre cualquier recurso.
Juntos, estos acuerdos añaden segmentación por tiempo a nuestro programador. Las cargas de trabajo en ejecución pueden eliminarse y reenviarse automáticamente, lo que permite que la distribución equitativa se acerque y disminuye los actos de ocupación no necesarios. También permiten que los hosts no saludables agoten sus cargas de trabajo al alcanzar su tiempo de ejecución mínimo, de modo que las actividades de reparación puedan ser completamente automatizadas. Este último punto fue más importante de lo que imaginábamos al planificar este trabajo. Redució los reparaciones que requerían una intervención humana en un 74%, lo que representó un ahorro enorme en el trabajo extra durante la atención.
Simulaciones
Sabemos que los cambios en la política de programación pueden tener consecuencias no deseadas. La naturaleza de suma cero del problema significa que dar tiempo a un investigador implica quitarlo a otro. Los usuarios que pierden este intercambio tienden a buscar nuevas soluciones. Antes de implementar el sistema basado en presupuesto, queríamos una forma rápida de predecir dónde podrían surgir esos tiempos de espera más largos y de probar los ajustes de configuración como la longitud de la ventana de retroceso o el valor máximo para permitir un tiempo de ejecución mínimo (elijimos 8 horas).
Construimos un pequeño entorno de simulación que recibe como entrada un conjunto de cargas de trabajo y su cronograma de presentación, y permite que el programador de horarios tome decisiones de preemisión y asignación de GPU. Con conocimiento del número de GPU solicitadas por cada carga de trabajo y el tiempo total de ejecución, el simulador puede avanzar hasta los momentos disponibles para ser programados y ofrecer análisis de tiempos de espera en la cola, eventos de preemisión y distribución del tiempo de GPU entre proyectos durante muchos días simulados en pocos segundos. Llevamos a cabo el simulador con los datos históricos de presentación y creamos escenarios que queríamos comprender mejor.
Una hipótesis que queríamos probar se refería a “cargas de trabajo de depuración”. Estas tareas requieren un pequeño número de GPU y un tiempo de ejecución mínimo de 15 minutos o menos, lo cual es suficiente para que el usuario vea si la tarea se inicia con éxito o se cae temprano debido a un error o una mala configuración. Queríamos saber si estas tareas tendrían un tiempo de espera en la cola más corto que las cargas de trabajo de entrenamiento más grandes, que a menudo necesitan muchos GPU y horas de ejecución para lograr progreso significativo. Intuitivamente, estos trabajos más pequeños deberían llegar al top de la cola, ya que un trabajo pequeño puede ocupar más espacio que uno grande. Pero la latencia precisa de la cola era importante. Una espera corta de un minuto o dos podría desbloquear una nueva práctica de desarrollo, pero una espera de diez minutos se vuelve imposible.
Nuestras simulaciones requirieron datos de casos de prueba construidos manualmente, ya que nuestro registro histórico no contenía un volumen suficiente de estas cargas de trabajo similares a las de depuración. Nuestros resultados respaldaron la hipótesis, mostrando que los tiempos de espera de la carga de trabajo de depuración p90 pasan de aproximadamente 6 horas a solo 5 minutos.
Visualización de simulador a menor escala del modelo base (izquierda) y el nuevo programador de “asignaciones” (derecha). Cada fila corresponde a una GPU; cada barra representa un trabajo, coloreado según la carga de trabajo parental, con un tono de color distinto para cada equipo; las líneas de sombra indican los momentos en que el trabajo puede ser interrumpido, y el borde rojo señala una preemisión. En el modelo base, los trabajos urgentes largos nunca se interrumpen, y hay menos preemisiones en el trabajo de menor prioridad. El nuevo programador tiene una mayor variedad de colores en cada GPU, lo que ilustra la rotación de ocupación entre los equipos.
Resultados
Con los resultados de la simulación en mano, comenzamos una implementación por clúster a finales de julio. Los resultados que nos importan son si las cargas de trabajo que elegimos para financiar recibieron el tiempo necesario, si el nuevo sistema mantuvo una ocupación completa y si los investigadores pudieron razonar sobre el programador para tomar decisiones informadas.
Desde el lanzamiento, hemos observado que los usuarios y los equipos reciben consistentemente el tiempo asignado de GPU. Contamos el tiempo que corresponde a un equipo según su asignación, limitada hora por hora según su demanda real. Durante el período de prueba de 30 días, los equipos recibieron el 98% de las horas de GPU que les correspondían, y 13 de las 15 asignaciones de equipo obtuvieron el 95% o más; en el caso más grave, fue del 90%. La ocupación en el clúster se mantuvo estable en el 98% antes y después del cambio, con la demanda superando la capacidad en 2-3 veces en ambos periodos. El 18% del tiempo de la GPU entregada estuvo no asignado, lo que permitió mantener una alta ocupación durante los períodos en que los casos de uso financiados no estaban listos para funcionar.
Los resultados del simulador demostraron ser precisos en su dirección, y los resultados reales superaron en rendimiento las predicciones realizadas. El tiempo de espera en la cola p90 de cargas de depuración disminuyó de 2 horas a 30 segundos bajo el nuevo programador, en comparación con una predicción simulada de 6 horas y 5 minutos basada en escenarios de prueba manuales. Cabe señalar que la menor tamaño de muestra de las cargas de depuración en el modelo base significó que había mayor variabilidad en esas mediciones. La latencia de la cola en general mejoró como efecto secundario del slicing de tiempo: en nuestro clúster H100 más grande, el tiempo medio de espera en la cola disminuyó de 5 minutos a 24 segundos, y el tiempo p90 de espera disminuyó aproximadamente un tercio (de 2.8 horas a 1.8 horas).
En comparación con los tres problemas que nos propusimos resolver:
Quieto en posición de cuatro patas: Las cargas de trabajo de depuración cortas comienzan en menos de un minuto, lo que reduce el valor del quieto en posición de cuatro patas. El costo de este comportamiento impone una carga en el presupuesto del que se ejerce el quieto, lo que impide que reciban tiempo cuando realmente lo necesitan.
Inflación de prioridad: Todavía permitimos que las cargas de trabajo declaren su prioridad, pero esto solo afecta el ordenamiento dentro de un equipo. Los gerentes están incentivados a monitorear la prioridad en todo el grupo para optimizar el uso de sus presupuestos.
Trabajo en espera: Los hosts no saludables se agotan automáticamente cuando las cargas de trabajo alcanzan el tiempo de funcionamiento mínimo. Las reparaciones que requieren intervención humana disminuyeron en un 74%.
Desafíos
La curva de aprendizaje fue más pronunciada de lo que habíamos supuesto. Implementamos el cambio de forma incremental, por lo que en los primeros días los investigadores experimentaron comportamientos variados dependiendo del grupo al que se dirigieran. Además, nuestras interfaces conservaron algunas terminologías antiguas (como la prioridad de carga) cuyos significados habían cambiado. La documentación por sí sola no resolvió la confusión. Lo que funcionó fue llevar a cabo sesiones explicativas en vivo, proporcionando un foro donde los investigadores pudieran hacer preguntas y donde el equipo de ingeniería pudiera ofrecer descripciones más detalladas de cómo y por qué el programador tomaba las decisiones de priorización utilizando ejemplos reales.
Esto fue un momento clave, ya que marcó un cambio de un período inicial de frustración y teorías populares hacia el modo actual, en el que los grupos de investigación se comunican con más frecuencia y de manera más amplia sobre las necesidades de GPU de sus experimentos. Los investigadores ahora participan en la discusión de presupuestos con un conocimiento más claro de los compromisos que se hacen para atender cualquier nueva solicitud.
Además de las sesiones en persona, introdujimos nuevas visualizaciones después del lanzamiento para que los usuarios tengan una mejor idea de qué tan cerca está el tiempo asignado a la GPU con respecto a las asignaciones esperadas, y para exponer directamente la métrica utilizada para ordenar la cola de carga de trabajo. Esto proporcionó un lugar sencillo donde buscar cuando una carga de trabajo se interrumpía para entender por qué. Estas visualizaciones también ayudaron a los propietarios de presupuestos, quienes podían ver cómo se utiliza el tiempo de GPU en varios proyectos bajo su gestión.
Ejemplo de visualización del uso de la asignación a lo largo del tiempo.
No todos los casos de uso mejoraron. Además del entrenamiento distribuido, nuestros investigadores inician sesiones interactivas en las que realizan el análisis de datos y prueban el código de entrenamiento mientras lo escriben. En el sistema antiguo, un investigador podía mantener una tal sesión durante hasta una semana. Con el desglose por tiempo, estaban sujetos al límite de 8 horas en el tiempo de ejecución protegido, después de lo cual una sesión se convierte en preemptible si supera su asignación. No apreciamos en qué medida los investigadores dependían del estado volátil de estas sesiones. El hecho de ser interrumpidos significaba esperar para obtener una nueva sesión y también reconstruir su estado manualmente. Después de entrevistar a los investigadores para entender la amplitud de este problema, creamos dos nuevos proyectos de roadmap. Estamos invirtiendo en un clúster únicamente con CPU junto a nuestro almacenamiento on-prem para las sesiones de desarrollo, centradas en tareas de preparación de datos. Esto preservará la capacidad del clúster de entrenamiento para las cargas de trabajo que realmente lo necesitan. Además, planeamos crear sesiones recuperables para estas cargas de trabajo que utilizan solo CPU. Esto nos permitirá continuar priorizando las cargas de trabajo al final de su tiempo de ejecución mínima para mantenimiento o segmentación de tiempo, y también podremos restablecer la sesión en otro lugar sin que el investigador tenga que reconstruirla. Podemos mantener los beneficios operativos y de programación de este nuevo sistema, al mismo tiempo que mejoramos la experiencia de usuario.
Seguimos en busca de problemas emergentes. Un problema potencial que estamos investigando es la fragmentación de la capacidad, lo que podría provocar un aumento en los tiempos de espera en la cola para las cargas de trabajo más grandes. Nuestra intuición es que se aplica una protección contra el tiempo de ejecución mínimo a los tipos de trabajos que antes dependían de mecanismos preemptibles para superar el límite simultáneo de GPU del equipo. Anteriormente, esos trabajos podían ser interrumpidos en cualquier momento, lo que podría desperdiciar ese tiempo, pero también facilitaba la programación de trabajos grandes. Ahora, el programador puede tener menos oportunidades para interrumpir muchos trabajos al mismo tiempo para colocar una carga de trabajo grande pendiente. Actualmente estamos utilizando nuestras herramientas de simulador para reproducir este problema, al mismo tiempo que midemos la verdad real en producción.
El futuro
Al mirar más allá del trabajo de programación descrito aquí, nos enfocamos en el objetivo final de la pirámide: la utilización. Necesitamos asegurarnos de que el bootstrapping, el checkpointing y las aplicaciones de entrenamiento se realicen de la manera más eficiente posible, maximizando el valor del tiempo programado que cada carga de trabajo recibe.
Si desea abordar desafíos como estos en colaboración estrecha con investigadores, le animamos a explorar roles de ingeniería abiertos en Ai2.