Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

Resumen Técnico: Más allá de lo agnóstico a la tarea: Agrupación consciente de la tarea para una inferencia de MoE eficiente en comunicación

1. Planteamiento del problema

Los modelos de Mezcla de Expertos (MoE) de activación dispersa escalan la capacidad de manera eficiente mediante computación condicional, pero enfrentan desafíos sistémicos significativos durante la inferencia distribuida, específicamente un excesivo overhead de comunicación entre GPUs y un severo desequilibrio de carga. Las estrategias existentes de colocación de expertos intentan mitigar estos costos mediante la co-localización de expertos que se co-activan frecuentemente, pero típicamente derivan un plan de despliegue global único a partir de trazas de enrutamiento agnósticas a la tarea (agregadas globalmente).

Los autores identifican un fallo crítico en este enfoque: la co-activación de expertos está fuertemente condicionada por la tarea. Como demuestra el análisis del artículo (Figura 1), los pares de expertos estrechamente acoplados en una familia de tareas (por ejemplo, generación de código) suelen estar no correlacionados o incluso compitiendo en otra (por ejemplo, razonamiento matemático). El promedio de estos patrones heterogéneos en un único grafo de afinidad colapsa los grupos de expertos específicos de cada tarea, dejando la localidad de comunicación sin explotar y fallando al adaptarse a los cambios de distribución en las cargas de trabajo en línea.

2. Metodología

El marco propuesto, Agrupación de Co-activación Consciente de la Tarea (TACG) junto con la Replicación Compartida de Expertos Genéricos (GESR), aborda estos problemas a través de un proceso de despliegue en dos etapas:

A. Agrupación de Co-activación Consciente de la Tarea (TACG)
TACG convierte las estadísticas de co-activación condicionadas por la tarea en un plan de colocación factible en capacidad:

1. Calificación de la preferencia de la familia de tareas: Utilizando trazas de calibración, el sistema calcula vectores de uso y matrices de co-activación para cada familia de tareas. Calcula una "puntuación de preferencia suave" ($s^{(f)}_{l,e}$) para cada experto relativo a familias de tareas específicas, normalizando estas puntuaciones para crear una distribución de probabilidad ($p^{(f)}_{l,e}$) que representa la afinidad de un experto por una familia.
2. Grafo de co-activación modulado por la tarea: En lugar de utilizar un grafo de co-activación agrupado, TACG construye un grafo modulado por la tarea ($\tilde{A}_l$). Este grafo repondera las aristas de co-activación estándar mediante el producto de las preferencias de familia de los expertos ($p^{(f)}_{l,e} p^{(f)}_{l,e'}$). Esto asegura que las aristas entre expertos que se co-activan y comparten una fuerte preferencia de familia sean amplificadas, mientras que las aristas entre distintas familias sean atenuadas.
3. Agrupación factible en capacidad: El marco particiona los expertos en $M$ grupos (uno por GPU) maximizando la afinidad dentro del grupo en el grafo modulado por la tarea bajo estrictas restricciones de capacidad. Esto se resuelve mediante un enfoque de agrupamiento espectral de dos etapas seguido de un paso de "reparación de capacidad" local para asegurar que se cumplan los límites exactos de carga de la GPU.

B. Replicación Compartida de Expertos Genéricos (GESR)
Para manejar el sesgo de carga causado por los expertos "genéricos" (aquellos activos en todas las familias y propensos a convertirse en puntos calientes o hotspots), GESR introduce un mecanismo acompañante ligero:

1. Identificación: Identifica expertos genéricos utilizando una puntuación que premia la centralidad global y la consistencia entre familias, mientras penaliza la desviación fuerte específica de una tarea.
2. Replicación: Estos expertos se replican en un pequeño conjunto de GPUs secundarias elegidas según su compatibilidad estructural (afinidad de co-activación).
3. Guardián en línea: Durante el servicio, una regla de selección consciente de la localidad y la carga elige la mejor GPU candidata para un experto genérico, utilizando señales de carga con decaimiento exponencial para prevenir puntos calientes sin alterar el enrutador del MoE.

3. Contribuciones Clave

• Perspectiva (Insight): El artículo establece que la co-activación de expertos en el servicio de MoE multi-tarea está fuertemente condicionada por la tarea, argumentando que una colocación efectiva debe ser impulsada por una co-activación consciente de la tarea en lugar de un único grafo de afinidad promediado globalmente.
• TACG: Una estrategia consciente de la tarea que repondera el grafo de co-activación utilizando las preferencias de experto-a-familia, produciendo una colocación que generaliza estrictamente el agrupamiento agnóstico a la tarea.
• GESR: Un módulo de replicación ligero que replica selectivamente expertos genéricos de diversas familias y realiza una selección en línea consciente de la carga para mantener el equilibrio bajo cambios de carga de trabajo.
• Validación Empírica: Evaluación extensa en tres modelos MoE de código abierto representativos (DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B) demostrando mejoras consistentes sobre bases sólidas.

4. Resultados Experimentales

El marco fue evaluado en una configuración de paralelismo de expertos de 16 GPUs a través de tres modelos con cuatro familias de tareas (Código, Consulta, Matemáticas, Razonamiento Complejo).

• Reducción de Comunicación: El marco completo (TACG + GESR) redujo el costo promedio de comunicación por token en un 31.39% en comparación con el baseline, superando a la estrategia de colocación previa más fuerte (ExFlow) por un 12.58%.
• Equilibrio de Carga: El método mantuvo un alto índice de equidad de Jain promedio de 0.9975 y una baja violación de carga máxima promedio (MaxVio) de 0.0736, igualando o excediendo el equilibrio de carga de los baselines mientras reducía significativamente la comunicación.
• Robustez: Bajo severos cambios en la distribución de tareas (cargas de trabajo de una sola tarea fuera de la distribución), el enfoque consciente de la tarea mantuvo reducciones de comunicación estables, mientras que las variantes agnósticas a la tarea sufrieron grandes oscilaciones de rendimiento y formación de puntos calientes (MaxVio aumentando ~7x en tareas de Código).
• Ablación: Eliminar la señal de co-activación modulada por la tarea redujo las ganancias de comunicación de 31.39% a 16.51%, confirmando que el condicionamiento por tarea es el factor dominante. Se demostró que GESR es esencial para controlar las violaciones de carga (MaxVio) sin sacrificar la eficiencia de comunicación.

5. Significancia y Reivindicaciones

El artículo afirma que su significancia radica en el cambio de paradigma de la colocación de expertos de MoE de una visión estática y agnóstica a la tarea hacia una perspectiva dinámica y consciente de la tarea. Al modelar explícitamente la heterogeneidad de la co-activación de expertos a través de diferentes familias de tareas, el marco desbloquea la localidad de comunicación que es invisible para los métodos de agregación global.

Los autores aseguran que este enfoque proporciona una mejora estricta en la frontera de equilibrio comunicación-carga. A diferencia de métodos anteriores que a menudo sacrifican el equilibrio de carga por ganancias de comunicación (o viceversa), TACG + GESR logra una reducción sustancial de la comunicación preservando un equilibrio de carga casi perfecto. El trabajo destaca que el servicio multi-tarea efectivo requiere estrategias de despliegue que respeten la estructura intrínseca de la colaboración de expertos específica de cada tarea, en lugar de promediarlas.

El artículo se mantiene modesto respecto a su alcance, señalando que se centra en la colocación estática offline y deja como trabajo futuro el reagrupamiento dinámico dentro de las ventanas de solicitud, el codiseño de memoria (prelectura/caché) y los expertos especializados en modelos particionados por dominio.