Efektívna a prenosná komunikácia mixture-of-experts

Predstavujeme vysoko výkonnú, prenosnú, open-source knižnicu pre komunikáciu Mixture-of-Experts (MoE), ktorá dosahuje 10-krát rýchlejší výkon v porovnaní so štandardnými komunikačnými primitívami All-to-All.

Naša implementácia prináša niekoľko kľúčových technických inovácií, ktoré zabezpečujú vynikajúcu efektivitu MoE komunikácie:

Komunikácia iniciovaná GPU (IBGDA): Podporuje priamu komunikáciu medzi GPU a NIC, čím výrazne znižuje latenciu obídením zapojenia CPU

Prekrývanie komunikácie a výpočtov: Architektúra deleného jadra so samostatnými fázami odosielania a prijímania umožňuje, aby výpočty prebiehali súčasne so sieťovými prenosmi

Najrýchlejší výkon na jednom uzle: 2,5-krát nižšia latencia než pri doteraz najrýchlejšej implementácii na konfiguráciách s jedným uzlom

Efektívny a prenosný výkon pre viaceré uzly: Naša implementácia dosahuje rýchlosť až 10-krát vyššiu než štandardná komunikácia all-to-all. Hoci je približne 2-krát pomalšia než vysoko špecializované implementácie, náš prístup ponúka lepšiu prenosnosť naprieč verziami NVSHMEM a sieťovými prostrediami (NVLink, CX-7 a EFA)

Knižnica je plne open-source a je k dispozícii na adrese https://github.com/ppl-ai/pplx-kernels.

V tomto článku skúmame výzvy paralelizmu expertov vo veľkých MoE modeloch a opisujeme náš prístup k efektivnemu smerovaniu a kombinovaniu tokenov v distribuovaných prostrediach GPU.

Úvod

Modely Mixture-of-Experts (MoE), ako sú DeepSeek R1 a Mixtral 8x7B, vylepšujú husté modely obmedzením počtu váh, ktoré sa aktivujú pre každý token. Namiesto dosahovania vysokého počtu parametrov zväčšovaním počtu vrstiev alebo zväčšovaním lineárnych vrstiev vo viacvrstvovom perceptróne (MLP) každej vrstvy dekodéra nahradzujú MoE modely tradičný MLP viacerými expertmi a smerovačom. Napríklad z 671 miliárd parametrov modelu DeepSeek R1 sa pri inferencii vynásobí s daným tokenom iba 37 miliárd. To výrazne znižuje množstvo výpočtov potrebných na dekódovanie tokenu a zlepšuje latenciu v porovnaní s hustým modelom.

Modely MoE predstavujú oproti hustým modelom pri inferencii niektoré ďalšie výzvy. Hoci samotní experti sú malé vrstvy MLP, každá dekóderová vrstva obsahuje router, ktorý rozhoduje, ku ktorým expertom sa token odošle, pričom každý token je odosielaný viacerým expertom. Router je zvyčajne malá lineárna vrstva, ktorá vytvára pravdepodobnostné rozdelenie. Zvyčajne sa vyberú experti s najvyššími K skóre a výsledná aktivácia sa vypočíta ako vážený priemer, pričom sa sčítajú výstupy expertov vynásobené váhou odvodenou z pravdepodobnostného rozdelenia routera.

Na minimalizáciu latencie v distribuovaných MoE systémoch možno paralelizmus využiť naprieč viacerými zariadeniami, čo však prináša komunikačné výzvy. Modely ako Mixtral 8x7B alebo Llama-70B sa zmestia do 8 zariadení v rámci jedného uzla, pričom využívajú rýchle prepojenia NVLink s nízkou latenciou (do 900 Gbps). Väčšie modely však vyžadujú rozdelenie (sharding) expertov do viacerých uzlov pomocou technológie InfiniBand (s vrcholom na 400 Gbps), čo prináša dodatočné problémy s latenciou.

Zväčšiť

V tomto článku skúmame problém paralelizmu expertov, pričom opisujeme našu implementáciu jadier GPU, ktoré efektívne smerujú tokeny do zariadení obsahujúcich priradených expertov a zbierajú ich späť výpočtom váženého súčtu výstupov expertov. Aj keď existujú agresívne optimalizované implementácie zamerané na konkrétny sieťový hardvér, predstavujeme vysoko prenosnú implementáciu, ktorá sa spolieha na minimálnu sadu primitív NVSHMEM, no napriek tomu poskytuje 10-krát rýchlejší výkon než štandardná komunikácia all-to-all. Naša implementácia dosahuje špičkový výkon na konfiguráciách s jedným uzlom a zároveň si zachováva vynikajúcu prenosnosť v rôznych sieťových prostrediach.

Sharding a paralelizmus

Na zabezpečenie efektivnej inferencie musia byť váhy modelu uložené v pamäti zariadenia, pričom musí zostať dostatočný priestor pre aktivácie, vyrovnávacie pamäte KV a ďalšie vyrovnávacie pamäte potrebné na prechod smerom dopredu (forward pass) modelom. Najvýkonnejšie modely presahujú kapacitu aj tých najvýkonnejších GPU, preto sa inferencia musí rozdeliť medzi viaceré zariadenia, ktoré dokážu spoločne uložiť váhy modelu. Na základe schém shardingu váh sa musia použiť rôzne komunikačné a výpočtové schémy na presun výpočtov a synchronizáciu zariadení.

Zväčšiť

Paralelizmus expertov (EP), ilustrovaný na prvom obrázku, paralelizuje iba výpočty expertov. Rôzni experti sú priradení k rôznym zariadeniam, ktoré uchovávajú ich váhy. Po smerovaní sa tokeny odoslú do príslušného zariadenia, pričom výsledky sa následne zhromaždia a akumulujú. Zložitosť smerovania závisí od stupňa paralelizmu v ostatných častiach modelu: replikácia iných vrstiev môže úplne eliminovať potrebu smerovania, pretože každá hodnosť (rank) môže vybrať tokeny z lokálne replikovanej smerovacej tabuľky. Ak však smerovanie vykonáva iba jedna hodnosť, vyžaduje sa vysielanie (broadcast) na odoslanie tokenov, indexov a váh príslušným expertom. Nakoniec operácia all-gather alebo vysielanie all-to-all synchronizuje výstupné tokeny s tou hodnosťou, ktorá pokračuje v spúšťaní modelu. Takáto implementácia je pomerne jednoduchá, pretože torch už poskytuje požadované primitíva, hoci určité výhody by sa dali získať zlúčením akumulácie s operáciou zhromažďovania, ktorá zbiera tokeny od expertov.

Zväčšiť

Paralelizmus zameraný len na expertov sa neškáluje ideálne, pretože uzly v klusteri môžu byť nečinné, zatiaľ čo model spúšťa vrstvy, ktoré nie sú spojené s expertmi, ako sú attention, norm a vzorkovanie. Výpočty týchto vrstiev, predovšetkým attention, však môžu ťažiť aj z paralelizmu tenzorov (TP). Väčšina modelov sa spolieha na multi-head attention, čo znamená, že hlavy attention a ich zodpovedajúce projekcie Q, K a V možno taktiež rozdeliť medzi zariadenia, pričom výrezy medzi rôznymi vrstvami sa replikujú alebo zhromažďujú. Ak je attention rozložená do všetkých zariadení, operácia all-gather môže synchronizovať aktivácie, čo umožňuje replikáciu smerovania, čo si na efektivnu implementáciu vyžaduje synchronizačné primitíva podobné prípadu paralelizmu expertov. Existujú však limity paralelizmu na tejto úrovni, pretože zníženie počtu hláv attention pod určitú hranicu povedie k zmenšujúcej sa návratnosti.

Zväčšiť

Na dosiahnutie čo najlepšieho využitia všetkých zariadení a podporu veľmi vysokého stupňa paralelizmu expertov (až 128 alebo 256 GPU) sa vyžaduje paralelizmus údajov (DP). V rámci tejto schémy sú zariadenia, medzi ktoré je model rozdelený, zoskupené na súbežné vybavovanie požiadaviek, pričom počítajú attention a udržiavajú vyrovnávacie pamäte KV rozdelené v rámci ich lokálnej skupiny. Viaceré inštancie týchto paralelých skupín spolupracujú pri vyhodnocovaní expertov, pričom každá hostí inú podmnožinu váh expertov. V závislosti od počtu hláv attention môže skupina typicky pokrývať až veľkosť celého uzla, pretože komunikácia v rámci uzla je rýchlejšia. Napríklad na obrázku vyššie jedna hodnosť DP nezávisle obsluhuje dve požiadavky, A a B, pričom spracováva attention, norm a akúkoľvek inú správu pre požiadavky. Druhá hodnosť DP spracováva odlišnú požiadavku C. Prvý uzol však hostí polovicu expertov a druhý uzol druhú polovicu, takže po smerovaní sa tokeny z A a B môžu odoslať do druhého uzla a naopak. To vedie k problému riedkej komunikácie: každé zariadenie môže odoslať iný počet tokenov do ľubovoľného iného cieľového umiestnenia. Existujúce primitíva z knižnice torch, predovšetkým all_to_all, nie sú obzvlášť vhodné, pretože môžu vyžadovať určitú formu dopĺňania (padding), synchronizáciu GPU s CPU a vysielanie metadát. Na efektivnu implementáciu komunikácie sa vyžadujú vlastné jadrá na spracovanie smerovacích informácií a iniciovanie komunikácie zo zariadení bez zásahu CPU. Po smerovaní musí dispečerské jadro odoslať tokeny do hodností, na ktoré boli smerované, zatiaľ čo na strane kombinovania sa musia zhromaždiť aktivácie patriace požiadavkám v aktuálnej skupine DP. Okrem toho sa práca musí vyvážiť v rámci skupiny DP, ktorá môže mať viacero zariadení, aby sa minimalizovalo vysielanie a správne replikovali tokeny na všetkých hodnostiach pre nasledujúce vrstvy.

NVSHMEM

NVSHMEM je implementácia OpenSHMEM špecifická pre spoločnosť NVIDIA, ktorá poskytuje prenosné komunikačné zariadenia medzi zariadeniami, abstrahujúce zložitosť základného hardvéru. Rozhranie API dokáže vyjadriť čítanie a zápis medzi zariadeniami, ktoré sa mapujú na primitíva jednotlivých transportných vrstiev. Podporované sú NVLink aj RDMA, čo poskytuje určitý stupeň prenosnosti a umožňuje jadrám NVSHMEM pracovať v rámci jedného uzla alebo naprieč klstrom viacerých uzlov. V našich jadrách používame operácie vzdialeného zápisu do pamäte a signalizácie na prenos údajov a synchronizáciu zúčastnených zariadení.

Operácie NVSHMEM sú postavené na koncepcii symetrickej pamäte: pracujú s vyrovnávacími pamäťami, ktoré boli alokované na všetkých zariadeniach zúčastnených na komunikácii medzi zariadeniami. Lokálne zariadenie môže použiť adresu odvodenú v rámci vlastnej lokálnej vyrovnávacej pamäte spolu s indexom cieľovej hodnosti na určenie miesta určenia operácie. Obrázok nižšie ilustruje tento koncept: obe GPU alokujú symetrické vyrovnávacie pamäte rovnakej veľkosti a uchovávajú na ne ukazovatele src. Prvá hodnosť chce odoslať 3 celé čísla druhej hodnosti a umiestniť ich na začiatok vyrovnávacej pamäte: nvshmem_int_put_nbi odvodzuje počiatočnú adresu z lokálnej vyrovnávacej pamäte a špecifikuje cieľové zariadenie. Druhá hodnosť odvodzuje posun zo svojej vlastnej vyrovnávacej pamäte, pričom odošle jeden prvok prvému zariadeniu s posunom o jeden. Zatiaľ čo cieľové adresy musia byť vždy symetrické vyrovnávacie pamäte alokované pomocou nvshmem_alloc, zdrojové vyrovnávacie pamäte môžu byť ľubovoľné oblasti pamäte zariadenia za predpokladu, že sú vopred zaregistrované v NVSHMEM.

Hoci NVSHMEM poskytuje širokú škálu primitív, naše jadrá sa spoliehajú iba na 3 funkcie, na ktorých stavajú všetku synchronizáciu a ohraničenie (fencing).

Zväčšiť

nvshmemx_putmem_signal_nbi_warp: Prenáša blok údajov z jedného zariadenia do druhého a zároveň nastavuje príznak na vzdialenom zariadení. Operácia buď nastaví (NVSHMEM_SIGNAL_SET), alebo inkrementuje (NVSHMEM_SIGNAL_ADD) 64-bitové umiestnenie. Príznak sa aktualizuje po prenose celého bloku pamäte. Ak vzdialené zariadenie zaznamená zmenu príznaku, môže bezpečne pristupovať k vyrovnávacej pamäti vo svojej vlastnej pamäti. Táto funkcia je užitočná na zlúčenie prenosu údajov a synchronizácie.

nvshmemx_signal_op pracuje s jedným pamäťovým umiestnením, zvyčajne 64-bitovým príznakom, pričom ho atomicky nastavuje alebo inkrementuje. Je užitočný pri odosielaní metadát a synchronizácii zariadení.

nvshmem_uint64_wait_until sa používa na strane príjmu signálu na cyklické dotytovanie (polling) príznaku, kým ho vzdialené zariadenie neaktualizuje.

Zabezpečenie akejkoľvek formy usporiadania medzi operáciami prostredníctvom nvshmem_quiet alebo nvshmem_fence je nákladné, preto naša implementácia jadier vyhýba použitiu akýchkoľvek iných primitív NVSHMEM a minimalizuje akékoľvek závislosti okrem jednotlivých párov odoslania a prijatia. Pri odosielaní údajov vždy používame neblokujúcu verziu funkcií bez čakania na to, kým budú údaje vôbec odoslané z lokálnej hodnosti. Iná, implicitná synchronizácia naprieč skupinou jadier zabezpečuje, že tieto vyrovnávacie pamäte nebudú deštruktívne prepísané, zatiaľ čo sieťové karty (NIC) potrebujú k nim prístup.

NVSHMEM využíva technológiu GPUDirect RDMA, aby sieťová karta (NIC) mohla priamo čítať a zapisovať do pamäte GPU na výmenu údajov bez zapojenia CPU. Okrem toho na sieťových kartách ConnectX podporuje NVSHMEM komunikáciu iniciovanú GPU (známu aj ako Infiniband GPUDirect Async alebo IBGDA), ktorá umožňuje GPU odosielať operácie RDMA priamo do NIC a priamo vyhodnocovať front dokončenia NIC. Na platformách, ktoré nepodporujú komunikáciu iniciovanú GPU, NVSHMEM spúšťa „proxy“ vlákno na hostiteľskom CPU na iniciovanie komunikácie a vyhodnocovanie dokončenia v mene GPU. Program NVSHMEM je prenosný bez ohľadu na to, či je komunikácia iniciovaná GPU podporovaná alebo nie. Komunikácia iniciovaná GPU však výrazne skracuje latenciu, pretože úplne obchádza obchádzkovú trasu cez CPU.

Prenosné jadrá

Komunikáciu MoE implementujeme prostredníctvom dvojice dispečerských a kombinovacích jadier. Dispečerské jadrá zodpovedajú za čítanie tokenov a informácií o smerovaní na každej hodnosti a ich odosielanie príslušným expertom. Kombinačné jadrá zhromažďujú aktivácie vytvorené expertmi a posielajú ich späť do zdrojových hodností a zároveň počítajú vážený priemer od vybraných expertov na základe váh vypočítaných smerovačom. Jadrá sú ďalej rozdelené na komponent odosielania a prijímania, aby sa umožnilo prekrývanie prenosov údajov s výpočtami. Jadrá odosielania sú neblokujúce a nesynchronizujúce: jednoducho odošlú všetky zápisy do vzdialených umiestnení. Na druhej strane, kombinačné jadrá čítajú iba z pamäte a čakajú, kým sa prenesú všetky potrebné údaje. Po odoslaní práce do NIC, zatiaľ čo sa údaje prenášajú asynchrónne po sieti, môžu GPU lokálne vykonávať inú užitočnú prácu, ako je aplikácia zdieľaných expertov alebo výpočet attention.

Každý pár jadier má vlastné symetrické pamäťové vyrovnávacie pamäte, pričom každá hodnosť alokuje súkromné úložisko na prijímanie údajov zo všetkých ostatných hodností. Maximálny počet tokenov, ktoré je možné odoslať zo skupiny DP, je pevne daný naprieč hodnosťami, čím sa stanovuje aj horná hranica, ktorú môže každá hodnosť prijať pre každého lokálneho experta z každej hodnosti DP. To umožňuje odosielajúcim hodnostiam odvodiť jedinečnú adresu na cieľovej hodnosti, do ktorej sa má zapisovať, bez potreby akejkoľvek synchronizácie medzi nimi. Po prijatí údajov jadrá preskupujú tokeny, aby ich usporiadali v pamäti v optimálnom formáte požadovanom prijímajúcimi jadrami. Hoci vyrovnávacie pamäte majú značnú veľkosť, opätovne sa používajú vo všetkých sériových vrstvách modelu.

Zväčšiť

Jedinou formou globálnej synchronizácie, ako je znázornené na obrázku vyššie, je implementácia v kombinačnom prijímacom jadre. Keď sú údaje odoslané z odosielania, hodnosť sa vráti a pokračuje v spúšťaní, keď prijme všetky tokeny, ktoré potrebuje od ostatných hodností, a vstúpi do kombinačných odosielacích jadier hneď, ako sú spracované. Bariéra v kombinačnom prijímacom jadre zabezpečuje, že žiadna hodnosť nemôže prebehnúť vopred a začať dispečerské odosielanie, zatiaľ čo iná hodnosť stále čaká na prijatie údajov, čo by mohlo potenciálne spôsobiť deštruktívne prekrývanie. Synchronizácia sa vykonáva v kombinačnom jadre kvôli jednoduchosti, pričom ako optimalizácia by sa tá istá bariéra mohla implementovať medzi dispečerským prijímacím a kombinačným odosielacím jadrom, čím by sa ďalej znížila latencia tým, že by sa sieťovým kartám poskytlo viac času na usporiadanie prenosov, zatiaľ čo GPU spúšťa iné jadrá.

Oba jadrá sú rozdelené medzi všetky dostupné SM zariadení: zatiaľ čo dispečerské odosielacie a kombinačné prijímacie jadrá musia paralelizovať naprieč maximálnym počtom tokenov na hodnosť (max_m), dispečerské prijímacie a kombinačné odosielacie jadrá musia spracovať toľko tokenov z každej hodnosti a lokálneho experta, pričom maximum je max_m * (num_experts // EP) * (EP // TP).

Dispečer

Dispečerské jadro zodpovedá za odosielanie tokenov a faktorov blokového škálovania partnerom, ktorí hostia expertov, ku ktorým boli tokeny smerované. Strana odosielateľa sa spolieha na špecializáciu warpov na paralelné spracovanie dvoch úloh: agregáciu smerovacích informácií na určenie, koľko tokenov sa odosiela z každej hodnosti, a balenie tokenov do správ, ktoré sa majú odoslať na vzdialené zariadenie. Strana prijímača najprv čaká na prijatie všetkých počtov tokenov, potom číta a rozbaľuje užitočné zaťaženia (payloads) zdieľanej pamäte do tenzorov, ktoré budú odovzdané do nadväzujúcich jadier Grouped GEMM. Tiež ukladá informácie do vyrovnávacích pamätí zdieľaných s kombinačnými jadrami, aby sa naznačilo, kam sa má každý token poslať späť. Táto informácia je potrebná, pretože prijímacie jadro preskupuje tokeny do súvislých vyrovnávacích pamätí v nedeterministickom poradí.

Zväčšiť

V časti odosielateľa je jeden warp s 32 vláknami na každom bloku zodpovedný za čítanie informácií o smerovaní a počítanie počtu tokenov, ktoré má každý cieľový expert prijať od tejto hodnosti, pričom odošle počet plus jeden pomocou funkcie nvshmem_signal_op. Počet sa inkrementuje o jeden, pretože prechod z nuly na nenulovú hodnotu na vzdialenom konci signalizuje prijatie počtov zo zdrojovej hodnosti. Zostávajúce warpy paralelne spolupracujú na kopírovaní tokenov do symetrickej pamäte naprieč všetkými blokmi, pričom balia aktivácie, faktory škálovania a svoj index na lokálnej hodnosti do súvislého bloku. Index je potrebný pre kombinačného odosielateľa na určenie adresy, na ktorú sa token zapíše. Ďalej po zabezpečení toho, že všetky údaje boli skopírované cez bariéru naprieč skupinami warpov, warpy opäť pracujú nezávisle, pričom každý odosiela rovnakú vyrovnávaciu pamäť inému expertovi paralelne. Tokeny sa odosielajú pomocou nvshmemx_putmem_signal_nbi_warp, čo tiež atomicky zvyšuje počet odoslaných tokenov z lokálnej hodnosti na vzdialenom zariadení. V rámci skupiny DP, keďže každá hodnosť vlastní repliku tokenu, ktorý sa má odoslať, je dispečing rovnomerne vyvážený, pričom každé zariadenie odosiela podmnožinu tokenov.

Na strane príjmu sa všetky jadrá najprv synchronizujú so zdrojovými hodnosťami čakaním na odoslanie celkových počtov tokenov. Potom všetky čakajú, kým sa atomicky inkrementované počty odoslaných tokenov ustália na celkových počtoch, čo naznačuje, že všetky užitočné zaťaženia zo zdrojových hodností boli tiež odoslané, vďaka sémantike volania putmem. Jadrá vykonávajú cyklické dotytovanie (polling) počtov pomocou nvshmem_uint64_wait_until, čím paralelizujú operáciu naprieč všetkými blokmi a vláknami. Následne bariéra medzi blokmi zabezpečuje, že žiadny blok nečíta vyrovnávacie pamäte, pokiaľ neboli všetky údaje správne prijaté. Rozložené naprieč blokmi a synchronizované pomocou atomicky inkrementovaného počítadla sa tokeny kopírujú zo symetrickej pamäte do bežných súvislých vyrovnávacích pamätí, ktoré sa stávajú vstupmi pre nadväzujúce jadrá implementujúce expertov. Zdrojová hodnosť, index experta a index tokenu sú uložené oddelene, čím presne určujú umiestnenie, kam musí kombinačné jadro odoslať aktivácie. Aj keď sa tokeny zo skupiny DP odosielajú z rôznych zariadení, všetky sú zoskupené dohromady, aby sa odovzdali príslušnému expertovi.

Kombinátor

Kombinačné jadrá sú opäť rozdelené na polovice odosielania a prijímania: odosielatelia kopírujú nekvantované 16-bitové aktivácie do vzdialených zariadení, pričom prijímacie jadrá čakajú na odoslanie údajov a lokálne počítajú vážený priemer všetkých príspevkov expertov. Okrem toho fungujú aj ako bariéra na synchronizáciu sekvencie odoslania a kombinovania: každá hodnosť nastaví príznak na každom partnerovi pri vstupe do odosielacích jadier, pričom prijímacím jadrám nie je dovolené vrátiť sa, pokiaľ nezistia, že príznak je nastavený. Latencia synchronizácie je minimálna, pretože sa prekrýva so samotnou komunikáciou a výpočtami.

Zväčšiť

Na strane odosielateľa jadrá paralelne prechádzajú zoznam tokenov priradených všetkým lokálnym expertom a zapisujú ich do vyrovnávacej pamäte na cieľovej hodnosti. Cieľová hodnosť, index experta a index tokenu sa čítajú z vyrovnávacích pamätí pre každý token, ktoré naplnili rozptyľovacie (scatter) jadrá po prijatí tokenov. Každý odosielateľ má svoju vlastnú súkromnú oblasť pamäte na každého experta, do ktorej zapisuje, ako je znázornené na obrázku vyššie, čím sa eliminuje potreba synchronizácie. Podobne ako pri dispečingu, kombinátor atomicky inkrementuje počítadlá pre každý token na cieľovej hodnosti, čím signalizuje prijatie údajov: keď počítadlo zodpovedá počtu expertov, ktorým bol token odoslaný, je možné pristupovať k obsahu tokenu.

V prijímacom jadre sa zoznam tokenov prechádza paralelne naprieč viacerými blokmi a čaká sa na príchod ich obsahu prostredníctvom dotytovania príznaku nastaveného signalizačnou operáciou. Po príchode sa užitočné zaťaženia načítajú zo súkromných vyrovnávacích pamätí, pričom smerovacia tabuľka určuje, z ktorej vyrovnávacej pamäte sa má čítať a akú váhu priradiť každému expertovi. Výsledky sa potom zapíšu do externe alokovaných tenzorov, pričom jadro ukončí vykonávanie, keď všetky zariadenia prejdú cez bariéru. V rámci skupiny DP dostávajú všetky hodnosti kópiu aktivácie každého experta na výpočet vlastných replík.

Porovnávacie testy (Benchmarks)

Naše jadrá vyhodnocujeme na klusteri 128 GPU rozložených medzi 16 uzlami prepojenými pomocou InfiniBand a sieťových kariet CX-7. Vyskúšali sme GPUDirect Async (IBGDA aj Reliable Connection (RC) s proxy pre CPU. Porovnávame ich s vysoko optimalizovanou implementáciou DeepSeek, ako aj s hustými primitívami poskytovanými systémom PyTorch (prostredníctvom NCCL) alebo NVSHMEM.

Zväčšiť

Hoci pri hustej operácii all-to-all v NVSHMEM je výkon IBRC a IBGDA podobný (6378 µs vs 6180 µs), IBGDA je pri riedkych jadrách podstatne rýchlejší. Po sčítaní dispečingu a kombinovania využíva IBGDA 902 µs, zatiaľ čo IBRC potrebuje 3223 µs – čo predstavuje 3,6-násobné zlepšenie latencie. Zatiaľ čo all-to-all je obmedzená šírkou pásma, riedke jadrá vysielajú o rády menej údajov a sú obmedzené latenciou. Spustením sieťových prenosov priamo z GPU bez potreby proxy pre CPU na koordináciu GPU a sieťovej karty sa celková latencia výrazne zníži.

Hoci naše prenosné jadrá sú približne 2-krát pomalšie než vysoko optimalizované jadrá DeepSeek, zlepšujú latenciu 10-krát v porovnaní s hustými jadrami (902 µs vs 9944 µs).

Rozdelenie na komponenty odosielateľa a prijímača navyše umožňuje skryť časť latencie, na rozdiel od primitív knižnice.

Zväčšiť

Na jednom uzle (EP8) využívame NVSHMEM na transport NVLINK, čím sa dosahuje nižšia latencia a vyššia priepustnosť než pri sieťovom prepojení medzi uzlami. Naše prenosné jadrá sú približne 2,5-krát rýchlejšie než DeepEP na jednom uzle (186 µs vs 481 µs).

Ďalšie vylepšenia

Tu opísané jadrá prekonávajú vstavané primitíva frameworkov ML a ponúkajú slušný výkon bez nadmernej špecializácie na konkrétne transporty medzi zariadeniami, ako sú NVLink alebo InfiniBand. Okrem už spomenutých možností je možné dosiahnuť ďalšie zvýšenie výkonu nahradením komunikačných primitív špecializovanejšími verziami. Napríklad v rámci NVLink by sa použitie symetrickej pamäte mohlo nahradiť vyrovnávacími pamäťami zdieľanými naprieč zariadeniami, čím by sa eliminovalo určité kopírovanie a umožnili by sa jemnejšie synchronizačné schémy pre jednotlivé tokeny namiesto dávok tokenov. V rámci InfiniBand by implementácia mohla pristupovať priamo k základným párom frontov, čím by sa eliminovala potreba prenášať počty očakávaných aj odoslaných tokenov a vyprázdňovať tokeny z frontov do výstupných vyrovnávacích pamätí pri ich príchode. Takéto implementácie sú však vykupované obetovaním prenosnosti, zatiaľ čo naše implementácie umožňujú vyhodnotiť klustery so špecifickými konfiguráciami pred vynaložením vývojového úsilia na optimalizáciu pre konkrétny hardvér.

Záver

Predstavili sme vysoko výkonnú prenosnú knižnicu pre komunikáciu MoE, ktorá dosahuje 10-krát rýchlejší výkon v porovnaní so štandardnou komunikáciou all-to-all pri zachovaní kompatibility naprieč rôznymi hardvérovými konfiguráciami. Pri nasadeniach s jedným uzlom s NVLink naše riešenie demonštruje 2,5-krát nižšiu latenciu než predchádzajúce implementácie.

Náš prístup vyvažuje výkon s prenosnosťou prostredníctvom kľúčových inovácií vrátane podpory komunikácie iniciovanej GPU, architektúry deleného jadra umožňujúcej prekrývanie výpočtov a komunikácie a efektivneho dispečingu tokenov s použitím minimálnych primitív NVSHMEM. Hoci je naša knižnica v porovnaní s vysoko špecializovanými implementáciami na nastaveniach s viacerými uzlami približne 2-krát pomalšia, ponúka vynikajúcu flexibilitu v rôznych sieťových prostrediach (NVLink, CX-7 a EFA).

Keďže sa MoE modely naďalej škálujú, efektívne komunikačné stratégie, akou je tá naša, budú čoraz dôležitejšie pre praktické nasadenie. Naša plne open-source implementácia je k dispozícii na adrese https://github.com/ppl-ai/pplx-kernels.

Odkazy

DeepEP

Technická správa DeepSeek-V3

NVIDIA NVSHMEM

Citácia

@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://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}