Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

Resumen Técnico: Compresión de Subespacios Aprendidos para el Paralelismo de Pipeline con Eficiencia de Comunicación

Planteamiento del Problema
El entrenamiento de modelos de lenguaje extensos (LLMs) a través de hardware heterogéneo y geográficamente distribuido enfrenta cuellos de botella significativos en la comunicación, particularmente en el paralelismo de pipeline. Mientras que el paralelismo de pipeline particiona los parámetros del modelo entre dispositivos para superar los límites de memoria, requiere el intercambio de activaciones intermedias entre etapas tanto durante las pasadas hacia adelante (forward) como hacia atrás (backward). En redes de bajo ancho de banda, esta comunicación entre etapas se convierte en el cuello de botella dominante.

Los enfoques existentes, como las Redes de Subespacio (SSNs), intentan mitigar esto proyectando las activaciones sobre un subespacio de bajo rango fijo y compartido definido por una matriz ortogonal. Sin embargo, este método requiere modificaciones arquitectónicas intrusivas, incluyendo la optimización restringida para mantener los pesos dentro de un subespacio compartido y el uso de optimizadores modificados. Empíricamente, estos subespacios globales fijos provocan una degradación significativa del rendimiento, especialmente a medida que aumentan las tasas de compresión, porque no logran adaptarse a la geometría evolutiva de las activaciones durante el entrenamiento. Además, el aprendizaje ingenuo de matrices de proyección sin restricciones provoca que estas escapen de la variedad de Stiefel (el conjunto de matrices ortogonales), destruyendo las propiedades isométricas requeridas para una compresión efectiva y provocando caídas severas en el rendimiento.

Metodología: Aprendizaje de Proyección Consciente de la Variedad (MAPL)
Los autores proponen el Aprendizaje de Proyección Consciente de la Variedad (MAPL), un método que trata la compresión entre etapas como una proyección geométrica aprendible en lugar de una restricción arquitectónica fija. Los componentes principales de MAPL son:

1. Proyectores Ortogonales Aprendibles por Etapa: A diferencia de las SSNs que imponen un único subespacio global, MAPL equipa cada límite de etapa de pipeline con su propio proyector ortogonal aprendible $A_p \in \mathbb{R}^{d \times r}$. Estos proyectores están restringidos a la variedad de Stiefel ($St(d, r)$), asegurando que $A_p^\top A_p = I_r$. Esto permite que cada etapa descubra y adapte continuamente su propio subespacio de compresión óptimo para la tarea.
2. Optimización Restringida a la Variedad (SPEL): Para evitar que los proyectores escapen de la variedad durante las actualizaciones de gradiente estándar, los autores emplean el Descenso de Gradiente Espectral en la Variedad de Stiefel (SPEL). Este algoritmo actualiza los proyectores utilizando proyecciones tangentes y operaciones de retracción (vía el método Polar Express) en cada paso de optimización, asegurando que la restricción de ortogonalidad se mantenga estrictamente.
3. Embeddings de Anclaje Factorizados: Para manejar los desplazamientos (offsets) de alto rango impulsados por tokens en las corrientes residuales sin agotar la capacidad del proyector, MAPL introduce embeddings de anclaje factorizados. El desplazamiento se descompone en una matriz de bajo rango entrenable $E_{small}^p$ y una matriz ortogonal congelada $P_p$. Solo se transmiten IDs de tokens enteros para reconstruir estos desplazamientos en la etapa receptora, lo que añade un costo de comunicación insignificante.
4. Cuantización Vectorial (VQ): Para una mayor compresión, el método soporta la cuantización vectorial residual (RVQ) sobre las representaciones proyectadas de bajo rango. Se utiliza un protocolo de sincronización de libro de códigos (codebook) en streaming para amortizar la comunicación del diccionario, permitiendo altas tasas de compresión (hasta 16×) con un costo de ancho de banda mínimo.

Contribuciones Clave

• Aprendizaje de Subespacio por Etapa: El artículo demuestra que imponer un subespacio global compartido degrada el aprendizaje en configuraciones de coincidencia de tokens. En su lugar, MAPL permite que cada etapa aprenda un subespacio de compresión distinto y restringido a la variedad de Stiefel, lo cual se alinea mejor con la estructura intrínseca de bajo rango de las activaciones en diferentes profundidades.
• Optimización Restringida a la Variedad: Los autores identifican el "escape de la variedad" como el principal modo de fallo del aprendizaje de proyectores ingenuo. Introducen SPEL para mantener la ortogonalidad, demostiendo que el aprendizaje sin restricciones es estrictamente peor que el uso de líneas base fijas.
• Reconstrucción Eficiente: Al combinar proyectores aprendibles con embeddings de anclaje factorizados, MAPL logra la reconstrucción de activaciones de rango completo con un costo de comunicación insignificante, eliminando la necesidad de las descomposiciones arquitectónicas intrusivas requeridas por las SSNs.
• Mejora de la Frontera de Pareto: El método traza la frontera de Pareto de la entropía cruzada de validación frente al costo de comunicación, logrando una alta compresión con una degradación de rendimiento insignificante en comparación con métodos previos.

Resultados Experimentales
Los autores evaluaron MAPL en modelos LLaMA de solo decodificador que van desde 150M hasta 1B de parámetros, entrenados en el corpus DCLM con configuraciones de pipeline de $P=4$ y $P=8$ etapas.

• Rendimiento vs. Compresión: MAPL supera consistentemente a las líneas base de SSN. Para un modelo de 150M con 4× compresión, MAPL logró una pérdida de validación dentro del 0.84% de la línea base sin comprimir, mientras que las líneas base de SSN se degradaron entre un 8–12%.
• Escalabilidad: A la escala de 1B de parámetros con 8× compresión, MAPL mantuvo una degradación de pérdida de solo 1.38–2.02%, mientras que las líneas base de SSN sufrieron degradaciones de hasta un 26%.
• Tareas de Aval (Downstream Tasks): En evaluaciones zero-shot (HellaSwag, PIQA, ARC), MAPL siguió de cerca a las líneas base sin comprimir (dentro de ~1 punto de precisión promedio en escalas de 500M/1B), mientras que las líneas base de SSN sufrieron caídas de precisión de hasta 9 puntos.
• Compresión Extrema: Cuando se combina con cuantización vectorial (MAPL+VQ), el método logró 16× de compresión con una degradación adicional modesta (por ejemplo, ~2.3% de aumento de pérdida a 1B de parámetros), aunque la precisión en tareas de aval mostró una caída más pronunciada en comparación con MAPL sin cuantización.

Significancia y Reivindicaciones
El artículo sostiene que las activaciones de los transformadores entre etapas admiten naturalmente estructuras de baja dimensión adaptables que pueden aprenderse conjuntamente con el objetivo de la tarea, en lugar de ser impuestas mediante subespacios globales fijos. Al tratar la compresión como una proyección geométrica aprendible en la variedad de Stiefel, MAPL permite una reducción agresiva de la comunicación (4× a 16×) preservando la calidad del entrenamiento y el rendimiento en tareas posteriores. El trabajo destaca el aprendizaje de representaciones restringido a la variedad como un mecanismo práctico para reducir los costos de comunicación en el entrenamiento de modelos fundacionales distribuidos, ofreciando un compromiso superior entre rendimiento y compresión en comparación con los enfoques basados en subespacios existentes.

Limitaciones
Los autores señalan que los experimentos están actualmente limitados a modelos de hasta 1B de parámetros, y se requiere una evaluación adicional a escalas mayores y bajo condiciones de red heterogéneas del mundo real. Además, aunque MAPL reduce significativamente la degradación, el rendimiento aún disminuye bajo tasas de compresión extremas, particularmente cuando se aplica la cuantización vectorial.