Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

1. Enunciado del Problema

El entrenamiento y la inferencia de modelos de lenguaje grandes (LLM) enfrentan dos cuellos de botella principales de memoria:

1. Memoria de Parámetros: El almacenamiento requerido para los pesos del modelo, los gradientes y los estados del optimizador.
2. Memoria de Activaciones: El almacenamiento requerido para las activaciones intermedias, que escala cuadráticamente con la longitud de la secuencia ($S^2$) en los mecanismos de atención.

Las estrategias de paralelismo existentes abordan estos problemas de forma aislada o con compensaciones:

• Paralelismo de Tensores (TP): Fragmenta los pesos del modelo entre dispositivos, reduciendo la memoria de parámetros pero dejando la memoria de activaciones completamente replicada. Se vuelve ineficiente en longitudes de secuencia largas debido al alto consumo de memoria de activaciones y al volumen de comunicación.
• Paralelismo de Secuencia (SP): Fragmenta la secuencia de entrada entre dispositivos, reduciendo la memoria de activaciones pero dejando los pesos del modelo completamente replicados. Es ineficaz cuando la memoria de parámetros es el cuello de botella.
• TP+SP Convencional: Combina ambos asignándolos a ejes ortogonales en una malla de dispositivos (por ejemplo, una cuadrícula $T \times \Sigma$). Aunque es efectivo, esto consume un gran número de dispositivos para el paralelismo de modelo ($T \times \Sigma$), dejando menos dispositivos para el paralelismo de datos. Además, a menudo obliga a que los grupos paralelos abarquen enlaces inter-nodo más lentos, subutilizando las interconexiones intra-nodo de alto ancho de banda (como NVLink o AMD Infinity Fabric).

El Desafío Central: ¿Cómo reducir simultáneamente tanto la memoria de parámetros como la de activaciones en un solo eje de dispositivo para maximizar la eficiencia de la memoria y liberar dispositivos para el paralelismo de datos, sin incurrir en costos de comunicación prohibitivos?

2. Metodología: Paralelismo de Tensores y Secuencia (TSP)

Los autores proponen TSP, una estrategia que "pliega" el paralelismo de tensores y de secuencia sobre un único eje de dispositivo. En lugar de asignar TP y SP a dimensiones ortogonales, TSP asigna a cada rango (GPU) tanto una fragmentación de los pesos del modelo como una fragmentación de la secuencia.

Componentes Clave del Diseño:

• Fragmentación Unificada: Cada dispositivo retiene $1/D$ de los parámetros del modelo y $1/D$ de los tokens de la secuencia. Esto reduce tanto la memoria proporcional a los parámetros como la memoria de activaciones por un factor de $D$ en un solo eje.
• Programaciones de Ejecución: Para manejar el hecho de que los pesos y las activaciones se fragmentan de manera diferente, TSP emplea programas de comunicación específicos para diferentes capas:
  • Capas de Atención: Utiliza un programa de difusión (broadcast).
    • Los rangos iteran sobre fragmentos de pesos (empaquetados $W_Q, W_K, W_V, W_O$).
    • En cada iteración, un rango difunde su fragmento de peso a todos los demás.
    • Todos los rangos calculan proyecciones locales en sus fragmentos de secuencia.
    • Una operación All-Gather reconstruye el contexto completo de Clave/Valor a través de la dimensión de la secuencia para realizar FlashAttention.
    • Se utiliza un esquema de particionamiento en zigzag para equilibrar la carga de trabajo de la atención causal (asegurando que los tokens posteriores no incurran en una carga de trabajo desproporcionadamente mayor).
  • Capas MLP con Puerta: Utiliza un programa de comunicación en anillo.
    • Los fragmentos de peso ($W_1, W_2, W_3$) circulan punto a punto alrededor del anillo.
    • Cada rango acumula salidas parciales localmente.
    • Esto elimina la necesidad de una operación All-Reduce (estándar en TP), ya que la secuencia permanece local y solo se mueven los pesos.
• Conciencia del Hardware: El diseño está optimizado para topologías como la MI300X de AMD (Infinity Fabric), donde el ancho de banda colectivo escala con el número de GPUs participantes. TSP asegura que los primitivos de comunicación (difusión y anillo) involucren al grupo completo o se solapen con el cómputo para maximizar la eficiencia.

3. Contribuciones Clave

1. Estrategia de Paralelismo Novel: Introducción de TSP, que colapsa dos ejes de paralelismo ortogonales en uno, permitiendo la reducción simultánea de la memoria de parámetros y de activaciones.
2. Análisis Teórico: Un modelo matemático riguroso que compara TSP frente a DP, TP, SP y TP+SP en cuanto a:
   • Memoria: TSP es el único esquema que reduce tanto la memoria proporcional a los pesos como la de activaciones por $1/D$ sin requerir $D^2$ dispositivos.
   • Comunicación: Aunque TSP introduce comunicación de movimiento de pesos (que TP+SP evita), los autores demuestran que para secuencias largas o tamaños de lote moderados, la comunicación proporcional a las activaciones domina, haciendo a TSP competitivo.
   • FLOPs: TSP mantiene la misma eficiencia computacional por dispositivo que TP y SP.
3. Implementación y Solapamiento: Implementación detallada de los bloques de atención y MLP de TSP con estrategias explícitas de solapamiento de comunicación/cómputo (pipeline de transferencias de pesos detrás de los cálculos GEMM).
4. Optimización de Topología: Demostración de que TSP permite que toda la réplica del modelo quepa dentro de un único nodo de alto ancho de banda (intra-nodo), evitando las penalizaciones de rendimiento de abarcar enlaces inter-nodo más lentos, a menudo requeridos por las mallas 2D TP+SP.

4. Resultados

Los autores evaluaron TSP en un clúster de hasta 1,024 GPUs AMD MI300X (128 nodos) con un modelo transformador denso de 7B de parámetros.

• Eficiencia de Memoria:
  • TSP logró consistentemente la memoria pico más baja en todas las longitudes de secuencia probadas (de 16K a 128K tokens).
  • En contextos cortos, TSP igualó a TP (dominado por la memoria de parámetros).
  • En contextos largos, TSP superó significativamente a TP y TP+SP (dominado por la memoria de activaciones), previniendo errores de Memoria Agotada (OOM) donde otras estrategias fallaron.
• Rendimiento (Throughput):
  • A pesar de un mayor volumen teórico de comunicación (debido al movimiento de pesos), TSP mantuvo un rendimiento competitivo con TP+SP.
  • La comunicación extra fue ocultada exitosamente detrás del cómputo mediante pipeline.
  • Las ventajas de rendimiento crecieron con el aumento de los grados de paralelismo ($D$) y las longitudes de secuencia.
• Escalabilidad:
  • TSP permitió micro-lotes más grandes debido a la huella de memoria reducida, mejorando aún más el rendimiento.
  • Permitió ajustar réplicas del modelo completamente dentro de un solo nodo (8 GPUs), liberando rangos para Paralelismo de Datos (DP) en clústeres más grandes.

5. Significado

• Escalado Consciente del Hardware: TSP proporciona una solución práctica para clústeres HPC modernos donde el ancho de banda intra-nodo es alto, pero el ancho de banda inter-nodo es un cuello de botella. Maximiza la utilidad de las interconexiones intra-nodo costosas.
• Viabilidad de Contexto Largo: A medida que los modelos se mueven hacia ventanas de contexto más largas (100k+ tokens), la memoria de activaciones se convierte en el cuello de botella principal. TSP está posicionado de manera única para manejar esto sin el sobrecosto masivo de dispositivos del paralelismo 2D.
• Flexibilidad: TSP puede componerse ortogonalmente con otros esquemas de paralelismo (Paralelismo de Pipeline, Paralelismo de Expertos, Paralelismo de Datos), ofreciendo un "eje" flexible para que los diseñadores de sistemas optimicen para topologías de hardware específicas.
• Eficiencia de Costos: Al reducir el número de dispositivos necesarios para el paralelismo de modelo, TSP permite que más dispositivos se utilicen para el paralelismo de datos, reduciendo potencialmente el tamaño total del clúster requerido para una tarea de entrenamiento dada o aumentando el tamaño de lote efectivo.

En conclusión, TSP representa una evolución significativa en las estrategias de entrenamiento distribuido, intercambiando un aumento manejable en el volumen de comunicación por ganancias sustanciales en la eficiencia de la memoria y la utilización del hardware, convirtiéndolo en una herramienta crítica para el entrenamiento de LLMs de próxima generación con contexto largo.