Olmo-core 3: Ai2 åpner MoE-treningsstakken som ga 2,7 ganger høyere gjennomstrømning

52 000 tokens per sekund per GPU. Det klarte en MoE-modell på 47 milliarder parametere på åtte Nvidia B300 i en foreløpig test med Olmo-core 3, mot 19 400 med den forrige implementasjonen. Ai2 slapp rammeverket 1. oktober, og versjon 3.0.0 ligger på GitHub under Apache 2.0. Det er treningsstakken bak neste generasjon Olmo, som skal bli en MoE-modell.

Hovedgrepet er å slutte å flytte vekter. Den gamle MoE-koden brukte FSDP, som samlet og fordelte modellvektene på nytt for hver lille batch. Olmo-core 3 bytter til DDP, holder ekspertene fast på GPU-ene og sender i stedet dataene dit. I en annen test økte Ai2 antall eksperter fra 8 til 128 med fire aktive per token. Aktive parametere lå fast på rundt 3,2 milliarder, total kapasitet vokste fra 4,6 til 47 milliarder, og gjennomstrømningen falt under 5 prosent.

Rammeverket støtter også MXFP8. På fire B300 ga det rundt 21 prosent høyere gjennomstrømning enn BF16, og toppminnet falt fra 103 til 95 GiB. Ai2 har kjørt konfigurasjoner opp til 1,2 billioner parametere på 512 GPU-er, men med tilfeldig ruting, så tallene sier noe om systemet og ingenting om modellkvalitet.

«Modellvekter er mer nyttige når infrastrukturen og treningsbeslutningene bak dem også er åpne.» — Ai2

For de fleste som bygger selv, ligger verdien i den tekniske rapporten. Ai2 beskriver blant annet «token gerrymandering»: en score som skal belønne jevn ruting, kan bli bedre mens arbeidsfordelingen mellom ekspertene i virkeligheten blir skjevere.

Hva bør du gjøre?

1. Les den tekniske rapporten før du trener MoE selv. Den dokumenterer også grep som ikke virket, som lavere læringsrate for ekspertene og overlappende kommunikasjon og beregning.
2. Trener du allerede MoE med Nvidias Megatron-Core, sammenlign mot Olmo-core 3 på din egen maskinvare før du bytter. Ai2s tall er målt på B300.