Olmo-core 3-tal skálázódik a trillió paraméteres AI a kutatók számára az Allen Institute-nél

Az Allen Institute (AI2) bemutatta az Olmo-core 3-at, egy új nyílt forráskódú rendszert, amely a nagy, Mixture-of-Experts (MoE) típusú nyelvi modellek betanítását hivatott felgyorsítani és költséghatékonyabbá tenni. A fejlesztők célja, hogy a trillió paraméteres tartományba skálázódó modellek fejlesztését elérhetővé tegyék a kutatók és kisebb laborok számára is. (Hugging Face)

A rendszer a korábbi, FSDP-alapú megközelítés helyett elosztott adatpárhuzamosságra (DDP) épít, így az expert-súlyokat nem kell minden kis batch-hez újra összegyűjteni és felosztani. Egy előzetes tesztben nyolc NVIDIA B300 GPU-n egy 47 milliárd paraméteres MoE modell 52 000 tokent dolgozott fel másodpercenként, ami 2,7-szeres gyorsulást jelent a korábbi implementációhoz képest — közölte az AllenAI blog.

Az Olmo-core 3 több technikát kombinál a nagy MoE-modellek GPU-fürtökön való elosztására és optimalizálására. Az expert párhuzamosság (expert parallelism) révén az expert-ek az egyes GPU-kra kerülnek, a folyamat párhuzamosság (folyamat parallelism) pedig a modell rétegeit osztja szét, csökkentve az egyes GPU-k memóriaterhelését. A distributed optimizer az optimizer állapotát is elosztja, így nem kell minden GPU-nak a teljes másolatot tárolnia.

A rendszer további optimalizációi közé tartozik a row-wise expert parallelism, amely közvetlenül az expert bemeneti pufferébe helyezi az adatokat, valamint a GPU-resident routing, amely a routing metadatát a GPU-kon tartja. A grouped GEMM pedig a kisebb expert-számításokat egyesíti a hatékonyabb végrehajtás érdekében. Az MXFP8 formátum támogatása további 21%-os növekedést hozott a tréning sebességében, miközben a csúcs aktív memória 103 GiB-ról 95 GiB-ra esett.

A teljesítménytesztek során egy 1,2 trillió paraméteres modellt teszteltek 512 GPU-n, amely 858 TFLOP/s/GPU átviteli sebességet ért el. Kísérleteztek DeepEP v2-vel is, elérve egy 2,38 trillió paraméteres konfigurációt, ami a rendszer skálázhatóságát demonstrálja. Az AI2 szerint ezek a fejlesztések csökkentik a modellfejlesztés költségeit és energiaigényét, így az akadémiai kutatók és kisebb laborok számára is elérhetővé válik a fejlett AI-modellek fejlesztése.