Hatékony és hordozható Mixture-of-Experts kommunikáció

Bemutatunk egy nagy teljesítményű, hordozható, nyílt forráskódú könyvtárat a Mixture-of-Experts (MoE) kommunikációhoz, amely a standard All-to-All kommunikációs primitívekhez képest 10-szer gyorsabb teljesítményt ér el.

Megvalósításunk számos kulcsfontosságú technikai innovációt tartalmaz, amelyek kiemelkedő MoE kommunikációs hatékonyságot biztosítanak:

GPU által kezdeményezett kommunikáció (IBGDA): Támogatja a közvetlen GPU–NIC kommunikációt, jelentősen csökkentve a késleltetést a CPU bevonásának megkerülésével

Kommunikáció és számítás átfedése: A különálló küldési és fogadási szakaszokkal rendelkező osztott kernel architektúra lehetővé teszi, hogy a számítás a hálózati átvitelekkel párhuzamosan haladjon

Leggyorsabb egycsomópontos teljesítmény: 2,5-ször alacsonyabb késleltetés, mint az egycsomópontos konfigurációkban eddig leggyorsabb megvalósításé

Hatékony és hordozható többfelhasználós vagy töbcsomópontos teljesítmény: Megvalósításunk akár 10-szer gyorsabb sebességet ér el, mint a standard all-to-all kommunikáció. Bár körülbelül 2-szer lassabb, mint a rendkívül specializált megvalósítások, a mi megközelítésünk jobb hordozhatóságot biztosít az NVSHMEM verziók és hálózati környezetek (NVLink, CX-7 és EFA) között

A könyvtár teljes mértékben nyílt forráskódú, és elérhető a https://github.com/ppl-ai/pplx-kernels címen.

Ebben a cikkben feltárjuk a szakértői párhuzamosság kihívásait a nagyméretű MoE modellekben, és bemutatjuk a hatékony token-diszpécselési és -kombinációs megközelítésünket a elosztott GPU-környezetekben.

A Mixture-of-Experts (MoE) modellek, mint például a DeepSeek R1 és a Mixtral 8x7B, a sűrű (dense) modelleket fejlesztik tovább azáltal, hogy korlátozzák az egyes tokenekhez aktivált súlyok számát. Ahelyett, hogy magas paraméterszámot érnének el a rétegek számának növelésével vagy a lineáris rétegek méretének növelésével az egyes dekóderrétegek többrétegű perceptronjában (MLP), a MoE modellek a hagyományos MLP-t több szakértővel és egy routerrel váltják fel. Például a DeepSeek R1 671 milliárd paraméteréből csak 37 milliárd szorzódik meg egy adott tokennel a következtetés (inference) során. Ez jelentősen csökkenti a token dekódolásához szükséges számítási mennyiséget, és a sűrű modellekhez képest javítja a késleltetést.

A MoE modellek a sűrű modellekhez képest néhány további kihívást jelentnek a következtetés során. Míg maguk a szakértők kis MLP rétegek, minden dekóderréteg tartalmaz egy routert, amely eldönti, hogy egy tokont melyik szakértőkhöz küldjön el, és minden token több szakértőhöz is elküldésre kerül. A router általában egy kis lineáris réteg, amely valószínűségi eloszlást hoz létre. Általában a top-K pontszámú szakértők kerülnek kiválasztásra, és a végső aktiválás súlyozott átlagként kerül kiszámításra, összegezve a szakértői kimeneteket, megorozva a router valószínűségi eloszlásából származó súllyal.

Az elosztott MoE rendszerek késleltetésének minimalizálása érdekében a párhuzamosság kihasználható több eszközön keresztül, de ez kommunikációs kihívásokat hoz magával. Az olyan modellek, mint a Mixtral 8x7B vagy a Llama-70B, 8 eszközön férnek el egyetlen csomóponton belül, kihasználva a gyors és alacsony késleltetésű NVLink összeköttetéseket (akár 900 Gbps). A nagyobb modellek azonban megkövetelik a szakértők shardingolását (több csomópontra osztását) InfiniBand használatával (amely 400 Gbps-on tetőzik), ami további késleltetési kihívásokat rejt magában.

Tokenirányítás, szakértői számítás és kimenetgyűjtés egy mixture-of-experts rétegben.

Ebben a cikkben feltárjuk a szakértői párhuzamosság problémáját, leírva azon GPU kernelek megvalósítását, amelyek hatékonyan küldik el a tokeneket a hozzárendelt szakértőket tartalmazó eszközökre, és gyűjtik össze őket a szakértői kimenetek súlyozott összegének kiszámításával. Noha léteznek agresszíven optimalizált, specifikus hálózati hardvereket célzó megvalósítások, mi egy rendkívül hordozható megvalósítást mutatunk be, amely az NVSHMEM primitívek minimális készletére támaszkodik, mégis a standard all-to-all kommunikációnál 10-szer gyorsabb teljesítményt nyújt. Megvalósításunk csúcsteljesítményt ér el egycsomópontos konfigurációkon, miközben kiváló hordozhatóságot biztosít a különböző hálózati környezetekben.

A hatékony következtetés érdekében a modell súlyait az eszköz memóriájában kell tartani, miközben elegendő tér marad az aktiválások, a KV-gyorsítótárak (cache) és a modell előremenő (forward) futtatásához szükséges egyéb pufferek számára is. A legképzettebb modellek még a legfejlettebb GPU-k kapacitását is meghaladják, így a következtetést meg kell osztani több eszköz között, amelyek együttesen képesek tárolni a modell súlyait. A súlyok shardingolási sémái alapján különböző kommunikációs és számítási sémákat kell alkalmazni a számítások kiszervezésére és az eszközök szinkronizálására.

A szakértői párhuzamosság (Expert Parallelism – EP), amelyet az első ábra mutat be, csak a szakértői számításokat párhuzamosítja. Különböző szakértők vannak hozzárendelve a különböző eszközökhöz, amelyek a súlyaikat tárolják. Az útválasztás (routing) után a tokenek a megfelelő eszközre kerülnek, majd az eredmények összegyűjtésre és akkumulációra kerülnek. Az útválasztás összeteksége a modell többi részének párhuzamossági fokától függ: más rétegek replikálása teljesen feleslegessé teheti az útválasztást, mivel az egyes rangok (rank) kiválaszthatják a tokeneket egy helyben replikált útválasztási táblázatból. Ha azonban csak az egyik rang futtatja az útválasztást, egy broadcast (szórás) szükséges a tokenek, indexek és súlyok elküldéséhez a saját szakértőiknek. Végül egy all-gather vagy all-to-all broadcast szinkronizálja a kimeneti tokeneket azzal a ranggal, amely folytatja a modell végrehajtását. Egy ilyen megvalósítás viszonylag egyszerű, mivel a torch már elérhetővé teszi a szükséges primitíveket, bár némi előny származhatna az akkumulációnak a szakértőktől a tokeneket gyűjtő gather művelettel való összevonásából.

Szakértői párhuzamosság és tenzorpárhuzamosság kombinálása két GPU-n.

A kizárólag szakértői alapú párhuzamosság nem skálázódik ideálisan, mivel a klaszter csomópontjai tétlenek lehetnek, amíg a modell nem szakértői rétegeket, például figyelmet (attention), normalizálást (norm) és mintavételezést futtat. Ezen rétegek – elsősorban a figyelem – számítása azonban szintén profitálhat a tenzorpárhuzamosságból (Tensor Parallelism – TP). A legtöbb modell többfejű figyelemre (multi-head attention) támaszkodik, ami azt jelenti, hogy a figyelmi fejek és a hozzájuk tartozó Q, K és V projekciók szintén szétoszthatók az eszközök között, replikálva vagy összegyűjtve a szeleteket a különböző rétegek között. Ha a figyelem szét van osztva az összes eszköz között, egy all-gather szinkronizálhatja az aktiválásokat, lehetővé téve az útválasztás replikálását, ami a hatékony megvalósításhoz a szakértői párhuzamossági esethez hasonló szinkronizációs primitíveket igényel. Az ezen a szinten történő párhuzamosságnak azonban megvannak a korlátai, mivel a figyelmi fejek számának egy bizonyos küszöbérték alá csökkentése csökkenő hozamot eredményez.

Szakértői párhuzamosság és adatszintű párhuzamosság kombinálása külön kérésekhez két GPU-n.

Az összes eszköz legjobb kihasználása és a szakértői párhuzamosság nagyon magas fokának – akár 128 vagy 256 GPU – támogatása érdekében adatpárhuzamosságra (Data Parallelism – DP) van szükség. E séma szerint azok az eszközök, amelyekre a modell szét van osztva, csoportokba rendeződnek, hogy egyidejűleg kezeljék a kéréseket, kiszámítva a figyelmet és fenntartva a helyi csoportjukon szétosztott KV-gyorsítótárakat. E párhuzamos csoportok több példánya működik együtt a szakértői értékelésben, mindegyik a szakértői súlyok egy-egy eltérő részhalmazát tárolja. A figyelmi fejek számától függően egy csoport tipikusan egy teljes csomópont méretét is elérheti, mivel a csomóponton belüli kommunikáció gyorsabb. Például a fenti ábrán az egyik DP rang függetlenül szolgál ki két kérést, az A-t és a B-t, kezelve a figyelmet, a normát és a kérések minden egyéb könyvelését. A másik DP rang egy különálló C kérést dolgoz fel. Az első csomópont azonban a szakértők felét tárolja, míg a másik csomópont a másik felét, így az A-ból és B-ből származó tokenek útválasztása után előfordulhat, hogy a második csomópontra küldődnek, és fordítva. Ez ritka (sparse) kommunikációs problémához vezet: előfordulhat, hogy minden egyes eszköz eltérő számú tokent küld bármely más célrangra. A torch meglévő primitívei, elsősorban az all_to_all, nem különösebben alkalmasak erre, mivel valamilyen kitöltést (padding) vagy GPU–CPU szinkronizációt és metaadat-broadcastot igényelnek. A kommunikáció hatékony megvalósításához egyedi kernelekre van szükség az útválasztási információk feldolgozásához és a kommunikáció eszközöktől történő, CPU beavatkozása nélküli kezdeményezéséhez. Az útválasztás után egy diszpécserkernelnek kell elküldenie a tokeneket azoknak a rangoknak, amelyekre útválasztódtak, míg a kombinációs oldalon összegyűjtésre kell kerülniük az aktuális DP csoportban lévő kérésekhez tartozó aktiválásoknak. Ezenkívül a munkát egyensúlyba kell hozni egy több eszközzel rendelkező DP csoporton belül, annak érdekében, hogy minimalizáljuk a broadcastot, és helyesen replikáljuk a tokeneket az összes rangon a következő rétegek számára.

Az NVSHMEM egy NVIDIA-specifikus OpenSHMEM megvalósítás, amely hordozható eszközök közötti kommunikációs lehetőségeket biztosít, elvonatkoztatva az alapul szolgáló hardver összetettégétől. Az API képes kifejezni az eszközök közötti olvasásokat és írásokat, amelyek leképeződnek az egyes átviteli rétegek primitíveire. Mind az NVLink, mind az RDMA támogatott, ami bizonyos mértékű hordozhatóságot biztosít, és lehetővé teszi, hogy az NVSHMEM kernelek egy csomóponton belül vagy több csomópontból álló klaszteren keresztül működjenek. Kerneleinkben távoli memóriaírásokat és jelzési (signaling) műveleteket használunk az adatok továbbítására és a részt vevő eszközök szinkronizálására.

Az NVSHMEM műveletek a szimmetrikus memória fogalmára épülnek: olyan puffereken operálnak, amelyek az eszközök közötti kommunikációban részt vevő összes eszközön lefoglalásra kerültek. Egy helyi eszköz a saját helyi pufferén belül származtatott címet használhat a célrang indexe mellett egy művelet célhelyének megadására. Az alábbi ábra ezt a fogalmat szemlélteti: mindkét GPU azonos méretű szimmetrikus puffereket foglal le, és src mutatókat tart fenn rájuk. Az első rang 3 egész számot szeretne elküldeni a másodiknak, a puffer elejére helyezve őket: a nvshmem_int_put_nbi a helyi pufferből származtatja a kezdőcímet, megadva a céleszközt. A második rang a saját pufferéből származtat egy eltolást (offsetet), egy elemet küldve az első eszköznek, egy eltolással. Míg a célcímeknek mindig nvshmem_alloc használatával lefoglalt szimmetrikus puffereknek kell lenniük, a forráspufferek az eszközmemória tetszőleges régiói lehetnek, feltéve, hogy előre regisztrálták őket az NVSHMEM-mel.

Míg az NVSHMEM a primitívek széles skáláját biztosítja, kerneleink mindössze 3 függvényre támaszkodnak, és minden szinkronizációt és korlátozást (fencing) ezekre építenek.

Szimmetrikus memória és távoli írások két GPU között NVSHMEM használatával.

nvshmemx_putmem_signal_nbi_warp: Adatblokkot továbbít az egyik eszközről a másikra, miközben beállít egy jelzőt (flaget) a távoli eszközön. A művelet vagy beállít (NVSHMEM_SIGNAL_SET), vagy megnövel (NVSHMEM_SIGNAL_ADD) egy 64 bites helyet. A jelző a teljes memóriablokk átvitele után frissül. Ha a távoli eszköz a jelző megváltozását észleli, biztonságosan hozzáférhet a saját memóriájában lévő pufferhez. Ez a függvény hasznos az adatátvitel és a szinkronizáció összevonásához.

A nvshmemx_signal_op egyetlen memóriarégión – tipikusan egy 64 bites jelzőn – operál, atomi módon beállítva vagy növelve azt. Hasznos metaadatok átvitelében és eszközök szinkronizálásában.

Az nvshmem_uint64_wait_until a jelzés fogadó oldalán kerül felhasználásra egy jelző leképezésére (polling), amíg a távoli fél azt nem frissíti.

A műveletek közötti bármilyen sorrendiség biztosítása az nvshmem_quiet vagy nvshmem_fence révén költséges, ezért a kernelek megvalósítása elkerüli minden más NVSHMEM primitív használatát, és az egyes küldés-fogadás párokon kívül minden függőséget minimalizál. Adatok küldésekor mindig a függvények nem blokkoló verzióját használjuk anélkül, hogy megvárnánk, amíg az adatok egyáltalán elhagyják a helyi rangot. A kernelek csoportján belüli egyéb, implicit szinkronizáció biztosítja, hogy ezek a pufferek ne íródjanak felül romboló módon, amíg a NIC-eknek hozzáférésre van szükségük hozzájuk.

Az NVSHMEM GPUDirect RDMA-t használ, így a NIC közvetlenül tudja olvasni és írni a GPU memóriát az adatcseréhez a CPU bevonása nélkül. Továbbá a ConnectX NIC-eken az NVSHMEM támogatja a GPU által kezdeményezett kommunikációt (más néven Infiniband GPUDirect Async vagy IBGDA), amely lehetővé teszi a GPU számára, hogy RDMA műveleteket közvetlenül a NIC-hez küldjön, és közvetlenül kérdezze le a NIC befejezési sorát. Azokon a platformokon, amelyek nem támogatják a GPU által kezdeményezett kommunikációt, az NVSHMEM egy „proxy” szálat indít a gazdagép (host) CPU-ján, hogy a GPU nevében kezdeményezze a kommunikációt és lekérdezze a befejezést. Az NVSHMEM program függetlenül attól, hogy a GPU által kezdeményezett kommunikáció támogatott-e vagy sem, hordozható. A GPU által azonban kezdeményezett kommunikáció jelentősen csökkenti a késleltetést, mivel teljesen megkerüli a CPU-hoz vezető detourt.

A MoE kommunikációt egy diszpécser- és egy kombinációs kernelekből álló páron keresztül valósítjuk meg. A diszpécserkernelek felelősek a tokenek és az útválasztási információk olvasásáért az egyes rangokon, és azok megfelelő szakértőkhöz történő elküldéséért. A kombinációs kernelek összegyűjtik a szakértők által előállított aktiválásokat, és visszaküldik azokat a forrásrangjaikra, miközben kiszámítják a súlyozott átlagot a kiválasztott szakértőktől a router által kiszámított súlyok alapján. A kernelek ráadásul egy küldési és egy fogadási komponensre oszlanak, hogy az adatátvitelek átfedésbe hozhatók legyenek a számításokkal. A küldési kernelek mind nem blokkolók és nem szinkronizálók: egyszerűen az összes írást elküldik a távoli feleknek. Ezzel szemben a kombinációs oldalon a kernelek csak a memóriából olvasnak, megvárva, amíg az összes szükséges adat átvitelre kerül. Miután elküldték a munkát a NIC-eknek, amíg az adatok aszinkron módon áramolnak a vezetéken, a GPU-k egyéb hasznos helyi munkát végezhetnek, például alkalmazhatják a megosztott szakértőket vagy kiszámíthatják a figyelmet.

A kernelek minden párja saját szimmetrikus memóriapufferekkel rendelkezik, és minden egyes rang privát tárolót foglal le az adatok fogadásához az összes többi rangtól. A DP csoportból legfeljebb elküldhető tokenek maximális száma rögzítve van a rangok között, ami egyben beállítja azt a felső határt is, amit az egyes rangok fogadhatnak az egyes helyi szakértőkhöz az egyes DP rangoktól. Ez lehetővé teszi a küldő rangok számára, hogy egyedi címet származtassanak a célrangon az íráshoz anélkül, hogy bármilyen szinkronizációra lenne szükség közöttük. Az adatok fogadása után a kernelek megkeverik a tokeneket, hogy a fogadó kernelek által megkövetelt optimális formátumban rendezzék el őket a memóriában. Míg a pufferek jelentős méretűek, a modell összes szekvenciális rétegén keresztül újra felhasználásra kerülnek.

Külön küldő és fogadó kernel a dispatch és combine műveletekhez.

A globális szinkronizáció egyetlen formája – amint azt a fenti ábra mutatja – a kombinációs-fogadási (combine-receive) kernelben valósul meg. Amint az adatok elküldésre kerülnek a küldéstől, egy rang visszatér, és folytatja a végrehajtást, amikor megkapja az összes szükséges tokent a többi rangtól, és belép a kombinációs-küldési kernelekbe, amint azok feldolgozásra kerülnek. A kombinációs-fogadási gát (barrier) biztosítja, hogy egyetlen rang sem haladhat előre, és nem kezdheti meg a diszpécser-küldést (dispatch-send), amíg bármely más rang még a adatok fogadására vár, ami potenciálisan romboló átfedést okozna. A szinkronizáció az egyszerűség kedvéért a kombinációs kernelben történik; optimalizálásként ugyanez a gát megvalósítható lenne a diszpécser-fogadás és a kombinációs-küldés kernelek között is, tovább csökkentve a késleltetést azzal, hogy több időt hagy a NIC-eknek az átvitelek letelepedésére, miközben a GPU más kerneleket futtat.

Mindkét kernel megoszlik az eszközök összes elérhető SM-je (streaming multiprocessor) között: míg a diszpécser-küldési és kombinációs-fogadási kerneleknek rangonkénti maximális tokenszám (max_m) szerint kell párhuzamosítaniuk, a diszpécser-fogadási és kombinációs-küldési kerneleknek ennyi tokent kell kezelniük minden egyes rangról és helyi szakértőtől, legfeljebb max_m * (num_experts // EP) * (EP // TP) méretig.

A diszpécserkernel felelős a tokenek és blokk-skálázási tényezők elküldéséért a társaknak, amelyek azokat a szakértőket üzemeltetik, amelyekhez a tokenek útválasztásra kerültek. A küldő oldal warp-specializációra támaszkodik két feladat párhuzamosításához: az útválasztási információk aggregálására annak meghatározására, hogy hány token kerül elküldésre az egyes rangokról, valamint a tokenek üzenetekbe csomagolására a távoli félhez történő elküldéshez. A fogadó oldal először megvárja az összes tokenszám megérkezését, majd beolvassa és kicsomagolja a megosztott memóriából a hasznos terheket (payloads) azokba a tenzorokba, amelyek a lenti Csoportos GEMM kernelekhez kerülnek továbbításra. Olyan információkat is tárol a kombinációs kernelekkel megosztott pufferekben, amelyek jelzik, hogy az egyes tokeneket hová kell visszaküldeni. Ez az információ azért szükséges, mert a fogadó kernel nem determinisztikus sorrendben keveri meg a tokeneket a folytonos pufferekben.

Tokenek továbbítása és érkezésük jelzése a forrás- és cél rankek között.

A küldő részben minden blokkon egyetlen, 32 szálból álló warp felelős az útválasztási információk olvasásáért és az egyes célcsoport-szakértők által ebből a rangból fogadni várt tokenek számának megszámlálásáért, elküldve a darabszámot plusz egyet az nvshmem_signal_op használatával. A darabszám egyel növekszik, mivel a nulláról a nem nullára való áttérés a távoli oldalon jelzi a forrásrangtól származó darabszámok megérkezését. Ezzel párhuzamosan a megmaradt warpok együttműködnek a tokeneknek a szimmetrikus memóriába történő másolásában az összes blokkon párhuzamosan, egyetlen folytonos darabba csomagolva az aktiválásokat, a skálázási tényezőket és a helyi rangon lévő indexüket. Az indexre a kombinációs küldőnek van szüksége annak a címnek a meghatározásához, ahová a token írásra kerül. Ezután, miután meggyőződtünk arról, hogy az összes adat átmásolásra került a warp-csoportokon átívelő gát révén, a warpok ismét függetlenül működnek, mindegyik ugyanazt a puffert küldi el egy másik szakértőnek párhuzamosan. A tokenek küldése az nvshmemx_putmem_signal_nbi_warp használatával történik, amely egyben atomi módon növeli a helyi rangról elküldött tokenek számát a távoli eszközön. Egy DP csoporton belül, mivel minden rang rendelkezik az elküldendő token egy replikájával, a diszpécserelés egyenletesen oszlik el, és minden eszköz a tokenek egy részhalmazát küldi ki.

A fogadó oldalon az összes kernel először szinkronizál a küldő rangokkal az összesített tokenszámok átérkezésének megvárásával. Ezt követően mindannyian megvárják, amíg az atomi módon növelt elküldött tokenszámok leülepednek a teljes darabszámra, jelezve, hogy a forrásrangoktól származó összes hasznos teher is át lett küldve a putmem hívás szemantikájának köszönhetően. A kernelek az nvshmem_uint64_wait_until használatával kérdezik le a számlálókat, párhuzamosítva a műveletet az összes blokkon és szállal. Ezt követően egy blokkok közötti gát biztosítja, hogy egyetlen blokk se olvassa a puffereket, amíg az összes adat helyesen meg nem érkezett. A blokkok között elosztva és egy atomi módon növelt számláló segítségével szinkronizálva a tokenek a szimmetrikus memóriából a szokásos folytonos pufferekbe kerülnek átmásolásra, amelyek a szakértőket megvalósító lenti kernelek bemeneteivé válnak. A forrásrang, a szakértői index és a tokenindex külön kerül tárolásra, pontosan megjelölve azt a helyet, ahová a kombinációs kernelnek küldenie kell az aktiválásokat. Mociaż a DP csoportból származó tokenek különböző eszközökről érkeznek, mindannyian össze vannak csoportosítva, hogy továbbíthatók legyenek a megfelelő szakértőnek.

A kombinációs kernelek ismét küldési és fogadási felekre oszlanak: a küldők átmásolják a nem kvantált 16 bites aktiválásokat a távoli eszközökre, míg a fogadó kernelek megvárják az adatok átküldését, és helyben kiszámítják az összes szakértői hozzájárulás súlyozott átlagát. Emellett gátként is szolgálnak a diszpécser-kombináció sorrendjének szinkronizálásához: minden rang beállít egy jelzőt minden társon a küldési kernelekbe való belépéskor, és a fogadó kernelek nem térhetnek vissza, amíg nem észlelik a jelző beállítását. A szinkronizáció késleltetése minimális, mivel átfedésben van a tényleges kommunikációval és számítással.

A szakértői kimenetek egyesítése a tokenekbe minden adatszintű párhuzamos ranken.

A küldő oldalon a kernelek párhuzamosan bejárják az összes helyi szakértőhöz rendelt tokenek listáját, és beírják azokat a célrangon lévő pufferbe. A célrangot, a szakértői indexet és a tokenindexet a tokenek megérkezésekor a szórási (scatter) kernelek által feltöltött, tokenenkénti pufferekből olvassák. Minden küldőnek saját privát memóriarégiója van szakértőként az íráshoz, amint azt a fenti ábra mutatja, elkerülve a szinkronizáció szükségességét. A diszpécsereléshez hasonlóan a kombináció is atomi módon növeli a tokenenkénti számlálókat a célrangon az adatok megérzésének jelzésére: amikor a számláló megegyezik azon szakértők számával, amelyekhez a token elküldésre került, a token tartalma elérhetővé válik.

A fogadó kernelben a tokenek listája több blokkon keresztül párhuzamosan van bejárva, megvárva a tartalmuk megérkezését a jelzési művelet által beállított jelző lekérdezésével. Megérkezéskor a hasznos terhek beolvasásra kerülnek a privát pufferekből, az útválasztási táblázat jelzi, hogy melyik puffert kell olvasni, és milyen súlyt kell rendelni az egyes szakértőkhöz. Az eredmények ezután külsőleg lefoglalt tenzorokba íródnak, és a kernel a végrehajtás végét jelzi, miután az összes eszköz átlépte a gátat. Egy DP csoporton belül az összes rang megkapja az egyes szakértői aktiválások egy-egy példányát a saját replikáik kiszámításához.

Kerneleinket egy 128 GPU-ból álló klaszteren értékeljük, amelyek InfiniBand és CX-7 NIC-ek útján összekapcsolt 16 csomóponton oszlanak meg. Kipróbáljuk a GPUDirect Async-ot (IBGDA) és a Megbízható Kapcsolatot (Reliable Connection – RC) is egy CPU proxyval. Ezeket összehasonlítjuk a rendkívül optimalizált DeepSeek megvalósítással, valamint a PyTorch (a NCCL-en keresztül) vagy az NVSHMEM által biztosított sűrű primitívekkel.

Kommunikációs késleltetés 128 token/GPU esetén többcsomópontos EP128 konfigurációban.

Bár a sűrű NVSHMEM all-to-all műveleten az IBRC és az IBGDA teljesítménye hasonló (6378 µs vs 6180 µs), az IBGDA lényegesen gyorsabb a ritka kernelekkel. Összeadva a Diszpécsert és a Kombinációt, az IBGDA 902 µs-ot használ, míg az IBRC 3223 µs-ot vesz igénybe – ez 3,6-szoros késleltetési javulás. Míg az all-to-all sávszélesség-korlátos, a ritka kernelek nagyságrenddel kevesebb adatot szórnak szét, így késleltetés-korlátosak. A hálózati átvitelek közvetlenül a GPU-ról történő indításával, anélkül, hogy CPU proxyra lenne szükség a GPU és a NIC koordinálásához, a végpontok közötti késleltetés jelentősen csökken.

Míg a hordozható kerneleink körülbelül kétszer olyan lassúak, mint a rendkívül optimalizált DeepSeek kernelek, a sűrű kernelekhez képest 10-szeresére javítják a késleltetést (902 µs vs 9944 µs).

Emellett a küldő-fogadó komponensekre való felosztás azt is lehetővé teszi, hogy a késleltetés egy része elrejthető legyen, ellentétben a könyvtári primitívekkel.

Kommunikációs késleltetés 128 token/GPU esetén egynode-os EP8 konfigurációban.

Egycsomópontos (EP8) környezetben az NVSHMEM az NVLINK-et használja az átvitelhez, alacsonyabb késleltetést és magasabb sávszélességet biztosítva, mint a csomópontok közötti hálózatépítés. Hordozható kerneleink körülbelül 2,5-szer gyorsabbak, mint a DeepEP egycsomópontos környezetben (186 µs vs 481 µs).

Az itt leírt kernelek felülmúlják a gépi tanulási keretrendszerek beépített primitíveit, és megfelelő teljesítményt nyújtanak anélkül, hogy túlságosan specializálódnának az egyes eszközök közötti átvitelekre, mint például az NVLink vagy az InfiniBand. A már említett lehetőségeken túl további teljesítménynövekedés érhető el a kommunikációs primitívek specializáltabb verziókkal való helyettesítésével. Például az NVLinken keresztül a szimmetrikus memória használata helyettesíthető az eszközök között megosztott pufferekkel, elhagyva bizonyos másolási műveleteket, és finomabb szemcsézettségű szinkronizációs sémákat téve lehetővé token-kötegek helyett az egyes tokeneken. Az InfiniBand-en keresztül egy megvalósítás közvetlenül hozzáférhet az alapul szolgáló sorpárokhoz, kiküszöbölve mind a várt, mind az elküldött tokenszámok átvitelének szükségességét, és a sorokból a kimeneti pufferekbe öblítve a tokeneket azok megérkezésekor. Az ilyen megvalósítások azonban a hordozhatóság kárára mennek, miközben a mi megvalósításunk lehetővé teszi a specifikus konfigurációkkal rendelkező klaszterek kiértékelését, mielőtt fejlesztési erőfeszítéseket tennénk egy adott hardver optimalizálására.

Bemutattunk egy nagy teljesítményű, hordozható könyvtárat a MoE kommunikációhoz, amely a standard all-to-all kommunikációhoz képest 10-szer gyorsabb teljesítményt ér el, miközben fenntartja a kompatibilitást a különböző hardverkonfigurációk között. Az NVLinkkel ellátott egycsomópontos kiépítéseknél a megoldásunk 2,5-ször alacsonyabb késleltetést mutat a korábbi megvalósításokhoz képest.

Megközelítésünk egyensúlyt teremt a teljesítmény és a hordozhatóság között olyan kulcsfontosságú innovációk révén, mint a GPU által kezdeményezett kommunikáció támogatása, a számítás és a kommunikáció átfedését lehetővé tevő osztott kernel architektúra, valamint a minimális NVSHMEM primitíveket használó hatékony tokendiszpécselés. Bár a töbcsomópontos beállításokon a rendkívül specializált megvalósításoknál körülbelül kétszer lassabb, könyvtárunk kiváló rugalmasságot biztosít a különböző hálózati környezetekben (NVLink, CX-7 és EFA).

Ahogy a MoE modellek tovább skálázódnak, a miénkhez hasonló hatékony kommunikációs stratégiák egyre fontosabbá válnak a gyakorlati üzembe helyezéshez. Teljes mértékben nyílt forráskódú megvalósításunk elérhető a https://github.com/ppl-ai/pplx-kernels címen.

DeepEP

DeepSeek-V3 műszaki jelentés

NVIDIA NVSHMEM

BibTeX

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}

Bevezetés

Sharding és párhuzamosság

NVSHMEM

Hordozható kernelek

Diszpécser

Kombináció

Teljesítménymérések (Benchmarks)

További fejlesztések

Összegzés

Hivatkozások

Idézet