Comunicație eficientă și portabilă Mixture-of-Experts

Vă prezentăm o bibliotecă open-source de înaltă performanță și portabilă pentru comunicația Mixture-of-Experts (MoE), care atinge o performanță de 10x mai rapidă în comparație cu primitivele standard de comunicație All-to-All.

Implementarea noastră include mai multe inovații tehnice cheie care asigură o eficiență superioară a comunicației MoE:

Comunicație inițiată de GPU (IBGDA): Suportă comunicația directă de la GPU la NIC, reducând semnificativ latența prin ocolirea implicării CPU

Suprapunerea comunicației și a calculului: Arhitectura de kernel împărțit cu etape separate de trimitere și primire permite ca calculele să se desfășoare în paralel cu transferurile în rețea

Cea mai rapidă performanță pe un singur nod: o latență de 2,5 ori mai mică decât cea mai rapidă implementare anterioară pe configurații cu un singur nod

Performanță eficientă și portabilă pe mai multe noduri: Implementarea noastră atinge viteze de până la 10x mai mari decât comunicația standard all-to-all. Deși este cu aproximativ 2x mai lentă decât implementările extrem de specializate, abordarea noastră oferă o portabilitate mai bună între versiunile NVSHMEM și mediile de rețea (NVLink, CX-7 și EFA)

Biblioteca este complet open-source și disponibilă la https://github.com/ppl-ai/pplx-kernels.

În acest articol, explorăm provocările paralelismului experților în modelele MoE la scară largă și descriem abordarea noastră pentru trimiterea și combinarea eficientă a tokenurilor în medii GPU distribuite.

Introducere

Modelele Mixture-of-Experts (MoE), cum ar aduce DeepSeek R1 și Mixtral 8x7B, aduc îmbunătățiri față de modelele dense prin limitarea numărului de ponderi care sunt activate pentru fiecare token. În loc să obțină un număr mare de parametri prin creșterea numărului de straturi sau mărimea straturilor liniare din perceptronul multi-strat (MLP) al fiecărui strat de decodor, modelele MoE înlocuiesc MLP-ul tradițional cu mai mulți experți și un ruter. De exemplu, din cei 671B parametri ai DeepSeek R1, doar 37B sunt înmulțiți cu un anumit token în timpul inferenței. Acest lucru reduce considerabil cantitatea de calcule necesară pentru decodarea unui token și îmbunătățește latența în comparație cu un model dens.

Modelele MoE prezintă unele provocări suplimentare pentru inferență în comparație cu modelele dense. În timp ce experții în sine sunt straturi MLP mici, fiecare strat de decodor include un ruter care decide către ce experți este expediat un token, fiecare token fiind expediat către mai mulți experți. Ruterul este de obicei un strat liniar mic care produce o distribuție de probabilitate. De obicei, sunt selectați experții cu scorurile top-K, iar activarea finală este calculată ca o medie ponderată, însumând ieșirile experților înmulțite cu o pondere derivată din distribuția de probabilitate a ruterului.

Pentru a minimiza latența în sistemele MoE distribuite, paralelismul poate fi exploatat pe mai multe dispozitive, dar acest lucru introduce provocări de comunicație. Modelele precum Mixtral 8x7B sau Llama-70B se potrivesc în 8 dispozitive pe un singur nod, beneficiind de interconexiuni NVLink rapide și cu latență redusă (de până la 900Gbps). Cu toate acestea, modelele mai mari necesită fragmentarea experților pe mai multe noduri folosind InfiniBand (cu un vârf de 400Gbps), ceea ce introduce provocări suplimentare legate de latență.

Extinde

În acest articol, explorăm problema paralelismului experților, descriind implementarea noastră de kerneluri GPU care trimit eficient tokenurile către dispozitivele care găzduiesc experții atribuiți și le colectează înapoi calculând suma ponderată a ieșirilor experților. Deși există implementări optimizate agresiv care vizează hardware de rețea specific, noi prezentăm o implementare extrem de portabilă care se bazează pe un set minim de primitive NVSHMEM, oferind totuși o performanță de 10x mai rapidă decât comunicația standard all-to-all. Implementarea noastră atinge o performanță de vârf pe configurații cu un singur nod, menținând în același timp o portabilitate excelentă în diferite medii de rețea.

Fragmentare și paralelism

Pentru o inferență eficientă, ponderile unui model trebuie să fie păstrate în memoria dispozitivului, lăsând în același timp un spațiu suficient pentru activări, cache-uri KV și alte buffere necesare pentru trecerea înainte prin model. Cele mai capabile modele depășesc capacitatea chiar și a celor mai capabile GPU-uri, prin urmare inferența trebuie distribuită pe mai multe dispozitive, care pot stoca colectiv ponderile modelului. Pe baza schemelor de fragmentare a ponderilor, trebuie utilizate diferite scheme de comunicație și calcul pentru a descărca calculele și a sincroniza dispozitivele.

Extinde

Paralelismul experților (EP), ilustrat în prima figură, paralelizează doar calculul experților. Diferiti experți sunt atribuiți unor dispozitive diferite, care le găzduiesc ponderile. După rutare, tokenurile sunt trimise către dispozitivul corespunzător, rezultatele fiind adunate și acumulate ulterior. Complexitatea rutării depinde de gradul de paralelism din celelalte părți ale modelului: replicarea altor straturi ar putea elimina complet nevoia de rutare, deoarece fiecare rang poate selecta tokenurile dintr-un tabel de rutare replicat local. Cu toate acestea, dacă doar un singur rang rulează rutarea, este necesară o difuzare (broadcast) pentru a trimite tokenurile, indicii și ponderile către experții lor respectivi. În cele din urmă, un all-gather sau o difuzare all-to-all sincronizează tokenurile de ieșire cu orice rang continuă execuția modelului. O astfel de implementare este relativ simplă, deoarece torch expune deja primitivele necesare, deși s-ar putea obține anumite beneficii prin fuzionarea acumulării cu operația de adunare care colectează tokenurile de la experți.

Extinde

Paralelismul bazat exclusiv pe experți nu sescalează în mod ideal, deoarece nodurile dintr-un cluster ar putea fi inactive în timp ce modelul rulează straturi care nu țin de experți, cum ar fi atenția (attention), norm și eșantionarea (sampling). Cu toate acestea, calculul acestor straturi, în principal atenția, poate beneficia și de paralelismul de tensor (TP). Majoritatea modelelor se bazează pe atenția cu capete multiple (multi-head attention), ceea ce înseamnă că acele capete de atenție și proiecțiile lor Q, K și V corespunzătoare pot fi, de asemenea, fragmentate pe dispozitive, replicând sau adunând secțiunile între diverse straturi. Dacă atenția este distribuită pe toate dispozitivele, un all-gather poate sincroniza activările, permițând replicarea rutării, ceea ce necesită primitive de sincronizare similare cazului paralelismului experților pentru o implementare eficientă. Cu toate acestea, există limite ale paralelismului la acest nivel, deoarece reducerea numărului de capete de atenție sub un anumit prag va genera beneficii descrescătoare.

Extinde

Pentru a utiliza optim toate dispozitivele și a susține un grad foarte ridicat de paralelism al experților, de până la 128 sau 256 de GPU-uri, este necesar paralelismul de date (DP). În cadrul acestei scheme, dispozitivele pe care este împărțit un model sunt grupate pentru a gestiona cererile în mod concurent, calculând atenția și menținând cache-urile KV fragmentate în grupul lor local. Mai multe instanțe ale acestor grupuri paralele colaborează la evaluarea experților, fiecare găzduind un subset diferit de ponderi ale experților. În funcție de numărul de capete de atenție, un grup se poate întinde de obicei până la dimensiunea unui întreg nod, deoarece comunicația intra-nod este mai rapidă. De exemplu, în figura de mai sus, un rang DP deservește independent două cereri, A și B, gestionând atenția, norm și orice altă evidență pentru cereri. Celălalt rang DP procesează o cerere distinctă C. Cu toate acestea, primul nod găzduiește jumătate dintre experți, în timp ce celălalt nod cealaltă jumătate; prin urmare, după rutare, tokenurile din A și B ar putea fi trimise către al doilea nod și invers. Acest lucru duce la o problemă de comunicație rară (sparse): fiecare dispozitiv poate trimite un număr diferit de tokenuri către orice alt rang de destinație. Primitivele existente din torch, în principal all_to_all, nu sunt deosebit de bine potrivite, deoarece ar putea necesita o formă de umplere (padding) sau sincronizare GPU-CPU și difuzare de metadate. Pentru a implementa eficient comunicația, sunt necesare kerneluri personalizate pentru a procesa informațiile de rutare și a iniția comunicația de pe dispozitive fără intervenția CPU-ului. După rutare, un kernel de expediere (dispatch) trebuie să trimită tokenurile către rangurile către care au fost rutate, în timp ce pe partea de combinare trebuie colectate activările aparținând cererilor din grupul DP curent. În plus, activitatea trebuie echilibrată în cadrul unui grup DP care poate avea mai multe dispozitive, pentru a minimiza difuzarea și a replica corect tokenurile pe toate rangurile pentru straturile ulterioare.

NVSHMEM

NVSHMEM este o implementare OpenSHMEM specifică NVIDIA, care oferă facilități de comunicație portabile între dispozitive, abstractizând complexitatea hardware-ului subiacent. API-ul poate exprima citiri și scrieri de la dispozitiv la dispozitiv, care sunt mapate la primitivele straturilor de transport individuale. Atât NVLink, cât și RDMA sunt acceptate, oferind un grad de portabilitate și permițând kernelurilor NVSHMEM să opereze în cadrul unui nod sau pe un cluster de mai multe noduri. În kernelurile noastre, utilizăm scrieri în memorie la distanță și operații de semnalizare pentru a transfera date și a sincroniza dispozitivele participante.

Operațiile NVSHMEM sunt construite în jurul conceptului de memorie simetrică: ele operează pe buffere care au fost alocate pe toate dispozitivele participante la comunicația între dispozitive. Un dispozitiv local poate utiliza o adresă derivată în propriul buffer local împreună cu indexul rangului țintă pentru a specifica destinația unei operații. Figura de mai jos ilustrează acest concept: ambele GPU-uri alocă buffere simetrice de aceeași dimensiune, păstrând pointeri src către acestea. Primul rang dorește să trimită 3 întregi către al doilea, plasându-le la începutul bufferului: nvshmem_int_put_nbi derivă adresa de pornire din bufferul local, specificând dispozitivul țintă. Al doilea rang derivă un offset din propriul buffer, trimițând un element către primul dispozitiv, aplicând un offset de unu. În timp ce adresele de destinație trebuie să fie întotdeauna buffere simetrice alocate folosind nvshmem_alloc, bufferele sursă pot fi regiuni arbitrare de memorie ale dispozitivului, cu condiția să fie preînregistrate cu NVSHMEM.

Deși NVSHMEM oferă o gamă largă de primitive, kernelurile noastre se bazează pe doar 3 funcții, construind toată sincronizarea și delimitarea (fencing) pe baza acestora.

Extinde

nvshmemx_putmem_signal_nbi_warp: Transferă un bloc de date de pe un dispozitiv pe altul, setând în același timp un semnalizator (flag) pe dispozitivul la distanță. Operația setează (NVSHMEM_SIGNAL_SET) sau incrementează (NVSHMEM_SIGNAL_ADD) o locație pe 64 de biți. Semnalizatorul este actualizat după ce întregul bloc de memorie este transferat. Dacă dispozitivul la distanță a observat o modificare a semnalizatorului, poate accesa în siguranță bufferul din propria memorie. Această funcție este utilă pentru comasarea transferului de date și a sincronizării.

nvshmemx_signal_op operează pe o singură locație de memorie, de obicei un semnalizator pe 64 de biți, setându-l sau incrementându-l atomic. Este utilă pentru transmiterea de metadate și sincronizarea dispozitivelor.

nvshmem_uint64_wait_until este utilizat la capătul receptor al unui semnal, pentru a interoga (poll) un semnalizator până când cel de la distanță îl actualizează.

Asigurarea oricărei forme de ordonare între operații prin nvshmem_quiet sau nvshmem_fence este costisitoare, prin urmare implementarea noastră a kernelurilor evită utilizarea altor primitive NVSHMEM și minimizează orice dependențe în afară de perechile individuale de trimitere-primire. Când trimitem date, utilizăm întotdeauna versiunea neblocantă a funcțiilor, fără a aștepta ca datele să fie măcar trimise în afara rangului local. O altă sincronizare implicită în cadrul grupului de kerneluri asigură că aceste buffere nu sunt suprascrise distructiv în timp ce plăcile NIC au nevoie de acces la ele.

NVSHMEM utilizează GPUDirect RDMA astfel încât placa NIC să poată citi și scrie direct în memoria GPU pentru schimbul de date, fără a implica CPU-urile. În plus, pe plăcile NIC ConnectX, NVSHMEM acceptă comunicația inițiată de GPU (cunoscută și sub denumirea de InfiniBand GPUDirect Async sau IBGDA), care permite GPU-ului să trimită operații RDMA direct către NIC și să interogueze direct coada de finalizare a NIC-ului. Pe platformele care nu acceptă comunicația inițiată de GPU, NVSHMEM generează un fir „proxy” pe CPU-ul gazdă pentru a iniția comunicația și a interoga finalizarea în numele GPU-ului. Programul NVSHMEM este portabil indiferent dacă comunicația inițiată de GPU este acceptată sau nu. Cu toate acestea, comunicația inițiată de GPU reduce semnificativ latența, deoarece ocolește complet ocolul prin CPU.

Kerneluri portabile

Implementăm comunicația MoE printr-o pereche de kerneluri de expediere (dispatch) și combinare (combine). Kernelurile de expediere sunt responsabile de citirea tokenurilor și a informațiilor de rutare pe fiecare rang, expediindu-le către experții adecvați. Kernelurile de combinare colectate activează produsele generate de experți și le trimit înapoi la rangurile sursă, calculând în același timp media ponderată de la experții selectați pe baza ponderilor calculate de ruter. Kernelurile sunt împărțite în continuare într-o componentă de trimitere și una de primire, pentru a permite suprapunerea transferurilor de date cu calculele. Kernelurile de trimitere sunt toate neblocante și nesincronizate: ele pur și simplu expediază toate scrierile către destinațiile la distanță. La celălalt capăt, kernelurile de combinare citesc doar din memorie, așteptând până când toate datele necesare au fost transferate. După expedierea activității către plăcile NIC, în timp ce datele sunt transferate asincron prin cablu, GPU-urile pot efectua alte sarcini utile la nivel local, cum ar fi aplicarea de experți partajați sau calcularea atenției.

Fiecare pereche de kerneluri are propriile buffere de memorie simetrică, fiecare rang alocând stocare privată pentru a primi date de la toate celelalte ranguri. Numărul maxim de tokenuri care pot fi expediate dintr-un grup DP este fix pe toate rangurile, ceea ce setează, de asemenea, limita superioară pe care fiecare rang o poate primi pentru fiecare expert local de la fiecare rang de date (DP). Acest lucru permite rangurilor expeditoare să deriveze o adresă unică pe rangul de destinație în care să scrie, fără a necesita vreo sincronizare între ele. După ce datele sunt primite, kernelurile rearanjează tokenurile pentru a le dispune în memorie în formatul optim cerut de kernelurile receptoare. Deși bufferele au o dimensiune considerabilă, ele sunt reutilizate în toate straturile secvențiale ale unui model.

Extinde

Singura formă de sincronizare globală, așa cum este ilustrată în figura de mai sus, este implementată în kernelul de combinare-primire (combine-receive). Odată ce datele sunt expediate de la trimitere, un rang se întoarce și își continuă execuția atunci când primește toate tokenurile necesare de la celelalte ranguri și intră în kernelurile de trimitere pentru combinare de îndată ce acestea sunt procesate. Bariera din combine-receive asigură că niciun rang nu poate avansa și începe expedierea-trimiterea în timp ce orice alt rang încă așteaptă să primească date, ceea ce ar putea provoca o suprapunere distructivă. Sincronizarea se face în kernelul de combinare pentru simplificare; ca optimizare, aceeași barieră ar putea fi implementată între kernelurile de primire a expedierii și de trimitere a combinării, reducând și mai mult latența prin oferirea unui timp suplimentar plăcilor NIC pentru a finaliza transferurile în timp ce GPU-ul rulează alte kerneluri.

Ambele kerneluri sunt împărțite pe toate SM-urile disponibile ale dispozitivului: în timp ce kernelurile de trimitere a expedierii și primire a combinării trebuie să se paralelizeze pe un număr maxim de tokenuri per rang (max_m), kernelurile de primire a expedierii și trimitere a combinării trebuie să gestioneze acel număr de tokenuri de la fiecare rang și expert local, pentru un maxim de max_m * (num_experts // EP) * (EP // TP).

Expediere (Dispatch)

Kernelul de expediere este responsabil pentru trimiterea tokenurilor și a factorilor de scalare a blocurilor către omologii (peers) care găzduiesc experții către care au fost rutate tokenurile. Partea de trimitere se bazează pe specializarea warp pentru a paralelizat două sarcini: agregarea informațiilor de rutare pentru a determina câte tokenuri sunt trimise de la fiecare rang și împachetarea tokenurilor în mesaje care urmează să fie expediate la distanță. Partea receptoare așteaptă mai întâi primirea tuturor numărului de tokenuri, apoi citește și despachetează sarcinile utile din memoria partajată în tensoarele care vor fi transmise către kernelurile Grouped GEMM din aval. De asemenea, stochează informații în buffere partajate cu kernelurile de combinare, pentru a indica unde ar trebui trimis înapoi fiecare token. Această informație este necesară deoarece kernelul de primire rearanjează tokenurile în buffere concontiguous într-o ordine nedeterministă.

Extinde

În partea de trimitere, un singur warp de 32 de fire pe fiecare bloc este responsabil pentru citirea informațiilor de rutare și numărarea numărului de tokenuri pe care se așteaptă le primească fiecare expert destinație de la acest rang, trimițând numărul plus unu utilizând nvshmem_signal_op. Numărul este incrementat cu unu, deoarece tranziția de la zero la diferit de zero de la distanță semnifică primirea numărului de la un rang sursă. În paralel, warp-urile rămase cooperează pentru a copia tokenurile în memoria simetrică pe toate blocurile în paralel, împachetând activările, factorii de scalare și indexul acestora de pe rangul local într-o bucată contiguă. Indexul este necesar de către expeditorul de combinare pentru a determina adresa în care va fi scris tokenul. În continuare, după ce ne asigurăm că toate datele au fost copiate printr-o barieră între grupurile warp, warp-urile operează din nou independent, fiecare trimițând același buffer către un expert diferit în paralel. Tokenurile sunt trimise folosind nvshmemx_putmem_signal_nbi_warp, care crește, de asemenea, atomic numărul de tokenuri trimise de la rangul local pe dispozitivul la distanță. În cadrul unui grup DP, deoarece fiecare rang deține o replică a tokenului de trimis, expedierea este echilibrată uniform, fiecare dispozitiv trimițând un subset de tokenuri.

La capătul de recepție, toate kernelurile se sincronizează mai întâi cu rangurile expeditoare așteptând ca numărul total de tokenuri să fie trimis. Ulterior, toate așteaptă ca numărătoarea de tokenuri trimise incrementată atomic să se stabilizeze la numărul total, indicând faptul că toate sarcinile utile de la rangurile sursă au fost de asemenea trimise, datorită semanticilor apelului putmem. Kernelurile interoghează numărătoarea folosind nvshmem_uint64_wait_until, paralelizând operația pe toate blocurile și firele de execuție. Ulterior, o barieră încrucișată între blocuri asigură că niciun bloc nu citește bufferele decât dacă toate datele au fost primite corect. Răspândite pe blocuri și sincronizate printr-un contor incrementat atomic, tokenurile sunt copiate din memoria simetrică în buffere contigue obișnuite care devin intrările pentru kernelurile din aval care implementează experții. Rangul sursă, indexul expertului și indexul tokenului sunt stocate separat, indicând exact locația în care kernelul de combinare trebuie să trimită activările. Chiar dacă tokenurile dintr-un grup DP sunt trimise de pe dispozitive diferite, ele sunt grupate împreună pentru a fi transmise expertului corespunzător.

Combinare (Combine)

Kernelurile de combinare sunt împărțite din nou în jumătăți de trimitere și de recepție: expeditorii copiază activările pe 16 biți necuantizate către dispozitivele la distanță, kernelurile de recepție așteptând ca datele să fie trimise și calculând media ponderată a tuturor contribuțiilor experților la nivel local. În plus, ele acționează și ca o barieră pentru a sincroniza secvența de expediere-combinare: fiecare rang setează un semnalizator pe fiecare omolog la intrarea în kernelurile de trimitere, kernelurile de recepție nefiind lăsate să revină decât dacă observă că semnalizatorul este setat. Latența de sincronizare este minimă, deoarece se suprapune peste comunicația și calculul efectiv.

Extinde

Pe partea de trimitere, kernelurile parcurg în paralel lista de tokenuri atribuite tuturor experților locali, scriindu-le într-un buffer de pe rangul de destinație. Rangul țintă, indexul expertului și indexul tokenului sunt citite din bufferele per-token populate de kernelurile de dispersie la primirea tokenurilor. Fiecare expeditor are propria regiune de memorie privată per expert în care să scrie, așa cum se indicând în figura de mai sus, evitând necesitatea sincronizării. Similar cu expedierea, combinarea incrementează atomic contoarele per-token de pe rangul de destinație pentru a indica primirea datelor: atunci când contorul se potrivește cu numărul de experți către care a fost expediat un token, se poate accesa conținutul tokenului.

În kernelul de recepție, lista de tokenuri este parcursă în paralel pe mai multe blocuri, așteptând ca conținutul acestora să sosească prin interogarea semnalizatorului setat de operația de semnalizare. La sosire, sarcinile utile sunt citite din bufferele private, tabelul de rutare indicând din ce buffer să se citească și ce pondere să se atribuie fiecărui expert. Rezultatele sunt apoi scrise în tensoare alocate extern, kernelul finalizând execuția odată ce toate dispozitivele au trecut de barieră. În cadrul unui grup DP, toate rangurile primesc o copie a fiecărei activări a experților pentru a-și calcula propriile replici.

Benchmarks

Evaluăm kernelurile noastre pe un cluster de 128 de GPU-uri răspândite pe 16 noduri conectate prin plăci NIC InfiniBand și CX-7. Încercăm atât GPUDirect Async (IBGDA), cât și Reliable Connection (RC) cu un proxy CPU. Le comparăm cu implementarea DeepSeek puternic optimizată, precum și cu primitivele dense furnizate de PyTorch (prin NCCL) sau NVSHMEM.

Extinde

Deși pe operația densă NVSHMEM all-to-all, performanța IBRC și IBGDA este similară (6378 µs vs 6180 µs), IBGDA este semnificativ mai rapidă cu kernelurile sparse. Adunând Dispatch și Combine, IBGDA utilizează 902 µs, în timp ce IBRC necesită 3223 µs - o îmbunătățire de 3,6x a latenței. În timp ce all-to-all este limitat de lățimea de bandă, kernelurile sparse difuzează ordine de mărime mai puține date, fiind limitate de latență. Prin declanșarea transferurilor de rețea direct de la un GPU, fără a necesita un proxy CPU pentru a coordona GPU-ul și placa NIC, latența end-to-end este redusă semnificativ.

În timp ce kernelurile noastre portabile sunt de aproximativ 2x mai lente decât kernelurile DeepSeek puternic optimizate, ele îmbunătățesc latența de 10x în comparație cu kernelurile dense (902 µs vs 9944 µs).

În plus, împărțirea în componente de trimitere-recepție permite, de asemenea, mascarea unei parte din latență, spre deosebire de primitivele de bibliotecă.

Extinde

Pe un singur nod (EP8), NVSHMEM utilizează NVLINK pentru transport, oferind o latență mai mică și un debit mai mare decât rețelele inter-nod. Kernelurile noastre portabile sunt de aproximativ 2,5x mai rapide decât DeepEP pe un singur nod (186 µs vs 481 µs).

Îmbunătățiri ulterioare

Kernelurile descrise aici depășesc performanțele primitivelor încorporate ale cadrelor (frameworks) ML și oferă o performanță decentă fără a fi supraspecializate pentru transporturi specifice între dispozitive, cum ar fi NVLink sau InfiniBand. Pe lângă oportunitățile deja menționate, se pot obține câștiguri suplimentare de performanță prin înlocuirea primitivelor de comunicație cu versiuni mai specializate. De exemplu, în NVLink, utilizarea memoriei simetrice ar putea fi înlocuită cu buffere partajate între dispozitive, eliminând o parte din copiere și permițând scheme de sincronizare mai fine, pe tokenuri individuale în loc de loturi de tokenuri. În InfiniBand, o implementare ar putea accesa direct perechile de cozi subiacente, eliminând necesitatea de a transfera atât numărul de tokenuri așteptate, cât și cele trimise, și de a goli tokenurile din cozi în bufferele de ieșire pe măsură ce sosesc. Cu toate acestea, astfel de implementări vin cu prețul portabilității, în timp ce implementarea noastră permite evaluarea clusterelor cu configurații specifice înainte de a aloca efort de dezvoltare pentru optimizarea pentru un anumit hardware.

Concluzie

Am prezentat o bibliotecă portabilă de înaltă performanță pentru comunicația MoE care obține o performanță de 10x mai rapidă comparativ cu comunicația standard all-to-all, menținând în același timp compatibilitatea între diverse configurații hardware. Pe implementările pe un singur nod cu NVLink, soluția noastră demonstrează o latență cu 2,5x mai mică decât implementările anterioare.

Abordarea noastră echilibrează performanța cu portabilitatea prin inovații cheie, inclusiv suport pentru comunicația inițiată de GPU, o arhitectură de kernel împărțit care permite suprapunerea calculului cu comunicația și expedierea eficientă a tokenurilor folosind primitive NVSHMEM minime. Deși este cu aproximativ 2x mai lentă decât implementările extrem de specializate pe configurații cu mai multe noduri, biblioteca noastră oferă o flexibilitate superioară în diferite medii de rețea (NVLink, CX-7 și EFA).

Pe măsură ce modelele MoE continuă să se dezvolte la scară, strategiile de comunicație eficiente precum a noastră vor deveni din ce în ce mai importante pentru implementarea practică. Implementarea noastră complet open-source este disponibilă la https://github.com/ppl-ai/pplx-kernels.

Referințe

DeepEP

Raport tehnic DeepSeek-V3

NVIDIA NVSHMEM

Citare

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