Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

Résumé Technique : Compression par Sous-Espace Appris pour le Parallélisme de Pipeline à Efficacité de Communication

Énoncé du Problème
L'entraînement de grands modèles de langage (LLM) sur du matériel hétérogène et géographiquement distribué fait face à d'importants goulots d'étranglement de communication, particulièrement dans le cadre du parallélisme de pipeline. Si le parallélisme de pipeline partitionne les paramètres du modèle entre les dispositifs pour surmonter les limites de mémoire, il nécessite l'échange d'activations intermédiaires entre les étapes lors des passes avant et arrière. Sur des réseaux à faible bande passante, cette communication inter-étapes devient le goulot d'étranglement dominant.

Les approches existantes, telles que les réseaux de sous-espaces (SSN), tentent d'atténuer ce problème en projetant les activations sur un sous-espace de faible rang partagé et fixe défini par une matrice orthogonale. Cependant, cette méthode nécessite des modifications architecturales intrusives, incluant une optimisation contrainte pour maintenir les poids dans un sous-espace partagé et l'utilisation d'optimiseurs modifiés. Empiriquement, ces sous-espaces globaux fixes entraînent une dégradation significative des performances, surtout à mesure que les taux de compression augmentent, car ils ne parviennent pas à s'adapter à la géométrie évolutive des activations pendant l'entraînement. De plus, l'apprentissage naïf de matrices de projection sans contraintes provoque leur sortie de la variété de Stiefel (l'ensemble des matrices orthogonales), détruisant les propriétés isométriques requises pour une compression efficace et entraînant de sévères chutes de performance.

Méthodologie : Apprentissage de Projection Sensible à la Variété (MAPL)
Les auteurs proposent le Manifold Aware Projection Learning (MAPL), une méthode qui traite la compression inter-étape comme une projection géométrique apprenable plutôt que comme une contrainte architecturale fixe. Les composantes clés de MAPL sont :

1. Projecteurs Orthogonaux Apprenables par Étape : Contrairement aux SSN qui imposent un sous-espace global unique, MAPL équipe chaque frontière d'étape de pipeline de son propre projecteur orthogonal apprenable $A_p \in \mathbb{R}^{d \times r}$. Ces projecteurs sont contraints à la variété de Stiefel ($St(d, r)$), garantissant que $A_p^\top A_p = I_r$. Cela permet à chaque étape de découvrir et d'adapter continuellement son propre sous-espace de compression optimal pour la tâche.
2. Optimisation Contrainte par la Variété (SPEL) : Pour empêcher les projecteurs de s'échapper de la variété lors des mises à jour de gradient standards, les auteurs emploient la Descente de Gradient Spectrale sur la Variété de Stiefel (SPEL). Cet algorithme met à jour les projecteurs en utilisant des projections tangentes et des opérations de rétraction (via la méthode Polar Express) à chaque étape d'optimisation, garantissant que la contrainte d'orthogonalité est strictement maintenue.
3. Embeddings d'Ancre Factorisés : Pour gérer les décalages (offsets) de haut rang induits par les tokens dans les flux résiduels sans épuiser la capacité du projecteur, MAPL introduit des embeddings d'ancre factorisés. Le décalage est décomposé en une matrice de faible rang entraînable $E_{small}^p$ et une matrice orthogonale gelée $P_p$. Seuls les identifiants de tokens entiers sont transmis pour reconstruire ces décalages à l'étape de réception, ajoutant une surcharge de communication négligeable.
4. Quantification Vectorielle (VQ) : Pour une compression supplémentaire, la méthode supporte la quantification vectorielle résiduelle (RVQ) sur les représentations projetées en faible rang. Un protocole de synchronisation de codebook en streaming est utilisé pour amortir la communication du dictionnaire, permettant des taux de compression élevés (jusqu'à 16×) avec un coût de bande passante minimal.

Contributions Clés

• Apprentissage de Sous-Espace par Étape : L'article démontre que l'imposition d'un sous-espace global partagé dégrade l'apprentissage dans les configurations de correspondance de tokens. Au lieu de cela, MAPL permet à chaque étape d'apprendre un sous-espace de compression distinct, contraint par la variété de Stiefel, ce qui s'aligne mieux avec la structure intrinsèque de faible rang des activations à différentes profondeurs.
• Optimisation Contrainte par la Variété : Les auteurs identifient l'« échappement de la variété » comme le principal mode de défaillance de l'apprentissage de projecteurs non contraints. Ils introduisent SPEL pour maintenir l'orthogonalité, montrant qu'un apprentissage non contraint est strictement moins performant que l'utilisation de bases fixes.
• Reconstruction Efficace : En combinant des projecteurs apprenables avec des embeddings d'ancre factorisés, MAPL atteint une reconstruction d'activation de plein rang avec une surcharge de communication négligeable, éliminant le besoin des décompositions architecturales intrusives requises par les SSN.
• Amélioration de la Frontière de Pareto : La méthode trace la frontière de Pareto de l'entropi croisée de validation par rapport au coût de communication, atteignant une compression élevée avec une dégradation de performance négligeable par rapport aux méthodes précédentes.

Résultats Expérimentaux
Les auteurs ont évalué MAPL sur des modèles LLaMA de type decoder-only allant de 150M à 1B de paramètres, entraînés sur le corpus DCLM avec des configurations de pipeline de $P=4$ et $P=8$ étapes.

• Performance vs Compression : MAPL surpasse systématiquement les bases de référence SSN. Pour un modèle de 150M à 4× de compression, MAPL a atteint une perte de validation à 0,84 % de la base non compressée, tandis que les bases SSN ont subi des dégradations de 8 à 12 %.
• Scalabilité : À l'échelle de 1 milliard de paramètres avec 8× de compression, MAPL a maintenu une dégradation de perte de seulement 1,38–2,02 %, alors que les bases SSN ont subi des dégradations allant jusqu'à 26 %.
• Tâches Aval : Lors des évaluations zero-shot (HellaSwag, PIQA, ARC), MAPL suit de près les bases non compressées (à environ 1 point de précision moyenne à l'échelle 500M/1B), tandis que les bases SSN ont subi des chutes de précision allant jusqu'à 9 points.
• Compression Extrême : Lorsqu'il est combiné avec la quantification vectorielle (MAPL+VQ), la méthode atteint une compression de 16× avec une dégradation supplémentaire modeste (par exemple, une augmentation de perte de ~2,3 % à 1 milliard de paramètres), bien que la précision des tâches aval subisse une chute plus prononcée par rapport à MAPL non quantifié.

Signification et Revendications
L'article affirme que les activations des transformers inter-étapes admettent naturellement des structures de faible dimension adaptatives qui peuvent être apprises conjointement avec l'objectif de la tâche, plutôt que d'être imposées via des sous-espaces globaux fixes. En traitant la compression comme une projection géométrique apprenable sur la variété de Stiefel, MAPL permet une réduction agressive de la communication (4× à 16×) tout en préservant la qualité de l'entraînement et les performances sur les tâches aval. Ce travail souligne que l'apprentissage de représentations contraint par une variété est un mécanisme pratique pour réduire les coûts de communication dans l'entraînement de modèles de fondation distribués, offrant un meilleur compromis entre performance et compression par rapport aux approches de sous-espaces existantes.

Limites
Les auteurs notent que les expériences sont actuellement limitées à des modèles allant jusqu'à 1 milliard de paramètres, et qu'une évaluation supplémentaire à plus grande échelle et dans des conditions de réseaux hétérogènes réelles est nécessaire. De plus, bien que MAPL réduise considérablement la dégradation, les performances déclinent toujours sous des taux de compression extrêmes, particulièrement lorsque la quantification vectorielle est appliquée.