AWS zrýchlilo tréning MoE modelov: EFA a DeepEP zvýšili priepustnosť o 40 %

AWS opísalo infraštruktúru na škálovanie posilňovacieho učenia modelov typu Mixture-of-Experts, ktorá kombinuje Amazon EKS, sieťové adaptéry Elastic Fabric Adapter, objektové úložisko Amazon S3 a open-source komunikačnú knižnicu DeepEP. Firma uvádza, že optimalizované prepojenie týchto vrstiev zvýšilo v jej teste súhrnnú priepustnosť generovania tréningových odpovedí o 40 %. Výsledok sa týka konkrétnej architektúry a pracovnej záťaže, nejde preto o všeobecne garantované zrýchlenie každého MoE modelu. Podstatná je však ukážka, že pri tréningu veľkých riedkych modelov už výkon neurčujú iba GPU. Rovnako dôležité sú komunikácia medzi uzlami, rozdelenie expertov a koordinácia medzi generovaním dát a aktualizáciou modelu.

Prečo sa riedke modely opierajú o sieť

Architektúra Mixture-of-Experts rozdeľuje časť neurónovej siete na skupinu špecializovaných expertov. Smerovací mechanizmus pre každý token aktivuje iba vybranú podmnožinu z nich. Model tak môže mať veľmi veľký celkový počet parametrov bez toho, aby pri každom kroku vykonával výpočet nad všetkými parametrami. Táto riedkosť pomáha obmedziť výpočtové náklady inferencie, no zároveň vytvára intenzívnu komunikáciu: tokeny sa musia presúvať k expertom umiestneným na rôznych akcelerátoroch a ich výstupy následne spojiť. Pri expert parallelism preto vzniká dynamická komunikácia typu all-to-all, ktorá je menej pravidelná než komunikácia pri dátovom, tenzorovom alebo pipeline paralelizme.

Problém sa prehlbuje pri post-tréningu pomocou RLHF, PPO alebo GRPO. Systém súčasne generuje veľké množstvo odpovedí, vyhodnocuje ich cez odmeňovacie modely či verifikátory a aktualizuje trénovanú politiku. PPO spravidla používa aj samostatný critic model, zatiaľ čo GRPO porovnáva relatívne odmeny v skupine a samostatný critic nepotrebuje. Z pohľadu infraštruktúry však oba prístupy kombinujú pružnú inferenciu s tesne previazaným distribuovaným tréningom. Ak sú tréningové kroky pomalé, čakajú generátory odpovedí. Ak generovanie nestíha, zostávajú tréningové GPU nevyužité. Optimalizácia jednej časti preto nemusí zlepšiť výkon celého systému.

EFA a DeepEP riešia inú časť toho istého úzkeho miesta

Elastic Fabric Adapter je sieťové rozhranie AWS určené pre tesne previazané HPC a AI úlohy. Podporuje nízkolatenčnú komunikáciu medzi podporovanými inštanciami EC2 a obchádza časť režijných nákladov tradičného sieťového zásobníka. EFA prevádzka však nemôže prechádzať medzi zónami dostupnosti ani medzi virtuálnymi sieťami VPC, čo priamo ovplyvňuje návrh klastra. Samotné zapnutie EFA navyše nestačí: aplikácia a komunikačné knižnice musia vedieť dostupnú sieťovú cestu efektívne využiť. AWS preto nasadilo EFA ako transportnú vrstvu pod komunikáciou expertov, nie ako izolovanú optimalizáciu.

DeepEP dopĺňa túto vrstvu o operácie určené špeciálne pre expert parallelism. Open-source projekt od DeepSeeku poskytuje vysoko priepustné aj nízkolatenčné GPU jadrá pre odosielanie tokenov expertom a následné spájanie výsledkov. Podporuje nízku presnosť vrátane FP8 a svoje jadrá kompiluje za behu. Aktuálna generácia knižnice zjednocuje komunikačné režimy v rozhraní ElasticBuffer a deklaruje podporu veľkých domén expert parallelism až po EP2048. Tieto vlastnosti nevytvárajú lepší model samy osebe; ich úlohou je znížiť čas, počas ktorého akcelerátory čakajú na presun tokenov medzi expertmi.

AWS spojilo DeepEP s EFA v prostredí Amazon EKS, kde Kubernetes koordinuje jednotlivé časti RL pipeline. Kontajnerová orchestrace umožňuje samostatne škálovať pracovníkov pre rollouty, tréning a ďalšie pomocné úlohy. S3 slúži ako spoločná vrstva pre checkpointy a výmenu artefaktov. Takéto rozdelenie je dôležité, pretože generovanie odpovedí a aktualizácia politiky nemajú rovnaký profil spotreby. Rollouty môžu potrebovať viac pružne pridávanej inferenčnej kapacity, kým distribuovaný tréning vyžaduje stabilnú skupinu uzlov s predvídateľnou a rýchlou komunikáciou.

Praktický dopad pre tímy trénujúce vlastné modely

Uvádzaných 40 % predstavuje zvýšenie súhrnnej priepustnosti rolloutov v meraní AWS po nasadení optimalizovaného komunikačného riešenia. Neznamená to automaticky o 40 % kratší celý tréning ani rovnakú úsporu nákladov. Koncový výsledok závisí od pomeru času stráveného generovaním, komunikáciou, výpočtom odmien, aktualizáciou politiky a ukladaním checkpointov. Ak je konkrétny systém limitovaný procesorom, úložiskom alebo verifikátorom, zrýchlenie siete sa môže v celkovom čase prejaviť menej. Meranie je preto vhodné čítať ako dôkaz významu komunikácie v jednej škálovanej konfigurácii, nie ako univerzálny benchmark EKS či DeepEP.

Pre infraštruktúrne tímy je dôležitejší samotný spôsob návrhu. AWS odporúča škálovať subsystémy nezávisle a vyvažovať výpočtový výkon, pamäť aj sieťovú priepustnosť spoločne. Pružnejšie časti pipeline možno prevádzkovať aj na Spot inštanciách, ak systém zvláda ich prerušenie, zatiaľ čo synchronizovaná tréningová skupina potrebuje opatrnejšiu stratégiu. Kubernetes môže zjednodušiť plánovanie a obnovu pracovných úloh, no nepravidelná komunikácia expertov stále vyžaduje vhodné rozmiestnenie podov, kompatibilné typy inštancií, správne sieťové rozhrania a sledovanie výkonu na úrovni jednotlivých uzlov.

Zaujímavý je aj vzťah medzi cloudovou službou a otvoreným softvérom. DeepEP je dostupné pod licenciou MIT a jeho základné komunikačné mechanizmy nie sú viazané výhradne na AWS. Integrácia s EFA je však výsledkom práce na konkrétnom transportnom a klastrovom prostredí. Organizácia, ktorá chce architektúru zopakovať, musí overiť kompatibilitu svojich GPU, ovládačov, verzie NCCL, sieťovej konfigurácie a tréningového frameworku. Prenos receptu medzi modelmi tiež nemusí byť priamy, pretože počet expertov, spôsob smerovania tokenov a nerovnomerné zaťaženie expertov menia objem aj charakter komunikácie.

Čo zatiaľ nevieme

AWS vo verejnom zhrnutí nedáva dostatok údajov na nezávislé porovnanie ceny za jeden tréningový token ani na oddelenie prínosu EFA od prínosu DeepEP. Na posúdenie prenositeľnosti výsledku by boli potrebné kompletné parametre modelu, počet a typ akcelerátorov, veľkosť dávok, topológia klastra, nastavenie expert parallelism a základná konfigurácia, voči ktorej vzniklo 40-percentné zlepšenie. Nie je tiež jasné, ako sa architektúra správa pri dlhodobom tréningu s výpadkami uzlov, pri výrazne nevyváženom smerovaní tokenov alebo pri súbehu viacerých experimentov v jednom klastri.

Napriek týmto neistotám výsledok vystihuje posun v infraštruktúre veľkých modelov. Pri hustých modeloch sa pozornosť často sústreďovala na výkon samotných GPU a klasické kolektívne operácie. Riedke MoE modely presúvajú väčšiu časť optimalizačného problému do siete, smerovania a koordinácie heterogénnych úloh. Pre európske a slovenské tímy, ktoré si nebudujú vlastné superpočítače, môže byť takýto cloudový návrh cestou k experimentom vo väčšom meradle. Pred produkčným nasadením však potrebujú vlastný benchmark založený na čase dokončenia úlohy, využití akcelerátorov, odolnosti a celkových nákladoch, nie iba na najvyššej nameranej priepustnosti.