Učinkovita i prijenosna komunikacija mješavine stručnjaka

Predstavljamo visokoučinkovitu, prijenosnu, open-source biblioteku za komunikaciju mješavine stručnjaka (MoE) koja ostvaruje 10 puta brže performanse u usporedbi sa standardnim All-to-All komunikacijskim primitivima.

Naša implementacija sadrži nekoliko ključnih tehničkih inovacija koje donose vrhunsku učinkovitost MoE komunikacije:

Komunikacija pokrenuta s GPU-a (IBGDA): Podržava izravnu komunikaciju između GPU-a i mrežne kartice (NIC), značajno smanjujući latenciju zaobilaženjem uključenosti CPU-a

Preklapanje komunikacije i računanja: Arhitektura podijeljene jezgre s odvojenim fazama slanja i primanja omogućuje da se računanje odvija istovremeno s mrežnim prijenosima

Najbrže performanse na jednom čvoru: 2,5 puta manja latencija od prethodno najbrže implementacije na konfiguracijama s jednim čvorom

Učinkovite i prijenosne performanse na više čvorova: Naša implementacija postiže brzine do 10 puta brže od standardne all-to-all komunikacije. Iako je otprilike 2 puta sporija od visoko specijaliziranih implementacija, naš pristup nudi bolju prenosivost kroz različite NVSHMEM verzije i mrežna okruženja (NVLink, CX-7 i EFA)

Biblioteka je u potpunosti open-source i dostupna na adresi https://github.com/ppl-ai/pplx-kernels.

U ovom članku istražujemo izazove paralelizma stručnjaka u velikim MoE modelima i opisujemo naš pristup učinkovitom slanju i kombiniranju tokena u distribuiranim GPU okruženjima.

Uvod

Modeli mješavine stručnjaka (MoE), kao što su DeepSeek R1 i Mixtral 8x7B, poboljšavaju guste modele ograničavanjem broja težina koje se aktiviraju za svaki token. Umjesto postizanja visokog broja parametara povećanjem broja slojeva ili povećanjem veličine linearnih slojeva u višeslojnom perceptronu (MLP) svakog sloja dekodera, MoE modeli zamjenjuju tradicionalni MLP s više stručnjaka i usmjerivačem (routerom). Na primjer, od 671 milijarde parametara modela DeepSeek R1, samo se 37 milijardi množi s određenim tokenom tijekom zaključivanja. Ovo uvelike smanjuje količinu računanja potrebnu za dekodiranje tokena i poboljšava latenciju u usporedbi s gustim modelom.

MoE modeli predstavljaju neke dodatne izazove za zaključivanje u usporedbi s gustim modelima. Dok su sami stručnjaci mali MLP slojevi, svaki sloj dekodera uključuje usmjerivač koji odlučuje na koje se stručnjake token šalje, pri čemu se svaki token šalje na više stručnjaka. Usmjerivač je obično mali linearni sloj koji proizvodi distribuciju vjerojatnosti. Obično se odabiru stručnjaci s najvišim K rezultatima, a konačna aktivacija se izračunava kao ponderirani prosjek, zbrajajući izlaze stručnjaka pomnožene s težinom izvedenom iz distribucije vjerojatnosti usmjerivača.

Kako bi se smanjila latencija u distribuiranim MoE sustavima, paralelizam se može iskoristiti na više uređaja, ali to uvodi komunikacijske izazove. Modeli poput Mixtral 8x7B ili Llama-70B staju unutar 8 uređaja na jednom čvoru, koristeći brze NVLink međusobne veze niske latencije (do 900 Gbps). Međutim, veći modeli zahtijevaju distribuciju (sharding) stručnjaka na više čvorova pomoću InfiniBanda (s vršnim vrijednostima od 400 Gbps), što donosi dodatne izazove s latencijom.

Povećaj

U ovom članku istražujemo problem paralelizma stručnjaka, opisujući našu implementaciju GPU jezgri koje učinkovito šalju tokene na uređaje koji drže dodijeljene stručnjake i ponovno ih prikupljaju izračunavanjem ponderiranog zbroja izlaza stručnjaka. Iako postoje agresivno optimizirane implementacije usmjerene na specifičan mrežni hardver, mi predstavljamo visoko prijenosnu implementaciju koja se oslanja na minimalan skup NVSHMEM primitiva, a ipak pruža performanse 10 puta brže od standardne all-to-all komunikacije. Naša implementacija postiže vrhunske performanse na konfiguracijama s jednim čvorom uz održavanje izvrsne prenosivosti u različitim mrežnim okruženjima.

Raspodjela i paralelizam

Za učinkovito zaključivanje, težine modela moraju se držati u memoriji uređaja, uz istovremeno ostavljanje dovoljno prostora za aktivacije, KV predmemorije i druge međuspremnike potrebne za prolazak prema naprijed kroz model. Najsposobniji modeli premašuju kapacitet čak i najsposobnijih GPU-ova, stoga se zaključivanje mora proširiti na više uređaja koji zajedno mogu pohraniti težine modela. Na temelju shema raspodjele težina, moraju se koristiti različite komunikacijske i računalne sheme za rasterećenje računanja i sinkronizaciju uređaja.

Povećaj

Paralelizam stručnjaka (EP), ilustriran na prvoj slici, paralelizira samo računanje stručnjaka. Različiti stručnjaci dodijeljeni su različitim uređajima koji drže njihove težine. Nakon usmjeravanja, tokeni se šalju na odgovarajući uređaj, s tim da se rezultati naknadno prikupljaju i akumuliraju. Složenost usmjeravanja ovisi o stupnju paralelizma u drugim dijelovima modela: repliciranje drugih slojeva moglo bi u potpunosti eliminirati potrebu za usmjeravanjem, budući da svaki rang može odabrati tokene iz lokalno replicirane tablice usmjeravanja. Međutim, ako samo jedan rang pokreće usmjeravanje, potreban je emitirani prijenos (broadcast) za slanje tokena, indeksa i težina njihovim odgovarajućim stručnjacima. Konačno, operacija all-gather ili all-to-all sinkronizira izlazne tokene s bilo kojim rangom koji nastavlja izvršavanje modela. Takva je implementacija relativno jednostavna jer torch već izlaže potrebne primitive, premda bi se određene koristi mogle postići spajanjem akumulacije s operacijom prikupljanja (gather) koja skuplja tokene od stručnjaka.

Povećaj

Paralelizam koji uključuje samo stručnjake ne skalira se idealno, jer čvorovi u klasteru mogu mirovati dok model pokreće slojeve koji nisu vezani uz stručnjake, kao što su pažnja (attention), normalizacija i uzorkovanje. Međutim, računanje ovih slojeva, prvenstveno pažnje, također može imati koristi od tenzorskog paralelizma (TP). Većina modela oslanja se na pažnju s više glava (multi-head attention), što znači da se glave pažnje i njihove odgovarajuće Q, K i V projekcije također mogu raspodijeliti po uređajima, replicirajući ili prikupljajući isječke između različitih slojeva. Ako je pažnja raspodijeljena po svim uređajima, operacija all-gather može sinkronizirati aktivacije, omogućujući repliciranje usmjeravanja, što zahtijeva primitivne sinkronizacije slične slučaju paralelizma stručnjaka za učinkovitu implementaciju. Ipak, postoje granice paralelizma na ovoj razini, budući da će smanjenje broja glava pažnje ispod određenog praga dati sve manjeg učinka.

Povećaj

Kako bi se najbolje iskoristili svi uređaji i podržao vrlo visok stupanj paralelizma stručnjaka, do 128 ili 256 GPU-ova, potreban je paralelizam podataka (DP). Prema ovoj shemi, uređaji na koje je model podijeljen grupirani su radi istovremenog obradivanja zahtjeva, izračunavanja pažnje i održavanja KV predmemorija raspodijeljenih po njihovoj lokalnoj grupi. Više instanci ovih paralelnih grupa surađuje u procjeni stručnjaka, pri čemu svaka hostira drugačiji podskup težina stručnjaka. Ovisno o broju glava pažnje, grupa obično može obuhvaćati do veličine cijelog čvora, budući da je komunikacija unutar čvora brža. Na primjer, na gornjoj slici, jedan DP rang neovisno opslužuje dva zahtjeva, A i B, obrađujući pažnju, normalizaciju i bilo koje drugo knjigovodstvo za zahtjeve. Drugi DP rang obrađuje zaseban zahtjev C. Međutim, prvi čvor ugošćuje polovinu stručnjaka, dok drugi ugošćuje drugu polovinu, stoga se nakon usmjeravanja tokeni iz A i B mogu poslati na drugi čvor i obrnuto. To dovodi do problema rijetke (sparse) komunikacije: svaki uređaj može poslati različit broj tokena bilo kojem drugom odredišnom rang-u. Postojeći primtivi iz torcha, prvenstveno all_to_all, nisu osobito prikladni jer mogu zahtijevati neki oblik dopunjavanja (padding) ili sinkronizacije GPU-a i CPU-a te emitiranja metapodataka. Za učinkovitu provedbu komunikacije potrebne su prilagođene jezgre za obradu informacija o usmjeravanju i pokretanje komunikacije s uređaja bez intervencije CPU-a. Nakon usmjeravanja, jezgra za slanje (dispatch) mora poslati tokene na rangove na koje su usmjereni, dok se na strani kombiniranja moraju prikupiti aktivacije koje pripadaju zahtjevima u trenutnoj DP grupi. Dodatno, rad se mora uravnotežiti unutar DP grupe koja može imati više uređaja, kako bi se smanjilo emitiranje i ispravno replicirali tokeni na svim rangovima za naknadne slojeve.

NVSHMEM

NVSHMEM je OpenSHMEM implementacija specifična za NVIDIA-u, koja pruža prenosive objekte za komunikaciju između uređaja koji apstrahiraju složenost temeljnog hardvera. API može izraziti čitanja i pisanja s uređaja na uređaj, koja se preslikavaju na primitive pojedinačnih transportnih slojeva. Podržani su i NVLink i RDMA, što daje određeni stupanj prenosivosti i omogućuje NVSHMEM jezgrama rad unutar čvora ili preko klastera s više čvorova. U našim jezgrama koristimo udaljeno pisanje u memoriju i operacije signalizacije kako bismo prenijeli podatke i sinkronizirali uređaje koji sudjeluju.

NVSHMEM operacije izgrađene su oko koncepta simetrične memorije: one rade na međuspremnicima koji su alocirani na svim uređajima koji sudjeluju u komunikaciji između uređaja. Lokalni uređaj može koristiti adresu izvedenu unutar vlastitog lokalnog međuspremnika zajedno s indeksom ciljnog ranga kako bi odredio odredište operacije. Slika u nastavku ilustrira ovaj koncept: oba GPU-a alociraju simetrične međuspremnike iste veličine, zadržavajući src pokazivače na njih. Prvi rang želi poslati 3 cijela broja drugome, stavljajući ih na početak međuspremnika: nvshmem_int_put_nbi izvodi početnu adresu iz lokalnog međuspremnika, određujući ciljni uređaj. Drugi rang izvodi pomak iz vlastitog međuspremnika, šaljući jedan element prvom uređaju, s pomakom za jedan. Dok odredišne adrese uvijek moraju biti simetrični međuspremnici alocirani pomoću nvshmem_alloc, izvorni međuspremnici mogu biti proizvoljna područja memorije uređaja, pod uvjetom da su unaprijed registrirani s NVSHMEM-om.

Iako NVSHMEM pruža širok raspon primitiva, naše se jezgre oslanjaju na samo 3 funkcije, gradeći svu sinkronizaciju i ograđivanje (fencing) na njima.

Povećaj

nvshmemx_putmem_signal_nbi_warp: Prenosi blok podataka s jednog uređaja na drugi, istovremeno postavljajući zastavicu na udaljenom uređaju. Operacija ili postavlja (NVSHMEM_SIGNAL_SET) ili uvećava (NVSHMEM_SIGNAL_ADD) 64-bitnu lokaciju. Zastavica se ažurira nakon što je cijeli blok memorije prenesen. Ako je udaljeni uređaj primijetio promjenu zastavice, može sigurno pristupiti međuspremniku u vlastitoj memoriji. Ova je funkcija korisna za spajanje prijenosa podataka i sinkronizacije.

nvshmemx_signal_op djeluje na jednoj memorijskoj lokaciji, obično 64-bitnoj zastavici, atomski je postavljajući ili uvećavajući. Koristan je za slanje metapodataka i sinkronizaciju uređaja.

nvshmem_uint64_wait_until koristi se na strani primanja signala, za anketiranje (polling) zastavice sve dok je udaljeni uređaj ne ažurira.

Osiguravanje bilo kakvog oblika redoslijeda između operacija putem nvshmem_quiet ili nvshmem_fence je skupo, stoga naša implementacija jezgri izbjegava korištenje bilo kojih drugih NVSHMEM primitiva i svodi sve ovisnosti na minimum osim pojedinačnih parova slanja i primanja. Prilikom slanja podataka uvijek koristimo verziju funkcija bez blokiranja, bez čekanja da podaci uopće budu poslani s lokalnog ranga. Druga, implicitna sinkronizacija preko grupe jezgri osigurava da se ovi međuspremnici ne prepišu na destruktivan način dok mrežnim karticama treba pristup njima.

NVSHMEM koristi GPUDirect RDMA tako da mrežna kartica može izravno čitati i pisati u GPU memoriju radi razmjene podataka bez uključivanja CPU-a. Nadalje, na ConnectX mrežnim karticama, NVSHMEM podržava komunikaciju pokrenutu s GPU-a (također poznatu kao InfiniBand GPUDirect Async ili IBGDA), što omogućuje GPU-u da šalje RDMA operacije izravno na mrežnu karticu i izravno provjerava red čekanja dovršetaka mrežne kartice. Na platformama koje ne podržavaju komunikaciju pokrenutu s GPU-om, NVSHMEM pokreće "proxy" nit na glavnom CPU-u radi pokretanja komunikacije i provjere dovršetaka u ime GPU-a. NVSHMEM program je prenosiv bez obzira na to je li komunikacija pokrenuta s GPU-a podržana ili nije. Međutim, komunikacija pokrenuta s GPU-om značajno smanjuje latenciju jer u potpunosti zaobilazi zaobilazni put do CPU-a.

Prijenosne jezgre

Implementiramo MoE komunikaciju kroz par jezgri za slanje (dispatch) i kombiniranje. Jezgre za slanje odgovorne su za čitanje tokena i informacija o usmjeravanju na svakom rangu, šaljući ih odgovarajućim stručnjacima. Jezgre za kombiniranje prikupljaju aktivacije koje proizvode stručnjaci i šalju ih natrag na njihove izvorne rangove, istovremeno izračunavajući ponderirani prosjek od odabranih stručnjaka na temelju težina koje je izračunao usmjerivač. Jezgre su dodatno podijeljene na komponentu slanja i primanja, kako bi se omogućilo preklapanje prijenosa podataka s računalom. Jezgre za slanje su neblokirajuće i nesinkronizirajuće: one jednostavno šalju sva pisanja na udaljene lokacije. S druge strane, jezgre za kombiniranje samo čitaju iz memorije, čekajući dok se svi potrebni podaci ne prenesu. Nakon slanja posla mrežnim karticama, dok se podaci prenose asinkrono preko mreže, GPU-ovi mogu obavljati druge korisne zadatke lokalno, kao što je primjena dijeljenih stručnjaka ili izračunavanje pažnje.

Svaki par jezgri ima svoje vlastite simetrične memorijske međuspremnike, pri čemu svaki rang alocira privatnu pohranu za primanje podataka sa svih drugih rangova. Maksimalni broj tokena koji se može poslati iz DP grupe fiksiran je na svim rangovima, što također postavlja gornju granicu koju svaki rang može primiti za svakog lokalnog stručnjaka iz svakog DP ranga. To omogućuje rangovima pošiljateljima da izvedu jedinstvenu adresu na odredišnom rangu u koju će pisati, bez potrebe za bilo kakvom sinkronizacijom među njima. Nakon što su podaci primljeni, jezgre premještaju tokene kako bi ih rasporedile u memoriji u optimalnom formatu koji zahtijevaju jezgre za primanje. Iako međuspremnici imaju znatnu dimenziju, ponovno se koriste kroz sve uzastopne slojeve modela.

Povećaj

Jedini oblik globalne sinkronizacije, kao što je ilustrirano na gornjoj slici, implementiran je u jezgri za kombiniranje i primanje (combine-receive). Jednom kada se podaci pošalju iz dijela za slanje, rang se vraća i nastavlja s izvršavanjem kada primi sve potrebne tokene s drugih rangova i ulazi u jezgre za kombiniranje i slanje čim se obrade. Barijera u dijelu za kombiniranje i primanje osigurava da nijedan rang ne može trčati naprijed i započeti slanje (dispatch-send) dok bilo koji drugi rang još uvijek čeka primitak podataka, što bi moglo potencijalno uzrokovati destruktivno preklapanje. Sinkronizacija se vrši u jezgri za kombiniranje radi jednostavnosti, a kao optimizacija, ista bi se barijera mogla implementirati između jezgri za primanje slanja (dispatch-receive) i slanja kombiniranja (combine-send), dodatno smanjujući latenciju dajući mrežnim karticama više vremena za rješavanje prijenosa dok GPU pokreće druge jezgre.

Obje su jezgre podijeljene po svim dostupnim SM jedinicama uređaja: dok jezgre za slanje (dispatch send) i primanje kombiniranja (combine receive) moraju paralelizirati preko maksimalnog broja tokena po rangu (max_m), jezgre za primanje slanja i slanje kombiniranja moraju obraditi toliko tokena sa svakog ranga i lokalnog stručnjaka, do maksimuma od max_m * (num_experts // EP) * (EP // TP).

Slanje (Dispatch)

Jezgra za slanje odgovorna je za slanje tokena i faktora skaliranja blokova vršnjacima koji ugošćuju stručnjake kojima su tokeni usmjereni. Strana pošiljatelja oslanja se na specijalizaciju warpova za paralelizaciju dvaju zadataka: agregiranje informacija o usmjeravanju kako bi se utvrdilo koliko se tokena šalje sa svakog ranga i pakiranje tokena u poruke koje se šalju na udaljenu lokaciju. Strana primatelja prvo čeka da se prime sve broji tokena, zatim čita i raspakira terete (payloads) iz dijeljene memorije u tenzore koji će se proslijediti nizvodnim Grouped GEMM jezgrama. Također pohranjuje informacije u međuspremnike koji se dijele s jezgrama za kombiniranje, kako bi označila gdje se svaki token treba poslati natrag. Te su informacije potrebne jer jezgra za primanje premješta tokene u susjedne međuspremnike u nedeterminističkom redoslijedu.

Povećaj

U dijelu pošiljatelja, jedan warp od 32 niti na svakom bloku odgovoran je za čitanje informacija o usmjeravanju i brojanje broja tokena za koje se očekuje da će svaki odredišni stručnjak primiti od ovog ranga, šaljući broj uvećan za jedan pomoću nvshmem_signal_op. Broj se uvećava za jedan, budući da prijelaz s nule na ne-nulu na udaljenom kraju označava primitak brojeva s izvornog ranga. Paralelno, preostali warpovi surađuju na kopiranju tokena u simetričnu memoriju paralelno preko svih blokova, pakirajući aktivacije, faktore skaliranja i njihov indeks na lokalnom rangu u jedan susjedni komad. Indeks je potreban pošiljatelju kombiniranja da bi odredio adresu na koju će token biti upisan. Zatim, nakon što se osigura da su svi podaci kopirani kroz barijeru preko grupa warpova, warpovi ponovno rade neovisno, od kojih svaki šalje isti međuspremnik različitom stručnjaku paralelno. Tokeni se šalju pomoću nvshmemx_putmem_signal_nbi_warp, što također atomski povećava broj poslanih tokena s lokalnog ranga na udaljenom uređaju. Unutar DP grupe, budući da svaki rang posjeduje repliku tokena koji se šalje, slanje je ravnomjerno uravnoteženo, pri čemu svaki uređaj šalje podskup tokena.

Na strani primanja, sve se jezgre prvo sinkroniziraju s rangovima pošiljateljima čekajući da se ukupni broj tokena pošalje. Nakon toga, svi oni čekaju da se atomski uvećani brojevi poslanih tokena slegnu na ukupne brojeve, što ukazuje na to da su i svi tereti s izvornih rangova poslani, zahvaljujući semantici putmem poziva. Jezgre provjeravaju brojeve pomoću nvshmem_uint64_wait_until, paralelizirajući operaciju preko svih blokova i niti. Naknadno, barijera između blokova osigurava da nijedan blok ne čita međuspremnike osim ako svi podaci nisu ispravno primljeni. Raspoređeni po blokovima i sinkronizirani putem atomski uvećanog brojača, tokeni se kopiraju iz simetrične memorije u obične susjedne međuspremnike koji postaju ulazi u nizvodne jezgre koje implementiraju stručnjake. Izvorni rang, indeks stručnjaka i indeks tokena pohranjuju se zasebno, točno označavajući lokaciju na koju jezgra za kombiniranje mora poslati aktivacije. Iako se tokeni unutar DP grupe šalju s različitih uređaja, oni su svi grupirani zajedno kako bi se proslijedili odgovarajućem stručnjaku.

Kombiniranje (Combine)

Jezgre za kombiniranje ponovno su podijeljene na polovicu slanja i primanja: pošiljatelji kopiraju nekvantizirane 16-bitne aktivacije na udaljene uređaje, pri čemu jezgre za primanje čekaju da se podaci pošalju i lokalno izračunavaju ponderirani prosjek svih doprinosa stručnjaka. Dodatno, one također djeluju kao barijera za sinkronizaciju slijeda slanja i kombiniranja: svaki rang postavlja zastavicu na svakom vršnjaku prilikom ulaska u jezgre za slanje, pri čemu jezgrama za primanje nije dopušteno da se vrate osim ako ne primijete da je zastavica postavljena. Latencija sinkronizacije je minimalna jer se preklapa s actualnom komunikacijom i računalom.

Povećaj

Na strani pošiljatelja, jezgre prolaze kroz popis tokena dodijeljenih svim lokalnim stručnjacima paralelno, upisujući ih u međuspremnik na odredišnom rangu. Ciljni rang, indeks stručnjaka i indeks tokena čitaju se iz međuspremnika po tokenu koje su popunile jezgre za raspršivanje (scatter) nakon primitka tokena. Svaki pošiljatelj ima svoje privatno memorijsko područje po stručnjaku u koje piše, kao što je prikazano na gornjoj slici, čime se izbjegava potreba za sinkronizacijom. Slično kao i kod slanja, kombiniranje atomski uvećava brojače po tokenu na odredišnom rangu kako bi označilo primitak podataka: kada brojač odgovara broju stručnjaka na koje je token poslan, može se pristupiti sadržaju tokena.

U jezgri za primanje, popis tokena se prolazi paralelno preko više blokova, čekajući da njihov sadržaj stigne anketiranjem zastavice postavljene operacijom signalizacije. Po dolasku, tereti se čitaju iz privatnih međuspremnika, pri čemu tablica usmjeravanja pokazuje iz kojeg međuspremnika treba čitati i koju težinu dodijeliti svakom stručnjaku. Rezultati se zatim upisuju u vanjski alocirane tenzore, pri čemu jezgra završava izvršavanje nakon što svi uređaji prođu barijeru. Preko DP grupe, svi rangovi primaju kopiju svake aktivacije stručnjaka kako bi izračunali vlastite replike.

Mjerenja performansi (Benchmarks)

Procjenjujemo naše jezgre na klasteru od 128 GPU-ova raspoređenih na 16 čvorova povezanih preko InfiniBand i CX-7 mrežnih kartica. Isprobavamo i GPUDirect Async (IBGDA) i pouzdanu vezu (Reliable Connection - RC) s CPU proxyjem. Uspoređujemo ih s visoko optimiziranom DeepSeek implementacijom, kao i s gustim primitivima koje nudi PyTorch (putem NCCL-a) ili NVSHMEM.

Povećaj

Iako je na gustom NVSHMEM all-to-all operaciji performansa IBRC-a i IBGDA-e slična (6378 µs vs 6180 µs), IBGDA je znatno brža s rijetkim (sparse) jezgrama. Kada se zbroje slanje i kombiniranje, IBGDA koristi 902 µs, dok IBRC troši 3223 µs – što je poboljšanje latencije 3,6 puta. Dok je all-to-all ograničen propusnošću (bandwidth bound), rijetke jezgre emitiraju redove veličine manje podataka, budući da su ograničene latencijom. Pokretanjem mrežnih prijenosa izravno s GPU-a, bez potrebe za CPU proxyjem koji bi koordinirao GPU i mrežnu karticu, latencija od kraja do kraja značajno je smanjena.

Iako su naše prijenosne jezgre otprilike 2 puta sporije od visoko optimiziranih DeepSeek jezgri, one poboljšavaju latenciju 10 puta u usporedbi s gustim jezgrama (902 µs vs 9944 µs).

Dodatno, podjela na komponente pošiljatelja i primatelja također omogućuje skrivanje dijela latencije, za razliku od primitiva biblioteke.

Povećaj

Na jednom čvoru (EP8), NVSHMEM koristi NVLINK za prijenos, pružajući manju latenciju i veću propusnost od umrežavanja između čvorova. Naše prijenosne jezgre otprilike su 2,5 puta brže od DeepEP-a na jednom čvoru (186 µs vs 481 µs).

Daljnja poboljšanja

Ovdje opisane jezgre nadmašuju ugrađene primitive ML okvira i nude pristojne performanse bez pretjerane specijalizacije za određene transporte između uređaja, kao što su NVLink ili InfiniBand. Osim već spomenutih prilika, daljnja poboljšanja performansi mogu se postići zamjenom komunikacijskih primitiva specijaliziranijim verzijama. Na primjer, preko NVLink-a, korištenje simetrične memorije moglo bi se zamijeniti međuspremnicima koji se dijele među uređajima, eliminirajući određeno kopiranje i omogućujući finije sheme sinkronizacije, preko pojedinačnih tokena umjesto serija tokena. Preko InfiniBanda, implementacija bi mogla izravno pristupiti temeljnim parovima redova čekanja, eliminirajući potrebu za prijenosom očekivanog i poslanog broja tokena te pražnjenjem tokena iz redova u izlazne međuspremnike kako pristižu. Međutim, takve implementacije dolaze po cijenu prenosivosti, dok naše implementacije omogućuju ocjenjivanje klastera s specifičnim konfiguracijama prije ulaganja razvojnog truda u optimizaciju za određeni hardver.

Zaključak

Predstavili smo visokoučinkovitu, prijenosnu biblioteku za MoE komunikaciju koja postiže 10 puta brže performanse u usporedbi sa standardnom all-to-all komunikacijom uz održavanje kompatibilnosti s različitim hardverskim konfiguracijama. Na primjenama na jednom čvoru s NVLink-om, naše rješenje pokazuje 2,5 puta manju latenciju od prethodnih implementacija.

Naš pristup balansira performanse s prenosivošću kroz ključne inovacije koje uključuju podršku za komunikaciju pokrenutu s GPU-a, arhitekturu podijeljene jezgre koja omogućuje preklapanje računanja i komunikacije te učinkovito slanje tokena pomoću minimalnih NVSHMEM primitiva. Iako je otprilike 2 puta sporija od visoko specijaliziranih implementacija na postavkama s više čvorova, naša biblioteka nudi vrhunsku fleksibilnost u različitim mrežnim okruženjima (NVLink, CX-7 i EFA).

Kako MoE modeli nastavljaju rasti, učinkovite komunikacijske strategije poput naše postajat će sve važnije za praktičnu primjenu. Naša u potpunosti open-source implementacija dostupna je na https://github.com/ppl-ai/pplx-kernels.

Reference

DeepEP

Tehničko izvješće za DeepSeek-V3

NVIDIA NVSHMEM

Citiranje

@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}
}