Resumen Técnico: Kareus – Reducción Conjunta de la Energía Dinámica y Estática en el Entrenamiento de Modelos Grandes

Planteamiento del Problema

El rápido crecimiento de la demanda de computación de IA ha superado la oferta de energía, convirtiendo a la energía en un recurso costoso y disputado para el entrenamiento de modelos grandes. Los esfuerzos de optimización actuales suelen abordar ya sea la energía dinámica (consumida por la computación activa, reducible mediante el escalado de frecuencia) o la energía estática (consumida por el hardware independientemente de la actividad, reducible mediante el acortamiento del tiempo de ejecución).

Las soluciones existentes sufren una limitación crítica: optimizan estos factores de forma aislada o mediante combinaciones ingenuas.

• Los enfoques de escalado de frecuencia (por ejemplo, Perseus) reducen la energía dinámica bajando las frecuencias de la GPU para las rutas no críticas, pero suelen seguir una ejecución secuencial de kernels, perdiendo oportunidades de solapar comunicación y computación.
• Los enfoques de programación de kernels a nivel fino (por ejemplo, Nanobatching) reducen la energía estática al solapar la comunicación y la computación, pero ignoran la energía dinámica y el escalado de frecuencia.

El artículo identifica que la asignación de SM, el tiempo de lanzamiento de kernels y la frecuencia de la GPU son interdependientes. Cambiar un factor desplaza la configuración óptima de los otros. Por ejemplo, reducir la frecuencia hace que los kernels sean relativamente más dependientes de la computación (compute-bound), alterando la estrategia óptima de contención de recursos y solapamiento. En consecuencia, optimizar subconjuntos de estos factores o combinarlos de forma ingenua impide alcanzar el límite teórico de la frontera de compromiso entre tiempo y energía.

Metodología

Los autores proponen Kareus, un sistema de entrenamiento diseñado para optimizar conjuntamente la asignación de SM, el tiempo de lanzamiento de kernels y la frecuencia de la GPU para encontrar la frontera de Pareto óptima entre tiempo y energía.

1. Modelo de Ejecución de Solapamiento Particionado

Para hacer que la optimización conjunta sea tratable, Kareus descompone el problema global en subproblemas locales utilizando un modelo de ejecución de solapamiento particionado:

• Partición: El grafo de ejecución de las pasadas hacia adelante (forward) y hacia atrás (backward) se divide en particiones fijas. Una partición consiste en un kernel de comunicación (por ejemplo, AllReduce) y la secuencia contigua más larga de kernels de computación de un microbatch diferente que no tiene dependencias de datos con el kernel de comunicación.
• Generalización: Este modelo generaliza las técnicas de nanobatching, añadiendo un control fino sobre la asignación de SM y el tiempo de lanzamiento para cada partición.
• Composición Jerárquica: El sistema optimiza cada partición de forma independiente y luego compone jerárquicamente estas fronteras locales en una frontera global de microbatch e iteración.

2. Optimización Bayesiana de Objetivos Múltiples y Pasadas Múltiples (MBO)

Dado que la búsqueda exhaustiva del espacio de configuración (que involucra conteos de SM, tiempos de lanzamiento y frecuencias) es prohibitivamente costosa (requiere miles de horas-GPU y causa interferencia térmica), Kareus emplea un algoritmo MBO especializado:

• Modelos Sustitutos: Utiliza dos modelos separados de Árboles de Decisión Potenciados por Gradiente (XGBoost) para predecir el tiempo de ejecución y la energía dinámica, respectivamente.
• Selección de Candidatos de Pasadas Múltiples: Para explorar la frontera de Pareto de manera efectiva, el algoritmo utiliza múltiples funciones de adquisición en paralelo:
  • Mejora del Hipervolumen (HVI): Tres funciones HVI distintas guían la búsqueda hacia una menor energía total, una menor energía dinámica (favoreciendo frecuencias más bajas) y una menor energía estática (favoreciendo una ejecución más rápida).
  • Exploración de Incertidumbre: Un conjunto bootstrap de modelos sustitutos estima la incertidumbre de la predicción, priorizando candidatos en regiones del espacio de búsqueda poco exploradas.
• Perfilado Térmicamente Estable: Para garantizar mediciones de energía precisas, el sistema ejecuta las particiones durante una ventana de 5 segundos e inserta un periodo de enfriamiento de 5 segundos entre configuraciones para evitar el sesgo térmico.

3. Implementación del Sistema

Kareus se implementa integrándose con Megatron-LM y Perseus.

• Utiliza MSCCL++ para el control fino de la asignación de SM.
• Emplea una función personalizada torch.autograd.Function para envolver las particiones, permitiendo esquemas distintos para las pasadas hacia adelante y hacia atrás.
• En tiempo de ejecución, selecciona un esquema de ejecución de la frontera precalculada basándose en las restricciones del usuario (por ejemplo, un plazo de tiempo o un presupuesto de energía) y ajusta dinámicamente las frecuencias de la GPU mediante un controlador adaptado de Perseus.

Resultados Clave

Kareus fue evaluado en 14 cargas de trabajo, incluyendo entrenamiento real en Llama 3.2 3B y Qwen 3 1.7B, y una emulación a gran escala de Llama 3.3 70B en 16 GPUs NVIDIA A100.

• Compromiso Tiempo-Energía: Comparado con la línea base de estado del arte (Megatron-LM + Perseus), Kareus logra:
  • Hasta un 28.3% de reducción en la energía de entrenamiento en el mismo tiempo de entrenamiento.
  • Hasta un 27.5% de reducción en el tiempo de entrenamiento con el mismo consumo de energía.
• Rendimiento de Máximo Rendimiento (Throughput): En modo de máximo rendimiento, Kareus reduce el tiempo de iteración hasta en un 14.9% y la energía hasta en un 22.1% en comparación con la línea base Megatron-LM.
• Estudios de Ablación: Eliminar ya sea el escalado de frecuencia o la programación de kernels resulta en una degradación significativa del rendimiento (por ejemplo, eliminar la programación de kernels conlleva un aumento del 10.8% en la energía), confirmando la necesidad de la optimización conjunta.
• Escalabilidad: Los resultados de emulación a gran escala para Llama 3.3 70B muestran tendencias consistentes con los resultados del banco de pruebas físico, demostrando la aplicabilidad del sistema a modelos masivos.

Significado y Reivindicaciones

El artículo afirma que Kareus empuja la frontera de compromiso tiempo-energía para el entrenamiento de modelos grandes al demostrar que la energía dinámica y la estática no pueden optimizarse de forma independiente. La contribución central es la idea de que los esquemas de ejecución (asignación de SM y tiempos de lanzamiento) y la frecuencia de la GPU están determinados de forma conjunta e interdependiente.

Los autores enfatizan que:

1. El Control Conjunto es Superior: Optimizar conjuntamente la frecuencia y la programación es demostrablemente más eficiente energéticamente que dejar el control de la frecuencia a los controladores de potencia del hardware o optimizar la programación sin conciencia de la frecuencia.
2. La Interdependencia es Clave: La asignación óptima de SM y el tiempo de lanzamiento son específicos de la frecuencia; un esquema óptimo a una frecuencia puede ser subóptimo en otra.
3. Eficiencia Práctica: Al descomponer el problema en particiones y utilizar MBO, Kareus logra estos avances sin el costo prohibitivo de una búsqueda exhaustiva, lo que lo convierte en una solución práctica para las cargas de trabajo de entrenamiento de IA actuales y futuras.