🤖 machine learning

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

Dieses Paper stellt DASH vor, eine signifikant schnellere Implementierung des Distributed-Shampoo-Optimierers, die 3D-Tensor-Stacking zur verbesserten GPU-Auslastung sowie neuartige Invertierungs-Wurzel-Solver (Newton-DB und Chebyshev-Approximationen) nutzt, um eine bis zu 5,6-fache Beschleunigung bei gleichbleibender oder verbesserter Konvergenzqualität zu erreichen.

Ursprüngliche Autoren: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Veröffentlicht 2026-06-26

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Ionut-Vlad Modoranu, Philip Zmushko, Erik Schultheis, Mher Safaryan, Dan Alistarh

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen, einem riesigen Roboter eine neue Sprache beizubringen. Um dies zu tun, benötigen Sie einen „Optimierer“ – einen klugen Coach, der das Gehirn des Roboters (seine Parameter) nach jeder Lektion anpasst, um Fehler zu minimieren.

Jahrelang war der beliebteste Coach Adam, eine Methode, die schnell ist, aber ein wenig „faul“ ist. Adam betrachtet nur, wie viel sich jedes einzelne Neuron ändern muss, und ignoriert dabei, wie Neuronen in Teams zusammenarbeiten.

Hier kommt Shampoo ins Spiel, ein viel klügerer Coach. Anstatt Neuronen einzeln zu betrachten, schaut Shampoo darauf, wie sie in Gruppen interagieren. Dies führt zu besserem Lernen und Modellen, die später leichter zu komprimieren sind. Es gibt jedoch einen Haken: Shampoo ist unglaublich langsam. Es ist wie ein genialer Coach, der so viel Zeit mit der Berechnung des perfekten Zuges verbringt, dass der Roboter kaum Zeit zum Üben bekommt.

Das Paper stellt DASH (Distributed Accelerated SHampoo) vor, ein neues System, das diesen genialen Coach mit der Geschwindigkeit eines Sprinters laufen lässt, ohne seine Intelligenz zu verlieren. So haben sie es geschafft, unter Verwendung einfacher Analogien:

1. Der „Stapeln“-Trick (Batched Block Preconditioning)

Das Problem:
Stellen Sie sich vor, Sie haben eine riesige Bibliothek von Büchern (die Daten), die organisiert werden müssen. Die alte Methode (Distributed Shampoo) bestand darin, ein Buch nach dem anderen zu nehmen, es zu organisieren, zurückzustellen, das nächste zu nehmen, zu organisieren und so weiter. Selbst wenn man 1.000 Arbeiter (GPUs) hatte, standen diese meistens nur herum und warteten darauf, dass die vorherige Person fertig wurde. Das ist ineffizient.

Die DASH-Lösung:
DASH ändert den Arbeitsablauf. Anstatt Bücher einzeln zu handhaben, stapelt es sie in ordentliche, einheitliche 3D-Türme. Jetzt können die Arbeiter einen ganzen Turm greifen und alle Bücher darin gleichzeitig organisieren.

Die Analogie: Es ist der Unterschied zwischen einem Kassierer, der Artikel einzeln scannt, und einem Förderband, das einen ganzen Karton mit Lebensmitteln sofort scannt.
Das Ergebnis: Dieses „Stapeln“ ermöglicht es den leistungsstarken Grafikchips (GPUs) des Computers, mit voller Kapazität zu arbeiten, wodurch die Optimierungsschritte bis zu 5,6-mal schneller werden.

2. Die „Abkürzungs“-Mathematik (Efficient Inverse-Root Solvers)

Das Problem:
Um seinen Job zu erledigen, muss Shampoo in jedem Schritt ein sehr schwieriges mathematisches Rätsel lösen: das Finden der „Inversen Quadratwurzel“ einer riesigen Matrix. Die alte Methode, dieses Rätsel zu lösen, war wie die Suche nach einer Nadel im Heuhaufen, indem man jeden einzelnen Halm einzeln überprüft (eine Methode namens Eigenwertzerlegung). Das ist genau, aber quälend langsam.

Die DASH-Lösung:
Die Autoren haben zwei neue „Abkürzungen“ eingeführt, um dieses Rätsel zu lösen:

Newton-DB: Eine clevere iterative Methode, die mit jeder Vermutung näher an die Antwort herankommt, ähnlich wie ein GPS, das die Route neu berechnet, während man fährt.
Chebyshev-Polynome: Eine mathematische Approximation, die die Antwort sehr schnell schätzt.
Die Analogie: Anstatt jeden Pfad in einem Labyrinth abzulaufen, um den Ausgang zu finden (die alte Methode), sind diese neuen Methoden wie eine Karte, die die allgemeine Richtung zeigt, sodass man direkt zum Ausgang sprinten kann.

3. Die „Lineal“-Kalibrierung (Matrix Scaling)

Das Problem:
Wenn man diese Abkürzungen verwendet, kann die Mathematik instabil werden, wenn die Zahlen zu groß oder zu klein sind. Die alte Methode verwendete ein „Lineal“ (Frobenius-Norm), das zu lang war, die Zahlen zu weit streckte und die Abkürzungen viele mehr Schritte zur Konvergenz benötigten. Es war, als würde man versuchen, einen winzigen Ameisen zu messen, mit einem 30-Meter-Maßband; die Präzision geht verloren.

Die DASH-Lösung:
Die Autoren erkannten, dass sie ein besseres Lineal brauchten. Sie entwickelten eine Technik namens Multi-Power-Iteration.

Die Analogie: Anstatt die Länge der Ameise mit einem riesigen Maßband zu schätzen, verwenden sie einen spezialisierten, hochpräzisen Messschieber, der perfekt zur Ameise passt. Dies stellt sicher, dass die Mathematik schnell konvergiert und nicht aufgrund numerischer Fehler abstürzt.

4. Die Ergebnisse

Das Paper testete DASH auf einem großen Sprachmodell (Llama mit 953 Millionen Parametern).

Geschwindigkeit: DASH schloss den „Denkpart“ des Trainingsschritts 5,6-mal schneller ab als die bisher beste Version.
Qualität: Trotz der viel höheren Geschwindigkeit waren die mit DASH trainierten Modelle genauso gut oder sogar etwas besser im Lernen. Tatsächlich erzeugte die neue „Newton-DB“-Abkürzung Modelle mit den niedrigsten Fehlerraten (Perplexity) aller getesteten Methoden.

Zusammenfassung

Betrachten Sie Shampoo als einen Ferrari-Motor, der im Stau feststeckte, weil der Fahrer die malerische Route nahm. DASH ist derselbe Ferrari-Motor, aber er hat nun:

Eine breitere Autobahn (Stacking Blocks), damit er schneller fahren kann.
Eine GPS-Abkürzung (Newton-DB), um Staus zu vermeiden.
Bessere Navigationswerkzeuge (Multi-Power-Iteration), um sicherzustellen, dass er sich nicht verirrt.

Das Ergebnis ist, dass der „schlaue“ Optimierer (Shampoo) nicht mehr zu langsam ist, um verwendet zu werden, was ihn zu einer praktischen Wahl für das Training der nächsten Generation von KI-Modellen macht.

Technisches Resümee: DASH – Schnelleres Shampoo durch gebatchte Block-Präkonditionierung und effiziente Invertierte-Wurzel-Solver

Problemstellung

Shampoo ist ein führender approximativer zweiter Ordnung bekannter Optimierer, der Modelle mit geringeren Aktivierungs-Ausreißern und verbesserter Komprimierbarkeit hervorbringt. Seine weite Verbreitung wird jedoch durch den erheblichen Rechenaufwand behindert, primär aufgrund der teuren Operation zur Berechnung der inversen Matrixwurzel ( $A^{-1/4}$ ) für Präkonditionierungsmatrizen. Während Distributed Shampoo versucht, dies durch die Aufteilung der Präkonditionierer in $B \times B$ Blöcke zu mildern, um die Komplexität von $O(n^3)$ auf $O(Bn^2)$ zu reduzieren, leiden bestehende Implementierungen unter kritischen Effizienzlücken:

Sequenzielle Verarbeitung: Standardimplementierungen verarbeiten Blöcke sequenziell, wodurch die moderne GPU-Parallelität nicht voll ausgeschöpft wird.
Ineffiziente Solver: Die Standardabhängigkeit von der Eigenwertzerlegung (EVD) ist auf GPUs schwer zu parallelisieren. Alternative iterative Methoden wie Coupled-Newton (CN) werden oft aufgrund von Bedenken hinsichtlich der numerischen Stabilität und Konvergenzgeschwindigkeit standardmäßig deaktiviert.
Suboptimale Skalierung: Aktuelle Methoden skalieren Matrizen oft mittels der Frobenius-Norm, was zwar rechentechnisch günstig ist, aber Eigenwerte gegen Null drücken kann, was wiederum signifikant mehr Iterationen erfordert, um eine Zielpräzision zu erreichen.

Methodik

Das Paper schlägt DASH (Distributed Accelerated SHampoo) vor, eine Hochleistungsimplementierung, die darauf ausgelegt ist, die Lücke zwischen theoretischer Effizienz und Hardware-Kapazitäten zu schließen. Die Methodik ruht auf zwei primären Säulen:

1. Gebatchte Block-Präkonditionierung (Systemseitige Optimierung)

Anstatt Präkonditionierer-Blöcke sequenziell zu verarbeiten, stapelt DASH diese in 3D-Tensoren.

3D-Tensor-Stapelung: Gradientenblöcke $G_{ij}$ werden in einem Tensor der Form $(N, B, B)$ gestapelt, wobei $N$ die Gesamtzahl der Blöcke ist.
Parallele Ausführung: Dies ermöglicht es dem Optimierer, die Verfahren zur inversen Wurzel (EVD, CN oder neue iterative Methoden) für die gesamte Gruppe von Blöcken gleichzeitig unter Verwendung von Batched-Matrix-Multiplikations-Primitiven (bmm) anzuwenden.
GPU-Auslastung: Diese architektonische Änderung maximiert den Durchsatz der GPU-Tensor-Cores, reduziert die Speicherfragmentierung und ermöglicht die effiziente Nutzung von Halbfachpräzision (FP16) für bestimmte Solver.

2. Effiziente Invertierte-Wurzel-Solver (Algorithmische Optimierung)

Das Paper führt fortgeschrittene linear-algebraische Ansätze ein und analystiert diese als Ersatz oder Ergänzung zur Standard-EVD und CN:

Newton-Denman-Beavers (NDB) Iteration: Ein neuartiger Ansatz für Shampoo, der sowohl die Quadratwurzel als auch die inverse Quadratwurzel berechnet. Die Autoren optimieren die erste Iteration, um redundante Matrixmultiplikationen zu eliminieren, wodurch der Overhead auf 3 Matmul pro Schritt reduziert wird (was in der Praxis effektiv weniger ist).
Chebyshev-Polynom-Approximation: Nutzung des Clenshaw-Algorithmus zur Approximation von Potenzfunktionen der Matrix, was einen alternativen iterativen Solver bietet.
Multi-Power-Iteration für die Skalierung: Um Konvergenzprobleme anzugehen, ersetzt das Paper die Standard-Frobenius-Norm-Skalierung durch eine Multi-Power-Iteration-Methode.
- Mechanismus: Es schätzt den Spektralradius ( $\lambda_{max}$ ) mithilfe eines Pools von Startvektoren (bis zu 32), die parallel via Power-Iteration verarbeitet werden.
- Vorteil: Dies liefert einen genaueren Skalierungsfaktor als die Frobenius-Norm, wodurch die Eigenwerte in einem Bereich gehalten werden, in dem iterative Solver schnell konvergieren (näher an 1), was die Anzahl der benötigten Iterationen reduziert und die numerische Stabilität verbessert.

Zentrale Beiträge

DASH-Implementierung: Ein GPU-effizientes Framework, das sequenzielle Schleifen durch parallele 3D-Tensor-Operationen ersetzt und bis zu 5,6× schnellere Optimiererschritte im Vergleich zu gut optimiertem Distributed Shampoo erreicht.
Neuartige Solver: Die Einführung der Newton-DB (NDB) Iteration und der Chebyshev-Polynom-Approximation als praktikable, schnellere Alternativen zu EVD und CN. Das Paper zeigt, dass NDB eine geringere Validierungs-Perplexität als CN und oft sogar EVD erzielt.
Konvergenzanalyse: Eine detaillierte Verhaltensanalyse, die zeigt, dass die Standard-Frobenius-Norm-Skalierung suboptimal ist und oft zu einer langsameren Konvergenz führt. Das Paper liefert intuitive Erklärungen für die unterschiedlichen Konvergenzverhaltensweisen von CN gegenüber NDB.
Multi-Power-Iteration: Eine effiziente Halbfachpräzisions-Implementierung der Power-Iteration, die robust den Spektralradius schätzt und somit eine optimale Matrixskalierung ermöglicht, welche die Konvergenzkriterien für NDB und CN mit weniger Schritten erfüllt.

Experimentelle Ergebnisse

Die Autoren evaluierten DASH auf einem Llama-Modell mit 953 Mio. Parametern (Embedding-Größe 2048, Sequenzlänge 1024), das auf dem C4-Datensatz trainiert wurde.

Geschwindigkeit:
- DASH reduziert die Laufzeit des Optimiererschritts um den Faktor 4× bis 5,6× im Vergleich zu Distributed Shampoo, abhängig von der inversen Wurzel-Methode und der verwendeten Präzision.
- Die signifikanteste Beschleunigung (5,6×) wurde beim Vergleich von Distributed Shampoos EVD (FP32) gegen DASHs CN (FP16) mit GPU-seitiger Stapelung beobachtet.
- Selbst bei häufigen Präkonditionierer-Updates ( $f=1$ ) behält DASH erhebliche Geschwindigkeitsvorteile bei.
Genauigkeit (Validierungs-Perplexität):
- DASH erreicht in fast allen stabilen Konfigurationen die Werte von Distributed Shampoo.
- NDB liefert konsistent eine geringere Validierungs-Perplexität als CN und übertrifft oft EVD.
- Die Verwendung der Power-Iteration-Skalierung mit NDB resultierte in der niedrigsten Validierungs-Perplexität (11,68) im Vergleich zur Frobenius-Skalierung (11,76), was die theoretische Analyse bestätigt.
Präzision:
- Die Implementierung unterstützt FP16 für CN und EVD, ohne die Perplexität zu verschlechtern, was die Laufzeit weiter reduziert.
- NDB in FP16/BF16 wurde in ihren Experimenten als numerisch instabil befunden, wobei die Autoren dies als Bereich für zukünftige Arbeiten markieren.

Bedeutung und Behauptungen

Das Paper behauptet, dass DASH die rechnerische Barriere für die Nutzung von Optimierern zweiter Ordnung wie Shampoo signifikant senkt und hochwertiges Preconditioning praktikabel für den breiten Einsatz in groß angelegtem Deep Learning macht.

Verbindung von Theorie und Hardware: Durch die Abstimmung des algorithmischen Designs (iterative Solver, Matrixskalierung) auf die Hardware-Kapazitäten (Tensor Cores, gebatchte Operationen) ermöglicht DASH es Forschern, Eigenschaften von Optimierern zu untersuchen, ohne prohibitive Laufzeiten in Kauf nehmen zu müssen.
Numerische Stabilität vs. Geschwindigkeit: Die Arbeit stellt die Annahme in Frage, dass iterative Methoden inhärent weniger stabil oder langsamer als EVD sind. Durch die Einführung besserer Skalierung (Multi-Power-Iteration) und optimierter Iterationen (NDB) demonstrieren die Autoren, dass iterative Methoden eine überlegene Konvergenz und geringere Perplexität erreichen können, während sie gleichzeitig signifikant schneller sind.
Zukünftige Richtungen: Die Autoren deuten an, dass die überlegene Leistung von Newton-DB Wege für eine dynamische Solver-Auswahl basierend auf den Konditionszahlen der Blöcke sowie für die Stabilisierung von NDB in niedrigeren Präzisionsformaten mittels stochastischer oder Fehlerkorrektur-Methoden eröffnet.

Das Paper kommt zu dem Schluss, dass der primäre Flaschenhals von Shampoo nicht die mathematische Komplexität war, sondern die fragmentierte Ausführung und die numerische Instabilität – Probleme, die DASH effektiv adressiert.