Optimierung von KI-Trainingsmethoden durch die Entwicklung von DASH

Optimierung von KI-Modellen: Wie DASH die Effizienz von Shampoo-Optimierern revolutioniert

Die kontinuierliche Entwicklung von Künstlicher Intelligenz und insbesondere von großen Sprachmodellen (LLMs) erfordert immer leistungsfähigere und effizientere Trainingsmethoden. Optimierer spielen dabei eine zentrale Rolle, indem sie die Anpassung der Modellparameter während des Lernprozesses steuern. Der Shampoo-Optimierer gilt als einer der führenden approximativen Optimierer zweiter Ordnung und hat sich in Wettbewerben wie MLCommons AlgoPerf bewährt. Er ist bekannt dafür, Modelle mit geringeren Aktivierungs-Ausreißern zu erzeugen, die leichter zu komprimieren sind. Trotz seiner Vorteile litt der Einsatz von Shampoo in der Praxis jedoch unter erheblichen Rechenverzögerungen, hauptsächlich aufgrund aufwendiger interner Operationen.

Ein Forschungsteam des ISTA-DASLab hat sich dieser Herausforderung angenommen und mit DASH (Distributed Accelerated SHampoo) eine innovative Lösung vorgestellt. DASH zielt darauf ab, die rechnerischen Engpässe des Distributed Shampoo zu beseitigen, ohne die numerische Präzision zu beeinträchtigen. Dieser Artikel beleuchtet die Kerninnovationen von DASH und deren Auswirkungen auf das Training von KI-Modellen.

Die Herausforderung: Rechenintensive Inverse-Wurzel-Berechnungen

Der Hauptgrund für die rechnerische Ineffizienz des Shampoo-Optimierers liegt in der Berechnung inverser Matrixwurzeln. Diese Operation skaliert typischerweise mit Θ(n³) für eine n×n-Preconditioner-Matrix, was bei großen Modellen schnell prohibitiv wird. Obwohl Distributed Shampoo diesen Engpass durch die Aufteilung der Preconditioner in B×B-Blöcke zu mildern versuchte, wurden diese Blöcke sequenziell verarbeitet, was die GPU-Auslastung limitierte.

DASH setzt genau hier an und führt zwei wesentliche Neuerungen ein, um diese Limitierungen zu überwinden:

1. Gebündelte Block-Vorkonditionierung: Anstatt die B×B-Preconditioner-Blöcke nacheinander zu verarbeiten, stapelt DASH diese zu 3D-Tensoren. Diese Architekturänderung ermöglicht batched matrix multiplications (bmm) auf GPUs und verbessert die Auslastung der Tensor-Kerne erheblich.
2. Effizientere Inverse-Wurzel-Löser: DASH integriert neue, fortschrittliche numerische Methoden zur Berechnung von Matrixpotenzen A^-1/p, darunter die Newton-Denman-Beavers (NDB)-Iteration und Chebyshev-Polynom-Approximationen.

Algorithmen im Detail: Newton-DB und Chebyshev-Polynome

DASH untersucht und optimiert verschiedene Ansätze zur Berechnung inverser Matrixwurzeln:

Coupled-Newton (CN)-Iteration

Die Coupled-Newton-Iteration, die bereits in früheren Shampoo-Implementierungen als schnellere Alternative zur Eigenwertzerlegung (EVD) verwendet wurde, wird in DASH in einer optimierten FP16-Implementierung eingesetzt. Diese Methode erfordert je nach Potenz (p=2 oder p=4) 3 bis 4 Matrixmultiplikationen pro Iteration.

Newton-Denman-Beavers (NDB)-Iteration

Die NDB-Iteration berechnet gleichzeitig die Quadrat- und Inverse-Quadratwurzel einer Matrix. Für die Berechnung der inversen vierten Wurzel werden zwei Aufrufe der NDB-Prozedur miteinander verkettet. Obwohl jede Iteration 3 Matrixmultiplikationen erfordert, kann die erste Iteration optimiert werden, um zwei Multiplikationen einzusparen. Experimente zeigen, dass NDB eine niedrigere Validierungs-Perplexität als CN und EVD erzielt.

Chebyshev-Polynom-Approximation

Diese Methode approximiert f(A) = A^-1/p mittels einer Chebyshev-Reihe. Die Koeffizienten werden einmalig für ein gewähltes Intervall bestimmt. Für Matrizen wird die Eingabematrix A auf ein Intervall [-1, 1] abgebildet, bevor der Clenshaw-Algorithmus angewendet wird. Eine optimierte Version des Algorithmus ermöglicht die Berechnung in d-1 Matrixmultiplikationen für ein Polynom vom Grad d. Diese Methode profitiert von FP16 für Multiplikationen mit FP32-Akkumulation.

Die Bedeutung der Matrixskalierung und Multi-Power-Iteration

Eine tiefgehende Analyse in der DASH-Studie zeigt, dass die Skalierung der Eingabematrix A durch ihre Frobenius-Norm, wie sie im Distributed Shampoo üblich ist, suboptimal ist. Da die Frobenius-Norm in der Regel größer oder gleich dem größten Eigenwert ist, verschiebt diese Skalierung die Eigenwerte in Richtung Null. Dies wiederum erfordert mehr Iterationen für Newton-basierte Methoden, um zu konvergieren.

DASH schlägt stattdessen eine Skalierung durch eine genaue Schätzung des größten Eigenwerts (λ_max(A)) vor, insbesondere durch 2λ_PI, wobei λ_PI eine Schätzung der Power-Iteration ist. Dies stellt sicher, dass die Eigenwerte in einem günstigen Intervall liegen, was eine schnellere Konvergenz ermöglicht. Um λ_max(A) robust zu schätzen, führt DASH die Multi-Power-Iteration ein. Diese Methode nutzt einen Pool von Startvektoren, um die Power-Iteration parallel auszuführen und den Eigenvektor mit dem größten Rayleigh-Quotienten auszuwählen. Durch die gebündelten Operationen ist dieser Ansatz rechnerisch effizient.

Praktische Ergebnisse und Auswirkungen

Experimente mit einem Llama-953M-Modell zeigen, dass DASH in den meisten stabilen Konfigurationen eine vergleichbare Validierungs-Perplexität wie Distributed Shampoo erreicht, während die Rechenzeit des Optimierungsschritts erheblich reduziert wird. Es konnten Beschleunigungen von bis zu 4,83-fach beobachtet werden. Insbesondere die DASH-CN-FP16-Konfiguration (Coupled-Newton in FP16 mit Frobenius-Normalisierung) erzielte die höchste Beschleunigung von 4,83x bei einer Blockgröße von 1024, wodurch die Optimierungsschrittzeit von 666 ms auf 138 ms reduziert wurde.

DASH-NDB (Newton-Denman-Beavers mit Power-Iteration-Skalierung in FP32) erreichte die niedrigste Validierungs-Perplexität (11,68) und übertraf damit EVD und CN, während es ebenfalls signifikante Beschleunigungen lieferte (z.B. von 740 ms auf 194 ms für B=1024). Die Verwendung der Power-Iteration-Skalierung führte konsistent zu einer niedrigeren Validierungs-Perplexität als die Frobenius-Norm-Skalierung, was die theoretische Analyse bestätigt.

Die Forschungsarbeit zu DASH bietet eine Hochleistungs-Implementierung des Shampoo-Optimierers, die die Lücke zwischen theoretischer Effizienz und praktischer GPU-Ausführung schließt. Durch Optimierungen in der numerischen Analyse und im Systemdesign (gebündelte Ausführung, neue Löser, verbesserte Skalierung) wird die Sekundärordnungsoptimierung für großskaliges Deep Learning zugänglicher und effizienter gestaltet.

Diese Fortschritte sind für Unternehmen im B2B-Bereich, die auf die Entwicklung und das Training großer KI-Modelle angewiesen sind, von großer Bedeutung. Eine schnellere und effizientere Modelloptimierung kann zu kürzeren Entwicklungszyklen, geringeren Rechenkosten und letztendlich zu leistungsfähigeren KI-Lösungen führen.

Ausblick

Die Arbeit an DASH eröffnet weitere Forschungsfelder. Beispielsweise könnte eine dynamische Auswahl des Solvers, basierend auf der Konditionszahl eines Blocks, weitere Effizienzsteigerungen ermöglichen. Auch die Stabilisierung der Newton-DB-Iteration bei geringerer Präzision durch Methoden wie stochastische oder Fehlerkorrekturverfahren stellt ein interessantes Forschungsfeld dar. Zudem ist die Validierung des Ansatzes bei noch größeren Modell- und Datenskalen sowie im Kontext des Tensor-Parallel-Trainings von Relevanz.

Bibliographie

- Modoranu, I.-V., Zmushko, P., Schultheis, E., Safaryan, M., & Alistarh, D. (2026). DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers. arXiv preprint arXiv:2602.02016. - TheMoonlight.io. (n.d.). [Literature Review] DASH: Faster Shampoo via Batched Block Preconditioning and Efficient Inverse-Root Solvers. Abgerufen von https://www.themoonligh.io/en/review/dash-faster-shampoo-via-batched-block-preconditioning-and-efficient-inverse-root-solvers - Hugging Face. (n.d.). Optimizers - a ISTA-DASLab Collection. Abgerufen von https://huggingface.co/collections/ISTA-DASLab/optimizers - ChatPaper.ai. (n.d.). DASH: Schnelleres Shampoo durch gebündelte Block-Vorkonditionierung und effiziente Inverse-Wurzel-Löser. Abgerufen von https://www.chatpaper.ai/de/dashboard/paper/ab1d89a9-de9a-4b50-b536-07832ac96d03