Resumen Técnico: Partición de Datos Adaptativa para el Aprendizaje Federado y Distribuido Energéticamente Eficiente en Sistemas Heterogéneos

1. Planteamiento del Problema

El entrenamiento de modelos distribuidos en sistemas heterogéneos se ve significativamente obstaculizado por las variaciones en el rendimiento del procesamiento, el comportamiento de la memoria y las condiciones operativas entre los dispositivos participantes. Cuando se aplica una división fija e igual de los datos de entrenamiento, los dispositivos más rápidos terminan su trabajo temprano y permanecen inactivos mientras esperan a los dispositivos más lentos ("stragglers") en los puntos de sincronización. Este efecto de rezago no solo aumenta el tiempo total de ejecución (makespan) del entrenamiento, sino que también eleleva el consumo energético de los dispositivos, ya que los participantes sincronizados deben permanecer activos hasta que el progreso colectivo pueda continuar.

Los enfoques existentes para mitigar este problema suelen basarse en la selección de participantes (excluyendo dispositivos lentos), la agregación asíncrona, cambios en los métodos de optimización o la adaptación del modelo. Si bien sistemas recientes han explorado la redistribución de los datos de entrenamiento, muchos dependen de controladores aprendidos, búsquedas previas, configuraciones de hardware específicas o alteraciones en la semántica de sincronización y los conjuntos de participación. Existe una brecha para un mecanismo de bucle cerrado común que pueda operar a través de diferentes fronteras de sincronización y propiedad de datos (como el Aprendizaje Federado y el Paralelismo de Datos Distribuido) manteniendo la configuración original, el modelo y las reglas de sincronización.

2. Metodología

Los autores proponen un controlador de asignación ligero, basado en mediciones, que reasigna un presupuesto fijo de muestras en los límites de cada ronda basándose en el tiempo de entrenamiento observado y la energía medida por muestra en cada dispositivo. El controlador opera sin modificar el modelo, la función de pérdida, el optimizador o el conjunto de dispositivos participantes.

Regímenes de Ejecución

El estudio evalúa el controlador a través de tres regímenes de ejecución distintos:

1. Promedio de Modelos Periódico (PMA): Los trabajadores entrenan réplicas de modelos independientes sobre datos compartidos dentro de un dominio confiable y se sincronizan únicamente en los límites de la ronda.
2. Paralelismo de Datos Distribuido (DDP): Los rangos sincronizan gradientes después de cada paso del optimizador. Este régimen requiere una estrategia específica de dimensionamiento de lotes (batch-sizing) para manejar cargas de trabajo desiguales por rango sin introducir relleno (padding) o lotes ficticios.
3. Aprendizaje Federado Jerárquico con DDP intra-silo (HFL-DDP): Los dispositivos se agrupan en "silos" fijos (dominios propietarios de datos). Los rangos dentro de un silo utilizan DDP, mientras que los modelos de los silos se agregan en el límite de la ronda federada. La reasignación se restringe a los rangos dentro del mismo silo para respetar la propiedad de los datos.

El Controlador Adaptativo

El controlador funciona de la siguiente manera:

• Medición: Al final de cada ronda $r$, cada participante $i$ reporta su tiempo de procesamiento ($t_i^{(r)}$), consumo de energía ($e_i^{(r)}$) y recuento de muestras ($n_i^{(r)}$).
• Cálculo de Costo: Los costos por muestra se derivan como $\tau_i^{(r)} = t_i^{(r)} / n_i^{(r)}$ y $\epsilon_i^{(r)} = e_i^{(r)} / n_i^{(r)}$. Estos se suavizan mediante una Media Móvil Exponencial (EMA) para reducir el ruido.
• División Proporcional Ponderada: El controlador calcula un costo combinado adimensional $c_i^{(r)}$ utilizando pesos seleccionados por el usuario $w_t$ (tiempo) y $w_e$ (energía).
  • $c_i^{(r)} = w_t \frac{\hat{\tau}_i^{(r)}}{\text{mediana}(\hat{\tau})} + w_e \frac{\hat{\epsilon}_i^{(r)}}{\text{mediana}(\hat{\epsilon})}$
• Asignación: Los participantes reciben muestras inversamente proporcionales a su costo. Un mecanismo de "banda muerta" (deadband) evita la reasignación innecesaria si las diferencias de costo son insignificantes.
• Especificaciones de DDP: Para DDP, el controlador asegura que todos los rangos ejecuten el mismo número de pasos del optimizador por ronda, construyendo esquemas de tamaño de lote por paso que redistribuyen las muestras locales, preservando el objetivo de la media de la muestra del gradiente.

Configuración Experimental

La evaluación se realizó en un banco de pruebas de Sistema en Chip (SoC) heterogéneo de once dispositivos que cuenta con arquitecturas Arm, RISC-V y x86. Se utilizaron tres tareas de aprendizaje supervisado: MNIST (clasificación de imágenes), ERA5 (regresión de series temporales) e IMDB (clasificación de sentimiento de texto). El estudio comparó las políticas adaptativas contra:

• Asignación igual: División uniforme de muestras.
• Perfilado estático: Una medición única seguida de una asignación fija.
• Selección tipo Oort: Un punto de referencia que excluye a los participantes lentos.
• FedAsync: PMA asíncrono.
• Tamaño de lote EMA: Un punto de referencia de rendimiento a nivel de paso para DDP.

3. Principales Contribuciones

El artículo presenta cuatro contribuciones primarias:

1. Un Controlador Basado en Mediciones: Un mecanismo ligero que reasigna un presupuesto fijo de muestras basado en el tiempo y la energía observados por muestra, preservando el conjunto completo de participación y la semántica de sincronización.
2. Diseño Unificado: Un diseño de asignación común aplicable a PMA, DRP y HFL-DDP, con restricciones específicas para las fronteras de propiedad de datos en entornos jerárquicos.
3. Objetivo Consciente de la Energía: Una función objetivo con pesos ajustables que expone la interacción entre la energía por muestra, el tiempo de ejecución (makespan) y la energía de los dispositivos sincronizados, desafiando la suposición de que maximizar el peso de la energía siempre minimiza la energía total.
4. Evaluación Exhaustiva: Un estudio replicado a través de tres tareas de aprendizaje y tres regímenes de sincronización en un banco de pruebas heterogéneo, conectando los cambios de asignación con el desequilibrio de cómputo, la exposición a la sincronización, el tiempo de ejecución y la energía medida.

4. Resultados

Los experimentos en el banco de pruebas heterogéneo arrojaron los siguientes hallazgos:

• Reducción de Makespan y Energía: La partición adaptativa (específicamente las políticas de "Solo Tiempo" y "Baja Energía") redujo consistentemente el tiempo de ejecución y la energía medida de los dispositivos en comparación con la asignación Igual. En PMA, la política de "Solo Tiempo" redujo la mediana del tiempo de ejecución aproximadamente un 71–80% y la energía entre un 28–40% en todas las tareas.
• Mitigación de Rezagados (Stragglers): El controlador redujo significativamente el desequilibrio de cómputo. En DDP, el coeficiente de variación para el tiempo del pase hacia adelante (forward-pass) cayó de ~0.8–1.1 (Igual) a ~0.02–0.18 (adaptativo), lo que se correlaciona directamente con la reducción de la espera de sincronización.
• Convergencia: Las políticas adaptativas alcanzaron regiones de bajo objetivo (pérdida/MSE de validación) más rápido que la asignación Igual, logrando a menudo una calidad final comparable en mucho menos tiempo.
• Respuesta de Energía No Monotónica: Aumentar el peso de la energía ($w_e$) no siempre redujo la energía total. Asignar más trabajo a dispositivos que son eficientes energéticamente pero más lentos podía extender la ventana de ejecución sincronizada, aumentando la energía acumulada de todos los participantes. Un pequeño peso de energía ($w_e = 0.25$) a menudo proporcionó un mejor compromiso que "Solo Energía" ($w_e = 1$).
• Sobrecarga (Overhead): El cálculo de la programación fue insignificante, consumiendo menos del 0.02% de la duración de la ronda.

5. Significado y Reivindicaciones

El artículo sostiene que la partición de datos basada en mediciones proporciona un método práctico para mitigar los rezagados en diversos regímenes de entrenamiento distribuido sin alterar la configuración de aprendizaje ni las reglas de sincronización.

• Practicidad: El método no requiere un perfilado previo, políticas aprendidas o cambios en el optimizador o la arquitectura del modelo. Opera como un controlador de bucle cerrado utilizando telemetría en tiempo de ejecución.
• Optimización Dependiente del Contexto: Los autores enfatizan que el equilibrio óptimo entre tiempo y energía depende de la tarea específica y el régimen de ejecución. Si bien "Solo Tiempo" es el valor predeterminado más fiable para minimizar el tiempo de ejecución, un pequeño peso de energía puede ser beneficioso en contextos específicos (por ejemplo, HFL-DDP en ciertas tareas).
• Distinción de Trabajos Previos: A diferencia de los métodos de selección de participantes (como Oort) que excluyen dispositivos, este enfoque retiene a todos los participantes provistos. A diferencia de los métodos asíncronos, preserva la semántica de sincronización del régimen objetivo.
• Limitaciones: Los autores señalan modestamente que los resultados son específicos para el banco de pruebas de once dispositivos evaluado y las condiciones de red fijas. No pretenden una superioridad universal sobre todos los métodos existentes, sino demostrar que una simple reasignación de presupuesto fijo basada en mediciones puede superar a las divisiones iguales estáticas y al perfilado de una sola vez en entornos heterogéneos.