Ai2 otvára Olmo-core 3: tréning MoE modelov škáluje nad bilión parametrov

Otvorený tréningový systém pre veľké MoE modely

Allen Institute for AI, známy ako Ai2, vydal Olmo-core 3, zásadnú aktualizáciu otvoreného frameworku, na ktorom vzniká rodina jazykových modelov Olmo. Nová verzia sa sústreďuje na tréning modelov typu mixture-of-experts, skrátene MoE, a podľa zverejnených výsledkov dokáže pracovať aj s konfiguráciami presahujúcimi bilión parametrov. Nejde pritom o nový hotový jazykový model určený na okamžité používanie. Olmo-core 3 je tréningová infraštruktúra: súbor nástrojov, paralelizačných mechanizmov a optimalizácií, ktoré majú vývojárom umožniť zostavovať a trénovať vlastné riedke modely vo veľkých GPU klastroch. Ai2 ho zároveň označuje za technologický základ ďalšej generácie Olmo, ktorá má prejsť z hustej architektúry na MoE.

Pri MoE modeli sa pre každý token aktivuje iba malá časť dostupných expertných blokov. Model tak môže mať veľmi veľkú celkovú kapacitu bez toho, aby pri každom vstupe vykonal výpočet so všetkými parametrami. Úspora však nie je automatická. Experti musia zostať rozdelení medzi GPU, router musí rýchlo rozhodnúť, kam jednotlivé tokeny poslať, a dáta sa následne presúvajú naprieč klastrom. Komunikácia, nerovnomerné zaťaženie a opakované preskupovanie tenzorov môžu teoretickú výhodu riedkej architektúry výrazne zmenšiť. Olmo-core 3 preto nerieši iba samotnú definíciu modelu, ale najmä pohyb dát a rozdelenie práce počas tréningu.

Ai2 uvádza experiment, v ktorom zvýšilo počet expertov z ôsmich na 128, pričom router naďalej vyberal štyroch expertov pre každý token. Počet aktívnych parametrov zostal približne na úrovni 3,2 miliardy, celková kapacita však narástla zo 4,6 na 47 miliárd parametrov. Tréningová priepustnosť podľa merania klesla o menej než päť percent. Tento výsledok naznačuje, že nový systém dokáže zväčšiť zásobu špecializovaných komponentov bez úmerného rastu výpočtu na token. Ide však o systémový benchmark zverejnený samotným tímom, nie o nezávislé porovnanie kvality výsledného modelu.

Prečo Ai2 nahradilo pôvodný spôsob distribúcie

Staršia implementácia MoE v Olmo-core využívala plne shardovaný dátový paralelizmus FSDP. Pri tomto prístupe sa časti váh podľa potreby zhromažďujú a po spracovaní opäť rozdeľujú. Olmo-core 3 pre expertnú časť prechádza na riešenie založené na distribuovanom dátovom paralelizme. Experti zostávajú uložené na pridelených GPU a systém k nim smeruje príslušné tokeny, namiesto opakovaného prenášania ich váh. Ai2 tvrdí, že pri predbežnom teste 47-miliardového MoE na ôsmich akcelerátoroch NVIDIA B300 dosiahla nová implementácia 52-tisíc tokenov za sekundu na GPU. Predchádzajúca verzia dosiahla 19 400 tokenov, takže namerané zrýchlenie bolo približne 2,7-násobné.

Framework kombinuje expert parallelism, pipeline parallelism a distribuovaný optimalizátor. Prvý mechanizmus rozkladá expertov medzi GPU, druhý rozdeľuje vrstvy modelu medzi skupiny zariadení a tretí rozptyľuje stav optimalizátora, aby ho nemusela v plnom rozsahu uchovávať každá karta. Olmo-core 3 navyše presúva údaje routera na GPU, zapisuje smerované tokeny priamo do vstupných bufferov expertov a spája viac menších maticových operácií pomocou grouped GEMM. Praktickým cieľom je obmedziť čakanie CPU, kopírovanie metadát a drobné výpočty, ktoré nedokážu efektívne využiť moderné akcelerátory.

Súčasťou vydania je podpora formátu MXFP8 s nižšou presnosťou. V kontrolovanom teste na štyroch GPU B300, pri rovnomernom rozdelení práce medzi expertov, nameralo Ai2 približne 21-percentné zvýšenie koncovej tréningovej priepustnosti oproti BF16. Špičková obsadená pamäť pritom klesla zo 103 na 95 GiB. Väčšina zisku podľa tímu nevznikla v mechanizme pozornosti, ale vo feed-forward výpočtoch a pri prenose dát medzi expertmi. To je dôležitý detail: nižšia presnosť pomáha iba vtedy, keď úsporu prenosu a výpočtu neprevážia konverzie formátov alebo iné časti pipeline.

Bilión parametrov ešte neznamená vytrénovaný model

Najväčší demonštrovaný test zahŕňal model s 1,2 bilióna celkových parametrov a 58,36 miliardy aktívnych parametrov na token na 512 GPU B300. Najvyšší pozorovaný výkon dosiahol 858 TFLOP/s na GPU. Ai2 zároveň skúšalo alternatívnu komunikáciu cez DeepEP v2 pri konfigurácii s 2,38 bilióna parametrov. Tieto čísla treba interpretovať opatrne. Test s 1,2 bilióna parametrov používal náhodné smerovanie a sledoval výkon infraštruktúry, nie jazykové schopnosti. Väčšia konfigurácia bola iba krátkym kapacitným testom, nie úplným a dlhodobo stabilným tréningom. Vydanie teda dokazuje dosiahnuteľný systémový rozsah, nie existenciu nového biliónového Olmo modelu.

Technická práca priniesla aj menej marketingové zistenia. Tím opisuje jav nazvaný token gerrymandering, pri ktorom sa skóre určené na podporu vyváženého routovania mohlo zlepšovať, hoci skutočné rozdelenie práce medzi expertov bolo horšie. Znižovanie rýchlosti učenia expertov podľa menšieho počtu spracovaných tokenov v testovanej rodine modelov nepomohlo. Výpočty na GPU sa tiež správali rozdielne pri iných hodnotách vstupov, aj keď rozmery matíc zostali rovnaké. A prekrývanie komunikácie s výpočtom na samostatných GPU streamoch v niektorých prípadoch systém namiesto zrýchlenia spomalilo. Tieto výsledky ukazujú, že optimalizácia MoE sa nedá zredukovať na pridávanie ďalších paralelných techník.

Praktický význam pre otvorený ekosystém

Olmo-core 3 je dostupný ako zdrojový kód pod licenciou Apache 2.0 a vydanie 3.0.0 je zverejnené aj ako balík ai2-olmo-core na PyPI. Repozitár obsahuje tréningové komponenty, testy a príklady; vydanie pridáva nový DDP tréningový modul pre expertov, export hybridných MoE konfigurácií do formátu Hugging Face a ochranné kontroly pre nepodporované kombinácie. Minimálnou podporovanou verziou je podľa poznámok k vydaniu PyTorch 2.10. To znižuje bariéru pri štúdiu implementácie, no nie finančnú bariéru veľkých experimentov: výsledky z B300 klastrov nemožno automaticky preniesť na menší akademický hardvér alebo odlišné prepojenie GPU.

Pre výskumné tímy je podstatné, že otvorenosť sa neposúva iba k váham hotového modelu. Verejný tréningový stack umožňuje skúmať routovanie, paralelizmus, numerickú presnosť aj zlyhania, ktoré sa z modelovej karty finálneho checkpointu zistiť nedajú. Menšie laboratóriá môžu využiť jednotlivé komponenty alebo overovať správanie na menších konfiguráciách, hoci biliónové experimenty zostanú dostupné iba organizáciám s rozsiahlym výpočtovým zázemím. Pre európske akademické a verejné iniciatívy môže byť zaujímavá najmä možnosť auditovať celý tréningový mechanizmus a prispôsobiť ho vlastným dátam či infraštruktúre bez závislosti od uzavretej služby.

Otvorenou otázkou zostáva, ako sa systém osvedčí pri viacmesačnom tréningu, pri prirodzene nevyváženom routovaní a na hardvéri mimo najnovšej generácie NVIDIA. Zverejnené merania zatiaľ pochádzajú od autorov a sú primárne systémové; nehovoria, či budúci MoE model Olmo prinesie lepšiu kvalitu, stabilitu alebo pomer ceny a výkonu než porovnateľné husté modely. Olmo-core 3 je preto presnejšie vnímať ako overiteľný technologický základ než ako hotový prielom v schopnostiach modelov. Jeho význam sa ukáže až po reprodukcii benchmarkov a po vydaní modelu, ktorý tento stack využije v plnom tréningu.