🤖 machine learning

DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers

Ce document présente DASH, une implémentation nettement plus rapide de l'optimiseur Distributed Shampoo qui exploite l'empilement de tenseurs 3D pour améliorer l'utilisation des GPU ainsi que de nouveaux solveurs d'inverse de racine (approximations Newton-DB et Tchebychev) pour atteindre une accélération allant jusqu'à 5,6x tout en maintenant ou en améliorant la qualité de la convergence.

Auteurs originaux : Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Publié 2026-06-26

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez d'apprendre à un robot géant à parler une nouvelle langue. Pour ce faire, vous avez besoin d'un « optimiseur » — un coach intelligent qui ajuste le cerveau du robot (ses paramètres) après chaque leçon pour minimiser les erreurs.

Pendant des années, le coach le plus populaire a été Adam, une méthode rapide mais un peu « paresseuse ». Elle ne regarde que la façon dont chaque neurone individuel doit changer, ignorant la manière dont les neurones travaillent ensemble en équipe.

Entrez en scène Shampoo, un coach bien plus intelligent. Au lieu de regarder les neurones un par un, Shampoo observe comment ils interagissent en groupes. Cela conduit à un meilleur apprentissage et à des modèles plus faciles à compresser par la suite. Cependant, il y a un hic : Shampoo est incroyablement lent. C'est comme un coach de génie qui passe tellement de temps à calculer le mouvement parfait que le robot ne parvient presque pas à s'entraîner.

Le papier présente DASH (Distributed Accelerated SHampoo), un nouveau système qui permet à ce coach de génie de courir à la vitesse d'un sprinteur sans perdre son intelligence. Voici comment ils ont fait, en utilisant des analogies simples :

1. L'astuce de l'« empilement » (Batched Block Preconditioning)

Le Problème :
Imaginez que vous avez une immense bibliothèque de livres (les données) qui doit être organisée. L'ancienne méthode (Distributed Shampoo) consistait à prendre un livre, l'organiser, le remettre en place, prendre le suivant, l'organiser, et ainsi de suite. Même si vous aviez 1 000 ouvriers (GPU), la plupart restaient là à attendre que la personne précédente ait terminé. C'est inefficace.

La Solution DASH :
DASH change le flux de travail. Au lieu de manipuler les livres un par un, il les empile en tours 3D nettes et uniformes. Désormais, les ouvriers peuvent saisir une tour entière et organiser tous les livres à l'intérieur en même temps.

L'Analogie : C'est la différence entre un caissier qui scanne les articles un par un et un tapis roulant qui alimente une machine scannant instantanément toute une boîte de courses.
Le Résultat : Cet « empilement » permet aux puces graphiques puissantes de l'ordinateur (GPU) de travailler à pleine capacité, rendant les étapes de l'optimiseur jusqu'à 5,6 fois plus rapides.

2. Les mathématiques de « raccourci » (Efficient Inverse-Root Solvers)

Le Problème :
Pour faire son travail, Shampoo doit résoudre un puzzle mathématique très difficile à chaque étape : trouver la « racine carrée inverse » d'une matrice géante. L'ancienne méthode pour résoudre cela revenait à chercher une aiguille dans une botte de foin en vérifiant chaque brin de paille un par un (une méthode appelée Décomposition en Valeurs Propres). C'est précis, mais terriblement lent.

La Solution DASH :
Les auteurs ont introduit deux nouveaux « raccourcis » pour résoudre ce puzzle :

Newton-DB : Une méthode itérative intelligente qui se rapproche de la réponse à chaque tentative, comme un GPS qui recalcule votre itinéraire pendant que vous conduisez.
Polynômes de Chebyshev : Une approximation mathématique qui devine la réponse très rapidement.
L'Analogie : Au lieu de parcourir chaque chemin d'un labyrinthe pour trouver la sortie (l'ancienne méthode), ces nouvelles méthodes sont comme avoir une carte qui montre la direction générale, vous permettant de sprinter vers la sortie.

3. Le calibrage de la « règle » (Matrix Scaling)

Le Problème :
Lorsque l'on utilise ces raccourcis, les mathématiques peuvent devenir instables si les nombres sont trop grands ou trop petits. L'ancienne méthode utilisait une « règle » (norme de Frobenius) trop longue, étirant les nombres et faisant en sorte que les raccourcis prennent beaucoup plus d'étapes pour converger. C'était comme essayer de mesurer une petite fourmi avec un ruban à mesurer de 30 mètres ; la précision est perdue.

La Solution DASH :
Les auteurs ont réalisé qu'ils avaient besoin d'une meilleure règle. Ils ont développé une technique appelée Multi-Power-Iteration.

L'Analogie : Au lieu de deviner la longueur de la fourmi avec un ruban à mesurer géant, ils utilisent un pied à coulisse spécialisé de haute précision qui s'adapte parfaitement à la fourmi. Cela garantit que les mathématiques convergent rapidement et ne plantent pas à cause d'erreurs numériques.

4. Les Résultats

Le papier a testé DASH sur un grand modèle de langage (Llama avec 953 millions de paramètres).

Vitesse : DASH a terminé la partie « réflexion » de l'étape d'entraînement 5,6 fois plus vite que la version précédente la plus performante.
Qualité : Malgré cette rapidité, les modèles entraînés avec DASH étaient aussi bons, voire légèrement meilleurs, pour apprendre. En fait, le raccourci « Newton-DB » a produit les modèles avec les taux d'erreur (perplexité) les plus bas de toutes les méthodes testées.

Résumé

Considérez Shampoo comme un moteur de Ferrari qui était coincé dans les embouteillages parce que le conducteur prenait le chemin le plus pittoresque. DASH est ce même moteur de Ferrari, mais qui possède désormais :

Une autoroute plus large (Empilement de blocs) pour pouvoir rouler plus vite.
Un raccourci GPS (Newton-DB) pour éviter les bouchons.
De meilleurs outils de navigation (Multi-Power-Iteration) pour s'assurer de ne pas se perdre.

Le résultat est que l'optimiseur « intelligent » (Shampoo) n'est plus trop lent pour être utilisé, ce qui en fait un choix pratique pour entraîner la prochaine génération de modèles d'IA.

Résumé Technique : DASH – Un Shampoo plus rapide via le préconditionnement par blocs par lots et des solveurs de racine inverse efficaces

Énoncé du Problème

Shampoo est un optimiseur de second ordre approché de premier plan, connu pour produire des modèles avec moins d'anomalies d'activation et une meilleure compressibilité. Cependant, son adoption généralisée est entravée par un surcoût de calcul important, principalement dû à l'opération coûteuse de calcul des racines de matrices inverses ( $A^{-1/4}$ ) pour les matrices de préconditionnement. Bien que le Shampoo Distribué tente d'atténuer cela en divisant les préconditionneurs en blocs $B \times B$ pour réduire la complexité de $O(n^3)$ à $O(Bn^2)$ , les implémentations existantes souffrent de lacunes d'efficacité critiques :

Traitement Séquentiel : Les implémentations standards traitent les blocs de manière séquentielle, ne parvenant pas à exploiter pleinement le parallélisme des GPU modernes.
Solveurs Inefficaces : La dépendance par défaut à la Décomposition en Valeurs Propres (EVD) est difficile à paralléliser sur les GPU. Les méthodes itératives alternatives comme le Coupled-Newton (CN) sont souvent désactivées par défaut en raison de préoccupations concernant la stabilité numérique et la vitesse de convergence.
Mise à l'Échelle Sous-optimale : Les méthodes actuelles mettent souvent à l'échelle les matrices en utilisant la norme de Frobenius, ce qui est peu coûteux en calcul mais peut pousser les valeurs propres vers zéro, nécessitant ainsi beaucoup plus d'itérations pour converger vers une précision cible.

Méthodologie

Le papier propose DASH (Distributed Accelerated SHampoo), une implémentation haute performance conçue pour combler le fossé entre l'efficacité théorique et les capacités matérielles. La méthodologie repose sur deux piliers principaux :

1. Préconditionnement par Blocs par Lots (Optimisation au Niveau Système)

Au lieu de traiter les blocs de préconditionnement séquentiellement, DASH empile les blocs indépendants dans des tenseurs 3D.

Empilement de Tenseurs 3D : Les blocs de gradient $G_{ij}$ sont empilés dans une matrice de forme $(N, B, B)$ , où $N$ est le nombre total de blocs.
Exécution Parallèle : Cela permet à l'optimiseur d'appliquer les procédures de racine inverse (EVD, CN, ou nouvelles méthodes itératives) à l'ensemble du lot de blocs simultanément en utilisant des primitives de multiplication de matrices par lots (bmm).
Utilisation du GPU : Ce changement architectural maximise le débit des Tensor Cores du GPU, réduisant la fragmentation de la mémoire et permettant une utilisation efficace de l'arithmétique en demi-précision (FP16) pour certains solveurs.

2. Solveurs de Racine Inverse Efficaces (Optimisation Algorithmique)

Le papier introduit et analyse des approches d'algèbre linéaire avancées pour remplacer ou compléter l'EVD et le CN standard :

Itération de Newton-Denman-Beavers (NDB) : Une approche novatrice pour Shampoo qui calcule à la fois les racines carrées et les racines carrées inverses. Les auteurs optimisent la première itération pour éliminer les multiplications de matrices redondantes, réduant la surcharge à 3 multiplications de matrices par étape (effectivement moins en pratique).
Approximations par Polynômes de Chebyshev : Utilisation de l'algorithme de Clenshaw pour approximer les puissances de matrices, offrant une alternative de solveur itératif.
Multi-Power-Iteration pour la Mise à l'Échelle : Pour remédier aux problèmes de convergence, le papier remplace la mise à l'échelle standard par la norme de Frobenius par une méthode de Multi-Power-Iteration.
- Mécanisme : Il estime le rayon spectral ( $\lambda_{max}$ ) en utilisant un pool de vecteurs de départ (jusqu'à 32) traités en parallèle via la Power-Iteration.
- Avantage : Cela fournit un facteur d'échelle plus précis que la norme de Frobenius, garantissant que les valeurs propres restent dans un régime où les solveurs itératifs convergent rapidement (proche de 1), réduisant ainsi le nombre d'itérations requises et améliorant la stabilité numérique.

Principales Contributions

Implémentation DASH : Un framework efficace pour GPU qui remplace les boucles séquentielles par des opérations de tenseurs 3D parallèles, atteignant jusqu'à 5,6× plus de rapidité par étape d'optimiseur par rapport au Shampoo Distribué bien optimisé.
Nouveaux Solveurs : L'introduction de l'itération Newton-DB (NDB) et des approximations par polynômes de Chebyshev comme alternatives plus rapides à l'EVD et au CN. Le papier démontre que le NDB atteint une perplexité de validation plus faible que le Coupled-Newton et l'EVD standard.
Analyse de Convergence : Une analyse comportementale détaillée montrant que la mise à l'échelle standard par la norme de Frobenius est sous-optimale, menant souvent à une convergence plus lente. Le papier fournit des explications intuitives pour les comportements de convergence distincts du CN par rapport au NDB.
Multi-Power-Iteration : Une implémentation efficace en demi-précision de la Power-Iteration qui estime de manière robuste les rayons spectraux, permettant une mise à l'échelle de matrice optimale qui satisfait les critères de convergence pour le NDB et le CN avec moins d'étapes.

Résultats Expérimentaux

Les auteurs ont évalué DASH sur un modèle Llama de 953M de paramètres (taille d'embedding 2048, longueur de séquence 1024) entraîné sur le dataset C4.

Vitesse :
- DASH réduit le temps d'exécution de l'étape d'optimiseur de 4× à 5,6× par rapport au Shampoo Distribué, selon la méthode de racine inverse et la précision utilisée.
- Le gain de vitesse le plus significatif (5,6×) a été observé en comparant l'EVD (FP32) du Shampoo Distribué contre le CN (FP16) de DASH avec l'empilement au niveau du GPU.
- Même avec des mises à jour fréquentes du préconditionneur ( $f=1$ ), DASH maintient des avantages de vitesse substantiels.
Précision (Perplexité de Validation) :
- DASH égale le Shampoo Distribué dans presque toutes les configurations stables.
- Le NDB produit systématiquement une perplexité de validation plus faible que le CN et surpasse souvent l'EVD.
- L'utilisation de la mise à l'échelle par Power-Iteration avec le NDB a donné la perplexité de validation la plus basse (11,68) par rapport à la mise à l'échelle par Frobenius (11,76), validant l'analyse théorique.
Précision :
- L'implémentation supporte le FP16 pour le CN et l'EVD sans dégrader la perplexité, réduisant davantage le temps d'exécution.
- Le NDB en FP16/BF16 s'est révélé numériquement instable dans leurs expériences, bien que les auteurs notent cela comme une zone de travail futur.

Signification et Revendications

Le papier affirme que DASH réduit considérablement la barrière de calcul pour l'utilisation d'optimiseurs de second ordre comme Shampoo, rendant le préconditionnement de haute qualité pratique pour une utilisation à grande échelle.

Pont entre Théorie et Matériel : En alignant la conception algorithmique (solveurs itératifs, mise à l'échelle de matrice) avec les capacités matérielles (Tensor Cores, opérations par lots), DASH permet aux chercheurs d'étudier les propriétés des optimiseurs sans temps d'exécution prohibitifs.
Stabilité Numérique vs Vitesse : Ce travail remet en question l'hypothèse selon laquelle les méthodes itératives sont intrinsèquement moins stables ou plus lentes que l'EVD. En introduisant une meilleure mise à l'échelle (Multi-Power-Iteration) et des itérations optimisées (NDB), les auteurs démontrent que les méthodes itératives peuvent atteindre une convergence supérieure et une perplexité plus faible tout en étant nettement plus rapides.
Directions Futures : Les auteurs suggèrent que la performance supérieure du Newton-DB ouvre des voies pour la sélection dynamique de solveurs basée sur les nombres de condition des blocs et le potentiel de stabilisation du NDB dans des formats de plus basse précision en utilisant des méthodes stochastiques ou de correction d'erreur.

Le papier conclut que le goulot d'étranglement principal de Shampoo n'était pas la complexité mathématique, mais l'exécution fragmentée et l'instabilité numérique, deux problèmes que DASH résout efficacement.