Resumen Técnico: Paralelismo Heterogéneo para el Entrenamiento de Modelos de Lenguaje Multimodales Grandes

1. Declaración del Problema

A medida que el entrenamiento de modelos fundacionales evoluciona desde transformadores de solo texto hacia sistemas multimodales (que abarcan imágenes, audio, video y documentos), el gráfico de entrenamiento se vuelve cada vez más heterogéneo en términos de profundidad, ancho y tamaño de contexto de los módulos. Sin embargo, las pilas de paralelismo actuales (Megatron-LM, DeepSpeed, etc.) suelen imponer una disposición paralela única y compartida (Paralelismo de Tensor, Contexto, Pipeline, Datos y Experto) en todo el gráfico de extremo a extremo.

Este enfoque de "talla única" genera ineficiencias significativas cuando los requisitos del codificador y del Modelo de Lenguaje Grande (LLM) divergen:

• Comunicación Innecesaria: Los LLM grandes a menudo requieren grados altos de Paralelismo de Tensor (TP). Obligar a los codificadores pequeños a heredar estos grados de TP introduce colectivas de all-reduce innecesarias alrededor de los cálculos locales del codificador, añadiendo sobrecarga que supera los ahorros de memoria.
• Desajuste del Paralelismo de Contexto: El Paralelismo de Contexto (CP) es esencial para que los LLM gestionen la memoria de atención sobre secuencias largas. Sin embargo, las longitudes de secuencia del codificador están acotadas por la resolución de la modalidad (por ejemplo, parches de imagen), y no por la ventana de contexto del LLM. Heredar las particiones de CP del LLM fija secuencias de tamaño fijo para el codificador sin proporcionar alivio de memoria, mientras que esos rangos podrían funcionar de manera más eficiente como réplicas independientes de Paralelismo de Datos (DP) para procesar más muestras.
• Restricciones del Pipeline: El Paralelismo de Pipeline (PP) a menudo confina a los codificadores a la primera etapa del pipeline del LLM. Esto obliga al codificador a heredar el grado de TP del LLM y puede sobrecargar rangos específicos, impidiendo que el LLM utilice una disposición de pipeline más eficiente.

El desafío central es que el punto de operación óptimo para el cómputo, la memoria y la comunicación cambia según los parámetros de la carga de trabajo (escala del codificador, longitud de contexto, densidad de tokens, tamaño del mundo). Una disposición estática y compartida no puede adaptarse a estos cambios.

2. Metodología

Los autores proponen Paralelismo Heterogéneo, una abstracción de tiempo de ejecución que permite a los módulos dentro de un único gráfico multimodal de extremo a extremo elegir disposiciones lógicas independientes y colocaciones físicas de rangos. El sistema admite dos modos de ejecución distintos:

A. Abstracción Central: Comunicadores de Frontera

Para desacoplar las disposiciones de los módulos mientras se preservan las semánticas de entrenamiento, el sistema introduce comunicadores de frontera. Estos son operadores conscientes de autograd adjuntos a las aristas intermodulares que manejan la transformación de tensores entre diferentes representaciones distribuidas.

• Paso Adelante: Materializa la activación de origen en la disposición del módulo de destino.
• Paso Atrás: Aplica la transformación inversa para enrutar los gradientes de vuelta a la disposición del módulo de origen.
• Corrección: El sistema garantiza que las semánticas de los tensores se preserven a través de cuadrículas independientes. Las asignaciones incorrectas se previenen definiendo estrictamente transformaciones lógicas de partición de lotes (DP Igual, Entrada en Abanico y Salida en Abanico) para cada arista.

B. Modos de Ejecución

1. Ejecución Colocada:

   • Mecanismo: Los módulos comparten las mismas GPUs físicas pero las interpretan bajo diferentes cuadrículas lógicas (por ejemplo, un codificador con alto DP y bajo TP compartiendo rangos con un LLM con bajo DP y alto TP).
   • Implementación: El comunicador actúa como un operador de disposición de rango compartido. Realiza operaciones colectivas locales (por ejemplo, all-gather o selección de intervalo) para reinterpretar la disposición de datos sin mover tensores entre conjuntos de rangos.
   • Programación: Utiliza una programación de tres fases para evitar conflictos entre las colectivas del codificador y el pipeline escalonado del LLM (1F1B).
     • Fase 1: Paso adelante del codificador + transformación colocada hacia adelante.
     • Fase 2: Ejecución del pipeline del LLM desvinculado (1F1B) sobre el tensor transformado.
     • Fase 3: Transformación colocada hacia atrás + paso atrás del codificador.
   • Beneficio: Permite descargar los parámetros del codificador durante la fase del pipeline del LLM, mejorando la eficiencia de la memoria.

2. Ejecución No Colocada:

   • Mecanismo: Los módulos ocupan conjuntos disjuntos de rangos (por ejemplo, una "isla de codificador" y una "isla de LLM").
   • Implementación: Utiliza comunicación explícita Punto a Punto (P2P) entre líderes de frontera deterministas (rangos con coordenadas canónicas).
     • Adelante: Los líderes de origen envían a los líderes de destino; las transmisiones locales materializan los datos.
     • Atrás: Los gradientes se enrutan de vuelta a través de las rutas inversas de los líderes.
   • Programación: Utiliza un despachador 1F1B consciente del gráfico que resuelve las aristas en un gráfico multietapa, manejando uniones y ramas donde múltiples flujos de codificador alimentan al LLM.
   • Beneficio: Evita que el estado del codificador (parámetros, activaciones, estados del optimizador) compita con el LLM por memoria en los mismos rangos, permitiendo que el LLM utilice una disposición de pipeline más eficiente.

3. Contribuciones Clave

1. Abstracción de Tiempo de Ejecución de Paralelismo Heterogéneo: Un sistema que permite disposiciones TP/CP/PP/DP/EP independientes y colocaciones físicas de rangos para módulos en un único gráfico de MLLM, admitiendo tanto la ejecución colocada como la no colocada.
2. Comunicadores de Frontera y Programación: La introducción de comunicadores de frontera conscientes de autograd que manejan transformaciones de disposición, acoplados con programadores especializados (tres fases para colocada, 1F1B consciente del gráfico para no colocada) para mantener la eficiencia del pipeline.
3. Evaluación del Régimen de Operación: Una evaluación exhaustiva que caracteriza cuándo la heterogeneidad es beneficiosa. El estudio identifica que la heterogeneidad colocada es óptima cuando hay margen de memoria suficiente pero las disposiciones lógicas difieren, mientras que la heterogeneidad no colocada se vuelve superior cuando el estado del codificador restringe la memoria del LLM o la disposición del pipeline.
4. Corrección y Código Abierto: Validación de la paridad en la convergencia de la pérdida frente a bases homogéneas mediante comprobaciones FP32 a nivel de paso y barridos de entrenamiento estilo LLaVA. El sistema se publica como una extensión de código abierto para Megatron-LM.

4. Resultados Experimentales

El sistema se evaluó en GPUs NVIDIA H100 de 80 GB en diversas cargas de trabajo multimodales (codificadores de 1B–12B, LLMs de 14B–120B) y longitudes de contexto (8K–128K).

• Heterogeneidad Colocada:
  • Mejora los TFLOPS/GPU hasta en un 49.3% (observado en escenarios de contexto largo, por ejemplo, 64K–128K, donde los LLM requieren CP pero los codificadores no).
  • Mejora el rendimiento en un 12.1%–21.6% en diversas configuraciones de codificador 1B/3B + LLM 14B/70B/120B al permitir que el codificador utilice un DP más alto y un TP más bajo.
• Heterogeneidad No Colocada:
  • Mejora el rendimiento agregado de tokens hasta en un 13.0% y los TFLOPS/GPU hasta en un 9.6% para cargas de trabajo de codificador grande (6B–12B) + LLM 120B.
  • Se convierte en el punto de operación preferido cuando la escala del codificador (por ejemplo, 6B+) crea presión de memoria que restringe la disposición del pipeline del LLM en configuraciones colocadas.
• Convergencia: Se validó que las disposiciones heterogéneas producen curvas de pérdida idénticas a las bases homogéneas, confirmando que los comunicadores de frontera preservan correctamente las semánticas de los gradientes.

5. Significado y Afirmaciones

El artículo afirma que el paralelismo heterogéneo es una evolución necesaria para el entrenamiento multimodal, alejándose de las restricciones rígidas de disposición a nivel de modelo. Al exponer grados de libertad adicionales, el sistema permite que la infraestructura de entrenamiento ajuste el punto de operación dinámicamente según las características de la carga de trabajo (tamaño del codificador, longitud de contexto, densidad de tokens).

Los autores enfatizan que esto no es meramente un cambio de configuración, sino un sustrato de tiempo de ejecución fundamental que preserva las semánticas de entrenamiento de extremo a extremo mientras habilita:

• Ejecución colocada para extraer el máximo rendimiento de un presupuesto fijo de GPU cuando los módulos pueden compartir hardware pero necesitan diferentes cuadrículas lógicas.
• Ejecución no colocada para desacoplar las presiones de memoria y programación cuando el estado del codificador de otro modo crearía un cuello de botella para el LLM.

El trabajo se presenta como una extensión de código abierto a Megatron-LM, proporcionando una base reutilizable para futuros marcos de trabajo MLLM modulares y sistemas de entrenamiento desagregados. El artículo no afirma resolver todos los desafíos multimodales, sino que aborda específicamente la ineficiencia causada por forzar módulos heterogéneos en una única disposición de paralelismo.