Aleph Alpha erklärt die Expertenbalancierung auf zwei Ebenen hinter seinem MoE-Modell Kolibri 1
Aleph Alpha hat einen Forschungsbeitrag zum Rezept der Lastverteilung veröffentlicht, mit dem Kolibri 1 vortrainiert wurde, ein spärlich aktiviertes Mixture-of-Experts-Modell. Die Methode kombiniert Exact Quantile Balancing (EQB) für die globale Balance mit Load-Error Injection (LEI) für die lokale Balance innerhalb jedes Microbatches.
- Kolibri 1 hat 78 Mrd. Parameter insgesamt und rund 3,5 Mrd. aktive pro Token, also 4,5 % der Gesamtzahl.
- Jede MoE-Schicht leitet jedes Token an 6 von 384 gerouteten Experten weiter.
- Exact Quantile Balancing legt Routing-Biases anhand exakter Quantile der Router-Logits über den globalen Batch fest.
- Load-Error Injection speist den Lastfehler jedes Experten im lokalen Microbatch in die Gradienten des Routers ein.
- Das Pre-Training von Kolibri nutzt EQB mit begrenztem LEI; Ablationen liefen an einem 7,5-Mrd.-Proxy-Modell über 100 Mrd. Tokens.
Ein sehr spärliches Modell, das auf ausgewogene Experten angewiesen ist
Kolibri 1 ist laut Aleph Alpha ein spärlich aktiviertes Mixture-of-Experts-Sprachmodell (MoE) mit 78 Mrd. Parametern insgesamt und rund 3,5 Mrd. aktiven Parametern pro Token, also 4,5 % der Gesamtzahl. In jeder MoE-Schicht wählt ein gelernter Router für jedes Token 6 von 384 geroutete Experten aus.
Das Unternehmen erklärt, dass das stabile Training eines derart spärlichen Modells im großen Maßstab stark von einem wirksamen Verfahren zur Lastverteilung zwischen den Experten abhängt. Der Router bestimmt, wie viele Trainingsdaten jeder Experte sieht und wie gleichmäßig die Arbeit auf die Geräte verteilt wird. Bevorzugt er dauerhaft nur wenige Experten, erhalten die übrigen zu wenige Tokens, um gut zu lernen, und ein Teil der Modellkapazität bleibt ungenutzt.
Hinzu kommt eine Hardwarefrage. Mit Expertenparallelität können Experten sowohl im Training als auch bei der Inferenz auf verschiedenen Geräten laufen. Gehen zu viele Tokens an ein einzelnes Gerät, müssen die anderen auf dieses warten. Aleph Alpha zeigt Daten aus einem einzelnen Pre-Training-Lauf, in denen ein höheres lokales Ungleichgewicht mit einem geringeren Trainingsdurchsatz einhergeht, was dazu passt, dass Geräte auf überlastete Experten warten.
Zwei Mechanismen für zwei Ebenen
Das Rezept zielt auf Balance auf zwei Ebenen. Die globale Balance erfasst alle Tokens eines vollständigen Optimierungsschritts, sodass jeder Experte genügend Gradientensignal erhält. Die lokale Balance betrifft jeden Microbatch, um Engpässe bei der expertenparallelen Ausführung zu vermeiden.
- Exact Quantile Balancing (EQB): erweitert Quantile Balancing (QB) und legt die Routing-Biases pro Experte anhand exakter Quantile der Router-Logits über den globalen Batch fest.
- Load-Error Injection (LEI): speist den Lastfehler jedes Experten im lokalen Microbatch direkt in die Gradienten des Routers ein.
Das Pre-Training von Kolibri nutzt EQB zusammen mit begrenztem LEI. Aleph Alpha zufolge ist LEI eine gradientenbasierte Korrektur, die dem Hilfsverlust von GShard eng verwandt ist.
Zur Messung des Ungleichgewichts verwendet der Beitrag die Maximal Violation (MaxVio), also die größte relative Überlastung unter den Experten. Ein Wert von 0 bedeutet perfekte Balance, 0,5 heißt, dass der am stärksten ausgelastete Experte das 1,5-Fache seines gleichmäßigen Anteils erhält.
Warum Quantile statt Updates mit fester Schrittweite
Verfahren zur Lastverteilung lassen sich grob in zwei Familien einteilen. Hilfsverluste wie der GShard-Verlust fügen dem Trainingsziel einen Balancierungsterm hinzu. Die Bias-Anpassung setzt die Routing-Biases pro Experte außerhalb des Gradientenabstiegs. Dazu zählen der hilfsverlustfreie Regler von DeepSeek, der jeden Bias um eine feste Schrittweite verschiebt, sowie QB, das Jianlin Su vorgeschlagen hat und das in MARIN und Kimi K3 zum Einsatz kommt.
Aleph Alpha beschreibt das Problem beim Top-K-Routing: Die Experten konkurrieren um die K Plätze jedes Tokens, sodass die Erhöhung des Bias eines Experten Zuweisungen anderer verdrängen kann. Der Regler von DeepSeek aktualisiert jeden Bias anhand des Lastfehlers des jeweiligen Experten und berücksichtigt verdrängte Zuweisungen nicht ausdrücklich. Eine kleine Update-Rate begrenzt diese Nebeneffekte, doch die Lasten nähern sich der Balance nur allmählich über viele Schritte.
QB braucht dagegen keine Update-Rate, weil es den Bias jedes Experten direkt aus dem Logit-Quantil bestimmt, das zu seiner Ziellast passt. Der Beitrag schreibt Top-K als Schwellenwertregel um, sodass sich die Last jedes Experten unabhängig einstellen lässt, sobald die Biases pro Token festgehalten werden. EQB berechnet diese globalen Quantile über alle Ränge hinweg exakt.
Was das bedeutet
Der Beitrag beschreibt ein Trainingsrezept und ist keine Produktveröffentlichung. Für Herleitungen verweist er auf ein vollständiges Paper und berichtet von Ablationen an einem Proxy-Modell mit 7,5 Mrd. Parametern, das mit 100 Mrd. Tokens trainiert wurde. Für Teams, die hochgradig spärliche MoE-Modelle bauen, dokumentiert er, wie ein Labor sowohl die Datenunterversorgung einzelner Experten als auch Stillstände auf Geräteebene in expertenparallelen Setups angeht.
