Resumen Técnico: DOPPLER – Aprendizaje de Política Dual para la Asignación de Dispositivos en Grafos de Flujo de Datos Asíncronos

1. Planteamiento del Problema

El artículo aborda el problema de la asignación de dispositivos en sistemas multi-GPU que ejecutan cargas de trabajo complejas de aprendizaje automático (por ejemplo, Transformers, Mixture-of-Experts) representadas como grafos de flujo de datos estáticos. El objetivo es mapear operaciones (vértices) a dispositivos para minimizar el tiempo total de ejecución dentro de un sistema de Trabajo Conservador (Work-Conserving, WC).

A diferencia de los marcos tradicionales de ejecución sincrónica por lotes (por ejemplo, PyTorch o TensorFlow estándar), donde todos los dispositivos esperan a la operación más lenta antes de proceder, los sistemas WC programan las operaciones de forma asíncrona tan pronto como se cumplen sus dependencias de datos. Aunque los sistemas WC ofrecen un potencial significativo de aceleración al superponer comunicación y computación, introducen desafíos únicos:

• Orden de Ejecución Estocástico: Sin una sincronización global, el orden de ejecución es dinámico y sensible a la heterogeneidad del hardware y a la contención de recursos.
• Optimización Compleja: Debido a esta estocasticidad, el tiempo de ejecución no puede expresarse como una función de forma cerrada de la asignación.
• Objetivos Contrapuestos: Una asignación efectiva debe equilibrar simultáneamente la carga de la GPU y minimizar la comunicación entre GPUs, una tarea dificultada porque el equilibrio de carga es inherentemente temporal en entornos asíncronos.

Los enfoques previos basados en aprendizaje están limitados por su dependencia de supuestos de sincronía por lotes, el uso de una única política de colocación sin modelar la dinámica del sistema y la falta de optimización durante el despliegue real.

2. Metodología: El Marco de Trabajo DOPPLER

Los autores proponen DOPPLER, un marco de aprendizaje por refuerzo (RL) que emplea un esquema de decisión secuencial de política dual y un pipeline de entrenamiento de tres etapas.

2.1 Formulación de Política Dual

En lugar de aprender una única política monolítica, DOPPLER descompone el proceso de asignación en dos políticas distintas, modeladas como un Proceso de Decisión de Markov (MDP) episódico:

1. SEL (Política de Selección): Selecciona la siguiente operación (vértice) para asignar del conjunto de nodos candidatos. Esta política recorre el grafo parcialmente asignado para aproximar el flujo no determinista del "tiempo", aprendiendo efectivamente un orden de ejecución.
2. PLC (Política de Colocación): Determina el dispositivo GPU específico para la operación seleccionada, con el fin de equilibrar la carga y minimizar la comunicación.

Arquitectura:

• Ambas políticas utilizan Redes Neuronales de Grafos (GNNs) con paso de mensajes para codificar la estructura del grafo de flujo de datos, las dependencias y las rutas críticas.
• SEL agrega información de las rutas críticas (rutas de nivel inferior y superior) y las características de los nodos para seleccionar el siguiente nodo.
• PLC combina las representaciones de los nodos, las características del dispositivo (por ejemplo, carga actual, tiempo de finalización) y el conjunto de nodos ya colocados para seleccionar el dispositivo óptimo.
• Optimización de Eficiencia: Para evitar costos de entrenamiento prohibitivos asociados con la reejecución del paso de mensajes en cada paso, DOPPLER realiza el paso de mensajes una vez por episodio y actualiza las características del dispositivo dinámicamente sin recomputar las incrustaciones (embeddings) del grafo en cada paso.

2.2 Pipeline de Entrenamiento de Tres Etapas

DOPPLER emplea un paradigma de "aprender haciendo" a través de tres etapas para asegurar la convergencia y la adaptabilidad:

1. Etapa I (Aprendizaje por Imitación / Offline): Las políticas se preentrenan utilizando aprendizaje supervisado para imitar a un maestro heurístico (específicamente, la heurística de Ruta Crítica). Esto acelera la convergencia al proporcionar una política inicial sólida antes de que comience la exploración.
2. Etapa II (RL basado en Simulación / Offline): Las políticas se refinan mediante aprendizaje por refuerzo dentro de un simulador de software (que implementa la lógica de ejecución WC). La recompensa se deriva del tiempo de ejecución simulado. Esta etapa permite una iteración rápida sin la latencia del hardware real.
3. Etapa III (RL en Sistema Real / Online): Las políticas entrenadas se despliegan en un sistema WC real. Se refinan continuamente mediante RL online utilizando los tiempos de ejecución reales como recompensas. Esta etapa adapta la política a las características específicas del hardware (por ejemplo, jitter de NVLink, contención) que los simuladores no pueden capturar perfectamente, sin sobrecarga adicional para el usuario.

3. Contribuciones Clave

1. Formulación del Problema: El artículo investiga la asignación de dispositivos específicamente dentro del contexto de sistemas multi-GPU asíncronos y de trabajo conservador, abordando la naturaleza estocástica del orden de ejecución.
2. Enfoque de Política Dual: Introduce una separación de responsabilidades donde una política aprende el orden de recorrido de los nodos y la otra gestiona la colocación de los dispositivos, capturando tanto la dinámica de ejecución como las restricciones del hardware.
3. Marco de Tres Etapas: DOPPLER propone una estrategia de entrenamiento híbrida que combina aprendizaje por imitación, RL basado en simulación y RL online en sistemas reales. Este enfoque acelera la convergencia y permite la optimización continua durante el despliegue.
4. Implementación Escalable de GNN: Los autores proponen una aproximación eficiente de paso de mensajes que reduce significativamente el tiempo de entrenamiento sin comprometer la convergencia.

4. Resultados Experimentales

Los experimentos se realizaron en cuatro GPUs NVIDIA Tesla P100 utilizando diversas cargas de trabajo: Multiplicaciones de Matrices en Cadena (CHAINMM), Redes Neuronales Feedforward (FFNN) y bloques/capas de transformadores Llama.

• Rendimiento vs. Baselines: DOPPLER (específicamente el DOPPLER-SYS totalmente entrenado) superó consistentemente a los modelos de referencia, incluyendo la heurística de Ruta Crítica, PLACETO (RL de política única), GDP (RL de incrustación de grafos) y un optimizador enumerativo personalizado.
  • DOPPLER logró tiempos de ejecución hasta un 52.7% menores en comparación con el mejor baseline (GDP).
  • Redujo el tiempo de ejecución hasta un 78.2% en comparación con la heurística de Ruta Crítica.
  • Superó al optimizador enumerativo personalizado hasta en un 13.8%.
• Estudios de Ablación:
  • Política Dual: La combinación de las políticas SEL y PLC produjo los mejores resultados, particularmente para modelos complejos, demostando la necesidad de aprender tanto el orden como la colocación.
  • Etapas de Entrenamiento: Entrenar únicamente en el sistema real llevó a una convergencia lenta e inestabilidad. La combinación de aprendizaje por imitación (Etapa I) y simulación (Etapa II) fue crucial para una convergencia estable y rápida antes del refinamiento online (Etapa III).
• Transferencia de Aprendizaje: DOPPLER demostró una fuerte capacidad de generalización. Las políticas entrenadas en grafos simples (FFNN) pudieron ajustarse con pocos episodios (2k–4k) para funcionar bien en arquitecturas no vistas (Llama) y diferentes configuraciones de hardware (por ejemplo, transferir de 4 P100s a 8 V100s).
• Escalabilidad: El entrenamiento y la inferencia de DOPPLER escalan linealmente con el tamaño del grafo de flujo de datos, superando a otros baselines de RL en eficiencia.

5. Significado y Reivindicaciones

El artículo afirma que DOPPLER representa un paso significativo hacia la optimización de las cargas de trabajo modernas de ML al ir más allá de los supuestos estáticos de sincronía por lotes. Su importancia radica en:

• Eficiencia Práctica: Al permitir la ejecución de trabajo conservador, DOPPLER puede ahorrar una cantidad sustancial de horas de GPU en despliegues a gran escala (por ejemplo, ahorrando millones de horas de GPU anualmente para cargas de trabajo de la escala de ChatGPT).
• Adaptabilidad: A diferencia de las heurísticas estáticas o los modelos puramente basados en simulación, la capacidad de aprendizaje online de DOPPLER le permite adaptarse a las realidades ruidosas de los entornos de hardware de producción.
• Innovación Metodológica: La descomposición de política dual ofrece un enfoque más matizado para la explosión combinatoria de los problemas de asignación de dispositivos, separando el "qué ejecutar a continuación" de "dónde ejecutarlo", lo cual se alinea mejor con la naturaleza temporal de la ejecución asíncrona.

Los autores enfatizan que su enfoque no requiere cambios en los marcos de trabajo de ML subyacentes, sino que optimiza la capa de programación (scheduling), lo que lo convierte en una solución práctica para las infraestructuras multi-GPU existentes.