Activarea modelelor cu mii de miliarde de parametri pe AWS EFA

La Perplexity, utilizăm cele mai bune modele pentru produsul nostru, pentru API-urile noastre și pentru echipele noastre de cercetare. Modelele mari open-source de tip Mixture-of-Experts, cum ar fi Kimi-K2, prezintă provocări particulare, deoarece cele mai mari noduri de inferență cu 8 unități GPU NVIDIA H200 nu le pot găzdui în mod eficient, fiind necesare implementări pe mai multe noduri. Vă prezentăm un set de nuclee pentru paralelismul experților care obțin latențe de ultimă generație pe ConnectX-7, depășind performanța DeepEP. Aceleași nuclee sunt, de asemenea, primele care obțin latențe Viabile pe AWS Elastic Fabric Adapter (EFA), permițând implementarea de modele cu mii de miliarde de parametri.

Încercați nucleele noastre pe GitHub și citiți lucrarea de cercetare completă pe arXiv.

Introducere

Arhitectura Mixture-of-Experts (MoE) a devenit arhitectura standard pentru scalarea la sute de miliarde sau mii de miliarde de parametri, oferind în același timp o latență rezonabilă în momentul inferenței. Modelele MoE înlocuiesc stratul dens al unui transformer cu un set de experți și un strat mic de rutare liniară care decide cu ce experți trebuie înmulțit un token. Prin gruparea tokenurilor asociate aceluiași expert, problema devine rară și ușor de paralelizat, deoarece diferenții experți pot fi acum găzduiți pe diferite unități GPU distribuite pe mai multe noduri.

Spre deosebire de alte forme de paralelism, cum ar fi paralelismul de tip Tensor (TP) și paralelismul de date (DP), care pot fi implementate cu ușurință folosind biblioteci de comunicare colectivă precum torch.distributed și NCCL, rutarea MoE implică o comunicare rară de tip peer-to-peer. Un nucleu de expediere trebuie să împartă setul de tokenuri primite și să le expedieze către rândurile care găzduiesc experții, în timp ce problema sa duală, nucleul de combinare, trebuie să primească omologii lor procesați și să calculeze media lor ponderată. Aceste transferuri nu se mapează la colective și necesită nuclee personalizate pentru a obține o latență redusă.

La Perplexity, am dezvoltat atât nuclee portabile între noduri, cât și nuclee specializate în cadrul aceluiași nod, prezentând lucrarea noastră în mai multă profunzime într-un articol anterior despre comunicarea MoE. De asemenea, am arătat într-o altă postare pe blog că implementarea pe mai multe noduri poate aduce o latență mai mică și un debit mai mare pentru modelele MoE mari.

În acest articol, analizăm în profunzime implementarea noilor noastre nuclee între noduri, care îmbunătățesc considerabil performanța lucrării noastre anterioare bazate pe NVSHMEM. Pe lângă prezentarea unor rezultate de ultimă generație pe ConnectX-7, demonstrăm, de asemenea, că în cadrul timpului nostru de execuție pentru inferență, nucleele pot fi utilizate pentru a deservi în mod eficient cele mai mari modele open-source.

Inferență prin InfiniBand

Paralelismul de tip tensor și paralelismul experților pot scala eficient în cadrul unui singur nod, deoarece NVLink oferă un debit ridicat (900 GB/s) la latențe de ordinul microsecundelor. În timp ce sistemele NVIDIA GB200 și GB300 sunt echipate cu memorie suficientă pentru a stoca ponderile și cache-urile KV ale modelelor cu mii de miliarde de parametri, implementările pe instanța AWS p5en cu până la 8 unități GPU H200 sunt limitate la 1120 GB de HBM, care trebuie împărțită între ponderile modelului și cache-urile KV. În consecință, anumite modele necesită implementări între noduri prin InfiniBand.

Deși InfiniBand este în continuare extrem de rapid, acesta oferă un debit de doar 400 Gbps (50 GB/s), adăugând zeci sau sute de microsecunde de latență pentru rutarea MoE. Din fericire, latențele suplimentare pot fi ascunse prin utilizarea unor experți partajați, micro-lotizare și suprapunerea calcul-comunicare pentru a oferi un debit viabil la un cost competitiv. DeepSeek a fost pionierul acestor tehnici prin intermediul DeepSeek-R1, implementat pe clustere de unități GPU H800 și adaptoare NVIDIA ConnectX-7 InfiniBand.

Pe clusterele AWS, situația este diferită din cauza utilizării adaptoarelor lor personalizate Elastic Fabric Adapters (EFA). Chiar dacă acestea oferă un debit maxim de 400 Gbps pentru colectivele utilizate de obicei în sarcinile de lucru de antrenament, ele sunt ușor sub nivelul ConnectX-7 în ceea ce privește dimensiunile mesajelor schimbate în timpul expedierii și combinării MoE. În plus, EFA nu acceptă GPUDirect Async, necesitând un fir proxy CPU care face legătura între unitatea GPU și placa de rețea (NIC) pentru a iniția transferurile. Pe lângă complexitatea adăugată, tranzacțiile PCIe suplimentare adaugă un cost suplimentar de ordinul microsecundelor la toate transferurile. Cadrele precum NVSHMEM le permit scriitorilor de nuclee să ofere implementări independente de dispozitiv, însă API-urile generice generează costuri suplimentare care fac ca implementările să fie ineficiente atât în ceea ce privește costul, cât și latența.

pplx-kernels bazate pe NVSHMEM

Nucleele noastre anterioare au abstractizat implementarea de transport subiacentă prin NVSHMEM. Ele s-au bazat pe nvshmemx_putmem_signal_nbi_warp pentru a transfera tokenurile unul câte unul, utilizând contoare atomice pentru a sincroniza nodurile partenere. Deși acest lucru a oferit o performanță rezonabilă pe plăcile de rețea ConnectX-7 cu IBGDA, implementarea ConnectX-7 IBRC bazată pe proxy a fost semnificativ mai lentă, în timp ce implementarea EFA nu a reușit să ofere latențe sub o milisecundă.

Am folosit nucleele ca bancă de testare pentru hardware-ul subiacent, stabilind în același timp cea mai potrivită interfață pe care să o expunem modelelor noastre. Am emis ipoteza că transferurile frecvente și interacțiunile dintre nuclee și plăcile de rețea sau firele proxy au afectat performanța. În plus, implementarea noastră nu s-a bazat pe nicio garanție din partea hardware-ului, în afara fiabilității, sugerând că am putea accelera lucrurile prin reducerea costului suplimentar al primitivelor de sincronizare.

În urma acestor observații, am conceput nucleele noastre bazate pe EFA care expun aceeași interfață, optimizând interacțiunile dintre nucleul dispozitivului și proxy-ul gazdă sub capotă pentru a obține o performanță maximă.

Rutare MoE all-to-all prin EFA

Extinde

Noile nuclee urmează o arhitectură hibridă CPU-GPU, în care nucleele GPU interacționează cu modelul de pe dispozitiv, iar un fir care rulează pe CPU-ul gazdă gestionează interacțiunile cu placa de rețea. Noile nuclee reutilizează același TransferEngine pe care l-am construit pentru transferurile de cache KV.

Rutarea MoE este efectuată de un nucleu de expediere, care primește de la model tokenurile și experții specifici fiecărui token pentru a le asigura rutarea, transmițându-le prin rețea către rândurile care găzduiesc experții. După ce nucleele GEMM grupate înmulțesc tokenurile cu ponderile experților, nucleul de combinare transferă tokenurile către rândurile lor sursă și calculează media lor ponderată înainte de a le returna modelului.

Atât nucleele de expediere, cât și cele de combinare sunt împărțite într-o jumătate de trimitere și o jumătate de recepție pentru a facilita suprapunerea și micro-lotizarea. Expeditorul amestecă tokenurile în bufferele de trimitere, astfel încât acestea să poată fi scrise către nodurile partenere printr-o singură operațiune de scriere RDMA. După executarea expeditorului, unitatea GPU poate efectua alte sarcini (cum ar fi înmulțirea experților partajați) în timp ce transferurile prin rețea sunt în așteptare în fundal. Nucleele receptoare se blochează până când toate transferurile au fost recepționate. În fundal, un fir de lucru de pe partea gazdă gestionează transferurile: este semnalat de unitatea GPU atunci când bufferele de trimitere au fost populate, după care expediază scrierile RDMA. Odată ce placa de rețea semnalează finalizarea, deblochează receptorul, care returnează tokenurile modelului.

Nucleele sunt strâns cuplate cu gazda, comunicând între ele atât prin memorie unificată, cât și prin GDRCopy. Memoria unificată CUDA permite ca memoria dispozitivului sau memoria gazdei să fie mapată în spațiile de adrese ale ambelor, gestionând în mod transparent comunicarea prin PCIe.

În timp ce memoria unificată este potrivită pentru transferuri în bloc, interogarea cu latență redusă se realizează prin intermediul GDRCopy. GDRCopy utilizează GPUDirect RDMA pentru a citi din și a scrie în memoria GPU la o latență de ordinul microsecundelor.

Pe lângă RDMA, în cadrul unui nod nucleele folosesc și NVLink pentru a face schimb de tokenuri. Acest lucru poate descărca o fracțiune semnificativă din trafic, până la un sfert pentru EP32 și până la o optime pentru EP64, obținând latențe generale mai mici.

Buffere de trimitere și recepție

Designul bufferelor de trimitere pentru nucleul de expediere este esențial pentru obținerea unor latențe viabile care implică proxy-ul gazdă. Dacă există N rânduri, fiecare expediind T tokenuri, fiecare rutat către R experți din E, un rând poate primi până la N * T * max(R, E / N) tokenuri în cel mai rău caz posibil, dacă toate sunt rutate către acesta. În plus, această limită minimă poate fi atinsă numai atunci când expeditorii își dispun scrierile în mod contiguu, deoarece bufferele de recepție private ar utiliza prea multă memorie. Pentru ca expeditorii să își poată dispune scrierile consecutiv fără conflicte, informațiile complete de rutare trebuie distribuite către fiecare rând. În consecință, nucleele noastre de expediere fac mai întâi schimb de numărăori de tokenuri per expert, după care fiecare poate determina unde merg scrierile lor și de la ce offset primesc tokenuri de la fiecare rând.

Simpla schimbare a informațiilor de rutare ar genera latențe ridicate inutil, deoarece fiecare rând ar trebui să aștepte ca scrierile să sosească înainte de a expedia propriile scrieri. În schimb, rezervăm o cantitate mică de spațiu privat per expeditor pe fiecare rând receptor și expediem o parte dintre tokenuri odată cu scrierile informațiilor de rutare. Acest lucru asigură utilizarea întregii lățimi de bandă pe parcursul celei mai mari părți a procesului de expediere. După ce rutele sunt schimbate, tokenurile rămase sunt transferate printr-o singură scriere per nod partener la distanță. Schimbul de informații de rutare și bufferele de dimensiune fixă, împreună cu transferurile care le populează, sunt ilustrate în figura de mai jos.

Extinde

Nucleele de combinare sunt substanțial mai simple, deoarece informațiile de rutare le sunt deja disponibile din etapa de expediere. Ele pot copia tokenurile printr-o singură scriere per nod partener, după care citesc din bufferul de recepție pentru a calcula media ponderată a tokenurilor în tesoarele de ieșire.

Pentru NVLink, bufferele de trimitere și recepție din cadrul unui nod sunt mapate în nodurile partenere pentru ca acestea să poată trimite sau prelua direct tokenuri de la rândurile învecinate.

Expediere și combinare

Figurile de mai jos detaliază sincronizarea și interacțiunile dintre unitatea GPU, CPU și placa de rețea, prezentând o iterare a expedierii și combinării intercalate cu nucleele modelului.

Lucrul începe în trimiterea de expediere, nucleul agregând informațiile de rutare și calculând numărătorile de tokenuri per expert prin acumulare atomică în memoria partajată dintr-un singur bloc. Odată ce numărătorile de tokenuri sunt disponibile, ele sunt copiate în memoria gazdei și proxy-ul este notificat imediat să înceapă expedierea informațiilor de rutare către toate celelalte rânduri partenere. În paralel, expedierea continuă să calculeze suma cumulativă a numărătorilor de tokenuri, identificând offseturile necesare pentru a împacheta tokenurile în buffere de trimitere contigue per rând. Odată ce offseturile sunt disponibile, toate tokenurile care urmează să fie trimise prin EFA sunt împachetate în bufferele de trimitere. Proxy-ul este notificat să posteze scrierile în bufferele de trimitere ale nodurilor partenere inter-nod. În cadrul aceluiași nod, nucleele împing, de asemenea, imediat tokenurile în buffere private din cadrul aceluiași nod prin NVLink.

Extinde

Extinde

Odată ce un rând primește informații de rutare de la toate nodurile partenere, le centralizează și calculează diverse offseturi necesare pentru identificarea tokenurilor pentru expediere și combinare. Această muncă se face pe CPU în timp ce transferurile sunt deja în așteptare, astfel încât nu este critică din punct de vedere al latenței. Ulterior, odată ce toate offseturile sunt cunoscute, scrierile pentru restul tokenurilor, care depășesc dimensiunea bufferului privat, sunt postate și faza de trimitere a expedierii este completă.

La capătul receptor, nucleul extrage mai întâi din același nod tokenurile care depășesc capacitatea bufferelor private prin NVLink. Apoi, se blochează până când proxy-ul gazdă semnalează finalizarea tuturor transferurilor. Odată ce semnalul este primit, tokenurile sunt amestecate din bufferele de recepție în tisorul de ieșire transmis către grupul GEMM. Deoarece bufferele sunt reutilizate între expediere și combinare, atât barierile NVLink, cât și cele RDMA se asigură că bufferele receptoare au fost citite înainte de a permite următorului nucleu să le suprascrie.

Nucleele de combinare urmează un model aproximativ similar, amestecând tokenurile, semnalând proxy-ul și așteptând finalizarea. Cu toate acestea, metadatele pentru postarea scrierilor lor sunt ușor disponibile și toate offseturile de la distanță sunt cunoscute, prin urmare toate tokenurile transferate în cadrul aceluiași nod sunt împinse prin NVLink către receptori, iar tokenurile inter-nod sunt transferate printr-o singură scriere. Înainte de a returna rezultatele, nucleul acumulează tokenurile din bufferele receptoare, scriindu-le în tesoarele de ieșire.

TransferEngine pentru rutarea MoE

Pentru a realiza prototipul implementării noastre, am reutilizat inițial motorul de transfer (Transfer Engine) construit pentru preumplerea și decodarea dezarhivate în vederea emiterii și urmăririi scrierilor RDMA. Cu toate acestea, spre deosebire de decodarea dezarhivată, unde fiecare nod postează un număr mare de scrieri către același nod partener, expedierea MoE postează transferuri mici către toate nodurile partenere într-o singură operațiune. Din cutie, performanța nu a fost optimă, iar implementarea noastră a rămas în urmă față de DeepEP pe ConnectX-7 cu aproximativ 20 us. În consecință, am specializat TransferEngine pentru a accepta două operațiuni, dispersare (scatter) și barieră (barrier), pentru a optimiza interacțiunile cu un număr mare de noduri partenere.

Atât operațiunile de dispersare, cât și cele de barieră expun grupuri de noduri partenere care pot fi pre-înregistrate pentru a pre-popula structurile de date utilizate de implementarea de transport subiacentă (șablonoare Work Request (WR)). Dispersarea (scatter) copiază diferite felii din același buffer sursă în diferite offseturi pe toți membrii unui grup de noduri partenere. Bariera semnalează toate nodurile partenere printr-un mesaj imediat. Încadrarea și gestionarea în masă a acestor operațiuni economisește microsecunde valoroase pe calea critică. Pe EFA, unde două plăci de rețea de 200 Gbps se agregoază pentru a oferi o lățime de bandă de 400 Gbps, fragmentăm grupurile de noduri partenere pe cele două plăci de rețea în loc să fragmentăm octeții fiecărui transfer.

Deși inițial am dezvoltat TransferEngine exclusiv pe EFA folosind libfabric, nucleele MoE funcționează pe ConnectX-7 fără a schimba o singură linie de cod după adăugarea suportului pentru ConnectX-7 în TransferEngine folosind libibverbs. Comparativ cu EFA bazat pe SRD (Scalable Reliable Datagram), ConnectX-7 necesită o muncă suplimentară pentru configurarea conexiunii și gestionarea nodurilor partenere. Folosim o pereche de cozi Unreliable Datagram (UD) pentru a inițializa conexiunea fiabilă (RC - Reliable Connection). Creăm două perechi de cozi pentru fiecare nod partener la distanță, una pentru SEND/RECV bidirecțional și una pentru WRITE_IMM unidirecțional, deoarece atât RECV, cât și WRITE_IMM consumă operațiunea RECV în ordine, în timp ce cele două operațiuni necesită dimensiuni diferite de buffer. În plus, relaxăm ordonarea tranzacțiilor PCIe între ConnectX-7 și memoria GPU. Similar cu EFA, implementăm optimizarea de șablonoare WR. În plus, atunci când trimitem mai multe mesaje mici către același nod partener, folosim înlănțuirea WR pentru a îmbunătăți rata mesajelor.

Deoarece lucrul pe partea GPU este identic pe ambele adaptoare, am împins implementarea de transport subiacentă la limitele sale pentru a depăși performanța DeepEP, evidențiind faptul că decalajul dintre rutarea MoE pe EFA și pe ConnectX-7 este determinat de capabilitățile hardware-ului subiacent.

Evaluare

Latența de expediere/combinare

Am evaluat latențele de expediere și combinare pe sisteme H200 echipate cu adaptoare ConnectX-7 și EFA pe configurația standard DeepSeek-V3 și Kimi-K2 cu o dimensiune ascunsă de 7168, expediere fp8 scalată pe blocuri și combinare bf16. Luăm în considerare o dimensiune de lot de 128 pentru expediere și de 4096 pentru combinare. Pe ConnectX-7, facem comparația cu DeepEP și cu nucleele noastre pplx anterioare bazate pe NVSHMEM, utilizând atât GPUDirect Async (IBGDA), cât și un proxy gazdă (IBRC). Obținem latențe combinate de 459 us, 582 us și 692 us pe EP16, EP32 și, respectiv, EP64. Depășim alte nuclee bazate pe EFA care se bazează pe proxy-uri generice, cum ar fi UCCL-EP, care obțin 519 us, 966 us și 1159 us.

Extinde

La decodare, în ciuda faptului că nucleele noastre sunt cu câteva microsecunde mai lente decât DeepEP, obținem latențe generale mai mici, deoarece nucleele noastre de combinare sunt substanțial mai rapide. Deși folosim un proxy gazdă, datorită transferurilor noastre în bloc, depășim performanța tuturor implementărilor fără proxy bazate pe IBGDA. De asemenea, arătăm că implementarea proxy-ului contează, deoarece proxy-ul nostru înalt specializat este cu un ordin de mreaddir mai rapid decât implementarea generică NVSHMEM. În ciuda faptului că EFA este de două ori mai lent decât ConnectX-7 la transferul de pachete de 256 KB, care sunt aproape de dimensiunea tipică a pachetului schimbată de rutarea MoE, decalajul dintre cele două adaptoare este per total mai puțin substanțial.

Extinde

Chiar dacă nucleele noastre optimizate pentru decodare pot scala la dimensiunile loturilor de preumplere fără nicio ajustare suplimentară, performanța rămâne în urmă față de DeepEP. Acest lucru se datorează în principal faptului că acumulăm tokenuri doar pe partea receptoare a combinării, în timp ce DeepEP acumulează în expeditor pentru a reduce cantitatea de date transferate. Cu toate acestea, această optimizare poate avea implicații asupra acurateței. În plus, deoarece nucleele transferă date într-o singură scriere, utilizăm mai multă memorie pentru bufferele de trimitere și recepție.

Debit Kimi-K2 / DeepSeek-V3

Pentru a demonstra viabilitatea nuclemelor noastre, facem un benchmark pentru performanța end-to-end a motorului nostru de inferență intern care deservește DeepSeek-V3 (671B parametri) și Kimi-K2 (1T parametri) pe multiple instanțe p5en echipate cu adaptoare Elastic Fabric Adapters (EFA) și unități GPU H200. Rulăm modelul în modul de verificare a ciornelor și raportăm latența presupunând o rată de acceptare de 80% folosind un singur strat MTP.

Pe DeepSeek-V3, facem o comparație cu o linie de bază EP=8 în cadrul aceluiași nod, care rulează nuclee de expediere și combinare eficiente exclusiv prin NVLink. Implementările între noduri egalează sau depășesc în general performanța implementării pe un singur nod, în special pe loturi de dimensiuni medii, demonstrând că scalarea între noduri aduce îmbunătățiri ale debitului general.

Pe Kimi-K2, nucleele fac posibilă servirea. Datorită dimensiunii sale, modelul nu se poate încadra în același nod folosind paralelismul de date (DP) pe un nod H200, prin urmare este necesară o implementare între noduri pentru a deservi modelul. Deoarece modelul are mai puține capete de atenție, obținem latențe mai bune decât DeepSeek-V3 și putem deservi modelul pe EFA la o latență viabilă.

Extinde

Extinde

De asemenea, am luat în considerare micro-lotizarea, deși aceasta a îmbunătățit latențele doar la dimensiuni foarte mari ale loturilor sau atunci când se utilizează mai mult de un token de ciornă. Pentru a beneficia de împărțirea problemelor GEMM grupate în două, dimensiunile loturilor de tokenuri trebuie să fie suficient de dense. Cea mai mare parte a timpului este petrecută în continuare pe expedierea și combinarea MoE, indicând faptul că orice îmbunătățiri ulterioare aduse nuclemelor vor îmbunătăți în consecință latențele și vor facilita micro-lotizarea.

Lucrări viitoare

Lucrăm strâns cu inginerii AWS pentru a îmbunătăți performanța pe EFA. Din partea noastră, vom urmări actualizările aduse libfabric care vor reduce costul suplimentar al planului de date. În plus, vom experimenta efa-direct pentru a reduce și mai mult costurile suplimentare ale spațiului de utilizator și pentru a interfața direct cu hardware-ul.

Referințe

Raport tehnic DeepSeek R1

Raport tehnic Kimi K2

DeepEP

UCCL-EP

Îmbunătățirea performanței rețelei sistemelor HPC folosind NVIDIA Magnum IO NVSHMEM și GPUDirect Async

Citare

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}