Efektyvus ir nešiojamas „Mixture-of-Experts“ ryšys

Pristatome našią, nešiojamą, atvirojo kodo „Mixture-of-Experts“ (MoE) ryšio biblioteką, kurios veikimo sparta yra 10 kartų didesnė, palyginti su standartinėmis „All-to-All“ ryšio primityvomis.

Mūsų įgyvendinime yra keletas pagrindinių techninių naujovių, užtikrinančių geresnį MoE ryšio efektyvumą:

GPU inicijuojamas ryšys (IBGDA): palaiko tiesioginį ryšį tarp GPU ir tinklo plokštės (NIC), o tai smarkiai sumažina delsimą, nes nereikia įtraukti CPU

Ryšio ir skaičiavimo persidengimas: padalinta branduolio architektūra su atskirais siuntimo ir gavimo etetais leidžia skaičiavimams vykti vienu metu su tinklo perkėlimais

Didžiausia vieno mazgo sparta: 2,5 karto mažesnis delsimas nei ankstesnės sparčiausios realizacijos vieno mazgo konfigūracijose

Efektyvus ir nešiojamas kelių mazgų našumas: Mūsų įgyvendinimas pasiekia iki 10 kartų didesnį greitį nei standartinis „all-to-all“ ryšys. Nors jis yra maždaug 2 kartus lėtesnis už itin specializuotas realizacijas, mūsų požiūris užtikrina geresnį perkeliamumą tarp „NVSHMEM“ versijų ir tinklo aplinkų („NVLink“, „CX-7“ ir „EFA“)

Biblioteka yra visiškai atvirojo kodo ir pasiekiama adresu https://github.com/ppl-ai/pplx-kernels.

Šiame straipsnyje nagrinėjame ekspertų paralelizavimo iššūkius didelio masto MoE modeliuose ir aprašome savo požiūrį į efektyvų žetonų paskirstymą bei sujungimą paskirstytose GPU aplinkose.

Įvadas

„Mixture-of-Experts“ (MoE) modeliai, tokie kaip „DeepSeek R1“ ir „Mixtral 8x7B“, patobulina tankius modelius apribodami svorių, kurie aktyvuojami kiekvienam žetonui, skaičių. Vietoj didelio parametrų skaičiaus didinant sluoksnių skaičių arba didinant tiesinių sluoksnių dydį kiekvieno dekoderio sluoksnio daugiásluoksniame perceptrone (MLP), MoE modeliai tradicinį MLP pakeičia keliais ekspertais ir maršrutizatoriumi. Pavyzdžiui, iš 671 mlrd. „DeepSeek R1“ parametrų su konkrečiu žetonu per inferenciją dauginami tik 37 mlrd. Tai smarkiai sumažina skaičiavimų, reikalingų žetonui dekoduoti, kiekį ir pagerina delsimą, palyginti su tankiuoju modeliu.

MoE modeliai sukelia papildomų inferencijos iššūkių, palyginti su tankiaisiais modeliais. Nors patys ekspertai yra maži MLP sluoksniai, kiekvienas dekoderio sluoksnis apima maršrutizatorių, kuris nustato, kurjems ekspertams išsiunčiamas žetonas (kiekvienas žetonas nusiunčiamas keliems ekspertams). Maršrutizatorius paprastai yra mažas tiesinis sluoksnis, sugeneruojantis tikimybių paskirstymą. Paprastai pasirenkami ekspertai, turintys didžiausius „Top-K“ balus, o galutinis aktyvinimas apskaičiuojamas kaip pasvertas vidurkis, sudedant ekspertų išvestis, padaugintas iš svorio, gauto iš maršrutizatoriaus tikimybių paskirstymo.

Siekiant sumažinti delsimą paskirstytose MoE sistemose, paralelizavimu galima naudotis per kelis įrenginius, tačiau tai kelia ryšio iššūkių. Tokie modeliai kaip „Mixtral 8x7B“ arba „Llama-70B“ telpa 8 įrenginiuose viename mazge, išnaudodami spartias ir mažo delsimas „NVLink“ jungtis (iki 900 Gbps). Tačiau didesniems modeliams reikia dalinti ekspertus per kelis mazgus naudojant „InfiniBand“ (maksimaliai pasiekiant 400 Gbps), o tai kelia papildomų delsimo iššūkių.

Padidinti

Šiame straipsnyje nagrinėjame ekspertų paralelizavimo problemą, aprašydami mūsų GPU branduolių realizaciją, kuri efektyviai paskirsto žetonus įrenginiams, laikantiems priskirtus ekspertus, ir surenka juos atgal apskaičiuodama svorinį ekspertų išvesčių sumą. Nors egzistuoja agresyviai optimizuotos realizacijos, orientuotos į specifinę tinklo aparatinę įrangą, mes pristatome itin nešiojamą realizaciją, kuri remiasi minimaliu „NVSHMEM“ primityvų rinkiniu, tačiau pasiekia 10 kartų spartesnį našumą nei standartinis „all-to-all“ ryšys. Mūsų realizacija pasiekia aukščiausio lygio našumą vieno mazgo konfigūracijose, išlaikydama puikų perkeliamumą įvairiose tinklo aplinkose.

Dalijimas ir paralelizavimas

Efektyviai inferencijai modelio svoriai turi būti saugomi įrenginio atmintyje, paliekant pakankamai erdvė aktyvinimams, KV podėliams ir kitiems buferiams, kurių reikalauja tiesioginis modelio perėjimas. Pajėgiausi modeliai viršija net galingiausių GPU talpą, todėl inferencija turi būti paskirstyta per kelis įrenginius, kurie gali kolektyviai saugoti modelio svorius. Remiantis svorių skaidymo schemomis, reikia naudoti skirtingas ryšio ir skaičiavimo schemas, kad būtų perkelti skaičiavimai ir sinchronizuoti įrenginiai.

Padidinti

Ekspertų paralelizavimas (EP), pavaizduotas pirmoje iliustracijoje, lygiagretina tik ekspertų skaičiavimus. Skirtingi ekspertai priskiriami skirtingiems įrenginiams, kuriuose saugomi jų svoriai. Po maršrutizavimo žetonai siunčiami į atitinkamą įrenginį, o rezultatai vėliau surenkami ir sumuojami. Maršrutizavimo sudėtingumas priklauso nuo paralelizavimo laipsnio kitose modelio dalyse: kitų sluoksnių replikavimas galėtų visiškai pašalinti maršrutizavimo poreikį, nes kiekvienas rangas gali pasirinkti žetonus iš vietoje replikuotos maršrutizavimo lentelės. Tačiau jei maršrutizavimą vykdo tik vienas rangas, žetonams, indeksams ir svoriams į atitinkamus ekspertus išsiųsti reikalingas transliavimas. Galiausiai „all-gather“ arba „all-to-all“ transliavimas sinchronizuoja išvesties žetonus su tuo rangu, kuris tęsia modelio vykdymą. Toks įgyvendinimas yra gana paprastas, nes „torch“ jau pateikia reikiamus primityvus, nors tam tikros naudos galima gauti sujungiant sumavimą su surinkimo operacija, surenkančia žetonus iš ekspertų.

Padidinti

Ekspertų paralelizavimas nepasiekia idealaus mastelio, nes klasterio mazgai gali būti nenaudojami, kai modelis vykdo ne ekspertų sluoksnius, tokius kaip dėmesys, normalizavimas ir atranka. Tačiau šių sluoksnių skaičiavimai, visų pirma dėmesio, taip pat gali būti naudingi iš tenzoriaus paralelizavimo (TP). Daugelis modelių remiasi daugiagalviu dėmesiu, o tai reiškia, kad dėmesio galvutės ir jų atitinkamos Q, K bei V projekcijos taip pat gali būti padalintos tarp įrenginių, replikuojant arba surenkant dalis tarp įvairių sluoksnių. Jei dėmesys paskirstytas visiems įrenginiams, „all-gather“ operacija gali sinchronizuoti aktyvinimus, leidžianti replikuoti maršrutizavimą ir reikalaujančią sinchronizavimo primityvų, panašių į ekspertų paralelizavimo atvejį, kad būtų pasiektas efektyvus įgyvendinimas. Tačiau šio lygio paralelizavimas turi ribas, nes dėmesio galvučių skaičiaus sumažinimas žemiau tam tikros ribos duos mažėjančią grąžą.

Padidinti

Norint kuo geriau išnaudoti visus įrenginius ir palaikyti labai didelį ekspertų paralelizavimo laipsnį (iki 128 arba 256 GPU), reikalingas duomenų paralelizavimas (DP). Pagal šią schemą įrenginiai, per kuriuos padalintas modelis, sugrupuojami taip, kad užklausos būtų apdorojamos vienu metu, apskaičiuojant dėmesį ir palaikant KV podėlius, padalintus jų vietinėje grupėje. Keli šių lygiagrečių grupių egzemplioriai bendradarbiauja vertindami ekspertus, o kiekvienas iš jų talpina vis kitokį ekspertų svorių posekį. Atsižvelgiant į dėmesio galvučių skaičių, grupė paprastai gali apimti visą mazgą, nes mazgo vidinis ryšys yra spartesnis. Pavyzdžiui, aukščiau pateiktoje iliustracijoje vienas DP rangas nepriklausomai aptarnauja dvi užklausas (A ir B), tvarkydamas dėmesį, normalizavimą ir bet kokį kitą užklausų apskaitą. Kitas DP rangas apdoroja atskirą užklausą C. Tačiau pirmame mazge yra pusė ekspertų, o kitame – kita pusė, todėl po maršrutizavimo žetonai iš A ir B gali būti nusiųsti į antrą mazgą ir atvirkščiai. Tai sukelia retą ryšio problemą: kiekvienas įrenginys į bet kurį kitą paskirties rangą gali nusiųsti skirtingą žetonų skaičių. Esami „torch“ primityvai, pirmiausia „all_to_all“, nėra itin tinkami, nes jiems gali prireikti tam tikro paminkštinimo arba GPU ir CPU sinchronizavimo bei metaduomenų transliavimo. Kad ryšys būtų įgyvendintas efektyviai, reikalingi pasirinktiniai branduoliai, galintys apdoroti maršrutizavimo informaciją ir inicijuoti ryšį iš įrenginių be CPU įsikišimo. Po maršrutizavimo išsiuntimo branduolys turi nusiųsti žetonus į rangus, į kuriuos jie buvo srautu nukreipti, o sujungimo pusėje turi būti surenkami aktyvinimai, priklausantys esamos DP grupės užklausoms. Be to, darbas turi būti subalansuotas DP grupėje, kuri gali turėti kelis įrenginius, siekiant sumažinti transliavimą ir teisingai replikuoti žetonus visuose rankuose vėlesniems sluoksniams.

NVSHMEM

„NVSHMEM“ yra „NVIDIA“ specifinė „OpenSHMEM“ realizacija, teikianti nešiojamas tarpįrenginio ryšio priemones, kurios abstrakčiu būdu paslepia pagrindinės aparatinės įrangos sudėtingumą. API gali išreikšti skaitymą ir rašymą tarp įrenginių, kurie susiejami su atskirų perdavimo sluoksnių primityvomis. Palaikomi tiek „NVLink“, tiek „RDMA“, suteikiantys tam tikrą perkeliamumą ir leidžiantys „NVSHMEM“ branduoliams veikti mazge arba kelių mazgų klasteryje. Savo branduoliuose naudojame nuotolinio atminties rašymo ir signalizavimo operacijas, kad perduotume duomenis ir sinchronizuotume dalyvaujančius įrenginius.

„NVSHMEM“ operacijos yra sukurtos aplink simetriškos atminties koncepciją: jos veikia su buferiais, kurie buvo paskirstyti visuose tarpįrenginio ryšiui dalyvaujančiuose įrenginiuose. Vietinis įrenginys gali naudoti adresą, išvestą jo paties vietiniame buferyje, kartu su tikslinio rango indeksu, kad nurodytų operacijos paskirtį. Žemiau pateiktoje iliustracijoje ši koncepcija paaiškinama: abu GPU paskirsto tokio paties dydžio simetriškus buferius, išlaikydami src rodykles į juos. Pirmasis rangas nori nusiųsti 3 sveikuosius skaičius antrajam, padėdamas juos buferio pradžioje: nvshmem_int_put_nbi išveda pradžios adresą iš vietinio buferio, nurodydamas tikslinį įrenginį. Antrasis rangas išveda poslinkį iš savo buferio, nusiųsdamas vieną elementą pirmam įrenginiui su vieneto poslinkiu. Nors paskirties adresai visada turi būti simetriški buferiai, paskirstyti naudojant nvshmem_alloc, šaltinio buferiai gali būti savavališki įrenginio atminties regionai, jei jie iš anksto užregistruoti „NVSHMEM“.

Nors „NVSHMEM“ siūlo platų primityvų asortimentą, mūsų branduoliai remiasi tik 3 funkcijomis, ant jų statydami visą sinchronizavimą ir aptvėrimą.

Padidinti

nvshmemx_putmem_signal_nbi_warp: Perkelia duomenų bloką iš vieno įrenginio į kitą ir nustato vėliavėlę nuotoliniame įrenginyje. Operacija arba nustato (NVSHMEM_SIGNAL_SET), arba padidina (NVSHMEM_SIGNAL_ADD) 64 bitų vietą. Vėliavėlė atnaujinama perkėlus visą atminties bloką. Jei nuotolinis įrenginys pastebėjo vėliavėlės pasikeitimą, jis gali saugiai pasiekti buferį savo atmintyje. Ši funkcija naudinga duomenų perdavimui ir sinchronizavimui apjungti.

nvshmemx_signal_op veikia su viena atminties vieta, paprastai 64 bitų vėliavėle, atomeškai ją nustatydama arba padidindama. Tai naudinga perduodant metaduomenis ir sinchronizuojant įrenginius.

nvshmem_uint64_wait_until naudojamas signalo gavimo pusėje, apklausos vėliavėlei, kol nuotolinis įrenginys ją atnaujins.

Užtikrinti bet kokią operacijų eiliškumo formą per nvshmem_quiet arba nvshmem_fence yra brangu, todėl mūsų branduolių įgyvendinimas vengia naudoti bet kokius kitus „NVSHMEM“ primityvus ir sumažina bet kokias priklausomybes, išskyrus atskiras siuntimo ir gavimo poras. Siųsdami duomenis, mes visada naudojame neblokuojančią funkcijos versiją, nelaukdami, kol duomenys netgi bus išsiųsti iš vietinio rango. Kitas netiesioginis sinchronizavimas branduolių grupėje užtikrina, kad šie buferiai nebūtų griaunamai perrašyti, kol tinklo plokštėms reikalinga prieiga prie jų.

„NVSHMEM“ naudoja „GPUDirect RDMA“, kad tinklo plokštė galėtų tiesiogiai skaityti ir rašyti GPU atmintį duomenų mainams, neįtraukdama CPU. Be to, „ConnectX“ tinklo plokštėse „NVSHMEM“ palaiko GPU inicijuojamą ryšį (taip pat žinomą kaip „Infiniband GPUDirect Async“ arba „IBGDA“), kuris leidžia GPU siųsti RDMA operacijas tiesiogiai į tinklo plokštę ir tiesiogiai apklausti tinklo plokštės užbaigimo eilę. Platformose, kurios nepalaiko GPU inicijuojamo ryšio, „NVSHMEM“ sukuria „įgaliotinio“ („proxy“) gija pagrindiniame CPU, kad inicijuotų ryšį ir apklaustų užbaigimą GPU vardu. „NVSHMEM“ programa yra nešiojama nepriklausomai nuo to, ar GPU inicijuojamas ryšys yra palaikomas, ar ne. Tačiau GPU inicijuojamas ryšys smarkiai sumažina delsimą, nes visiškai aplenkia apėjimus per CPU.

Nešiojami branduoliai

MoE ryšį įgyvendiname per siuntimo ir gavimo branduolių porą. Siuntimo branduoliai yra atsakingi už žetonų ir maršrutizavimo informacijos skaitymą kiekviename ranke, paskirstant juos atitinkamiems ekspertams. Sujungimo branduoliai surenka ekspertų sugeneruotus aktyvinimus ir nusiunčia juos atgal į jų šaltinio rangus, kartu apskaičiuodami pasvertą vidurkį iš pasirinktų ekspertų, remdamiesi maršrutizatoriaus apskaičiuotais svoriais. Branduoliai papildomai padalijami į siuntimo ir gavimo komponentus, kad duomenų perkėlimai galėtų persidengti su skaičiavimais. Siuntimo branduoliai yra neblokuojantys ir nesinchronizuojantys: jie tiesiog išsiunčia visus įrašus į nuotolinius įrenginius. Kitoje pusėje sujungimo branduoliai skaito tik iš atminties, laukdami, kol bus perkelti visi reikalingi duomenys. Paskirstę darbą tinklo plokštėms, kol duomenys perkeliami asinkroniškai laidu, GPU gali atlikti kitus naudingus darbus vietoje, pavyzdžiui, pritaikydami bendrus ekspertus arba skaičiuodami dėmesį.

Kiekviena branduolių pora turi savo simetriškos atminties buferius, o kiekvienas rangas paskirsto privačią saugyklą duomenims gauti iš visų kitų rangų. Didžiausias žetonų skaičius, kurį galima išsiųsti iš DP grupės, yra fiksuotas tarp rangų, o tai taip pat nustato viršutinę ribą, kurią kiekvienas rangas gali gauti kiekvienam vietiniam ekspertui iš kiekvieno DP rango. Tai leidžia siuntėjo rangams išvesti unikalų adresą paskirties ranke, į kurį reikia rašyti, nereikalaujant jokio sinchronizavimo tarp jų. Gavus duomenis, branduoliai perkelia žetonus, kad išdėstytų juos atmintyje optimaliu formatu, kurio reikalauja gaunantys branduoliai. Nors buferiai yra nemažo dydžio, jie yra pakartotinai naudojami visuose nuosekliuose modelio sluoksniuose.

Padidinti

Vienintelė pasaulinio sinchronizavimo forma, kaip parodyta aukščiau esančioje iliustracijoje, įgyvendinta sujungimo-gavimo branduolyje. Kai duomenys išsiunčiami iš siuntimo, rangas grįžta ir tęsia vykdymą, kai gauna visus reikalingus žetonus iš kitų rangų, ir įeina į sujungimo siuntimo branduolius, kai tik jie apdorojami. Barjeras sujungimo-gavimo metu užtikrina, kad joks rangas negalėtų aplenkti ir pradėti paskirstymo-siuntimo, kol bet kuris kitas rangas vis dar laukia duomenų gavimo, o tai gali sukelti destruktyvų persidengimą. Sinchronizavimas atliekamas sujungimo branduolyje siekiant paprastumo, o kaip optimizavimas tas pats barjeras galėtų būti įgyvendintas tarp paskirstymo-gavimo ir sujungimo-siuntimo branduolių, dar labiau sumažinant delsimą, suteikiant tinklo plokštėms daugiau laiko atlikti perkėlimus, kol GPU vykdo kitus branduolius.

Abudu branduoliai yra padalinti per visus pasiekiamus įrenginių SM: kol paskirstymo siuntimo ir sujungimo gavimo branduoliai turi paralelizuoti pagal vieno rango maksimalų žetonų skaičių (max_m), paskirstymo gavimo ir sujungimo siuntimo branduoliai turi apdoroti tiek žetonų iš kiekvieno rango ir vietinio eksperto, neviršijant daugiausia max_m * (num_experts // EP) * (EP // TP).

Paskirstymas

Paskirstymo branduolys yra atsakingas už žetonų ir bloko mastelio keitimo faktorių paskirstymą partneriams, kurie talpina ekspertus, į kuriuos buvo nukreipti žetonai. Siuntėjo pusė remiasi „warp“ specializacija, kad paralelizuotų dvi užduotis: maršrutizavimo informacijos agregavimą, siekiant nustatyti, kiek žetonų siunčiama iš kiekvieno rango, ir žetonų supakavimą į pranešimus, kurie bus išsiųsti į nuotolinį įrenginį. Gavėjo pusė pirmiausia laukia, kol bus gauti visi žetonų kiekiai, po to skaito ir išpakuoja naudingąją apkrovą iš bendros atminties į tenzorius, kurie bus perduodami tolesniems sugrupuotiems „Grouped GEMM“ branduoliams. Jis taip pat saugo informaciją buferiuose, bendrinamuose su sujungimo branduoliais, kad nurodytų, kur kiekvienas žetonas turi būti išsiųstas atgal. Ši informacija reikalinga, nes gavimo branduolys perkelia žetonus gretutiniuose buferiuose nedeterministine tvarka.

Padidinti

Siuntėjo dalyje viena 32 gijos gija kiekviename bloke yra atsakinga už maršrutizavimo informacijos skaitymą ir žetonų, kuriuos kiekvienas paskirties ekspertas tikisi gauti iš šio rango, skaičiavimą, siunčiant kiekį plius vienetą naudojant nvshmem_signal_op. Kiekis padidinamas vienetu, nes perėjimas nuo nulio prie nenuulinio skaičiaus nuotoliniame gale reiškia skaičių gavimą iš šaltinio rango. Lygiagrečiai likusios gijos bendradarbiauja nukopijuodamos žetonus į simetrišką atmintį visuose blokuose lygiagrečiai, supakuodamos aktyvinimus, mastelio keitimo faktorius ir jų indeksą vietiniame ranke į vientisą gabalą. Indeksas reikalingas sujungimo siuntėjui, kad būtų nustatytas adresas, į kurį bus įrašytas žetonas. Toliau, užtikrinus, kad visi duomenys buvo nukopijuoti per barjerą tarp gijų grupių, gijos vėl veikia nepriklausomai, kiekviena siųsdama tą patį buferį skirtingam ekspertui lygiagrečiai. Žetonai siunčiami naudojant nvshmemx_putmem_signal_nbi_warp, kuris taip pat atomeškai padidina išsiųstų žetonų iš vietinio rango skaičių nuotoliniame įrenginyje. DP grupėje, kadangi kiekvienas rangas turi siunčiamo žetono repliką, paskirstymas subalansuojamas tolygiai, o kiekvienas įrenginys išsiunčia žetonų posekį.

Gavimo gale visi branduoliai pirmiausia sinchronizuojasi su siuntėjo rangais, laukdami, kol bus atsiųsti bendri žetonų kiekiai. Po to jie visi laukia, kol atomeškai padidinti išsiųstų žetonų kiekiai susistovės iki bendrų kiekių, nurodydami, kad visos naudingosios apkrovos iš šaltinio rangų taip pat buvo išsiųstos, dėka putmem iškvietos semantikos. Branduoliai apklausia kiekius naudodami nvshmem_uint64_wait_until, paralelizuodami operaciją visuose blokuose ir gijos. Vėliau tarpblokinis barjeras užtikrina, kad joks blokas neskaitytų buferių, nebent visi duomenys buvo teisingai gauti. Paskirstyti per blokus ir sinchronizuoti per atomeškai padidintą skaitliuką, žetonai nukopijuojami iš simetriškos atminties į įprastus gretutinius buferius, kurie tampa įvestimis tolesniems branduoliams, įgyvendinantiems ekspertus. Šaltinio rangas, eksperto indeksas ir žetono indeksas saugomi atskirai, tiksliai nurodant vietą, kurioje sujungimo branduolys turi išsiųsti aktyvinimus. Nors žetonai iš DP grupės siunčiami iš skirtingų įrenginių, jie visi sugrupuojami kartu, kad būtų perduoti atitinkamam ekspertui.

Sujungimas

Sujungimo branduoliai vėlgi padalijami į siuntimo ir gavimo puses: siuntėjai nukopijuoja nekantuotus 16 bitų aktyvinimus į nuotolinius įrenginius, o gavimo branduoliai laukia, kol duomenys bus atsiųsti, ir vietoje apskaičiuoja visų ekspertų indėlių pasvertą vidurkį. Be to, jie taip pat veikia kaip barjeras, sinchronizuodami paskirstymo-sujungimo seką: kiekvienas rangas nustato vėliavėlę kiekviename partnerio įrenginyje įvedimo į siuntimo branduolius metu, o gavimo branduoliams neleidžiama grįžti, nebent jie pastebi, kad vėliavėlė nustatyta. Sinchronizavimo delsimas yra minimalus, nes jis persidengia su tikruoju ryšiu ir skaičiavimais.

Padidinti

Siuntėjo pusėje branduoliai lygiagrečiai pereina per visiems vietiniams ekspertams priskirtų žetonų sąrašą, įrašydami juos į buferį paskirties ranke. Tikslinis rangas, eksperto indeksas ir žetono indeksas nuskaitomi iš žetonų buferių, kuriuos išsklaidymo branduoliai užpildė gavus žetonus. Kiekvienas siuntėjas turi savo privačią atminties sritį kiekvienam ekspertui, į kurią rašo, kaip parodyta aukščiau esančioje iliustracijoje, išvengiant sinchronizavimo poreikio. Panašiai kaip ir paskirstymo atveju, sujungimas atomeškai padidina skaitliukus žetonui paskirties ranke, kad nurodytų duomenų gavimą: kai skaitliukas sutampa su ekspertų, į kuriuos buvo nukreiptas žetonas, skaičiumi, galima pasiekti žetono turinį.

Gavimo branduolyje žetonų sąrašas pereinamas lygiagrečiai per kelis blokus, laukiant, kol bus gautas jų turinys, apklausiant signalizavimo operacijos nustatytą vėliavėlę. Atvykus, naudingosios apkrovos nuskaitomos iš privačių buferių, o maršrutizavimo lentelė nurodo, kurį buferį skaityti ir kokį svorį priskirti kiekvienam ekspertui. Tuomet rezultatai įrašomi į išoriškai paskirstytus tenzorius, o branduolys baigia vykdymą, kai visi įrenginiai praeina barjerą. Visoje DP grupėje visi rangai gauna kiekvieno eksperto aktyvinimo kopiją, kad galėtų apskaičiuoti savo replikas.

Skačiavimo testai

Įvertiname savo branduolius 128 GPU klasteryje, paskirstytame per 16 mazgų, sujungtų per „InfiniBand“ ir „CX-7“ tinklo plokštes. Išbandome tiek „GPUDirect Async“ (IBGDA), tiek patikimą ryšį (RC) su CPU įgaliotiniu. Lyginame juos su itin optimizuota „DeepSeek“ realizacija, taip pat su tankiaisiais primityvais, kuriuos teikia „PyTorch“ (per NCCL) arba „NVSHMEM“.

Padidinti

Nors atliekant tankiąją „NVSHMEM“ „all-to-all“ operaciją „IBRC“ ir „IBGDA“ našumas yra panašus (6378 µs palyginti su 6180 µs), „IBGDA“ veikia žymiai greičiau su retaisiais branduoliais. Sudėjus paskirstymą ir sujungimą, „IBGDA“ sunaudoja 902 µs, o „IBRC“ užtrunka 3223 µs – tai 3,6 karto geresnis delsimas. Nors „all-to-all“ ribojasi pralaidumu, retieji branduoliai perduoda daug mažiau duomenų ir ribojasi delsimu. Tiesiogiai inicijuojant tinklo perkėlimus iš GPU, nereikalaujant CPU įgaliotinio koordinuoti GPU ir tinklo plokštės, galutinis delsimas smarkiai sumažėja.

Nors mūsų nešiojami branduoliai yra maždaug 2 kartus lėtesni už itin optimizuotus „DeepSeek“ branduolius, jie pagerina delsimą 10 kartų, palyginti su tankiaisiais branduoliais (902 µs palyginti su 9944 µs).

Be to, padalijimas į siuntėjo ir gavėjo komponentus taip pat leidžia paslėpti dalį delsimo, skirtingai nei bibliotekos primityvai.

Padidinti

Vieno mazgo konfigūracijoje (EP8) „NVSHMEM“ naudoja „NVLINK“ perdavimui, užtikrindama mažesnį delsimą ir didesnį pralaidumą nei tarpmazginis tinklas. Mūsų nešiojami branduoliai yra maždaug 2,5 karto greitesni nei „DeepEP“ viename mazge (186 µs palyginti su 481 µs).

Tolimesni patobulinimai

Čia aprašyti branduoliai pranoksta įmontuotus mašininio mokymosi karkasų primityvus ir siūlo padorų našumą nepergaliant konkretiems tarpįrenginių transportams, tokiams kaip „NVLink“ arba „InfiniBand“. Be jau minėtų galimybių, tolesnio našumo padidėjimo galima pasiekti pakeičiant ryšio primityvus specializuotesnėmis versijomis. Pavyzdžiui, per „NVLink“ simetriškos atminties naudojimas galėtų būti pakeistas tarp įrenginių bendrinamais buferiais, pašalinant kai kurį kopijavimą ir leidžiant taikyti smulkesnes sinchronizavimo schemas, atskiriems žetonams, o ne jų paketiškiems blokams. Per „InfiniBand“ realizacija galėtų tiesiogiai pasiekti pagrindines eilių poras, pašalinant poreikį perduoti tiek laukiamą, tiek išsiųstą žetonų kiekį ir išvalyti žetonus iš eilių į išvesties buferius jiems atvykus. Tačiau tokios realizacijos reikalauja perkeliamumo aukos, o mūsų realizacija leidžia įvertinti specifinių konfigūracijų klasterius prieš skiriant kūrimo pastangas optimizavimui konkrečiai aparatinei įrangai.

Išvada

Pristatėme našią, nešiojamą MoE ryšio biblioteką, kuri užtikrina 10 kartų spartesnį našumą, palyginti su standartiniu „all-to-all“ ryšiu, išlaikydama suderinamumą tarp įvairių aparatinės įrangos konfigūracijų. Vieno mazgo diegimuose su „NVLink“ mūsų sprendimas demonstruoja 2,5 karto mažesnį delsimą nei ankstesnės realizacijos.

Mūsų požiūris subalansuoja našumą su perkeliamumu pasitelkiant pagrindines naujoves, įskaitant GPU inicijuojamo ryšio palaikymą, padalintą branduolių architektūrą, įgalinančią skaičiavimų ir ryšio persidengimą, bei veiksmingą žetonų paskirstymą naudojant minimalius „NVSHMEM“ primityvus. Nors daugiaskylių sąrankose jis yra maždaug 2 kartus lėtesnis už itin specializuotas realizacijas, mūsų biblioteka pasižymi geresniu lankstumu įvairiose tinklo aplinkose („NVLink“, „CX-7“ ir „EFA“).

Augant MoE modelių masteliui, efektyvios ryšio strategijos, tokios kaip mūsiškė, taps vis svarbesnės praktiniam diegimui. Mūsų visiškai atvirojo kodo realizacija pasiekiama adresu https://github.com/ppl-ai/pplx-kernels.

Literatūra

DeepEP

„DeepSeek-V3“ techninė ataskaita

„NVIDIA NVSHMEM“

Citata

@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://perplexity.coolkeypoint.com/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}