Aktivering av modeller med billioner av parametere på AWS EFA

Hos Perplexity bruker vi de beste modellene for produktet vårt, våre API-er og forskningsteamene våre. Store åpne Mixture-of-Experts-modeller, som Kimi-K2, byr på spesielle utfordringer, ettersom de største inferensnoden med 8x NVIDIA H200-GPUer ikke effektivt kan romme dem, noe som gjør distribusjoner på tvers av noder nødvendig. Vi presenterer et sett med kjerner for ekspertparallellisme som oppnår toppmoderne ventetider på ConnectX-7, og som overgår ytelsen til DeepEP. De samme kjernene er også de første som oppnår levedyktige ventetider på AWS Elastic Fabric Adapter (EFA), noe som muliggjør modelldistribusjoner på billioner av parametere.

Prøv kjernene våre på GitHub og les hele forskningsartikkelen på arXiv.

Mixture-of-Experts (MoE) har dukket opp som standardarkitekturen for skalering til hundrevis av milliarder eller billioner av parametere, samtidig som den tilbyr rimelig ventetid ved inferenstid. MoE-modeller erstatter det tette laget i en transformator med et sett med eksperter og et lite lineært rutingslag som avgjør hvilken ekspert et token skal multipliseres med. Ved å gruppere tokenene som er knyttet til den samme eksperten, blir problemet glissent og trivielt parallelliserbart, siden forskjellige eksperter nå kan være vertskap på forskjellige GPUer spredt over flere noder.

I motsetning til andre former for parallellisme, slik som tensorparallellisme (TP) og dataparallellisme (DP) som enkelt kan implementeres ved hjelp av kollektive kommunikasjonsbiblioteker som torch.distributed og NCCL, innebærer MoE-routing glissen peker-til-peker-kommunikasjon. En utsendelseskjerne må dele opp settet med innkommende tokener og sende dem til rangeringene som er vert for ekspertene, mens dens duale problem, kombinasjonskjernen, må motta de behandlede motpartene og beregne deres vektede gjennomsnitt. Disse overføringene kartlegges ikke til kollektiver og krever tilpassede kjerner for å oppnå lav ventetid.

Hos Perplexity har vi utviklet både bærbare kjerner på tvers av noder og spesialiserte kjerner innenfor noder, og vi presenterer arbeidet vårt mer dypgående i en tidligere artikkel om MoE-kommunikasjon. Vi viste også i et annet blogginnlegg at distribusjon på tvers av noder kan gi lavere ventetid og høyere gjennomstrømning for store MoE-modeller.

I denne artikkelen dykker vi dypt ned i implementeringen av de nye kjernene våre på tvers av noder, noe som i stor grad forbedrer ytelsen til vårt NVSHMEM-baserte tidligere arbeid. I tillegg til å vise frem toppmoderne resultater på ConnectX-7, demonstrerer vi også at kjernene i vår inferenskjøretid kan brukes til effektivt å betjene de største åpen kildekode-modellene.

Tensorparallellisme og ekspertparallellisme kan skalere effektivt innenfor en enkelt node, ettersom NVLink tilbyr høy gjennomstrømning (900 GB/s) ved mikrosekunder med ventetid. Selv om NVIDIA GB200- og GB300-systemene er utstyrt med tilstrekkelig minne til å holde vektene og KV-mellomlagrene for modeller med billioner av parametere, er distribusjoner på AWS p5en-forekomster med opptil 8 H200 GPUer begrenset til 1120 GB HBM, som må deles mellom både modellvekter og KV-mellomlagre. Følgelig krever visse modeller distribusjoner på tvers av noder over InfiniBand.

Selv om InfiniBand fremdeles er lynraskt, tilbyr det bare 400 Gbps (50 GB/s) gjennomstrømning, noe som legger til titalls til hundrevis av mikrosekunder med ventetid for MoE-routing. Heldigvis kan de ekstra ventetidene skjules ved hjelp av delte eksperter, mikrobatching og overlapping av beregning og kommunikasjon for å levere levedyktig gjennomstrømning til en konkurransedyktig kostnad. DeepSeek var pioner for disse teknikkene via DeepSeek-R1, distribuert på klynger av H800 GPUer og NVIDIA ConnectX-7 InfiniBand-adaptere.

På AWS-klynger er situasjonen en annen på grunn av bruken av deres egne Elastic Fabric Adapters (EFA). Selv om disse leverer en toppgjennomstrømning på 400 Gbps på kollektivene som vanligvis brukes i treningsarbeidsbelastninger, faller de litt til kort for ConnectX-7 på meldingsstørrelsene som utveksles under MoE-utsendelse og -kombinasjon. I tillegg støtter ikke EFA GPUDirect Async, noe som krever en CPU-mellommannstråd som presterer en bro mellom GPU-en og nettverkskortet for å initiere overføringer. I tillegg til den ekstra kompleksiteten legger de ekstra PCIe-transaksjonene til mikrosekunder med overhead på alle overføringer. Rammeverk som NVSHMEM tillater kjerneforfattere å tilby enhetsuavhengige implementeringer, men de generiske API-ene medfører overhead som gjør distribusjoner ineffektive når det gjelder både kostnad og ventetid.

De forrige kjernene våre abstraherte bort den underliggende transportimplementeringen gjennom NVSHMEM. De stolte på nvshmemx_putmem_signal_nbi_warp for å overføre tokener én etter én, ved å bruke atomiske tellere for å synkronisere partnere. Selv om dette tilbød rimelig ytelse på ConnectX-7-nettverkskort med IBGDA, var den proxy-baserte ConnectX-7 IBRC-implementeringen betydelig tregere, mens EFA-implementeringen ikke klarte å levere ventetider under et millisekund.

Vi brukte kjernene som en testbenk for den underliggende maskinvaren, samtidig som vi fastslo det mest egnede grensesnittet å eksponere for modellene våre. Vi satte opp en hypotese om at hyppige overføringer og interaksjoner mellom kjernene og nettverkskortene eller mellommannstrådene hindret ytelsen. I tillegg var implementeringen vår ikke avhengig av andre garantier fra maskinvaren enn pålitelighet, noe som tyder på at vi kunne få fortgang i ting ved å redusere overheaden til synkroniseringsprimitiver.

Etter disse observasjonene utarbeidet vi våre EFA-baserte kjerner som eksponerer det samme grensesnittet, og optimaliserer enhetskjernen og verts-mellommannens interaksjoner under panseret for topp ytelse.

RDMA- og NVLink-hybridkommunikasjon med en hostproxy, som viser initierende og mottakende rader, proxyarbeidere og NIC-er.

De nye kjernene følger en hybrid CPU-GPU-arkitektur, der GPU-kjerner samhandler med modellen på enheten, og en mellommannstråd som kjører på vertens CPU, håndterer interaksjoner med nettverkskortet. De nye kjernene gjenbruker den samme TransferEngine som vi bygde for KV-mellomlageroverføringer.

MoE-routing utføres av en utsendelseskjerne, som mottar tokenene og per-token-ekspertene for å rute dem fra modellen, og sender dem over nettverket til rangeringene som er vert for ekspertene. Etter at de grupperte GEMM-kjernene multipliserer tokenene med ekspertvektene, overfører kombinasjonskjernen tokenene til kilderangeringene deres og beregner deres vektede gjennomsnitt før de returneres til modellen.

Både utsendelses- og kombinasjonskjernene er delt inn i en sender- og en mottakerhalvdel for å legge til rette for overlapping og mikrobatching. Senderen stokker tokenene inn i utsendelsesbuffere slik at de kan skrives til partnere gjennom en enkelt RDMA-skriveoperasjon. Etter å ha kjørt senderen, kan GPU-en utføre annet arbeid (som delt ekspertmultiplikasjon) mens overføringene over nettverket venter i bakgrunnen. Mottakerkjernene blokkeres til alle overføringer er mottatt. I bakgrunnen håndterer en arbeiderstråd på vertsiden overføringene: den blir varslet av GPU-en når utsendelsesbufferene er fylt ut, hvorefter den sender ut RDMA-skrivingene. Når nettverkskortet signaliserer fullføring, opphever den blokkering av mottakeren, som returnerer tokenene til modellen.

Kjernene er tett koblet til verten, og kommuniserer med hverandre gjennom både unificert minne og GDRCopy. CUDA unified memory gjør det mulig å tilordne enhetsminne eller vertsminne til adressarommet til begge, og håndterer kommunikasjon transparent via PCIe.

Selv om unificert minne er egnet for bulkoverføringer, utføres avstemming med lav ventetid gjennom GDRCopy. GDRCopy bruker GPUDirect RDMA til å lese fra og skrive til GPU-minne med mikrosekunds ventetid.

I tillegg til RDMA utnytter kjernene også NVLink innenfor en node for å utveksle tokener. Dette kan avlaste en betydelig brøkdel av trafikken, opptil en fjerdedel for EP32 og opptil en åttendedel for EP64, noe som gir totalt lavere ventetider.

Utformingen av utsendelsesbufferne for utsendelseskjernen er avgjørende for å oppnå levedyktige ventetider som involverer mellommannen til verten. Hvis det er N rangeringer som hver sender T tokener, som hver rutes til R eksperter ut av E, kan en rangering motta opptil N * T * max(R, E / N) tokener i det verste mulige tilfellet, hvis alle rutes til den. Videre kan denne minimumsgrensen bare oppnås når senderne legger ut skrivingene sine sammenhengende, siden private mottaksbuffere ville bruke for mye minne. For at sendere skal kunne legge ut skrivingene sine etter hverandre uten konflikter, må full kunnskap om rutinginformasjonen distribueres til hver rangering. Følgelig utveksler utsendelseskjernene våre først per-ekspert-tokentellinger, hvorefter de hver kan bestemme hvor skrivingene deres går og hvilket forskyvning de mottar tokener fra hver rangering fra.

Bare å utveksle rutinginformasjon ville resultere i unødvendig høye ventetider, ettersom hver rangering måtte vente på at skrivingene skulle ankomme før de sendte ut noen av sine egne skrivinger. I stedet reserverer vi en liten mengde privat per-sender-rom på hver mottakerrangering og sender ut noen av tokenene ved siden av skrivingene av rutinginformasjonen. Dette sikrer full båndbreddeutnyttelse gjennom mesteparten av utsendelsesprosessen. Etter at ruter er utvekslet, overføres de resterende tokenene via en enkelt skriving per ekstern partner. Utvekslingen av rutinginformasjon og bufferne med fast størrelse, sammen med overføringene som fyller dem ut, er illustrert i figuren nedenfor.

Bufferoppsett for MoE-distribusjon, som viser hvordan tokendata og rutingsmetadata er organisert og utvekslet mellom to rader.

Kombinasjonskjernene er vesentlig enklere, ettersom rutinginformasjon allerede er tilgjengelig for dem fra utsendelsesfasen. De kan kopiere tokenene over i en enkelt skriving per partner, hvorefter de leser fra mottaksbufferen for å beregne det vektede gjennomsnittet av tokenene inn i utgangstensorene.

For NVLink er utsendelses- og mottaksbufferne innenfor en node tilordnet partnerne slik at de direkte kan skyve eller trekke tokener fra omliggende rangeringer.

Figurene nedenfor detaljerer tidsberegningen og interaksjonene mellom GPU-en, CPU-en og nettverkskortet, og viser én iterering av utsendelse og kombinasjon flettet inn med modellens kjerner.

Arbeidet starter i utsendelsessending, der kjernen aggregerer rutinginformasjonen og beregner per-ekspert-tokentellinger via atomisk akkumulering i delt minne innenfor en enkelt blokk. Når tokentellingene er tilgjengelige, kopieres de til vertsminnet, og mellommannen blir umiddelbart varslet om å begynne å sende ut rutinginformasjonen til alle andre partnerrangeringer. Parallelt fortsetter utsendelsen med å beregne den kumulative summen av tokentellingene, og identifiserer forskyvningene som kreves for å pakke tokener inn i sammenhengende per-per-utsendelsesbuffere. Når forskyvningene er tilgjengelige, pakkes alle tokener som skal sendes via EFA inn i utsendelsesbufferne. Mellommannen blir varslet om å legge inn skrivingene i utsendelsesbufferne for partnere på tvers av noder. Internt i noden skyver kjernene også umiddelbart tokener inn i private buffere innenfor den samme noden gjennom NVLink.

Timing og koordinering for MoE-distribusjon, som viser trinnvise GPU-, CPU- og NIC-operasjoner og synkroniseringen deres.

Timing og koordinering for MoE-kombinering, som viser GPU-, CPU- og NIC-operasjoner og synkroniseringen deres.

Når en rangering mottar rutinginformasjon fra alle partnere, sentraliserer den den og beregner forskjellige forskyvninger som er nødvendige for å identifisere tokener for utsendelse og kombinasjon. Dette arbeidet utføres på CPU-en mens overføringer allerede venter, og dermed er det ikke kritisk for ventetiden. Deretter, når alle forskyvninger er kjent, legges skrivinger for resten av tokenene, som overstiger den private bufferstørrelsen, inn, og utsendelsesfasen er fullført.

På mottakersiden trekker kjernen først tokener fra den samme noden som er utenfor kapasiteten til private buffere gjennom NVLink. Deretter blokkeres den til vertsmellommannen signaliserer fullføring av alle overføringer. Når signalet er mottatt, stokkes tokener fra mottaksbufferne inn i utgangstensoren som sendes til gruppe-GEMM. Siden buffere gjenbrukes mellom utsendelse og kombinasjon, sikrer både NVLink- og RDMA-barrierer at mottaksbufferne har blitt lest før den neste kjernen får overskrive dem.

Kombinasjonskjernene følger et omtrent liknende mønster, der de stokker tokener, signaliserer mellommannen og venter på fullføring. Imidlertid er metadataene for å legge inn skrivingene lett tilgjengelige og alle eksterne forskyvninger er kjent, dermed skyves alle tokener som overføres internt i noden via NVLink til mottakerne, og tokenene på tvers av noder overføres via en enkelt skriving. Før resultatene returneres, akkumulerer kjernen tokenene fra mottaksbufferne og skriver dem inn i utgangstensorene.

For å lage en prototype av implementeringen vår gjenbrukte vi opprinnelig Transfer Engine som ble bygget for disaggregert prefill og dekoding for å utstede og spore RDMA-skrivinger. I motsetning til disaggregert dekoding der hver node legger inn et stort antall skrivinger til den samme partneren, legger imidlertid MoE-utsendelse inn små overføringer til alle partnerne i én operasjon. Utfra esken var ytelsen ikke optimal, og implementeringen vår lå omtrent 20 us etter DeepEP på ConnectX-7. Følgelig spesialiserte vi TransferEngine til å støtte to operasjoner, spredning og barriere, for å optimalisere interaksjoner med et stort antall partnere.

Både sprednings- og barriereoperasjoner eksponerer partnergrupper som kan forhåndsregistreres for å forhåndsutfylle datastrukturene som brukes av den underliggende transportimplementeringen (Work Request (WR)-maler). Spredning kopierer forskjellige utsnitt fra den samme kildebufferen til forskjellige forskyvninger på alle medlemmene i en partnergruppe. Barriere signaliserer alle partnerne gjennom en umiddelbar. Å sette i kø og håndtere disse operasjonene i bulk sparer verdifulle mikrosekunder på den kritiske banen. På EFA, der to 200 Gbps nettverkskort aggregerer for å tilby 400 Gbps båndbredde, sjarderer vi gruppedeltakere på tvers av de to nettverkskortene i stedet for å sjardere byte i hver overføring.

Selv om vi opprinnelig utviklet TransferEngine utelukkende på EFA ved hjelp av libfabric, fungerer MoE-kjernene på ConnectX-7 uten å endre en eneste linje med kode etter å ha lagt til ConnectX-7-støtte i TransferEngine ved hjelp av libibverbs. Sammenlignet med Scalable Reliable Datagram (SRD)-basert EFA, krever ConnectX-7 ekstra arbeid for tilkoblingsoppsett og partneradministrasjon. Vi bruker et Unreliable Datagram (UD)-køpar for å starte Reliable Connection (RC). Vi oppretter to køparer for hver eksterne partner, én for tosidig SEND/RECV og én for ensidig WRITE_IMM, fordi både RECV og WRITE_IMM forbruker RECV-operasjon i rekkefølge, mens de to operasjonene krever forskjellige bufferstørrelser. Dessuten slakker vi på rekkefølgen av PCIe-transaksjoner mellom ConnectX-7 og GPU-minne. I likhet med EFA implementerer vi WR-maloptimaliseringen. I tillegg, når vi sender flere små meldinger til den samme partneren, bruker vi WR-kjedekobling for å forbedre meldingsraten.

Siden arbeidet på GPU-siden er identisk på tvers av begge adaptere, presset vi den underliggende transportimplementeringen til grensene for å overgå ytelsen til DeepEP for å fremheve at gapet mellom EFA og ConnectX-7 MoE-routing bestemmes av egenskapene til den underliggende maskinvaren.

Vi evaluerte ventetider for utsendelse og kombinasjon på H200-systemer utstyrt med ConnectX-7- og EFA-adaptere på standard DeepSeek-V3- og Kimi-K2-konfigurasjon med en skjult dimensjon på 7168, blokkskalert fp8-utsendelse og bf16-kombinasjon. Vi vurderer en batchstørrelse på 128 for utsendelse og 4096 for kombinasjon. På ConnectX-7 sammenligner vi med DeepEP og våre forrige NVSHMEM-baserte pplx-kjerner, ved å bruke både GPUDirect Async (IBGDA) og en vertsmellommann (IBRC). Vi oppnår kombinerte ventetider på henholdsvis 459 us, 582 us og 692 us på EP16, EP32 og EP64. Vi utkonkurrerer andre EFA-baserte kjerner som stoler på generiske mellommenn, slik som UCCL-EP, som oppnår 519 us, 966 us og 1159 us.

Distribusjons- og kombineringslatens for EP64 MoE-dekoding, som sammenligner implementasjonene og kommunikasjonsbanene som vises.

Ved dekoding, til tross for at kjernene våre er noen mikrosekunder tregere enn DeepEP, oppnår vi generelt lavere ventetider ettersom kombinasjonskjernene våre er vesentlig raskere. Selv om vi bruker en vertsmellommann, overgår vi på grunn av våre bulkoverføringer ytelsen til alle IBGDA-baserte mellommannsløse implementeringer. Vi viser også at mellommannsimplementering betyr noe, ettersom vår svært spesialiserte mellommann er en størrelsesorden raskere enn den generiske NVSHMEM-implementeringen. Til tross for at EFA er dobbelt så treg som ConnectX-7 ved overføring av 256 KB-pakker, som ligger nær den typiske pakkestørrelsen som utveksles av MoE-routing, er gapet mellom de to adaptere generelt mindre vesentlig.

Distribusjons- og kombineringslatens for EP64 MoE-prefylling, som sammenligner pplx- og DeepEP-kommunikasjonsbanene.

Selv om de dekodingsoptimaliserte kjernene våre kan skalere til prefill-batchstørrelser uten noen ekstra justeringer, henger ytelsen etter DeepEP. Dette skyldes først og fremst det faktum at vi akkumulerer tokener bare på mottakersiden av kombinasjon, mens DeepEP akkumulerer i senderen for å redusere mengden data som overføres. Imidlertid kan denne optimaliseringen ha konsekvenser for nøyaktigheten. I tillegg, siden kjernene overfører data i én skriving, utnytter vi mer minne for utsendelses- og mottaksbuffere.

For å bevise levedyktigheten til kjernene våre, tester vi den helhetlige ytelsen til vår interne inferensmotor som betjener DeepSeek-V3 (671B parametere) og Kimi-K2 (1T parametere) på flere p5en-forekomster utstyrt med Elastic Fabric Adapters (EFA) og H200 GPUer. Vi kjører modellen i utkastverifiseringsmodus og rapporterer ventetid med en antakelse om en akseptrate på 80 % ved bruk av et enkelt MTP-lag.

På DeepSeek-V3 sammenligner vi mot en EP=8 internodegrunnlinje som kjører effektive NVLink-enkle utsendelses- og kombinasjonskjerner. Distribusjoner på tvers av noder samsvarer generelt med eller overgår ytelsen til distribusjonen på en enkelt node, spesielt på mellomstore batcher, noe som viser at skalering på tvers av noder gir forbedringer i den generelle gjennomstrømningen.

På Kimi-K2 gjør kjernene betjening mulig. På grunn av sin størrelse kan ikke modellen få plass innenfor den samme noden ved hjelp av dataparallellisme (DP) på en H200-node, dermed kreves en distribusjon på tvers av noder for å betjene modellen. Siden modellen har færre oppmerksomhetshoder, oppnår vi bedre ventetider enn DeepSeek-V3, og vi kan betjene modellen på EFA til en levedyktig ventetid.

DeepSeek-V3 gjennomstrømning per GPU sammenlignet med utgangshastighet for EP32-, EP16- og EP8-konfigurasjoner.

Kimi-K2 gjennomstrømning per GPU sammenlignet med utgangshastighet for EP32- og EP16-konfigurasjoner.

Vi har også vurdert mikrobatching, selv om det bare forbedret ventetider ved svært høye batchstørrelser eller når du bruker mer enn én utkasttokener. For å dra nytte av å dele opp grupperte GEMM-problemer i to, må tokenbatchstørrelsene være tilstrekkelig tette. Mesteparten av tiden brukes fremdeles på MoE-utsendelse og -kombinasjon, noe som indikerer at eventuelle ytterligere forbedringer av kjernene følgelig vil forbedre ventetider og legge til rette for mikrobatching.

Vi jobber tett med AWS-ingeniører for å forbedre ytelsen på EFA. På vår side vil vi følge opp oppdateringer til libfabric som vil redusere dataplanets overhead. I tillegg vil vi eksperimentere med efa-direct for ytterligere å redusere overhead i brukerspace og grensesnitt direkte med maskinvaren.

DeepSeek R1 teknisk rapport

Kimi K2 teknisk rapport

DeepEP

UCCL-EP

Forbedring av nettverksytelsen til HPC-systemer ved hjelp av NVIDIA Magnum IO NVSHMEM og GPUDirect Async

BibTeX

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

Introduksjon

Inferens over InfiniBand

NVSHMEM-baserte pplx-kjerner

MoE all-to-all over EFA

Utsendelses- og mottaksbuffere

Utsendelse og kombinasjon

TransferEngine for MoE-routing

Evaluering

Ventetid for utsendelse/kombinasjon

Kimi-K2 / DeepSeek-V3-gjennomstrømning

Fremtidig arbeid

Referanser

Sitat