Resumen Técnico: DisagMoE

1. Declaración del Problema

Las arquitecturas de Mezcla de Expertos (MoE) a gran escala permiten el entrenamiento de Modelos de Lenguaje Grande (LLM) con billones de parámetros activando de forma dispersa un subconjunto de expertos. Sin embargo, el entrenamiento de estos modelos enfrenta cuellos de botella severos debido a la comunicación de todos-a-todos requerida para el paralelismo de expertos (EP). A medida que los modelos se escalan a través de múltiples nodos de GPU, el ancho de banda limitado de la red inter-nodo (por ejemplo, InfiniBand) exacerba la latencia de comunicación, lo cual puede consumir hasta el 50% del tiempo de un paso de entrenamiento.

Las estrategias de optimización existentes se centran principalmente en solapar la comunicación con el cálculo a nivel de operador (por ejemplo, fragmentando los cálculos de la FFN) o utilizando la canalización de micro-lotes. Sin embargo, estos enfoques sufren de dos limitaciones fundamentales:

1. Paradas Residuales: Los desequilibrios inherentes en las proporciones de cálculo-comunicación entre las capas de atención y FFN impiden un solapamiento perfecto. El cálculo de atención escala cuadráticamente con la longitud de la secuencia, mientras que la FFN y la comunicación escalan linealmente. Esto crea "colas de comunicación" que no pueden ocultarse, especialmente a medida que aumentan las longitudes de secuencia.
2. Ineficiencia de Recursos: Los enfoques desagregados anteriores (separar la atención y la FFN en hardware diferente) a menudo requieren que cada grupo mantenga los parámetros completos del modelo, lo que conduce a errores de memoria insuficiente (OOM) a escala, o no logran adaptar la asignación de recursos a las características específicas de cálculo-comunicación de cada componente.

2. Metodología: DisagMoE

DisagMoE es un sistema de entrenamiento diseñado para mitigar los cuellos de botella de todos-a-todos mediante la optimización conjunta de la colocación del modelo y la programación a través de una arquitectura desagregada.

2.1 Colocación Desagregada de Componentes

A diferencia del Paralelismo de Expertos tradicional, donde las capas de atención y FFN coexisten en el mismo grupo de GPU, DisagMoE particiona el modelo por tipo de componente:

• Grupos de Atención (A-Workers): Estos grupos contienen capas de auto-atención de múltiples capas transformadoras. Replican componentes densos y distribuyen los datos de entrenamiento mediante Paralelismo de Datos (DP).
• Grupos de FFN (F-Workers): Estos grupos contienen capas FFN (expertos) de múltiples capas transformadoras. Los expertos se fragmentan entre las GPU dentro del grupo mediante Paralelismo de Expertos (EP).
• Asignación de Capas: Las capas se asignan a los grupos de manera intercalada (por ejemplo, el Grupo 0 maneja las capas 0, $p$, $2p$; el Grupo 1 maneja las capas 1, $p+1$, etc.), maximizando la utilización del dispositivo bajo restricciones de memoria.

2.2 AF-Pipe: Canalización Multi-etapa con Solapamiento

DisagMoE introduce AF-Pipe, un esquema de canalización que trata la comunicación como una etapa de primer nivel en lugar de un efecto secundario.

• Comunicación Muchos-a-Muchos (M2N): En lugar de los intercambios estándar de todos-a-todos dentro de un solo grupo de nodos, AF-Pipe establece intercambios de tokens unidireccionales y muchos-a-muchos entre los grupos de Atención y FFN.
• Alineación de Etapas: La canalización alinea los límites entre grupos, fusionando las operaciones tradicionales de "despacho" y "combinación" en una única etapa de comunicación unificada. Esto reduce la sobrecarga de comunicación en aproximadamente $1/k$ (donde $k$ son los expertos top-$k$) y elimina transferencias de datos redundantes.
• Ejecución Asíncrona: El sistema utiliza tres flujos coordinados (hacia adelante, hacia atrás y comunicación) dentro de cada trabajador. Esto permite que el grupo FFN envíe gradientes (N2M) mientras simultáneamente calcula el paso hacia adelante en los datos disponibles, asegurando una ejecución no bloqueante y una utilización casi continua de la GPU.

2.3 Asignación Adaptativa de Trabajadores

Para abordar el desequilibrio entre la atención (limitada por cálculo) y la FFN (limitada por comunicación), DisagMoE emplea un asignador adaptativo guiado por un Modelo de Techo de Cálculo-Comunicación.

• Modelado: El sistema modela el punto de inflexión ($\hat{I}$) donde un componente cambia de estar limitado por la red a estar limitado por el cálculo. Las capas de atención alcanzan el techo de cálculo más rápido que las capas FFN a medida que aumenta la longitud de la secuencia.
• Optimización: El sistema formula la partición de GPU y Tarjetas de Interfaz de Red (NIC) como un problema de Programación Lineal Entera Mixta (MILP). Asigna dinámicamente más ancho de banda de red a los grupos FFN (para elevar su techo de comunicación) y más recursos de cálculo a los grupos de Atención, equilibrando las etapas de la canalización para minimizar las burbujas.
• Estrategia de Búsqueda: Un algoritmo utiliza el modelo de techo para generar una configuración inicial, la cual se refina luego mediante búsqueda local utilizando los tiempos medidos por paso para tener en cuenta sobrecargas no modeladas.

3. Contribuciones Clave

1. Colocación Desagregada de Componentes: Una arquitectura novedosa que separa las capas de atención y FFN en grupos de trabajadores disjuntos, permitiendo un escalado independiente y una asignación de recursos.
2. Esquema AF-Pipe: Una canalización multi-etapa que reemplaza el todos-a-todos convencional con intercambios M2N inter-grupo, solapando sistemáticamente la comunicación tanto con los cálculos de atención como con los de FFN.
3. Asignación Adaptativa de Recursos: Un mecanismo guiado por el modelo de techo que dimensiona conjuntamente los recursos de GPU y NIC para los grupos de atención y FFN, equilibrando efectivamente la proporción cálculo-comunicación a través de diversas longitudes de secuencia y configuraciones de modelos.

4. Resultados Experimentales

Los autores evaluaron DisagMoE en clústeres de 16 nodos (128× GPUs H800) utilizando tres modelos MoE: DeepSeek-MoE, GPT-OSS-120B y Qwen3, con longitudes de secuencia que oscilan entre 4K y 32K.

• Mejora del Rendimiento: DisagMoE superó consistentemente a las líneas base sólidas, incluyendo Megatron-LM (1F1B intercalado), Tutel, Comet y DualPipe.
  • Logró una aceleración de hasta 1.81× sobre Megatron-1F1B.
  • Logró una aceleración de hasta 1.34× sobre los sistemas de solapamiento MoE más avanzados.
• Reducción de Comunicación: En comparación con las líneas base, DisagMoE redujo el tiempo de comunicación no solapada hasta en un 88% (vs. Tutel), 75% (vs. Comet) y 45% (vs. DualPipe).
• Escalabilidad: El sistema demostró efectividad a través de diferentes valores de top-$k$ y tamaños de Paralelismo de Expertos (EP).
• Sensibilidad a Recursos: Los experimentos confirmaron que las proporciones óptimas de GPU entre los grupos de atención y FFN varían con la longitud de la secuencia (por ejemplo, secuencias más largas favorecen más recursos de atención), validando la necesidad de una asignación adaptativa.

5. Significado y Afirmaciones

El artículo afirma que DisagMoE aborda una limitación fundamental en el entrenamiento actual de MoE: la incapacidad de solapar perfectamente la comunicación debido al desequilibrio estructural entre las capas de atención y FFN. Al desagregar estos componentes y asignar recursos adaptativamente basándose en un modelo de techo, el sistema expande la ventana de solapamiento efectiva más allá de lo que puede lograr un aprovisionamiento simétrico de ancho de banda.

Los autores enfatizan que su enfoque es particularmente significativo para cargas de trabajo de pre-entrenamiento a gran escala con longitudes de secuencia fijas. Señalan que, aunque su diseño actual asume formas estables (longitud de secuencia fija y micro-lote), la idea central de que el hardware y el software futuros deben apuntar a proporciones favorables de cálculo-comunicación entre componentes en lugar de simplemente escalar el ancho de banda crudo, ofrece una nueva dirección para escalar el entrenamiento de MoE de manera eficiente. El sistema se implementa sobre Megatron-LM y se muestra como una solución viable para entrenar modelos con cientos de miles de millones de parámetros sin las restricciones de memoria que plagaron los intentos anteriores de entrenamiento desagregado.