Az Ai2 megnyitotta az Olmo-core 3 MoE-tréningrendszerét

A nagy MoE-modellek egyik fő problémáját célozza

Az Ai2 október 1-jén tette közzé az Olmo-core 3-at. A keretrendszer az Olmo modellek következő generációjának egyik alapvető rendszere, és az Ai2 azon törekvésének része, hogy a modellek mögötti eszközöket és tréninginfrastruktúrát is nyíltan hozzáférhetővé tegye.

A szakértőkeverékes, azaz mixture-of-experts, MoE-modellek sok tanult komponenst, vagyis paramétert tartalmazhatnak úgy, hogy egy-egy bemenet feldolgozásakor csak a szakértők egy részét használják. A teljes modellt ugyanakkor a GPU-memóriában kell tartani és frissíteni, a bemenetek megfelelő szakértőkhöz irányítása pedig kommunikációs és koordinációs költségekkel jár. Az Ai2 szerint ezek a költségek a modellek növekedésével csökkenthetik az MoE-architektúra számítási előnyét.

Az egyik benchmarkban a szakértők számát 8-ról 128-ra emelték, miközben tokenenként továbbra is csak négy szakértőt választottak ki. Az aktív paraméterek száma tokenenként nagyjából 3,2 milliárd maradt, a teljes kapacitás viszont 4,6 milliárd paraméterről 47 milliárdra nőtt. A tréning áteresztőképessége eközben kevesebb mint 5 százalékkal csökkent. Az infrastruktúrát összesen több mint egybillió paraméteres konfiguráción is tesztelték.

Új elosztási rendszer és nagyobb áteresztőképesség

Az Olmo-core korábbi MoE-megoldása teljesen particionált adatpárhuzamosságot, vagyis FSDP-t használt. Ez minden kisebb tréningköteg előtt összegyűjtötte, majd újra felosztotta a modell súlyait. Az Olmo-core 3 ehelyett elosztott adatpárhuzamosságra, DDP-re épül. A szakértők így a GPU-kon maradnak, a rendszer pedig a hozzájuk tartozó adatokat irányítja, elkerülve a súlyok ismételt összegyűjtését.

Az Ai2 előzetes tesztjében egy 47 milliárd paraméteres MoE-modell nyolc NVIDIA B300 GPU-n másodpercenként és GPU-nként 52 000 tokent dolgozott fel az új rendszerrel. A korábbi, FSDP-alapú megoldás 19 400 tokenes eredményt ért el, ami körülbelül 2,7-szeres áteresztőképességet jelent.

Az Olmo-core 3 több elosztási technikát kombinál. A szakértőpárhuzamosság a GPU-k között osztja szét a szakértőket, a csővezeték-párhuzamosság a modell rétegeit, az elosztott optimalizáló pedig a tréninghez szükséges optimalizálói állapotot. Így egyetlen GPU-nak nem kell a teljes modellt és annak teljes tréningállapotát tárolnia.

Billió paraméteres tesztek, nyílt hozzáférés

A rendszer további optimalizációkat is tartalmaz az adatok szakértőkhöz irányítására és a számítások végrehajtására. A rowwise szakértőpárhuzamosság közvetlenül a szakértők bemeneti puffereibe helyezi az adatokat, a GPU-n futó útválasztás pedig a GPU-kon tartja az útválasztási információkat. A grouped GEMM több kisebb szakértői számítást egyesít a hatékonyabb GPU-s végrehajtás érdekében.

Az MXFP8 alacsonyabb pontosságú számformátum támogatásával egy négy NVIDIA B300 GPU-n végzett tesztben a tréning áteresztőképessége körülbelül 21 százalékkal nőtt a BF16-hoz képest. A csúcsként mért aktív memóriahasználat 103 GiB-ről 95 GiB-re csökkent.

Az Ai2 512 GPU-n egy 1,2 billió paraméteres modellt is vizsgált, amelyből tokenenként 58,36 milliárd paraméter volt aktív. A legmagasabb mért érték 858 TFLOP/s/GPU lett. Ezek a tesztek véletlenszerű útválasztást használtak, ezért a rendszer teljesítményét mérték, nem egy betanított modell minőségét. Egy másik kísérletben 2,38 billió teljes paraméterig jutottak, ez azonban rövid kapacitásteszt volt, nem teljes tréningfutás.

Az Olmo-core 3 nyíltan használható: kutatók és fejlesztők saját MoE-modelljeik betanítására, eltérő hardverekhez való igazítására, valamint útválasztási és párhuzamosítási kísérletekre alkalmazhatják. Az Ai2 következő generációs Olmo-modellje szintén MoE-architektúrát használ majd, és a szervezet szerint a korábbi MoE-munkánál nagyobb léptékű tréningre épül.