Olmo-core 3 má škálovat trénink expertních modelů až na biliony parametrů

Více expertů bez výrazného zpomalení

Olmo-core 3 je nová verze otevřeného frameworku, který Ai2 používá při vývoji jazykových modelů Olmo. Jeho hlavní novinkou je přepracovaný systém pro trénování architektury mixture of experts, zkráceně MoE. Tyto modely obsahují specializované části zvané experti, ale pro zpracování jednotlivých tokenů aktivují jen jejich vybraný podíl. Díky tomu mohou mít vysokou celkovou kapacitu, aniž by se při každém kroku počítalo se všemi parametry.

Rozšiřování MoE však přináší vlastní náklady: experty je třeba uložit do paměti GPU a data mezi nimi efektivně přesouvat. Ai2 uvádí test, v němž zvýšilo počet expertů z 8 na 128, přičemž pro každý token dál vybíralo čtyři. Počet aktivních parametrů na token zůstal přibližně 3,2 miliardy, celková kapacita modelu ale vzrostla ze 4,6 na 47 miliard parametrů. Propustnost trénování přitom klesla o méně než 5 %.

Nové rozdělení práce mezi GPU

Dřívější implementace MoE v Olmo-core využívala plně shardované paralelní zpracování dat (FSDP). Při každé menší dávce tréninkových dat shromažďovala váhy modelu a následně je znovu rozdělovala. Olmo-core 3 přechází na distribuované paralelní zpracování dat (DDP), které ponechává experty na GPU a posílá k nim příslušná data. Tím se vyhne opakovanému přesouvání vah.

V předběžném testu na osmi GPU NVIDIA B300 zpracoval 47miliardový MoE model s novým systémem 52 000 tokenů za sekundu na každé GPU. Předchozí implementace zvládla 19 400 tokenů, takže naměřená propustnost byla přibližně 2,7krát vyšší. Výsledek se týká výkonu trénovací infrastruktury, nikoli kvality natrénovaného modelu.

Paralelismus a nižší přesnost

Olmo-core 3 kombinuje několik způsobů rozdělení výpočtů. Expertový paralelismus rozmisťuje experty mezi GPU, pipeline paralelismus rozděluje vrstvy modelu mezi skupiny GPU a distribuovaný optimalizátor rozděluje mezi zařízení také stav potřebný k aktualizaci modelu. Systém dále omezuje režii při směrování dat k expertům a slučuje jejich výpočty do větších operací, které lze na GPU provádět efektivněji.

Framework podporuje také formát MXFP8 s nižší přesností. V kontrolovaném testu na čtyřech GPU NVIDIA B300 zvýšil při vhodném použití propustnost trénování přibližně o 21 % oproti formátu BF16. Špičková aktivní paměť přitom klesla ze 103 na 95 GiB. Ai2 připisuje většinu zrychlení výpočtům dopředných vrstev a přesunu dat mezi experty, nikoli samotné pozornosti.

Testy až na bilionové modely

Ai2 testovalo Olmo-core 3 také na konfiguraci s 1,2 bilionu parametrů, z nichž bylo na jeden token aktivních 58,36 miliardy. Model běžel na 512 GPU a nejvyšší naměřený výkon dosáhl 858 TFLOP/s na GPU. Testy využívaly náhodné směrování dat, takže ověřovaly výkon systému, nikoli schopnosti natrénovaného modelu. S alternativním způsobem komunikace mezi experty DeepEP v2 organizace vyzkoušela konfiguraci s 2,38 bilionu parametrů; šlo však o krátký test kapacity, ne o úplný trénink.

Olmo-core 3 má sloužit jako základ příští generace modelů Olmo, která bude využívat architekturu MoE. Ai2 zveřejňuje infrastrukturu, aby ji další výzkumníci a vývojáři mohli použít pro vlastní modely a upravovat ji podle hardwaru či zvoleného způsobu trénování. Podrobnosti o návrhu a experimentech organizace uvádí v technické zprávě a kód nabízí na GitHubu.