Memory-Efficient LLM Training with Dynamic Sparsity: From Stability to Practical Scaling

Résumé Technique : Entraînement de LLM à Efficacité Mémoire avec Sparsité Dynamique

Énoncé du Problème
La Sparsité Dynamique d'Entraînement (DST - Dynamic Sparse Training) offre un paradigme pour améliorer l'efficacité des réseaux de neurones profonds en maintenant une topologie parcieuse tout au long de l'entraînement. Cependant, l'application de la DST aux Grands Modèles de Langage (LLM) s'est avérée difficile en raison de l'instabilité de l'optimisation. Les auteurs identifient que l'application naïve d'optimiseurs standards basés sur Adam à la DST entraîne de graves « pics de perte » (loss spikes) immédiatement après les mises à jour de la topologie (élagage et régénération).

La cause profonde est identifiée comme un effet de « démarrage à froid » (cold-start) : les paramètres nouvellement régénérés manquent d'états d'optimiseur accumulés (estimations des premiers et seconds moments). Dans Adam, les termes de correction de biais pour ces nouveaux paramètres entraînent des magnitudes de mise à jour disproportionnées par rapport aux poids matures. Cela perturbe la trajectoire d'entraînement, provoquant des pics de perte transitoires et entravant la convergence, particulièrement lorsque l'échelle du modèle augmente ou que les mises à jour de topologie deviennent plus agressives. De plus, les implémentations conventionnelles de la DST échouent souvent à réaliser des économies de mémoire car elles allouent des états d'optimiseur de taille complète pour l'ensemble de l'espace de paramètres, même si seule une sous-partie des paramètres est active.

Méthodologie : SMET (Sparse Memory-Efficient Training)
Pour relever ces défis, les auteurs proposent SMET, un cadre conçu pour stabiliser la DST et réduire la consommation de mémoire pour les LLM. La méthode comprend trois composantes clés :

1. Warm-up de l'État de l'Optimiseur (Optimizer-State Warm-up) : Pour atténuer l'effet de démarrage à froid, SMET emploie une stratégie de chauffe (warm-up) en deux étapes pour les paramètres nouvellement régénérés :

   • Warm-up Implicite (Réinitialisation du pas de temps) : Le compteur interne de pas de temps ($t_i$) pour les paramètres régénérés est réinitialisé. Cela réactive le mécanisme de correction de biais d'Adam, gonflant le dénominateur de la règle de mise à jour et contraignant naturellement la magnitude initiale de la mise à jour.
   • Warm-up Explicite (Mise à l'échelle du taux d'apprentissage) : Un warm-up linéaire local est appliqué où le taux d'apprentissage pour les paramètres régénérés augmente linéairement de zéro jusqu'au taux cible sur $W$ étapes.
   • Ces mécanismes empêchent conjointement les explosions de mise à jour, lissant la trajectoire d'optimisation immédiatement après les changements de topologie.

2. Mise à l'échelle du Taux d'Apprentissage Sensible à la Densité (Density-Aware Learning Rate Scaling) : Pour compenser les mises à jour conservatrices introduites par le warm-up et pour tenir compte de la réduction du fan-in effectif des couches creuses, SMET ajuste le taux d'apprentissage. Inspiré par la Paramétrisation de Mise à Jour Maximale ($\mu$P), le taux d'apprentissage est mis à l'échelle par $1/\sqrt{d}$, où $d$ est la densité actuelle de la sparsité. Cela garantit que les magnitudes de mise à jour restent comparables à celles des régimes d'entraînement denses.

3. Optimisation à Efficacité Mémoire (Memory-Efficient Optimization) : SMET implémente une stratégie de stockage parcieuse où les gradients et les états d'optimiseur (estimations de moments $m$ et $v$) sont stockés uniquement pour les paramètres actifs (non nuls).

   • Les gradients sont conservés exclusivement aux indices actifs lors de la rétropropagation.
   • Pour gérer le surcoût lié au suivi des indices actifs et des pas de temps, une stratégie de stockage par blocs est utilisée, regroupant les paramètres en blocs de taille $B$ pour amortir les coûts de métadonnées.
   • Cette approche préserve la sparsité tout au long de l'intégralité du cycle d'entraînement, contrairement aux méthodes qui compressent les états mais produisent des modèles denses.

Principales Contributions

1. Analyse de l'Instabilité : L'article identifie et analyse théoriquement le phénomène de pics de perte dans le pré-entraînement des LLM avec la DST, l'attribuant à l'effet de démarrage à froid des poids nouvellement régénérés dans les optimiseurs Adam.
2. Cadre de Stabilisation : Les auteurs proposent SMET, qui intègre le warm-up de l'état de l'optimiseur et la mise à l'échelle du taux d'apprentissage sensible à la densité pour stabiliser les mises à jour de topologie et améliorer les performances d'entraînement.
3. Implémentation à Efficacité Mémoire : Une implémentation d'optimisation parcieuse est développée qui stocke les gradients et les états d'optimiseur uniquement pour les paramètres actifs, supportant à la fois des motifs de sparsité non structurés et par blocs, tout en maintenant la sparsité durant l'entraînement.
4. Évaluation Systématique : Ce travail fournit une évaluation complète de la DST à travers diverses échelles de modèles et niveaux de sparsité, démontrant que SMET surpasse les méthodes existantes de DST et se rapproche des performances de l'entraînement dense.

Résultats Expérimentaux
Des expériences ont été menées sur des modèles basés sur LLaMA (allant de 60M à 1B de paramètres) pré-entraînés sur les jeux de données C4 et OpenWebText.

• Performance : SMET surpasse systématiquement les autres bases de référence économes en mémoire (LoRA, ReLoRA, GaLore) et les méthodes DST existantes (SET, RigL). À des échelles de modèles plus grandes (ex: 1B de paramètres), SMET atteint une perplexité proche de l'entraînement dense, même à des niveaux de sparsité élevés (ex: densité $d=0.1$).
• Stabilité : La méthode élimine efficacement les pics de perte observés dans la version vanilla de la DST, résultant en des courbes d'entraînement lisses.
• Efficacité Mémoire : En stockant les états uniquement pour les paramètres actifs, SMET réduit considérablement l'empreinte mémoire des gradients et des états d'optimiseur. Pour un modèle de 1B à une densité $d=0.4$, des économies de mémoire substantielles sont observées.
• Scalabilité : L'écart de performance entre l'entraînement parcieux et dense se réduit à mesure que la taille du modèle augmente, suggérant que l'entraînement parcieux est particulièrement viable pour les modèles à grande échelle.
• Généralisation : Les bénéfices de SMET se généralisent à travers différentes stratégies de régénération (aléatoires et basées sur le gradient) et différentes architectures de modèles (LLaMA et Qwen).

Signification et Revendications
L'article affirme que SMET ouvre la voie à la sparsité comme une alternative pratique à l'entraînement dense pour les LLM. En résolvant l'instabilité d'optimisation inhérente à la sparsité dynamique d'entraînement et en réalisant de réelles économies de mémoire grâce au stockage d'états parcieux, SMET permet un pré-entraînement stable, évolutif et efficace en mémoire. Les auteurs notent que bien que SMET améliore la stabilité de l'entraînement et réduise la mémoire de l'optimiseur, des défis subsistent concernant la mémoire d'activation et la nécessité d'un support matériel efficace pour la sparsité à grain fin afin de pleinement réaliser les gains de vitesse d'inférence. Néanmoins, ce travail démontre qu'avec une stabilisation et une gestion de la mémoire appropriées, la sparsité dynamique peut être étendue avec succès aux modèles de fondation.