Triljoni parameetriga mudelite lubamine AWS EFA-l
Muutke triljoni parameetriga mudelid pilveplatvormi teisaldatavusega kättesaadavaks
Ettevõttes Perplexity kasutame parimaid mudeleid oma toodete, API-de ja uurimisrühmade jaoks. Suured avatud lähtekoodiga eksperdasegud (Mixture-of-Experts) mudelid, nagu Kimi-K2, esitavad erilisi väljakutseid, kuna suurimad 8x NVIDIA H200 graafikaprotsessoriga (GPU) järeldussõlmed ei suuda neid tõhusalt mahutada, mistõttu on vaja sõlmeüleseid juurutusi. Esitame eksperdaparalleelsuse jaoks kernelite komplekti, mis saavutab ConnectX-7-l tipptasemel latentsusaja, ületades DeepEP jõudlust. Samad tuumad on ka esimesed, mis saavutavad elujõulise latentsusaja AWS Elastic Fabric Adapteril (EFA), võimaldades triljoni parameetriga mudelite juurutamist.
Proovige meie kerneleid GitHubis ja lugege täielikku teadustööd arXiv-is.
Sissejuhatus
Eksperdasegude (Mixture-of-Experts, MoE) arhitektuurist on saanud standardne lähenemine sadade miljardite või triljonite parameetritega mastaapseks laiendamiseks, pakkudes samal ajal järelduse tegemisel mõistlikku latentsusaega. MoE mudelid asendavad transfömeri tiheda kihi ekspertide kogumi ja väikese lineaarse ruutingukihiga, mis otsustab, millise eksperdiga märki korrutada. Rühmitades sama eksperdiga seotud märgid, muutub probleem hõredaks ja hõlpsasti paralleelsemaks muudetavaks, kuna erinevaid eksperte saab nüüd majutada mitme sõlme vahel jagatud erinevatel GPU-del.
Erinevalt teistest paralleelsuse vormidest, nagu tensorparalleelsus (TP) ja andmeparalleelsus (DP), mida saab hõlpsasti rakendada kollektiivse suhtluse teekide (näiteks torch.distributed ja NCCL) abil, hõlmab MoE ruuting hõredat punktist-punkti suhtlust. Edastustuum (dispatch kernel) peab sissetulevate märkide komplekti tükeldama ja need ekspertide majandajate järkudesse saatma, samas kui selle duaalprobleem, ühendamistuum (combine kernel), peab saama nende töödeldud vasted ja arvutama nende kaalutud keskmise. Need edastused ei vasta kollektiivsetele toimingutele ja nõuavad madala latentsusaja saavutamiseks kohandatud kerneleid.
Ettevõttes Perplexity oleme arendanud nii kaasaskantavaid sõlmevahelisi kerneleid kui ka spetsiaalseid sõlmesiseseid kerneleid, tutvustades oma tööd põhjalikumalt varasemas MoE kommunikatsiooni teemalises artiklis. Samuti näitasime teises blogipostituses, et sõlmeüülene juurutus võib suurte MoE mudelite jaoks tagada madalama latentsusaja ja suurema läbilaskevõime.
Selles artiklis käsitleme põhjalikult oma uute sõlmevaheliste kernelite juurutust, mis parandab oluliselt meie NVSHMEM-põhise eelmise töö jõudlust. Lisaks ConnectX-7 tipptasemel tulemuste esitlemisele demonstreerime ka seda, et meie järelduste käituskeskkonnas saab kerneleid kasutada suurimate avatud lähtekoodiga mudelite tõhusaks teenindamiseks.
Järelduste tegemine InfiniBandi kaudu
Tensorparalleelsus ja eksperdaparalleelsus saavad ühes sõlmes tõhusalt mastaapseks muutuda, kuna NVLink pakub kõrget läbilaskevõimet (900 GB/s) mikrosekundilise latentsusega. Kuigi NVIDIA GB200 ja GB300 süsteemid on varustatud piisava mäluga triljoni parameetriga mudelite kaalude ja KV-vahemälude hoidmiseks, on juurutused AWS-i p5en eksemplaril (kuni 8 H200 GPU-ga) piiratud 1120 GB HBM-mäluga, mis tuleb jagada nii mudeli kaalude kui ka KV-vahemälude vahel. Seetõttu nõuavad teatud mudelid sõlmevahelisi juurutusi InfiniBandi kaudu.
Ehkki InfiniBand on endiselt äärmiselt kiire, pakub see vaid 400 Gbps (50 GB/s) läbilaskevõimet, lisades MoE ruutingule kümneid kuni sadu mikrosekundeid latentsusaega. Õnneks saab lisandunud latentsusaega peita jagatud ekspertide, mikropakettimise ja arvutus-suhtlus-kattuvuse abil, et tagada elujõulise läbilaskevõime konkurentsivõimelise hinnaga. DeepSeek oli nende tehnikate teerajajaks DeepSeek-R1 kaudu, mis on juurutatud H800 GPU-de ja NVIDIA ConnectX-7 InfiniBandi adapterite klastritesse.
AWS-i klastrites on olukord teistsugune nende kohandatud Elastic Fabric Adapterite (EFA) kasutamise tõttu. Kuigi need tagavad treeningtöökoormustes tavaliselt kasutatavatel kollektiividel maksimaalselt 400 Gbps läbilaskevõime, jäävad nad MoE edastamise ja kombineerimise ajal vahetatavate sõnumisuuruste osas ConnectX-7-st veidi alla. Lisaks ei toeta EFA funktsiooni GPUDirect Async, mis nõuab edastuste algatamiseks protsessori (CPU) puhversõlme lõime, mis loob silla GPU ja võrgukaardi (NIC) vahel. Lisaks keerukusele lisavad täiendavad PCIe tehingud kõigile edastustele mikrosekundite jagu lisakoormust. Raamistikud, nagu NVSHMEM, võimaldavad kernelite loojatel pakkuda seadmesõltumatuid teostusi, kuid üldised API-d toovad kaasa lisakoormuse, mis muudab juurutused ebatõhusaks nii kulude kui ka latentsusaja osas.
NVSHMEM-põhised pplx-kernelid
Meie eelmised tuumad abstraheerisid aluseks oleva transporditeostuse NVSHMEM-i kaudu. Need tuginesid funktsioonile nvshmemx_putmem_signal_nbi_warp, et edastada märke ükshaaval, kasutades partnereerijate sünkroonimiseks aatomloendureid. Kuigi see pakkus IBGDA-ga ConnectX-7 võrgukaartidel mõistlikku jõudlust, oli puhversõlmel põhinev ConnectX-7 IBRC teostus märgatavalt aeglasem, samas kui EFA teostus ei suutnud pakkuda alla millisekundilist latentsusaega.
Kasutasime kerneleid alusriistvara katseplatvormina, määrates samal ajal kindlaks meie mudelitele eksponeerimiseks sobivaima liidese. Me oletasime, et sagedased edastused ja interaktsioonid kernelite ning võrgukaartide või puhversõlme lõimede vahel takistasid jõudlust. Lisaks ei tuginenud meie teostus ühelegi muule riistvaralisele garantiile peale töökindluse, mis viitab sellele, et saaksime asju kiirendada, vähendades sünkroonimisprimitiivide lisakoormust.
Nende tähelepanekute põhjal koostasime oma EFA-põhised tuumad, mis pakuvad sama liidese ja optimeerivad parima jõudluse saavutamiseks seadme tuuma ja hosti puhversõlme vahelist suhtlust kapoti all.
MoE all-to-all EFA kaudu
Uued tuumad järgivad hübriidset CPU-GPU arhitektuuri, kus GPU tuumad suhtlevad seadmes oleva mudeliga ja hosti CPU-s töötav puhversõlme lõim käsitleb suhtlust võrgukaardiga. Uued tuumad taaskasutavad sama TransferEngine'i, mille ehitasime KV vahemälu edastuste jaoks.
MoE ruutingut teostab edastustuum (dispatch kernel), mis võtab mudelilt vastu märgid ja märgipõhised eksperdid, et need ruutida, saates need üle võrgu eksperte majandavatesse järkudesse. Pärast seda, kui rühmitatud GEMM-tuumad korrutavad märgid ekspertide kaaludega, kannab kombineerimistuum (combine kernel) märgid üle nende lähtejärkudesse ja arvutab nende kaalutud keskmise enne mudelile tagastamist.
Nii edastus- kui ka kombineerimistuumad on kattuvuse ja mikropakettimise hõlbustamiseks jagatud saatja- ja vastuvõtupooleks. Saatja paigutab märgid saatmispuhvritesse, et neid saaks ühe RDMA-kirjutustoiminguga partneritele kirjutada. Pärast saatja käivitamist saab GPU teha muud tööd (näiteks jagatud eksperdi korrutamist), kuni võrguülesed edastused on taustal ootel. Vastuvõtja tuumad blokeerivad töö seni, kuni kõik edastused on vastu võetud. Taustal haldab edastusi hostipoolne töölõim: GPU annab sellele märku, kui saatmispuhvrid on täidetud, mis järel see lähetab RDMA kirjutamised. Kui võrgukaart annab märku valmimisest, vabastab see vastuvõtja blokaadi, mis tagastab märgid mudelile.
Tuumad on tihedalt seotud hostiga ja suhtlevad üksteisega nii unifitseeritud mälu (unified memory) kui ka GDRCopy kaudu. CUDA unifitseeritud mälu võimaldab seadme mälu või hosti mälu vastendada mõlema aadressiruumidesse, hallates sidet läbipaistvalt PCIe kaudu.
Kuigi unifitseeritud mälu sobib hulgiedastusteks, tehakse madala latentsusega pärimist GDRCopy abil. GDRCopy kasutab GPUDirect RDMA-d mikrosekundilise latentsusega GPU mälust lugemiseks ja sinna kirjutamiseks.
Lisaks RDMA-le kasutavad tuumad märkide vahetamiseks ära ka NVLinki sõlmesiseselt. See võib maha laadida märkimisväärse osa liiklusest – kuni veerandi EP32 ja kuni kaheksandiku EP64 puhul –, saavutades üldiselt madalama latentsusaja.
Saatmis- ja vastuvõtupuhvrid
Edastustuuma saatmispuhvrite disain on hosti puhversõlme kaasava elujõulise latentsusaja saavutamiseks ülioluline. Kui on N järku, millest igaüks edastab T märki ning igaüks on ruuditud R eksperdile E-st, võib järk halvimateralisel juhul vastu võtta kuni N * T * max(R, E / N) märki, kui need kõik on sellele ruuditud. Lisaks saab selle minimaalse piiri saavutada ainult siis, kui saatjad paigutavad oma kirjutised külgnevalt, kuna privaatsed vastuvõtupuhvrid kulutaksid liiga palju mälu. Selleks, et saatjad saaksid oma kirjutised järjestikku ja ilma konfliktideta paigutada, tuleb ruutinguteave täielikult igale järgule edastada. Sellest tulenevalt vahetavad meie edastustuumad esmalt eksperdipõhised märgiloendid, mille järel igaüks saab kindlaks teha, kuhu nende kirjutised lähevad ja millisest nihkest nad igalt järgult märke saavad.
Pelgalt ruutinguteabe vahetamine tooks kaasa tarbetult suure latentsusaja, kuna iga järk peaks ootama kirjutiste saabumist, enne kui saab ise ühtegi kirjutist lähetada. Selle asemel reserveerime igal vastuvõtva järgul väikese hulga privaatset saatjaspõhist ruumi ja edastame mõned märgid koos ruutinguteabe kirjutistega. See tagab ribalaiuse täieliku ärakasutamise suurema osa edastusprotsessi vältel. Pärast marsruutide vahetamist kantakse ülejäänud märgid üle ühe kirjutise kaudu kaughalduspartneri kohta. Ruutinguteabe vahetamist ja fikseeritud suurusega puhvreid koos neid täitvate edastustega on illustreeritud alloleval joonisel.
Kombineerimistuumad on oluliselt lihtsamad, kuna ruutinguteave on neile edastusetapist juba saadaval. Nad saavad märgid kopeerida ühe kirjutisega partneri kohta, mille järel loevad nad vastuvõtupuhvrist, et arvutada märkide kaalutud keskmine väljundtensoritesse.
NVLinki puhul vastendatakse sõlmesisesed saatmis- ja vastuvõtupuhvrid partneritesse, et nad saaksid naaberjärkudelt otse märke tõmmata või neisse lükata.
Edastamine ja kombineerimine
Allolevad joonised selgitavad üksikasjalikult GPU, CPU ja võrgukaardi ajastust ja interaktsioone, näidates ühte edastuse ja kombineerimise iteratsiooni, mis on põimitud mudeli tuumadega.
Töö algab edastuse saatmises, kus tuum koondab ruutinguteabe ja arvutab eksperdipõhised märgiloendid aatomilise akumulatsiooni teel ühes plokis olevas ühismälus. Kui märgiloendid on saadaval, kopeeritakse need hosti mällu ja puhversõlme teavitatakse viivitamatult ruutinguteabe saatmise alustamisest kõigile teistele partnerijärkudele. Paralleelselt jätkab edastus märgiloendite kumulatiivse summa arvutamist, tuvastades nihked, mis on vajalikud märkide pakkimiseks külgnevatesse saatmispuhvritesse järgu kohta. Kui nihked on saadaval, pakitakse kõik EFA kaudu saadetavad märgid saatmispuhvritesse. Puhversõlme teavitatakse kirjutiste postitamisest sõlmevaheliste partnerite saatmispuhvritesse. Sõlmesiseselt lükkavad tuumad märgid ka kohe NVLinki kaudu sama sõlme privaatsetesse puhvritesse.
Kui järk saab ruutinguteabe kõigilt partneritelt, tsentraliseerib see selle ja arvutab erinevad nihked, mis on vajalikud edastamiseks ja kombineerimiseks mõeldud märkide tuvastamiseks. See töö tehakse protsessoril (CPU), samal ajal kui edastused on juba ootel, mistõttu see ei ole latentsuskriitiline. Seejärel, kui kõik nihked on teada, postitatakse ülejäänud märkide (mis ületavad privaatse puhvri suurust) kirjutised ja edastuse saatmisfaas on lõppenud.
Vastuvõtja poolel tõmbab tuum esmalt NVLinki kaudu samast sõlmest märgid, mis ületavad privaatsete puhvrite mahutavust. Järgmisena blokeerib see töö seni, kuni hosti puhversõlm annab märku kõigi edastuste lõpuleviimisest. Kui signaal on saabunud, tõstetakse märgid vastuvõtupuhvritest ümber grupi GEMM-ile edastatud väljundtensorisse. Kuna puhvreid taaskasutatakse edastamise ja kombineerimise vahel, tagavad nii NVLinki kui ka RDMA tõkked, et vastuvõtupuhvreid on loetud, enne kui järgmistel tuumadel lubatakse neid üle kirjutada.
Kombineerimistuumad järgivad ligikaudu sarnast mustrit: segavad märke, annavad puhversõlmele märku ja ootavad valmimist. Kuid nende kirjutiste postitamise metaandmed on kohe saadaval ja kõik kaugnihked on teada, mistõttu kõik sõlmesiseselt edastatud märgid lükatakse NVLinki kaudu vastuvõtjatele ja sõlmevahelised märgid edastatakse ühe kirjutisega. Enne tulemuste tagastamist akumuleerib tuum märgid vastuvõtupuhvritest, kirjutades need väljundtensoritesse.
TransferEngine MoE ruutinguks
Oma teostuse prototüübi loomiseks taaskasutasime algselt eraldatud eel- ja dekodeerimise jaoks ehitatud Transfer Engine'i RDMA kirjutiste esitamiseks ja jälgimiseks. Erinevalt eraldatud dekodeerimisest, kus iga sõlm postitab samale partnerile suure hulga kirjutisi, postitab MoE edastus aga väikseid edastusi kõigile partneritele ühe toiminguga. Kohe karbist välja tulles ei olnud jõudlus optimaalne ja meie teostus jäi ConnectX-7-l DeepEP-ist umbes 21 mikrosekundit maha. Sellest tulenevalt spetsialiseerisime TransferEngine'i toetamaks kahte toimingut: hajutamist (scatter) ja tõket (barrier), et optimeerida suhtlust suure hulga partneritega.
Nii hajutamise kui ka tõkke toimingud paljastavad partnerigruppe, mida saab eelnevalt registreerida, et eeltäita aluseks oleva transporditeostuse kasutatavaid andmestruktuure (töötaotluse (WR) mallimine). Hajutamine kopeerib erinevad viilud samast lähtepuhvrist partnerigrupi kõigi liikmete erinevatesse nihketesse. Tõke annab signaali kõigile partneritele kiirtoimingu kaudu. Nende toimingute hulgisseminek (enqueue) ja käsitlemine säästab kriitilisel teel väärtuslikke mikrosekundeid. EFA-l, kus kaks 200 Gbps võrgukaarti tagavad koos 400 Gbps ribalaiuse, jagame grupi partnerid kahe võrgukaardi vahel kildudeks (sharding), selle asemel et iga edastuse baite killustada.
Kuigi algselt arendasime TransferEngine'i ainult EFA-l libfabrici abil, töötavad MoE tuumad ConnectX-7-l ilma ühtegi koodirida muutmata pärast ConnectX-7 toe lisamist TransferEngine'i libibverbsi abil. Võrreldes skaleeritaval usaldusväärsel datagrammil (SRD) põhineva EFA-ga, nõuab ConnectX-7 ühenduse häälestamiseks ja partnerite haldamiseks täiendavat tööd. Kasutame usaldusväärse ühenduse (RC) käivitamiseks ebausaldusväärse datagrammi (UD) järjekorrapaari. Loome kummagi kaugpartneri jaoks kaks järjekorrapaari: ühe kahepoolse SEND/RECV jaoks ja teise ühepoolse WRITE_IMM jaoks, kuna nii RECV kui ka WRITE_IMM tarbivad RECV toimingut järjestikku, samas kui need kaks toimingut nõuavad erinevat puhvrisuurust. Lisaks lõdvestame PCIe tehingute järjekorda ConnectX-7 ja GPU mälu vahel. Sarnaselt EFA-le rakendame WR mallimise optimeerimist. Lisaks kasutame samale partnerile mitme väikese sõnumi saatmisel sõnumisageduse parandamiseks WR aheldamist.
Kuna GPU-poolne töö on mõlema adapteri puhul identne, viisime aluseks oleva transporditeostuse selle piirideni, et ületada DeepEP jõudlust ja rõhutada, et EFA ja ConnectX-7 MoE ruutingue erinevus on määratud alusriistvara võimetega.
Hindamine
Edastuse/kombineerimise latentsusaeg
Hindasime edastus- ja kombineerimislatentsust H200 süsteemides, mis on varustatud ConnectX-7 ja EFA adapteritega standardse DeepSeek-V3 ja Kimi-K2 konfiguratsiooni korral (varjatud dimensioon 7168, blokkskaalutud fp8 edastus ja bf16 kombineerimine). Arvestame partii suurusega 128 edastuse ja 4096 kombineerimise puhul. ConnectX-7-l võrdleme DeepEP-i ja meie eelmiste NVSHMEM-põhiste pplx-kernelitega, kasutades nii GPUDirect Async-i (IBGDA) kui ka hosti puhversõlme (IBRC). Saavutame kombineeritud latentsusajaks vastavalt 459 us, 582 us ja 692 us EP16, EP32 ja EP64 puhul. Oleme paremad teistest üldistel puhversõlmedel põhinevatest EFA-põhistest kernelitest, nagu UCCL-EP, mis saavutavad 519 us, 966 us ja 1159 us.
Dekodeerimisel saavutame vaatamata sellele, et meie tuumad on DeepEP-ist mõne mikrosekundi võrra aeglasemad, üldiselt madalama latentsusaja, kuna meie kombineerimistuumad on märgatavalt kiiremad. Kuigi kasutame hulgiedastuste tõttu hosti puhversõlme, ületame kõigi IBGDA-põhiste puhversõlmeta teostuste jõudlust. Samuti näitame, et puhversõlme teostusel on tähtsus, kuna meie kõrgelt spetsialiseerunud puhversõlm on üldisest NVSHMEM-i teostusest suurusjärgu võrra kiirem. Vaatamata asjaolule, et EFA on 256 KB pakettide edastamisel (mis on lähedased MoE ruutingu poolt vahetatavale tüüpilisele pakisuurusele) ConnectX-7-st kaks korda aeglasem, on kahe adapteri vaheline erinevus üldiselt vähem oluline.
Isegi kui meie dekodeerimiseks optimeeritud tuumad saavad mastaapseks muutuda eel-laadimise partiisuurusteni ilma täiendavate muudatusteta, jääb jõudlus DeepEP-ile alla. See tuleneb peamiselt asjaolust, et akumuleerime märke ainult kombineerimise vastuvõtjapoolses otsas, samas kui DeepEP akumuleerib saatjas, et vähendada edastatava andme mahtu. Kuid sellel optimeerimisel võivad olla täpsusmõjud. Lisaks kasutame andmete ühekordse kirjutamisega edastamisel saatmis- ja vastuvõtupuhvrite jaoks rohkem mälu.
Kimi-K2 / DeepSeek-V3 läbilaskevõime
Oma kernelite elujõulise olemuse tõestamiseks testime DeepSeek-V3 (671B parameetrit) ja Kimi-K2 (1T parameetrit) teenindava ettevõttesisese järeldusmootori lõpp-punkti jõudlust mitmel p5en eksemplaril, mis on varustatud Elastic Fabric Adapterite (EFA) ja H200 GPU-dega. Käitame mudelit mustandi verifitseerimise režiimis ja esitame latentsusaja, eeldades 80% aktsepteerimismäära ühe MTP-kihi abil.
DeepSeek-V3 puhul võrdleme tulemusi EP=8 sõlmesisese võrdlusalusega, mis käitab tõhusaid ainult NVLink-põhiseid edastus- ja kombineerimistuumi. Sõlmeülesed juurutused kas vastavad ühesõlmelise juurutuse jõudlusele või ületavad seda üldiselt, eriti keskmise suurusega partiide puhul, mis näitab, et sõlmede vahel mastaapne laiendamine toob üldisele läbilaskevõimele tõepoolest parandusi.
Kimi-K2 puhul võimaldavad tuumad teenindamist. Oma suuruse tõttu ei mahu mudel H200 sõlmes Data Parallelism (DP) abil samasse sõlmesse, mistõttu on mudeli teenindamiseks vaja sõlmeüülest juurutust. Kuna mudelil on vähem tähelepanupäid (attention heads), saavutame DeepSeek-V3-st paremad latentsusajad ja saame mudelit EFA-l elujõulise latentsusega teenindada.
Oleme kaalunud ka mikropakettimist, ehkki see parandas latentsusaega ainult väga suurte partiisuuruste korral või rohkema kui ühe mustandimärgi kasutamisel. Rühmitatud GEMM-probleemide kaheks jagamisest kasu saamiseks peavad märkide partiisuurused olema piisavalt tihedad. Suurem osa ajast kulub endiselt MoE edastamisele ja kombineerimisele, mis viitab sellele, et kernelite edasised täiustused parandavad sellest tulenevalt latentsusaega ja hõlbustavad mikropakettimist.
Tulevane töö
Teeme tihedat koostööd AWS-i inseneridega, et parandada EFA jõudlust. Omalt poolt jälgime libfabrici värskendusi, mis vähendavad andmetasandi lisakoormust. Lisaks katsetame efa-directiga, et vähendada veelgi kasutajaruumi lisakoormust ja suhelda otse riistvaraga.
Viited
HPC-süsteemide võrgujõudluse parandamine NVIDIA Magnum IO NVSHMEM ja GPUDirect Async abil
Tsitaat
@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}
}