Ai2s Olmo-core 3 trainiert MoE-Modelle mit 2,7-fachem Durchsatz
Serguey Asael Shinder2 Min. LesezeitDas Allen Institute for AI hat Olmo-core 3 veröffentlicht, eine neue Version des offenen Frameworks, mit dem es seine Olmo-Sprachmodelle trainiert. Die wichtigste Änderung ist ein neu gestaltetes System zum Training von Mixture-of-Experts-Modellen, das laut Ai2 bis in den Bereich von Billionen Parametern skalieren soll und eine der Grundlagen der nächsten Olmo-Generation ist.
Mixture-of-Experts-Modelle enthalten viele spezialisierte Komponenten, die Experten, doch jedes Token nutzt nur wenige davon. Der Haken, so der Beitrag: Das gesamte Modell muss trotzdem im GPU-Speicher liegen und aktualisiert werden, und das Weiterleiten von Tokens an die richtigen Experten im Cluster verursacht Kommunikationskosten, die den Vorteil aufzehren können.
Die zentrale Änderung ist die Abkehr vom Fully Sharded Data Parallelism, das die Gewichte für jeden kleinen Batch sammelte und neu verteilte, hin zu einem Ansatz auf Basis von Distributed Data Parallelism, der die Experten dauerhaft auf den GPUs hält und die Daten zu ihnen schickt. Dazu kommen Expert Parallelism, das die Experten auf GPUs verteilt, Pipeline Parallelism, das Schichten auf GPU-Gruppen aufteilt, und ein verteilter Optimizer, der den Optimizer-Zustand verteilt, statt ihn auf jede GPU zu kopieren. Routing-Metadaten bleiben auf der GPU, sodass die CPU Arbeit einreihen kann, ohne auf deren Rückkopie zu warten.
Ai2 nennt zwei Messungen. Wurde der Expertenpool von 8 auf 128 vergrößert, bei weiterhin vier Experten pro Token, blieben die aktiven Parameter bei etwa 3,2 Milliarden, die Gesamtparameter stiegen von 4,6 auf 47 Milliarden, und der Trainingsdurchsatz sank um weniger als 5 %. In einem vorläufigen Test auf acht NVIDIA-B300-GPUs verarbeitete ein Modell mit 47 Milliarden Parametern 52.000 Tokens pro Sekunde und GPU, gegenüber 19.400 mit der früheren Implementierung – etwa das 2,7-Fache. Getestet wurde der Stack auch mit mehr als einer Billion Gesamtparametern.
Als etablierte Option für großes MoE-Training nennt der Beitrag NVIDIAs Megatron-Core; Olmo-core 3 bringt eine integrierte Alternative in Ai2s eigenes offenes Framework.
Warum das wichtig ist
Trainingscode für große dünn besetzte Modelle lag bisher meist bei wenigen Unternehmen oder in Herstellerstacks. Ein offenes Framework mit veröffentlichten Durchsatzzahlen gibt Hochschulen und kleineren Laboren die Möglichkeit, MoE-Modelle zu trainieren, ohne die Parallelisierungsschicht selbst neu zu bauen.