Omogućavanje modela s bilijunima parametara na AWS EFA

U tvrtki Perplexity koristimo najbolje modele za naš proizvod, naše API-je i naše istraživačke timove. Veliki open-source modeli s mješavinom stručnjaka (MoE - Mixture-of-Experts), kao što je Kimi-K2, predstavljaju posebne izazove jer najveći inferencijski čvorovi s 8x NVIDIA H200 GPU-ova ih ne mogu učinkovito smjestiti, što zahtijeva implementacije na više čvorova. Predstavljamo skup jezgri za paralelizam stručnjaka koje postižu latencije na vrhunskoj razini na ConnectX-7, nadmašujući performanse DeepEP-a. Iste su jezgre također prve koje postižu održive latencije na AWS Elastic Fabric Adapteru (EFA), omogućujući implementacije modela s bilijunima parametara.

Isprobajte naše jezgre na GitHubu i pročitajte cijeli istraživački rad na arXivu.

Uvod

Arhitektura s mješavinom stručnjaka (MoE - Mixture-of-Experts) nametnula se kao standardna arhitektura za skaliranje na stotine milijardi ili bilijune parametara, istovremeno nudeći razumnu latenciju u vrijeme inferencije. MoE modeli zamjenjuju gusti sloj transformatora skupom stručnjaka i malim linearnim slojem usmjeravanja koji odlučuje s kojim stručnjakom treba pomnožiti token. Grupiranjem tokena povezanih s istim stručnjakom, problem postaje rijedak i trivijalno paralelizabilan, budući da različiti stručnjaci sada mogu biti smješteni na različitim GPU-ovima raspoređenima na više čvorova.

Za razliku od drugih oblika paralelizma, kao što su paralelizam tenzora (TP) i paralelizam podataka (DP) koji se mogu lako implementirati pomoću biblioteka kolektivne komunikacije kao što su torch.distributed i NCCL, MoE usmjeravanje uključuje rijetku peer-to-peer komunikaciju. Jezgra za slanje (dispatch kernel) mora podijeliti skup dolaznih tokena i poslati ih na rangove koji ugošćuju stručnjake, dok njen dualni problem, jezgra za kombiniranje (combine kernel), mora primiti njihove obrađene pandane i izračunati njihov ponderirani prosjek. Ovi prijenosi ne odgovaraju kolektivima i zahtijevaju prilagođene jezgre za postizanje niske latencije.

U tvrtki Perplexity razvijali smo i prijenosne jezgre između čvorova i specijalizirane jezgre unutar čvorova, a naš smo rad detaljnije predstavili u prethodnom članku o MoE komunikaciji. U drugom smo blog postu također pokazali da implementacija na više čvorova može donijeti manju latenciju i veći protok za velike MoE modele.

U ovom članku detaljno proučavamo implementaciju naših novih jezgri između čvorova, što uvelike poboljšava performanse našeg prethodnog rada temeljenog na NVSHMEM-u. Osim što prikazujemo vrhunske rezultate na ConnectX-7, također pokazujemo da se unutar našeg izvođenja inferencije jezgre mogu koristiti za učinkovito posluživanje najvećih open-source modela.

Inferencija preko InfiniBanda

Paralelizam tenzora i paralelizam stručnjaka mogu se učinkovito skalirati unutar jednog čvora, budući da NVLink nudi visok protok (900 GB/s) uz latenciju u mikrosekundama. Iako su NVIDIA GB200 i GB300 sustavi opremljeni s dovoljno memorije za pohranu težina i KV predmemorija modela s bilijunima parametara, implementacije na AWS p5en instanci s do 8 H200 GPU-ova ograničene su na 1120 GB HBM-a, koji se mora podijeliti između težina modela i KV predmemorija. Posljedično, određeni modeli zahtijevaju implementacije između čvorova preko InfiniBanda.

Iako je InfiniBand i dalje izuzetno brz, nudi protok od samo 400 Gbps (50 GB/s), dodajući desetke do stotine mikrosekundi latencije za MoE usmjeravanje. Srećom, dodane se latencije mogu sakriti korištenjem dijeljenih stručnjaka, mikro-serijskom obradom i preklapanjem izračuna i komunikacije kako bi se isporučio održiv protok po konkurentnoj cijeni. DeepSeek je bio pionir ovih tehnika putem DeepSeek-R1, implementiranog na klasterima H800 GPU-ova i NVIDIA ConnectX-7 InfiniBand adaptera.

Na AWS klasterima situacija je drugačija zbog korištenja njihovih prilagođenih Elastic Fabric Adaptera (EFA). Iako oni osiguravaju vršni protok od 400 Gbps na kolektivima koji se obično koriste u radnim opterećenjima za treniranje, oni neznatno zaostaju za ConnectX-7 u pogledu veličina poruka razmijenjenih tijekom MoE slanja i kombiniranja. Dodatno, EFA ne podržava GPUDirect Async, što zahtijeva pomoćnu nit procesora (CPU proxy nit) koja premošćuje GPU na mrežnu karticu (NIC) kako bi se pokrenuli prijenosi. Osim dodane složenosti, dodatne PCIe transakcije dodaju režijske troškove u vrijednosti mikrosekundi svim prijenosima. Okviri poput NVSHMEM-a omogućuju autorima jezgri pružanje implementacija neovisnih o uređaju, međutim generički API-ji uzrokuju režijske troškove koji čine implementacije neučinkovitima u pogledu i cijene i latencije.

pplx-jezgre temeljene na NVSHMEM-u

Naše prethodne jezgre apstrahirale su temeljnu implementaciju prijenosa putem NVSHMEM-a. Oslanjale su se na nvshmemx_putmem_signal_nbi_warp za prijenos tokena jedan po jedan, koristeći atomske brojače za sinkronizaciju partnera (peerova). Iako je to nudilo razumnu učinkovitost na ConnectX-7 NIC-ovima s IBGDA-om, ConnectX-7 IBRC implementacija temeljena na proxyju bila je znatno sporija, dok EFA implementacija nije uspjela isporučiti latencije ispod milisekunde.

Koristili smo jezgre kao testno okruženje za temeljni hardver, a istovremeno smo definirali najprikladnije sučelje za izlaganje našim modelima. Pretpostavili smo da česti prijenosi i interakcije između jezgri i mrežnih kartica ili proxy niti ometaju performanse. Osim toga, naša implementacija nije se oslanjala ni na kakva jamstva hardvera osim pouzdanosti, što sugerira da bismo mogli ubrzati stvari smanjenjem režijskih troškova primitiva sinkronizacije.

Prateći ova zapažanja, osmislili smo naše jezgre temeljene na EFA-i koje izlažu isto sučelje, optimizirajući interakcije između jezgre uređaja i host proxyja "ispod haube" radi vršnih performansi.

MoE all-to-all preko EFA-e

Povećaj

Nove jezgre slijede hibridnu CPU-GPU arhitekturu, gdje GPU jezgre komuniciraju s modelom na uređaju, a proxy nit koja se izvodi na host CPU-u upravlja interakcijama s mrežnom karticom (NIC). Nove jezgre ponovno koriste isti TransferEngine koji smo izgradili za prijenose KV predmemorije.

MoE usmjeravanje se izvodi pomoću jezgre za slanje (dispatch kernel), koja prima tokene i stručnjake po tokenu za usmjeravanje od modela, šaljući ih preko mreže do rangova koji ugošćuju stručnjake. Nakon što grupirane GEMM jezgre pomnože tokene s težinama stručnjaka, jezgra za kombiniranje prenosi tokene natrag do njihovih izvorišnih rangova i izračunava njihov ponderirani prosjek prije nego što ih vrati modelu.

I jezgra za slanje i jezgra za kombiniranje podijeljene su na pošiljateljsku i primateljsku polovicu kako bi se olakšalo preklapanje i mikro-serijska obrada. Pošiljatelj raspoređuje tokene u međuspremnike za slanje tako da se mogu zapisati partnerima putem jedne RDMA operacije pisanja. Nakon izvršavanja pošiljatelja, GPU može obavljati druge poslove (kao što je množenje s dijeljenim stručnjakom) dok su prijenosi preko mreže na čekanju u pozadini. Jezgre primatelja blokiraju se dok se svi prijenosi ne prime. U pozadini, radna nit na strani hosta (host-side worker nit) upravlja prijenosima: GPU joj šalje signal kada su međuspremnici za slanje popunjeni, nakon čega ona pokreće RDMA pisanja. Jednom kada mrežna kartica signalizira dovršetak, deblokira primatelja, koji vraća tokene modelu.

Jezgre su čvrsto povezane s hostom, komunicirajući međusobno putem jedinstvene memorije (unified memory) i GDRCopyja. CUDA jedinstvena memorija omogućuje mapiranje memorije uređaja ili memorije hosta u adresne prostore oba entiteta, transparentno upravljajući komunikacijom putem PCIe-a.

Iako je jedinstvena memorija prikladna za masovne prijenose, ispitivanje stanja (polling) s niskom latencijom vrši se putem GDRCopyja. GDRCopy koristi GPUDirect RDMA za čitanje iz i pisanje u GPU memoriju uz latenciju u mikrosekundama.

Kao nadopuna RDMA-u, unutar čvora jezgre također iskorštavaju NVLink za razmjenu tokena. To može rasteretiti značajan dio prometa, do četvrtine za EP32 i do osmine za EP64, postižući općenito niže latencije.

Međuspremnici za slanje i primanje

Dizajn međuspremnika za slanje za jezgru za slanje ključan je za postizanje održivih latencija koje uključuju host proxy. Ako postoji N rangova od kojih svaki šalje T tokena, pri čemu je svaki usmjeren na R stručnjaka od ukupno E, rang može primiti do N * T * max(R, E / N) tokena u najgorem mogućem slučaju, ako su svi usmjereni na njega. Nadalje, ova se minimalna granica može postići samo kada pošiljatelji poslože svoja pisanja uzastopno (kontikurno), budući da bi privatni međuspremnici za primanje trošili previše memorije. Kako bi pošiljatelji mogli uzastopno posložiti svoja pisanja bez sukoba, potpune informacije o usmjeravanju moraju biti distribuirane svakom ranku. Posljedično, naše jezgre za slanje prvo razmjenjuju broj tokena po stručnjaku, nakon čega svaka može odrediti kamo idu njihova pisanja i s kojeg pomaka primaju tokene od svakog ranga.

Samo razmjenjivanje informacija o usmjeravanju rezultiralo bi nepotrebno visokim latencijama, budući da bi svaki rang morao čekati da pisanja stignu prije slanja bilo kakvih vlastitih pisanja. Umjesto toga, rezerviramo malu količinu privatnog prostora po pošiljatelju na svakom rangu primatelju i šaljemo neke od tokena zajedno s pisanjima informacija o usmjeravanju. To osigurava iskorištenje pune propusnosti tijekom većeg dijela procesa slanja. Nakon što se rute razmijene, preostali se tokeni prenose putem jednog pisanja po udaljenom partneru. Razmjena informacija o usmjeravanju i međuspremnici fiksne veličine, zajedno s prijenosima koji ih pune, prikazani su na slici u nastavku.

Povećaj

Jezgre za kombiniranje znatno su jednostavnije jer su im informacije o usmjeravanju već dostupne iz faze slanja. One mogu kopirati tokene u jednom pisanju po partneru, nakon čega čitaju iz međuspremnika za primanje kako bi izračunale ponderirani prosjek tokena u izlazne tenzore.

Za NVLink, međuspremnici za slanje i primanje unutar čvora mapirani su u partnere kako bi oni mogli izravno gurati ili povlačiti tokene iz susjednih rangova.

Slanje i kombiniranje (Dispatch i Combine)

Slike u nastavku detaljno prikazuju vremenski raspored i interakcije između GPU-a, CPU-a i mrežne kartice, prikazujući jednu iteraciju slanja i kombiniranja isprepletenu s jezgrama modela.

Rad započinje u slanju (dispatch send), pri čemu jezgra agregira informacije o usmjeravanju i izračunava broj tokena po stručnjaku putem atomske akumulacije u dijeljenoj memoriji unutar jednog bloka. Jednom kada su brojevi tokena dostupni, kopiraju se u memoriju hosta i proxy se odmah obavještava da započne s odašiljanjem informacija o usmjeravanju svim ostalim rangovima partnerima. Paralelno s tim, slanje nastavlja izračunavati kumulativnu sumu broja tokena, identificirajući pomake potrebne za pakiranje tokena u uzastopne međuspremnike za slanje po partneru. Jednom kada su pomaci dostupni, svi tokeni koji se šalju preko EFA-e pakiraju se u međuspremnike za slanje. Proxy je obaviješten da objavi pisanja u međuspremnike za slanje partnera između čvorova. Unutar čvora, jezgre također odmah guraju tokene u privatne međuspremnike unutar istog čvora putem NVLinka.

Povećaj

Povećaj

Jednom kada rang primi informacije o usmjeravanju od svih partnera, on ih centralizira i izračunava različite pomake nužne za identificiranje tokena za slanje i kombiniranje. Ovaj se posao obavlja na CPU-u dok su prijenosi već na čekanju, stoga nije kritičan za latenciju. Nakon toga, kada su poznati svi pomaci, objavljuju se pisanja za preostali dio tokena, koji premašuju veličinu privatnog međuspremnika, i faza slanja je završena.

Na strani primatelja, jezgra prvo povlači tokene iz istog čvora koji premašuju kapacitet privatnih međuspremnika putem NVLinka. Zatim se blokira dok host proxy ne signalizira dovršetak svih prijenosa. Jednom kada se primi signal, tokeni se premještaju iz međuspremnika za primanje u izlazni tenzor proslijeđen grupiranom GEMM-u. Budući da se međuspremnici ponovno koriste između slanja i kombiniranja, i NVLink i RDMA barijere osiguravaju da su međuspremnici primatelja pročitani prije nego što se sljedećoj jezgri dopusti da ih prespiše.

Jezgre za kombiniranje slijede otprilike sličan uzorak, premještajući tokene, signalizirajući proxyju i čekajući dovršetak. Međutim, metapodaci za objavljivanje njihovih pisanja odmah su dostupni i svi udaljeni pomaci su poznati, stoga se svi tokeni preneseni unutar čvora guraju putem NVLinka primateljima, a tokeni između čvorova prenose se putem jednog pisanja. Prije vraćanja rezultata, jezgra akumulira tokene iz međuspremnika primatelja, zapisujući ih u izlazne tenzore.

TransferEngine za MoE usmjeravanje

Kako bismo izradili prototip naše implementacije, inicijalno smo ponovno upotrijebili Transfer Engine izgrađen za razdvojeno predpunjenje (prefill) i dekodiranje kako bismo inicirali i pratili RDMA pisanja. Međutim, za razliku od razdvojenog dekodiranja gdje svaki čvor objavljuje veliki broj pisanja istom partneru, MoE slanje objavljuje mala slanja svim partnerima u jednoj operaciji. Izvan okvira, performanse nisu bile optimalne i naša je implementacija zaostajala za DeepEP-om na ConnectX-7 za oko 20 us. Posljedično, specijalizirali smo TransferEngine da podržava dvije operacije, raspršivanje (scatter) i barijeru (barrier), kako bismo optimizirali interakcije s velikim brojem partnera.

I operacije raspršivanja i barijere izlažu grupe partnera koje se mogu unaprijed registrirati radi unaprijed popunjavanja struktura podataka koje koristi temeljna implementacija prijenosa (predlošci radnih zahtjeva - WR templating). Raspršivanje kopira različite odsječke iz istog međuspremnika izvora u različite pomake na svim članovima grupe partnera. Barijera signalizira svim partnerima putem neposredne poruke (immediate). Redanje i rukovanje ovim operacijama u masi štedi dragocjene mikrosekunde na kritičnoj putanji. Na EFA-i, gdje se dvije mrežne kartice od 200 Gbps udružuju kako bi pružile propusnost od 400 Gbps, dijelimo partnere grupe na dvije mrežne kartice umjesto da dijelimo bajtove svakog prijenosa.

Iako smo inicijalno razvili TransferEngine isključivo na EFA-i koristeći libfabric, MoE jezgre rade na ConnectX-7 bez promjene ijednog retka koda nakon dodavanja podrške za ConnectX-7 u TransferEngine korištenjem libibverbs. U usporedbi s EFA-om temeljenom na pouzdanom datagramu koji se može skalirati (SRD), ConnectX-7 zahtijeva dodatni rad za uspostavu veze i upravljanje partnerima. Koristimo par redova čekanja s nepouzdanim datagramom (UD) za pokretanje pouzdane veze (RC). Stvaramo dva para redova čekanja za svakog udaljenog partnera, jedan za dvosmjerni SEND/RECV i jedan za jednosmjerni WRITE_IMM, budući da i RECV i WRITE_IMM troše RECV operaciju redom, dok dvije operacije zahtijevaju različite veličine međuspremnika. Osim toga, ublažavamo redoslijed PCIe transakcija između ConnectX-7 i GPU memorije. Slično kao kod EFA-e, primjenjujemo optimizaciju predložaka radnih zahtjeva (WR templating). Dodatno, prilikom slanja više malih poruka istom partneru, koristimo uvezivanje radnih zahtjeva (WR chaining) radi poboljšanja brzine poruka.

Budući da je rad na strani GPU-a identičan na oba adaptera, gurnuli smo temeljnu implementaciju prijenosa do njezinih granica kako bismo nadmašili performanse DeepEP-a i istaknuli da je jaz između EFA i ConnectX-7 MoE usmjeravanja određen mogućnostima temeljnog hardvera.

Evaluacija

Latencija slanja i kombiniranja

Procijenili smo latencije slanja i kombiniranja na H200 sustavima opremljenima ConnectX-7 i EFA adapterima na standardnoj DeepSeek-V3 i Kimi-K2 konfiguraciji sa skrivenom dimenzijom od 7168, slanjem skaliranim po blokovima fp8 i kombiniranjem bf16. Razmatramo veličinu serije od 128 za slanje i 4096 za kombiniranje. Na ConnectX-7 uspoređujemo s DeepEP-om i našim prethodnim pplx-jezgrama temeljenim na NVSHMEM-u, koristeći i GPUDirect Async (IBGDA) i host proxy (IBRC). Postižemo kombinirane latencije od 459 us, 582 us i 692 us na EP16, EP32 odnosno EP64. Nadmašujemo druge jezgre temeljene na EFA-i koje se oslanjaju na generičke proxyje, kao što je UCCL-EP, koje postižu 519 us, 966 us i 1159 us.

Povećaj

Pri dekodiranju, unatoč tome što su naše jezgre nekoliko mikrosekundi sporije od DeepEP-a, postižemo općenito niže latencije jer su naše jezgre za kombiniranje znatno brže. Iako koristimo host proxy, zbog naših masovnih prijenosnih operacija nadmašujemo performanse svih implementacija bez proxyja temeljenih na IBGDA-i. Također pokazujemo da je implementacija proxyja važna, budući da je naš visoko specijalizirani proxy za red veličine brži od generičke NVSHMEM implementacije. Unatoč činjenici da je EFA dvostruko sporija od ConnectX-7 prilikom prijenosa paketa od 256 KB, koji su blizu tipičnoj veličini paketa razmijenjenoj MoE usmjeravanjem, jaz između dvaju adaptera općenito je manje značajan.

Povećaj

Iako se naše jezgre optimizirane za dekodiranje mogu skalirati na veličine serija predpunjenja (prefill batch size) bez ikakvih dodatnih prilagodbi, performanse zaostaju za DeepEP-om. To je prvenstveno zbog činjenice da mi akumuliramo tokene samo na strani primatelja kombiniranja, dok DeepEP akumulira kod pošiljatelja kako bi se smanjila količina prenesenih podataka. Međutim, ova optimizacija može imati implikacije na točnost. Dodatno, budući da jezgre prenose podatke u jednom pisanju, koristimo više memorije za međuspremnike slanja i primanja.

Protok za Kimi-K2 / DeepSeek-V3

Kako bismo dokazali održivost naših jezgri, mjerimo performanse end-to-end našeg vlastitog inferencijskog stroja koji poslužuje DeepSeek-V3 (671B parametara) i Kimi-K2 (1T parametara) na više p5en instanci opremljenih Elastic Fabric Adapterima (EFA) i H200 GPU-ovima. Izvodimo model u načinu verifikacije skice (draft verification mode) i izvještavamo o latenciji pretpostavljajući stopu prihvaćanja od 80% koristeći jedan MTP sloj.

Na DeepSeek-V3 uspoređujemo se s baznom linijom unutar čvora EP=8 koja pokreće učinkovite jezgre slanja i kombiniranja samo s NVLink-om. Implementacije između čvorova općenito ili izjednačavaju ili premašuju performanse implementacije s jednim čvorom, osobito na serijama srednje veličine, što pokazuje da skaliranje preko čvorova donosi poboljšanja ukupnog protoka.

Na Kimi-K2, jezgre omogućuju posluživanje. Zbog svoje veličine, model se ne može smjestiti unutar istog čvora koristeći paralelizam podataka (DP) na H200 čvoru, stoga je potrebna implementacija između čvorova za posluživanje modela. Kako model ima manje glava pažnje (attention heads), postižemo bolje latencije nego kod DeepSeek-V3 i možemo posluživati model na EFA-i uz održivu latenciju.

Povećaj

Povećaj

Također smo razmotrili mikro-serijsku obradu, premda je ona poboljšala latencije samo pri vrlo velikim veličinama serija ili kada se koristi više od jednog tokena skice. Da bi se imalo koristi od dijeljenja grupiranih GEMM problema na dva dijela, veličine serija tokena moraju biti dovoljno guste. Većina vremena se i dalje troši na MoE slanje i kombiniranje, što ukazuje na to da će bilo kakva daljnja poboljšanja jezgri posljedično poboljšati latencije i olakšati mikro-serijsku obradu.

Budući rad

Blisko surađujemo s AWS inženjerima na poboljšanju performansi na EFA-i. S naše strane, pratit ćemo ažuriranja libfabric biblioteke koja će smanjiti režijske troškove podatkovne ravnine. Dodatno, eksperimentirat ćemo s efa-direct tehnologijom kako bismo dodatno smanjili režijske troškove korisničkog prostora i izravno komunicirali s hardverom.

Reference

Tehničko izvješće za DeepSeek R1

Tehničko izvješće za Kimi K2

DeepEP

UCCL-EP

Poboljšanje mrežnih performansi HPC sustava korištenjem NVIDIA Magnum IO NVSHMEM-a i GPUDirect Async-a

Citiranje

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