SkewAdam reduziert den Speicherbedarf beim Training von Mixture-of-Experts um mehr als 60 % und liefert einen spürbaren Genauigkeitssprung, wodurch Entwickler ein 6,78 Milliarden Parameter umfassendes MoE-Modell auf einer einzigen 40-GB-GPU unterbringen können.

Warum das MoE-Training an eine Speicherbarriere stößt

Mixture-of-Experts-Architekturen leiten jeden Input durch eine kleine Teilmenge von „Experten“-Subnetzwerken, während ein dichter Backbone beibehalten wird. Der Vorteil ist ein Modell, das auf Milliarden von Parametern skalieren kann, ohne dass der Rechenaufwand proportional ansteigt. In der Praxis verbraucht der Optimizer – insbesondere die AdamW-Variante, die die meisten Teams verwenden – den Großteil des GPU-RAMs. Bei einem Modell mit 6,78 Mrd. Parametern beanspruchen allein die Momentum- und Varianz-Tensoren des Optimizers etwa 50 GB. Rechnet man Aktivierungen und Modellgewichte hinzu, erreicht der Spitzenwert des Speichers 81,4 GB, was Multi-GPU-Setups oder langsamere Trainingspläne erforderlich macht.

Was SkewAdam anders macht

SkewAdam erfindet keine neue Lernregel. Es ordnet lediglich neu an, wo die Momente von Adam gespeichert werden:

  • Der dichte Backbone behält das Momentum in voller Präzision bei, um die für dichte Schichten erforderlichen glatten Updates zu gewährleisten.
  • Die Experten-Bank speichert eine faktorisierte Approximation der Varianz, wodurch die für jeden Experten gespeicherten Daten reduziert werden.
  • Der Router, der entscheidet, welche Experten aktiviert werden, behält eine exakte Schätzung des zweiten Moments.

Indem der Optimizer diese drei Komponenten als separate „Ebenen“ (tiers) behandelt, verzichtet er auf redundante Präzision dort, wo sie weniger wichtig ist, und behält sie dort, wo sie am wichtigsten ist.

Messbare Auswirkungen

Das Ausführen desselben 6,78 Mrd. Parameter umfassenden MoE-Modells mit SkewAdam liefert:

  • Spitzen-GPU-Speicher: 31,3 GB (statt 81,4 GB)
  • Speicherbedarf des Optimizer-States: 1,29 GB (statt 50 GB)

Der reduzierte Zustand passt problemlos in einen einzelnen 40-GB-Beschleuniger.

Genauigkeitsgewinne, nicht nur Einsparungen

Speicherreduzierungen beeinträchtigen oft die Modellqualität, aber SkewAdam verbessert die Perplexity – eine Standardmetrik für Sprachmodelle – von 126,8 (AdamW) auf 108,4. Es übertrifft zudem Muon (120,2) und Lion (393,7) bei derselben Aufgabe. Die Autoren geben an, dass der Schub durch die Beibehaltung des Momentums über alle drei Ebenen hinweg zustande kommt; Methoden, die das Momentum vollständig verwerfen, wie etwa Adafactor, bleiben dahinter zurück.

Offene Fragen

Die Ergebnisse von SkewAdam werden an einem Modell mit 6,78 Mrd. Parametern demonstriert. Es bleibt unklar, ob dieselben Speicher-Zustands-Verhältnisse auch für Modelle einer Größenordnung höher oder für Aufgaben außerhalb der Sprachmodellierung gelten.

Worauf man achten sollte

  • Benchmarks für größere MoE-Konfigurationen (Dutzende von Milliarden Parametern).
  • Übernahme durch Open-Source-Frameworks und Integration in populäre Trainings-Skripte.
  • Feedback der Community zu versteckten Kompromissen, wie etwa der Konvergenzgeschwindigkeit oder der Stabilität unter verschiedenen Lernraten-Schedules.

Quelle: Entwickler-Post, der das Design des Optimizers und die empirischen Ergebnisse detailliert beschreibt.