Effektiv och bärbar kommunikation för mixture-of-experts

Vi presenterar ett högpresterande, bärbart, öppen källkods-bibliotek för Mixture-of-Experts (MoE)-kommunikation som uppnår 10x snabbare prestanda jämfört med vanliga All-to-All-kommunikationsprimitiver.

Vår implementation har flera viktiga tekniska innovationer som ger överlägsen effektivitet för MoE-kommunikation:

GPU-initierad kommunikation (IBGDA): Stöder direkt kommunikation mellan GPU och NIC, vilket avsevärt minskar latensen genom att kringgå CPU-inblandning

Överlappning av kommunikation och beräkning: Delad kärnarkitektur med separata sänd- och mottagningssteg gör att beräkningar kan fortgå samtidigt som nätverksöverföringar

Snabbaste prestanda för enskild nod: 2,5x lägre latens än den tidigare snabbaste implementationen på konfigurationer med en enskild nod

Effektiv och bärbar prestanda för flera noder: Vår implementation uppnår hastigheter upp till 10x snabbare än standard all-to-all-kommunikation. Även om den är ungefär 2x långsammare än mycket specialiserade implementationer erbjuder vår approach bättre portabilitet mellan NVSHMEM-versioner och nätverksmiljöer (NVLink, CX-7 och EFA)

Biblioteket är helt öppen källkod och tillgängligt på https://github.com/ppl-ai/pplx-kernels.

I den här artikeln utforskar vi utmaningarna med expertparallelism i storskaliga MoE-modeller och beskriver vår approach för effektiv token-dispatch och -kombination i distribuerade GPU-miljöer.

Introduktion

Mixture-of-Experts (MoE)-modeller, såsom DeepSeek R1 och Mixtral 8x7B, förbättrar täta modeller genom att begränsa antalet vikter som aktiveras för varje token. Istället för att uppnå höga parameterantal genom att öka antalet lager eller öka storleken på de linjära lagren i Multi-Layer Perceptron (MLP) i varje avkodarlager, ersätter MoE-modeller den traditionella MLP:n med flera experter och en router. Av DeepSeek R1:s 671B parametrar multipliceras till exempel endast 37B med ett givet token under inferens. Detta minskar avsevärt mängden beräkningar som krävs för att avkoda ett token och förbättrar latensen jämfört med en tät modell.

MoE-modeller medför ytterligare utmaningar för inferens jämfört med täta modeller. Medan själva experterna är små MLP-lager, innehåller varje avkodarlager en router som avgör vilka experter ett token skickas till, där varje token skickas till flera experter. Routern är typiskt ett litet linjärt lager som producerar en sannolikhetsfördelning. Vanligtvis väljs de experter med de högsta K-poängen och den slutliga aktiveringen beräknas som ett vägt genomsnitt genom att summera expertutdata multiplicerat med en vikt som härleds från routerns sannolikhetsfördelning.

För att minimera latensen i distribuerade MoE-system kan parallelism utnyttjas över flera enheter, men detta introducerar kommunikationsutmaningar. Modeller som Mixtral 8x7B eller Llama-70B ryms inom 8 enheter över en enda nod och drar nytta av snabba NVLink-sammankopplingar med låg latens (upp till 900 Gbps). Större modeller kräver dock att experter delas upp (sharding) över flera noder med InfiniBand (med en topp på 400 Gbps), vilket introducerar ytterligare latensutmaningar.

Förstora

I den här artikeln utforskar vi problemet med expertparallelism och beskriver vår implementation av GPU-kärnor som effektivt skickar tokens till enheter som hyser de tilldelade experterna och samlar ihop dem igen genom att beräkna den viktade summan av expertutdata. Även om det finns aggressivt optimerade implementationer inriktade på specifik nätverkshårdvara presenterar vi en mycket bärbar implementation som förlitar sig på en minimal uppsättning NVSHMEM-primitiver men ändå levererar prestanda som är 10x snabbare än standard all-to-all-kommunikation. Vår implementation uppnår toppmodern prestanda på konfigurationer med en enskild nod samtidigt som den bibehåller utmärkt portabilitet över olika nätverksmiljöer.

Sharding och parallelism

För effektiv inferens måste en modells vikter hållas i enhetens minne, samtidigt som det lämnas tillräckligt utrymme för aktiveringar, KV-cacher och andra buffertar som krävs för framåtriktad passering (forward pass) genom modellen. De mest kapabla modellerna överstiger kapaciteten hos även de mest kapabla GPU:erna, och därför måste inferensen spridas över flera enheter som kollektivt kan lagra modellens vikter. Baserat på sharding-scheman för vikter måste olika kommunikations- och beräkningsscheman användas för att avlasta beräkningar och synkronisera enheter.

Förstora

Expert Parallelism (EP), som illustreras i den första figuren, parallelliserar endast expertberäkningen. Olika experter tilldelas olika enheter som lagrar deras vikter. Efter routning skickas tokens till motsvarande enhet, och resultaten samlas in och ackumuleras efteråt. Routingens komplexitet beror på graden av parallelism i modellens övriga delar: replikering av andra lager kan eliminera behovet av routning helt och hållet eftersom varje rang kan välja tokens från en lokalt replikerad routingtabell. Om endast en rang utför routning krävs dock en broadcast för att skicka tokens, index och vikter till respektive experter. Slutligen synkroniserar en all-gather eller en all-to-all-broadcast utdatatokens med den rang som fortsätter exekveringen av modellen. En sådan implementation är relativ eftersom torch redan exponerar de nödvändiga primitiverna, även om vissa fördelar kan uppnås genom att fusioinera ackumulering med gather-operationen som samlar in tokens från experterna.

Förstora

Expert-endimensionell parallelism skalas inte idealiskt, eftersom noder i ett kluster kan vara inaktiva medan modellen kör icke-expertlager som attention, norm och sampling. Beräkningen av dessa lager, i synnerhet attention, kan dock också dra nytta av Tensor Parallelism (TP). De flesta modeller förlitar sig på multi-head attention, vilket innebär att attention-huvudena och deras motsvarande Q-, K- och V-projektioner också kan delas upp över enheter genom sharding, där skivorna mellan olika lager replikeras eller samlas in. Om attention sprids över alla enheter kan en all-gather synkronisera aktiveringarna, vilket gör att routning kan replikeras och kräver synkroniseringsprimitiver som liknar fallet med expertparallelism för en effektiv implementation. Det finns dock gränser för parallelism på denna nivå, eftersom en minskning av antalet attention-huvuden under ett visst tröskelvärde ger avtagande avkastning.

Förstora

För att bäst utnyttja alla enheter och stödja en mycket hög grad av expertparallelism på upp till 128 eller 256 GPU:er krävs Data Parallelism (DP). Enligt detta schema grupperas de enheter som en modell delas upp över för att hantera förfrågningar samtidigt, där attention beräknas och KV-cacher som har delats upp via sharding i den lokala gruppen underhålls. Flera instanser av dessa parallella grupper samarbetar vid expertutvärdering, där var och en är värd för en differentierad delmängd av expertvikterna. Beroende på antalet attention-huvuden kan en grupp vanligtvis omfatta upp till storleken av en hel nod, eftersom intranodkommunikation är snabbare. Till exempel betjänar en DP-rang i figuren ovan oberoende två förfrågningar, A och B, och hanterar attention, norm och all annan bokföring för förfrågningarna. Den andra DP-rangen bearbetar en separat förfrågan C. Den första noden är dock värd för hälften av experterna, medan den andra noden är värd för den andra hälften, och därför kan tokens från A och B efter routning skickas till den andra noden och vice versa. Detta leder till ett glest kommunikationsproblem: varje enhet kan skicka ett olika antal tokens till vilken annan destinationsrang som helst. Befintliga primitiver från torch, i synnerhet all_to_all, är inte särskilt väl lämpade eftersom de kan kräva någon form av utfyllnad (padding) eller synkronisering mellan GPU och CPU samt sändning av metadata. För att implementera kommunikation effektivt krävs anpassade kärnor för att bearbeta routningsinformationen och initiera kommunikation från enheterna utan CPU-ingripande. Efter routning måste en dispatch-kärna skicka tokens till de ranger de routades till, medan aktiveringarna som hör till förfrågningarna i den aktuella DP-ruppen måste samlas in på kombinationssidan. Dessutom måste arbetet balanseras inom en DP-grupp som kan ha flera enheter för att minimera broadcast och korrekt replikera tokens på alla ranger för de efterföljande lagren.

NVSHMEM

NVSHMEM är en NVIDIA-specifik OpenSHMEM-implementation som tillhandahåller bärbara kommunikationsfaciliteter mellan enheter som abstraherar bort komplexiteten hos den underliggande hårdvaran. API:et kan uttrycka läsningar och skrivningar mellan enheter, vilka mappas till primitiverna i enskilda transportlager. Både NVLink och RDMA stöds, vilket ger en viss grad av portabilitet och gör att NVSHMEM-kärnor kan arbeta inom en nod eller över ett kluster med flera noder. I våra kärnor använder vi fjärrminneskrivningar och signaleringsoperationer för att överföra data och synkronisera de deltagande enheterna.

NVSHMEM-operationer är uppbyggda kring konceptet symmetriskt minne: de arbetar på buffertar som har allokerats på alla enheter som deltar i kommunikation mellan enheter. En lokal enhet kan använda en adress som härleds inom dess egen lokala buffert tillsammans med indexet för målrangen för att ange destinationen för en operation. Figuren nedan illustrerar detta koncept: båda GPU:erna allokerar symmetriska buffertar av samma storlek och behåller src-pekare till dem. Den första rangen vill skicka 3 heltal till den andra och placera dem i början av bufferten: nvshmem_int_put_nbi härleder startadressen från den lokala bufferten och anger målenheten. Den andra rangen härleder en förskjutning från sin egen buffert och skickar ett element till den första enheten med en förskjutning på ett. Även om destinationsadresser alltid måste vara symmetriska buffertar som allokerats med nvshmem_alloc, kan källbuffertar vara godtyckliga regioner i enhetsminnet förutsatt att de är förregistrerade hos NVSHMEM.

Även om NVSHMEM tillhandahåller ett brett utbud av primitiver förlitar sig våra kärnor på endast 3 funktioner och bygger all synkronisering och stängsel (fencing) på dem.

Förstora

nvshmemx_putmem_signal_nbi_warp: Överför ett datablock från en enhet till en annan samtidigt som en flagga sätts på fjärrenheten. Operationen antingen sätter (NVSHMEM_SIGNAL_SET) eller inkrementerar (NVSHMEM_SIGNAL_ADD) en 64-bitars plats. Flaggan uppdateras efter att hela minnesblocket har överförts. Om fjärrenheten observerar en ändring i flaggan kan den säkert komma åt bufferten i sitt eget minne. Denna funktion är användbar för att sammanfoga dataöverföring och synkronisering.

nvshmemx_signal_op verkar på en enskild minnesplats, vanligtvis en 64-bitars flagga, och sätter eller inkrementerar den atomiskt. Den är användbar för att skicka över metadata och synkronisera enheter.

nvshmem_uint64_wait_until används på mottagarsidan av en signal för att avläsa (poll) en flagga tills fjärrenheten uppdaterar den.

Att säkerställa någon form av ordning mellan operationer genom nvshmem_quiet eller nvshmem_fence är kostsamt, och därför undviker vår implementation av kärnorna användningen av andra NVSHMEM-primitiver och minimerar alla beroenden utöver enskilda sänd-mottagnings-par. När vi skickar data använder vi alltid den icke-blockerande versionen av funktioner utan att vänta på att data ens ska skickas ut från den lokala rangen. Andra, implicita synkroniseringar över gruppen av kärnor säkerställer att dessa buffertar inte skrivs över destruktivt medan NIC:arna behöver åtkomst till dem.

NVSHMEM använder GPUDirect RDMA så att NIC kan läsa och skriva GPU-minne direkt för datautbyte utan att blanda in CPU:er. Vidare stöder NVSHMEM på ConnectX NIC GPU-initierad kommunikation (även känd som Infiniband GPUDirect Async eller IBGDA), vilket gör att GPU:n kan posta RDMA-operationer direkt till NIC och läsa av NIC:s slutförandekö (completion queue) direkt. På plattformar som inte stöder GPU-initierad kommunikation startar NVSHMEM en "proxy"-tråd på värd-CPU:n för att initiera kommunikation och avläsa slutförande för GPU:ns räkning. NVSHMEM-programmet är bärbart oavsett om GPU-initierad kommunikation stöds eller inte. GPU-initierad kommunikation minskar dock latensen avsevärt eftersom den kringgår omvägen till CPU:n helt och hållet.

Bärbara kärnor

Vi implementerar MoE-kommunikation genom ett par dispatch- och combine-kärnor. Dispatch-kärnorna ansvarar för att läsa tokens och routningsinformation på varje rang och skicka dem till lämpliga experter. Combine-kärnorna samlar in aktiveringarna som produceras av experterna och skickar tillbaka dem till sina källranger, samtidigt som de beräknar det vägda genomsnittet från de valda experterna baserat på de vikter som beräknats av routern. Kärnorna delas dessutom upp i en sänd- och en mottagningskomponent för att göra det möjligt att överlappa dataöverföringar med beräkningar. Sändkärnorna är helt icke-blockerande och icke-synkroniserande: de skickar helt enkelt ut alla skrivningar till fjärrenheterna. I den andra änden läser combine-kärnorna endast från minnet och väntar tills all nödvändig data har överförts. Efter att ha skickat arbete till NIC:arna, medan data överförs asynkront över ledningen, kan GPU:erna utföra annat användbart arbete lokalt, såsom att applicera delade experter eller beräkna attention.

Varje par av kärnor har sina egna symmetriska minnesbuffertar, där varje rang allokerar privat lagring för att ta emot data från alla andra ranger. Det maximala antalet tokens som kan skickas från en DP-grupp är fast över ranger, vilket också sätter den övre gräns som varje rang kan ta emot för varje lokal expert från varje DP-rang. Detta gör att sändarranger kan härleda en unik adress på destinationsrangen att skriva till, utan att kräva någon synkronisering dem emellan. Efter att data har tagits emot flyttar kärnorna runt tokens för att placera dem i minnet i det optimala format som krävs av mottagarkärnorna. Även om buffertarna har en avsevärd dimension återanvänds de över alla sekventiella lager i en modell.

Förstora

Den enda formen av global synkronisering, som illustreras i figuren ovan, implementeras i combine-receive-kärnan. När data har skickats från sändsidan återvänder en rang och fortsätter exekveringen när den tar emot alla tokens den behöver från de andra rangerna, och går in i combine-send-kärnorna så snart de har bearbetats. Barriären i combine-receive säkerställer att ingen rang kan springa före och starta dispatch-send medan någon annan rang fortfarande väntar på att ta emot data, vilket potentiellt skulle kunna orsaka destruktiv överlappning. Synkronisering görs i combine-kärnan för enkelhetens skull; som en optimering skulle samma barriär kunna implementeras mellan dispatch-receive- och combine-send-kärnorna, vilket ytterligare minskar latensen genom att ge NIC:arna mer tid att slutföra överföringarna medan GPU:n kör andra kärnor.

Båda kärnorna är uppdelade över alla tillgängliga SM på enheterna: medan dispatch-send- och combine-receive-kärnorna måste parallellisera över ett maximalt antal tokens per rang (max_m), måste dispatch-receive- och combine-send-kärnorna hantera det antalet tokens från varje rang och lokal expert, upp till maximalt max_m * (num_experts // EP) * (EP // TP).

Dispatch

Dispatch-kärnan ansvarar för att skicka tokens och blockskalningsfaktorer till parter som är värdar för de experter som tokens routades till. Sändarsidan förlitar sig på warp-specialisering för att parallellisera två uppgifter: att aggregera routningsinformation för att fastställa hur många tokens som skickas från varje rang och att packa in tokens i meddelanden som ska skickas till fjärrenheten. Mottagarsidan väntar först på att alla tokental ska tas emot, varefter den läser och packar upp nyttolasterna (payloads) från det delade minnet till de tensorer som kommer att skickas vidare till de nedströms Grouped GEMM-kärnorna. Den lagrar också information i buffertar som delas med combine-kärnorna för att indikera var varje token ska skickas tillbaka. Denna information krävs eftersom mottagarkärnan flyttar runt tokens i sammanhängande buffertar i en icke-deterministisk ordning.

Förstora

I sändardelen ansvarar en enskild warp med 32 trådar på varje block för att läsa routningsinformationen och räkna antalet tokens som varje destinationsexpert förväntas ta emot från denna rang, och skickar antalet plus ett med hjälp av nvshmem_signal_op. Antalet inkrementeras med ett, eftersom övergången från noll till icke-noll i den andra änden signalerar mottagningen av antalen från en källrang. Parallellt samarbetar de återstående warparna för att kopiera tokens till symmetriskt minne över alla block parallellt, och packar aktiveringar, skalningsfaktorer och deras index på den lokala rangen till ett sammanhängande stycke. Indexet krävs av combine-sändaren för att bestämma den adress där token kommer att skrivas. Därefter, efter att ha säkerställt att all data har kopierats genom en barriär över warp-grupperna, arbetar warparna återigen oberoende, där var och en skickar samma buffert till en differentierad expert parallellt. Tokens skickas med hjälp av nvshmemx_putmem_signal_nbi_warp, vilket också atomiskt ökar antalet skickade tokens från den lokala rangen på fjärrenheten. Eftersom varje rang inom en DP-grupp äger en replik av det token som ska skickas, balanseras dispatch jämnt, där varje enhet skickar ut en delmängd av tokens.

På mottagarsidan synkroniserar alla kärnor först med sändarrangerna genom att vänta på att de totala tokentalen ska skickas över. Därefter väntar alla på att de atomiskt inkrementerade skickade tokentalen ska stabiliseras till de totala antalen, vilket indikerar att alla nyttolaster från källrangerna också har skickats över, tack vare semantiken i putmem-anropet. Kärnorna avläser (poll) antalen med hjälp av nvshmem_uint64_wait_until och parallelliserar operationen över alla block och trådar. Därefter säkerställer en barriär över blocken att inget block läser buffertarna om inte all data har tagits emot korrekt. Utspridda över block och synkroniserade via en atomiskt inkrementerad räknare kopieras tokens från det symmetriska minnet till vanliga sammanhängande buffertar som blir indata till de nedströms kärnor som implementerar experterna. Källrangen, expertindexet och tokenindexet lagras separat och pekar exakt ut den plats dit combine-kärnan måste skicka aktiveringarna. Även om tokens från en DP-grupp skickas från olika enheter grupperas de alla tillsammans för att skickas vidare till motsvarande expert.

Combine

Combine-kärnorna delas återigen upp i sändar- och mottagarhalvor: sändarna kopierar över de okvantiserade 16-bitarsaktiveringarna till fjärrenheterna, medan mottagarkärnorna väntar på att data ska skickas över och beräknar det vägda genomsnittet av alla expertbidrag lokalt. Dessutom fungerar de som en barriär för att synkronisera dispatch-combine-sekvensen: varje rang sätter en flagga på varje part vid inträde till sändkärnorna, och mottagarkärnorna tillåts inte att återvända om de inte observerar att flaggan är satt. Synkroniseringens latens är minimal eftersom den överlappar med den faktiska kommunikationen och beräkningen.

Förstora

På sändarsidan traverserar kärnorna listan över tokens som tilldelats alla lokala experter parallellt och skriver dem till en buffert på destinationsrangen. Målrang, expertindex och tokenindex läses från de per-token-buffertar som fyllts i av scater-kärnorna vid mottagandet av tokens. Varje sändare har sin egen privata minnesregion per expert att skriva till, vilket visas i figuren ovan, vilket undviker behovet av synkronisering. På samma sätt som för dispatch inkrementerar combine atomiskt räknare per token på destinationsrangen för att indikera mottagandet av data: när räknaren matchar antalet experter som ett token skickats till kan tokenets innehåll kommas åt.

I mottagarkärnan traverseras listan över tokens parallellt över flera noder, och man väntar på att deras innehåll ska anlända genom att avläsa flaggan som satts av signaleringsoperationen. Vid ankomst läses nyttolasterna från de privata buffertarna, där routingtabellen anger vilken buffert som ska läsas från och vilken vikt som ska tilldelas varje expert. Resultaten skrivs sedan till externt allokerade tensorer, och kärnan avslutar exekveringen när alla enheter har passerat barriären. Över en DP-grupp tar alla ranger emot en kopia av varje expertaktivering för att beräkna sina egna repliker.

Benchmarks

Vi utvärderar våra kärnor på ett kluster med 128 GPU:er utspridda över 16 noder anslutna via InfiniBand- och CX-7-NIC:ar. Vi provar både GPUDirect Async (IBGDA) och Reliable Connection (RC) med en CPU-proxy. Vi jämför dem med den highly optimerade DeepSeek-implementationen samt de täta primitiver som tillhandahålls av PyTorch (genom NCCL) eller NVSHMEM.

Förstora

Även om prestandan för IBRC och IBGDA är likartad vid den täta NVSHMEM all-to-all-operationen (6378 µs vs 6180 µs), är IBGDA betydligt snabbare med de glesa kärnorna. Om man adderar Dispatch och Combine använder IBGDA 902 µs medan IBRC tar 3223 µs – en förbättring av latensen med 3,6x. Medan all-to-all är bandbreddsbundet skickar de glesa kärnorna storleksordningar mindre data och är latensbundna. Genom att utlösa nätverksöverföringar direkt från en GPU utan att kräva en CPU-proxy för att koordinera GPU:n och NIC:en reduceras end-to-end-latensen avsevärt.

Även om våra bärbara kärnor är ungefär 2x långsammare än de högt optimerade DeepSeek-kärnorna förbättrar de latensen med 10x jämfört med de täta kärnorna (902 µs vs 9944 µs).

Dessutom gör uppdelningen i sändar- och mottagarkomponenter att en del av latensen kan döljas, till skillnad från bibliotekets primitiver.

Förstora

På enskild nod (EP8) använder NVSHMEM NVLINK för transport, vilket ger lägre latens och högre genomströmning än nätverkskommunikation mellan noder. Våra bärbara kärnor är ungefär 2,5x snabbare än DeepEP på enskild nod (186 µs vs 481 µs).

Ytterligare förbättringar

De kärnor som beskrivs här presterar bättre än ML-ramverkens inbyggda primitiver och erbjuder anständig prestanda utan att överpecialisera för särskilda transporter mellan enheter, såsom NVLink eller InfiniBand. Utöver de redan nämnda möjligheterna kan ytterligare prestandavinster uppnås genom att ersätta kommunikationsprimitiverna med mer specialiserade versioner. Över NVLink kan till exempel användningen av symmetriskt minne ersättas med buffertar som delas över enheterna, vilket eliminerar viss kopiering och möjliggör finkornigare synkroniseringsscheman över enskilda tokens istället för token-satser. Över InfiniBand kan en implementation komma åt de underliggande köparen (queue pairs) direkt, vilket eliminerar behovet av att överföra både förväntade och skickade tokental samt tömma tokens från köerna till utdatabuffertar när de anländer. Sådana implementationer kommer dock till priset av portabilitet, medan våra implementationer gör det möjligt att utvärdera kluster med specifika konfigurationer innan utvecklingsansträngningar läggs på att optimera för specifik hårdvara.

Slutsats

Vi har presenterat ett högpresterande, bärbart bibliotek för MoE-kommunikation som uppnår 10x snabbare prestanda jämfört med standard all-to-all-kommunikation samtidigt som kompatibiliteten mellan olika hårdvarukonfigurationer bibehålls. Vid distribution på enskild nod med NVLink uppvisar vår lösning 2,5x lägre latens än tidigare implementationer.

Vår approach balanserar prestanda med portabilitet genom viktiga innovationer, inklusive stöd för GPU-initierad kommunikation, en delad kärnarkitektur som möjliggör överlappning mellan beräkning och kommunikation, samt effektiv token-dispatch med minimala NVSHMEM-primitiver. Även om det är ungefär 2x långsammare än mycket specialiserade implementationer på flernodsinstallationer erbjuder vårt bibliotek överlägsen flexibilitet i olika nätverksmiljöer (NVLink, CX-7 och EFA).

Allt eftersom MoE-modeller fortsätter att skalas upp kommer effektiva kommunikationsstrategier som vår att bli allt viktigare för praktisk implementering. Vår implementation med helt öppen källkod finns tillgänglig på https://github.com/ppl-ai/pplx-kernels.

Referenser

DeepEP

Teknisk rapport för DeepSeek-V3

NVIDIA NVSHMEM

Citering

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