Efektīva un portatīva ekspertu maisījuma saziņa

Mēs piedāvājam augstas veiktspējas, portatīvu, atvērtā koda bibliotēku Ekspertu maisījuma (MoE) saziņai, kas nodrošina 10 reizes ātrāku veiktspēju salīdzinājumā ar standarta All-to-All saziņas primitīviem.

Mūsu implementācijā ir vairāki galvenie tehniskie jauninājumi, kas nodrošina izcilu MoE saziņas efektivitāti:

GPU iniciēta saziņa (IBGDA): atbalsta tiešu GPU-uz-NIC saziņu, ievērojami samazinot latentumu, novēršot CPU iesaisti

Saziņas un skaitļošanas pārklāšanās: dalīta kodola arhitektūra ar atsevišķiem nosūtīšanas un saņemšanas posmiem ļauj veikt skaitļošanu vienlaikus ar tīkla pārsūtīšanu

Ātrākā viena mezgla veiktspēja: 2,5 reizes mazāks latentums nekā iepriekš ātrākajai implementācijai viena mezgla konfigurācijās

Efektīva un portatīva vairāku mezglu veiktspēja: Mūsu implementācija sasniedz līdz pat 10 reizēm lielāku ātrumu nekā standarta all-to-all saziņa. Lai gan tā ir aptuveni 2 reizes lēnāka nekā ļoti specializētas implementācijas, mūsu pieeja piedāvā labāku pārnesamību starp NVSHMEM versijām un tīkla vidēm (NVLink, CX-7 un EFA)

Bibliotēka ir pilnībā atvērtā koda un pieejama vietnē https://github.com/ppl-ai/pplx-kernels.

Šajā rakstā mēs pētām ekspertu paralelisma izaicinājumus liela mēroga MoE modeļos un aprakstām mūsu pieeju efektīvai žetonu (token) nosūtīšanai un kombinēšanai izkliedētās GPU vidēs.

Ekspertu maisījuma (MoE) modeļi, piemēram, DeepSeek R1 un Mixtral 8x7B, uzlabo blīvos modeļus, ierobežojot katram žetonam aktivizēto svaru skaitu. Tā vietā, lai sasniegtu lielu parametru skaitu, palielinot slāņu skaitu vai palielinot lineāro slāņu izmēru katra dekodētāja slāņa daudzslāņu perceptronā (MLP), MoE modeļi aizstāj tradicionālo MLP ar vairākiem ekspertiem un maršrutētāju. Piemēram, no DeepSeek R1 671B parametriem tikai 37B tiek reizināti ar konkrētu žetonu secināšanas (inference) laikā. Tas ievērojami samazina žetona dekodēšanai nepieciešamo skaitļošanas apjomu un uzlabo latentumu salīdzinājumā ar blīvu modeli.

MoE modeļi rada dažus papildu izaicinājumus secināšanai salīdzinājumā ar blīviem modeļiem. Lai gan paši eksperti ir mazi MLP slāņi, katrs dekodētāja slānis ietver maršrutētāju, kas izlemj, uz kuriem ekspertiem žetons tiek nosūtīts, turklāt katrs žetons tiek nosūtīts vairākiem ekspertiem. Maršrutētājs parasti ir mazs lineārs slānis, kas ģenerē varbūtību sadalījumu. Parasti tiek izvēlēti eksperti ar augstākajiem K rādītājiem, un galīgā aktivizācija tiek aprēķināta kā svērtais vidējais, summējot ekspertu izvades, kas reizinātas ar svaru, kurš atvasināts no maršrutētāja varbūtību sadalījuma.

Lai minimizētu latentumu izkliedētās MoE sistēmās, paralelismu var izmantot vairākās ierīcēs, taču tas rada saziņas izaicinājumus. Tādi modeļi kā Mixtral 8x7B vai Llama-70B iekļaujas 8 ierīcēs vienā mezglā, gūstot labumu no ātriem un zema latentuma NVLink savienojumiem (līdz 900 Gb/s). Tomēr lielākiem modeļiem ir nepieciešams sadalīt (shard) ekspertus vairākos mezglos, izmantojot InfiniBand (ar pīķa ātrumu 400 Gb/s), kas rada papildu latentuma izaicinājumus.

Tokenu maršrutēšana, ekspertu aprēķini un izvades apkopošana mixture-of-experts slānī.

Šajā rakstā mēs pētām ekspertu paralelisma problēmu, aprakstot mūsu GPU kodolu implementāciju, kas efektīvi nosūta žetonus uz ierīcēm, kurās atrodas piešķirtie eksperti, un savāc tos atpakaļ, aprēķinot ekspertu izvadu svērto summu. Lai gan pastāv agresīvi optimizētas implementācijas, kas orientētas uz noteiktu tīkla aparatūru, mēs piedāvājam ļoti portatīvu implementāciju, kas balstās uz minimālu NVSHMEM primitīvu kopu, tomēr nodrošina 10 reizes ātrāku veiktspēju nekā standarta all-to-all saziņa. Mūsu implementācija sasniedz augstākā līmeņa veiktspēju viena mezgla konfigurācijās, vienlaikus saglabājot lielisku pārnesamību dažādās tīkla vidēs.

Efektīvai secināšanai modeļa svariem jāatrodas ierīces atmiņā, vienlaikus atstājot pietiekami daudz telpa aktivizācijām, KV kešatmiņām un citiem buferiem, kas nepieciešami caurlaidei caur modeli. Jaudīgākie modeļi pārsniedz pat visjaudīgāko GPU ietilpību, tāpēc secināšana ir jāsadala starp vairākām ierīcēm, kuras kopīgi var uzglabāt modeļa svarus. Pamatojoties uz svaru sadalīšanas shēmām, ir jāizmanto dažādas saziņas un skaitļošanas shēmas, lai atslogotu skaitļošanu un sinhronizētu ierīces.

Ekspertu paralelisms (EP), kas parādīts pirmajā attēlā, paralelizē tikai ekspertu skaitļošanu. Dažādi eksperti tiek piešķirti dažādām ierīcēm, kurās glabājas to svari. Pēc maršrutēšanas žetoni tiek nosūtīti uz attiecīgo ierīci, un pēc tam rezultāti tiek apkopoti un uzkrāti. Maršrutēšanas sarežģītība ir atkarīga no paralelisma pakāpes pārējās modeļa daļās: citu slāņu replicēšana varētu pilnībā novērst maršrutēšanas nepieciešamību, jo katrs rangs var atlasīt žetonus no lokāli replicētas maršrutēšanas tabulas. Tomēr, ja maršrutēšanu veic tikai viens rangs, ir nepieciešama apraide (broadcast), lai nosūtītu žetonus, indeksus un svarus uz to attiecīgajiem ekspertiem. Visbeidzot, visu apkopojoša (all-gather) vai all-to-all apraide sinhronizē izvades žetonus ar to rangu, kas turpina modeļa izpildi. Šāda implementācija ir salīdzinoši vienkārša, jo torch jau nodrošina nepieciešamos primitīvus, lai gan zināmus ieguvumus varētu gūt, sapludinot uzkrāšanu ar savākšanas darbību, kas apkopo žetonus no ekspertiem.

Ekspertu paralēlisma apvienošana ar tensoru paralēlismu divos GPU.

Ekspertu paralelisms vien neizmērojas ideāli, jo klastera mezgli var būt dīkstāvē, kamēr modelis izpilda ne-ekspertu slāņus, piemēram, uzmanības (attention), normas un partveršanas slāņus. Tomēr šo slāņu, galvenokārt uzmanības, skaitļošana var gūt labumu arī no tensora paralelisma (TP). Lielākā daļa modeļu balstās uz vairāku galviņu uzmanību (multi-head attention), kas nozīmē, ka uzmanības galviņas un to atbilstošās Q, K un V projekcijas var arī sadalīt starp ierīcēm, replicējot vai apkopojot šķēles starp dažādiem slāņiem. Ja uzmanība ir sadalīta starp visām ierīcēm, all-gather var sinhronizēt aktivizācijas, ļaujot replicēt maršrutēšanu, kas efektīvai implementācijai prasa sinhronizācijas primitīvus, līdzīgus ekspertu paralelisma gadījumam. Tomēr pastāv paralelisma ierobežojumi šajā līmenī, jo uzmanības galviņu skaita samazināšana zem noteikta sliekšņa dos arvien mazāku atdevi.

Ekspertu paralēlisma apvienošana ar datu paralēlismu atsevišķiem pieprasījumiem divos GPU.

Lai vislabāk izmantotu visas ierīces un atbalstītu ļoti augstu ekspertu paralelisma pakāpi (līdz 128 vai 256 GPU), ir nepieciešams datu paralelisms (DP). Saskaņā ar šo shēmu ierīces, starp kurām modelis ir sadalīts, tiek sagrupētas, lai vienlaikus apstrādātu pieprasījumus, aprēķinot uzmanību un uzturot KV kešatmiņas, kas ir sadalītas to lokālajā grupā. Vairākas šo paralēlo grupu instances sadarbojas ekspertu novērtēšanā, un katra no tām mitina citu ekspertu svaru apakškopu. Pamatojoties uz uzmanības galviņu skaitu, grupa parasti var aptvert vesela mezgla lielumu, jo saziņa mezglā ir ātrāka. Piemēram, iepriekš redzamajā attēlā viens DP rangs neatkarīgi apkalpo divus pieprasījumus (A un B), apstrādājot uzmanību, normu un jebkādu citu grāmatvedību šiem pieprasījumiem. Otrs DP rangs apstrādā atsevišķu pieprasījumu C. Tomēr pirmais mezgls mitina pusi ekspertu, bet otrs – otru pusi, tāpēc pēc maršrutēšanas žetoni no A un B var tikt nosūtīti uz otro mezglu un otrādi. Tas rada retas saziņas problēmu: katra ierīce var nosūtīt atšķirīgu žetonu skaitu uz jebkuru citu mērķa rangu. Esošie primitīvi no torch, galvenokārt all_to_all, nav īpaši piemēroti, jo tiem var būt nepieciešama kāda veida polsterēšana (padding) vai GPU-CPU sinhronizācija un metadatu apraide. Lai efektīvi ieviestu saziņu, ir nepieciešami pielāgoti kodoli, lai apstrādātu maršrutēšanas informāciju un iniciētu saziņu no ierīcēm bez CPU iejaukšanās. Pēc maršrutēšanas nosūtīšanas kodolam (dispatch kernel) ir jānosūta žetoni uz rangu, uz kuru tie tika maršrutēti, savukārt kombinēšanas pusē ir jāsavāc aktivizācijas, kas pieder pieprasījumiem pašreizējā DP grupā. Turklāt darbs ir jālīdzsvaro DP grupā, kurā var būt vairākas ierīces, lai minimizētu apraidi un pareizi replicētu žetonus visos rangos nākamajiem slāņiem.

NVSHMEM ir NVIDIA specifiska OpenSHMEM implementācija, kas nodrošina portatīvas starp-ierīču saziņas iespējas, kuras abstrahē pamata aparatūras sarežģītību. API var izteikt ierīču-starp-ierīcēm nolasīšanu un ierakstīšanu, kas tiek kartētas uz atsevišķu transporta slāņu primitīviem. Tiek atbalstīti gan NVLink, gan RDMA, nodrošinot noteiktu pārnesamības pakāpi un ļaujot NVSHMEM kodoliem darboties mezglā vai vairāku mezglu klasterī. Mūsu kodolos mēs izmantojam attālinātās atmiņas ierakstīšanu un signalizācijas darbības, lai pārsūtītu datus un sinhronizētu iesaistītās ierīces.

NVSHMEM operācijas ir balstītas uz simetriskās atmiņas koncepciju: tās darbojas ar buferiem, kas ir piešķirti visās ierīcēs, kuras piedalās starp-ierīču saziņā. Lokāla ierīce var izmantot adresi, kas atvasināta tās pašas lokālajā buferī kopā ar mērķa ranga indeksu, lai norādītu operācijas galamērķi. Tālāk redzamajā attēlā ir ilustrēta šī koncepcija: abi GPU piešķir vienāda izmēra simetriskus buferus, saglabājot src rādītājus uz tiem. Pirmais rangs vēlas nosūtīt 3 veselus skaitļus otrajam, novietojot tos bufera sākumā: nvshmem_int_put_nbi atvasina sākuma adresi no lokālā bufera, norādot mērķa ierīci. Otrais rangs atvasina nobīdi no sava bufera, nosūtot vienu elementu pirmajai ierīcei, nobīdot par vienu. Lai gan mērķa adresēm vienmēr ir jābūt simetriskiem buferiem, kas piešķirti, izmantojot nvshmem_alloc, avota buferi var būt patvaļīgi ierīces atmiņas reģioni, ja tie ir iepriekš reģistrēti NVSHMEM.

Lai gan NVSHMEM nodrošina plašu primitīvu klāstu, mūsu kodoli balstās uz tikai 3 funkcijām, veidojot uz tām visu sinhronizāciju un nožogojumu (fencing).

Simetriskā atmiņa un attālināti ieraksti starp diviem GPU, izmantojot NVSHMEM.

nvshmemx_putmem_signal_nbi_warp: pārsūta datu bloku no vienas ierīces uz citu, vienlaikus iestatot karodziņu attālinātajā ierīcē. Operācija vai nu iestata (NVSHMEM_SIGNAL_SET), vai palielina (NVSHMEM_SIGNAL_ADD) 64 bitu atrašanās vietu. Karodziņš tiek atjaunināts pēc tam, kad viss atmiņas bloks ir pārsūtīts. Ja attālinātā ierīce konstatē karodziņa izmaiņas, tā var droši piekļūt buferim savā atmiņā. Šī funkcija ir noderīga datu pārsūtīšanas un sinhronizācijas apvienošanai.

nvshmemx_signal_op darbojas ar vienu atmiņas vietu, parasti 64 bitu karodziņu, atomiski to iestatot vai palielinot. Tas ir noderīgi metadatu nosūtīšanā un ierīču sinhronizācijā.

nvshmem_uint64_wait_until tiek izmantots signāla saņemšanas galā, lai aptaujātu karodziņu, līdz attālinātā ierīce to atjaunina.

Jebkādas secības nodrošināšana starp operācijām, izmantojot nvshmem_quiet vai nvshmem_fence, ir dārga, tāpēc mūsu kodolu implementācija izvairās no jebkādu citu NVSHMEM primitīvu izmantošanas un minimizē jebkādas atkarības, izņemot atsevišķus nosūtīšanas-saņemšanas pārus. Nosūtot datus, mēs vienmēr izmantojam funkciju nebloķējošo versiju, negaidot, kamēr dati vispār tiek nosūtīti no lokālā ranga. Cita, netieša sinhronizācija visā kodolu grupā nodrošina, ka šie buferi netiek destruktīvi pārrakstīti, kamēr NIC ir nepieciešama piekļuve tiem.

NVSHMEM izmanto GPUDirect RDMA, lai NIC varētu tieši nolasīt un ierakstīt GPU atmiņu datu apmaiņai bez CPU iesaistīšanas. Turklāt ConnectX NIC NVSHMEM atbalsta GPU iniciētu saziņu (pazīstama arī kā Infiniband GPUDirect Async jeb IBGDA), kas ļauj GPU iesniegt RDMA operācijas tieši NIC un tieši aptaujāt NIC pabeigšanas rindu. Platformās, kas neatbalsta GPU iniciētu saziņu, NVSHMEM izveido “starpniekservera” pavediens resursdatora CPU, lai iniciētu saziņu un aptaujātu pabeigšanu GPU vārdā. NVSHMEM programma ir portatīva neatkarīgi no tā, vai GPU iniciēta saziņa tiek atbalstīta vai nē. Tomēr GPU iniciēta saziņa ievērojami samazina latentumu, jo tā pilnībā novērš apiešanu caur CPU.

Mēs implementējam MoE saziņu, izmantojot nosūtīšanas un kombinēšanas kodolu pāri. Nosūtīšanas kodoli ir atbildīgi par žetonu un maršrutēšanas informācijas nolasīšanu katrā rangā, nosūtot tos attiecīgajiem ekspertiem. Kombinēšanas kodoli apkopo ekspertu ģenerētās aktivizācijas un nosūta tās atpakaļ uz to avota rangiem, vienlaikus aprēķinot svērto vidējo no izvēlētajiem ekspertiem, pamatojoties uz svariem, ko aprēķinājis maršrutētājs. Kodoli ir tālāk sadalīti nosūtīšanas un saņemšanas komponentē, lai ļautu datu pārsūtīšanai pārklāties ar skaitļošanu. Nosūtīšanas kodoli visi ir nebloķējoši un nesinhronizējoši: tie vienkārši nosūta visus ierakstus uz attālinātajām ierīcēm. Savukārt kombinēšanas kodoli tikai nolasa no atmiņas, gaidot, kamēr visi nepieciešamie dati ir pārsūtīti. Pēc darba nosūtīšanas uz NIC, kamēr dati tiek pārsūtīti asinhroni pa vadu, GPU var veikt citu noderīgu darbu lokāli, piemēram, lietojot koplietojamos ekspertus vai aprēķinot uzmanību.

Katram kodolu pārim ir savi simetriskās atmiņas buferi, un katrs rangs piešķir privātu krātuvi, lai saņemtu datus no visiem pārējiem rangiem. Maksimālais žetonu skaits, ko var nosūtīt no DP grupas, ir fiksēts visos rangos, kas arī nosaka augšējo robežu, ko katrs rangs var saņemt katram lokālajam ekspertam no katra DP ranga. Tas ļauj sūtītāja rangiem atvasināt unikālu adresi mērķa rangā, kurā ierakstīt, neprasot nekādu sinhronizāciju starp tiem. Pēc datu saņemšanas kodoli sajauc žetonus, lai izkārtotu tos atmiņā optimālā formātā, kas nepieciešams saņemšanas kodoliem. Lai gan buferiem ir ievērojams izmērs, tie tiek atkārtoti izmantoti visos modeļa secīgajos slāņos.

Atsevišķi nosūtīšanas un saņemšanas kodoli dispatch un combine operācijām.

Vienīgā globālās sinhronizācijas forma, kā redzams iepriekšējā attēlā, ir ieviesta kombinēšanas-saņemšanas kodolā. Kad dati ir nosūtīti no sūtīšanas, rangs atgriežas un turpina izpildi, kad tas saņem visus tam nepieciešamos žetonus no pārējiem rangiem, un nonāk kombinēšanas sūtīšanas kodolos, tiklīdz tie ir apstrādāti. Barjera kombinēšanas-saņemšanas kodolā nodrošina, ka neviens rangs nevar aizsteigties priekšā un sākt nosūtīšanas-sūtīšanu, kamēr kāds cits rangs joprojām gaida datu saņemšanu, kas varētu izraisīt destruktīvu pārklāšanos. Sinhronizācija vienkāršības labad tiek veikta kombinēšanas kodolā; kā optimizāciju to pašu barjeru varētu ieviest starp nosūtīšanas saņemšanas un kombinēšanas sūtīšanas kodoliem, vēl vairāk samazinot latentumu, dodot NIC vairāk laika pārsūtīšanas nostabilizēšanai, kamēr GPU izpilda citus kodolus.

Abi kodoli ir sadalīti visos pieejamajos ierīču SM: kamēr nosūtīšanas sūtīšanas un kombinēšanas saņemšanas kodoliem ir jāparalelizējas ar maksimālo žetonu skaitu uz rangu (max_m), nosūtīšanas saņemšanas un kombinēšanas sūtīšanas kodoliem ir jāapstrādā tik daudz žetonu no katra ranga un lokālā eksperta, maksimāli max_m * (num_experts // EP) * (EP // TP).

Nosūtīšanas kodols ir atbildīgs par žetonu un bloka mērogošanas faktoru nosūtīšanu vienaudžiem, kuri mitina ekspertus, uz kuriem žetoni tika maršrutēti. Sūtītāja puse balstās uz pavediena specializāciju, lai paralelizētu divus uzdevumus: maršrutēšanas informācijas apkopošanu, lai noteiktu, cik žetoni tiek nosūtīti no katra ranga, un žetonu iepakošanu ziņojumos, kas nosūtāmi uz attālināto ierīci. Saņēmēja puse vispirms gaida visu žetonu skaita saņemšanu, pēc tam nolasa un izpako lietderīgo slodzi no koplietojamās atmiņas tensoros, kas tiks nodoti tālākajiem grupētajiem GEMM kodoliem. Tā arī saglabā informāciju buferos, kas tiek koplietoti ar kombinēšanas kodoliem, lai norādītu, kur katrs žetons ir jānosūta atpakaļ. Šī informācija ir nepieciešama, jo saņemšanas kodols sajauc žetonus nepārtrauktos buferos nedeterministiskā secībā.

Tokenu nosūtīšana un to ierašanās signalizēšana starp avota un mērķa rangiem.

Sūtītāja daļā viens 32 pavediens pavediens katrā blokā ir atbildīgs par maršrutēšanas informācijas nolasīšanu un to žetonu skaita skaitīšanu, ko katram mērķa ekspertam ir paredzēts saņemt no šī ranga, nosūtot skaitu plus viens, izmantojot nvshmem_signal_op. Skaits tiek palielināts par viens, jo pāreja no nulles uz nenulli attālinātajā galā nozīmē skaitu saņemšanu no avota ranga. Paralēli atlikušie pavedieni sadarbojas, lai kopētu žetonus simetriskajā atmiņā visos blokos paralēli, iepakojot aktivizācijas, mērogošanas faktorus un to indeksu lokālajā rangā nepārtrauktā blokā. Indekss ir nepieciešams kombinēšanas sūtītājam, lai noteiktu adresi, kurā žetons tiks ierakstīts. Pēc tam, kad caur barjeru starp pavedienu grupām ir nodrošināta visu datu kopēšana, pavedieni atkal darbojas neatkarīgi, katrs nosūtot to pašu buferi citam ekspertam paralēli. Žetoni tiek nosūtīti, izmantojot nvshmemx_putmem_signal_nbi_warp, kas arī atomiski palielina nosūtīto žetonu skaitu no lokālā ranga attālinātajā ierīcē. DP grupā, tā ka katram rangam pieder nosūtāmā žetona kopija, nosūtīšana ir vienmērīgi līdzsvarota, katrai ierīcei nosūtot žetonu apakškopu.

Saņemšanas galā visi kodoli vispirms sinhronizējas ar sūtītāja rangiem, gaidot kopējā žetonu skaita nosūtīšanu. Pēc tam tie visi gaida, līdz atomiski palielinātais nosūtīto žetonu skaits nostabilizējas līdz kopējam skaitam, norādot, ka visa lietderīgā slodze no avota rangiem arī ir nosūtīta, pateicoties putmem izsaukuma semantikai. Kodoli aptaujā skaitu, izmantojot nvshmem_uint64_wait_until, paralelizējot operāciju visos blokos un pavedienos. Pēc tam starp-bloku barjera nodrošina, ka neviens bloks nenolasa buferus, ja vien visi dati nav pareizi saņemti. Sadalīti starp blokiem un sinhronizēti, izmantojot atomiski palielinātu skaitītāju, žetoni tiek kopēti no simetriskās atmiņas parastos nepārtrauktos buferos, kas kļūst par ievadēm tālākajiem kodoliem, kuri implementē ekspertus. Avota rangs, eksperta indekss un žetona indekss tiek uzglabāti atsevišķi, precīzi norādot vietu, kur kombinēšanas kodolam ir jānosūta aktivizācijas. Lai gan žetoni no DP grupas tiek nosūtīti no dažādām ierīcēm, tie visi tiek sagrupēti kopā, lai tos nodotu tālāk atbilstošajam ekspertam.

Kombinēšanas kodoli atkal ir sadalīti sūtīšanas un saņemšanas pusēs: sūtītāji kopē nekvantētās 16 bitu aktivizācijas uz attālinātajām ierīcēm, bet saņemšanas kodoli gaida datu nosūtīšanu un lokāli aprēķina visu ekspertu ieguldījumu svērto vidējo. Turklāt tie darbojas arī kā barjera, lai sinhronizētu nosūtīšanas-kombinēšanas secību: katrs rangs iestata karodziņu katram vienaudzim, ieejot sūtīšanas kodolos, un saņemšanas kodoliem nav ļauts atgriezties, ja vien tie nekonstatē karodziņa iestatīšanu. Sinhronizācijas latentums ir minimāls, jo tas pārklājas ar faktisko saziņu un skaitļošanu.

Ekspertu izvadu apvienošana tokens katrā datu paralēlajā rangā.

Sūtītāja pusē kodoli paralēli šķērso visu lokālo ekspertu žetonu sarakstu, ierakstot tos buferī galamērķa rangā. Mērķa rangs, eksperta indekss un žetona indekss tiek nolasīti no katra žetona buferiem, ko skatīšanas kodoli aizpildīja pēc žetonu saņemšanas. Katram sūtītājam ir savs privātās atmiņas reģions katram ekspertam, kurā rakstīt, kā norādīts iepriekšējā attēlā, novēršot nepieciešamību pēc sinhronizācijas. Līdzīgi kā nosūtīšanā, kombinēšana atomiski palielina katra žetona skaitītājus galamērķa rangā, lai norādītu datu saņemšanu: kad skaitītājs sakrīt ar to ekspertu skaitu, uz kuriem žetons tika nosūtīts, var piekļūt žetona saturam.

Saņemšanas kodolā žetonu saraksts tiek šķērsots paralēli vairākos blokos, gaidot to satura pienākšanu, aptaujājot signalizācijas operācijas iestatīto karodziņu. Pēc pienākšanas lietderīgā slodze tiek nolasīta no privātajiem buferiem, un maršrutēšanas tabula norāda, kuru buferi nolasīt un kādu svaru piešķirt katram ekspertam. Rezultāti pēc tam tiek ierakstīti ārēji piešķirtajos tensoros, un kodols pabeidz izpildi, kad visas ierīces ir šķērsojušas barjeru. DP grupā visi rangi saņem katra eksperta aktivizācijas kopiju, lai aprēķinātu savas kopijas.

Mēs novērtējam savus kodolus 128 GPU klasterī, kas sadalīts 16 mezglos, kuri savienoti, izmantojot InfiniBand un CX-7 NIC. Mēs izmēģinām gan GPUDirect Async (IBGDA), gan uzticamo savienojumu (RC) ar CPU starpniekserveri. Mēs salīdzinām tos ar augsti optimizēto DeepSeek implementāciju, kā arī ar blīvajiem primitīviem, ko nodrošina PyTorch (izmantojot NCCL) vai NVSHMEM.

Saziņas latentums 128 tokeniem uz vienu GPU multi-node EP128 konfigurācijā.

Lai gan blīvajai NVSHMEM all-to-all operācijai IBRC un IBGDA veiktspēja ir līdzīga (6378 µs pret 6180 µs), IBGDA ir ievērojami ātrāka ar retajiem (sparse) kodoliem. Saskaitot kopā nosūtīšanu un kombinēšanu, IBGDA izmanto 902 µs, savukārt IBRC aizņem 3223 µs — tas ir 3,6 reižu latentuma uzlabojums. Lai gan all-to-all ir joslas platuma ierobežots, retais kodols pārraida par vairākām pakāpēm mazāk datu, un to ierobežo latentums. Tieši aktivizējot tīkla pārsūtīšanu no GPU, neprasot CPU starpniekserverim koordinēt GPU un NIC, gala-līdz-gala latentums tiek ievērojami samazināts.

Lai gan mūsu portatīvie kodoli ir aptuveni 2 reizes lēnāki nekā augsti optimizētie DeepSeek kodoli, tie uzlabo latentumu 10 reizes salīdzinājumā ar blīvajiem kodoliem (902 µs pret 9944 µs).

Turklāt sadalījums sūtītāja-saņēmēja komponentēs arī ļauj paslēpt daļu latentuma, atšķirībā no bibliotēkas primitīviem.

Saziņas latentums 128 tokeniem uz vienu GPU single-node EP8 konfigurācijā.

Viena mezgla (EP8) gadījumā NVSHMEM transportēšanai izmanto NVLINK, nodrošinot mazāku latentumu un lielāku Caurlaides spēja nekā starp-mezglu tīklojums. Mūsu portatīvie kodoli ir aptuveni 2,5 reizes ātrāki nekā DeepEP viena mezgla režīmā (186 µs pret 481 µs).

Šeit aprakstītie kodoli pārspēj ML sistēmu iebūvētos primitīvus un piedāvā pienācīgu veiktspēju, īpaši nespecializējoties konkrētiem starp-ierīču transportiem, piemēram, NVLink vai InfiniBand. Papildus jau minētajām iespējām, turpmākus veiktspējas ieguvumus var sasniegt, aizstājot saziņas primitīvus ar specializētākām versijām. Piemēram, NVLink tīklā simetriskās atmiņas izmantošanu varētu aizstāt ar buferiem, kas tiek koplietoti starp ierīcēm, novēršot daļu kopēšanas un ļaujot izmantot smalkākas sinhronizācijas shēmas starp atsevišķiem žetoniem, nevis žetonu paketēm. InfiniBand tīklā implementācija varētu piekļūt pamata rindas pāriem tieši, novēršot nepieciešamību pārsūtīt gan gaidīto, gan nosūtīto žetonu skaitu un izskalot žetonus no rindām izvades buferos, tiklīdz tie pienāk. Tomēr šādas implementācijas prasa kompromisu ar portabilitāti, kamēr mūsu implementācija ļauj novērtēt klasterus ar specifiskām konfigurācijām pirms izstrādes pūļu veltīšanas optimizācijai konkrētai aparatūrai.

Mēs esam prezentējuši augstas veiktspējas, portatīvu bibliotēku MoE saziņai, kas sasniedz 10 reizes ātrāku veiktspēju salīdzinājumā ar standarta all-to-all saziņu, vienlaikus saglabājot saderību dažādās aparatūras konfigurācijās. Viena mezgla izvietojumos ar NVLink mūsu risinājums uzrāda 2,5 reizes zemāku latentumu nekā iepriekšējās implementācijas.

Mūsu pieeja līdzsvaro veiktspēju ar portabilitāti, izmantojot tādus galvenos jauninājumus kā GPU iniciētas saziņas atbalsts, dalītu kodolu arhitektūra, kas nodrošina skaitļošanas un saziņas pārklāšanos, un efektīvu žetonu nosūtīšanu, izmantojot minimālus NVSHMEM primitīvus. Lai gan tā ir aptuveni 2 reizes lēnāka nekā ļoti specializētas implementācijas vairāku mezglu iestatījumos, mūsu bibliotēka piedāvā izcilu elastību dažādās tīkla vidēs (NVLink, CX-7 un EFA).

Tā kā MoE modeļi turpina mērogoties, efektīvas saziņas stratēģijas, piemēram, mūsējā, kļūs arvien svarīgākas praktiskai izvietošanai. Mūsu pilnībā atvērtā koda implementācija ir pieejama vietnē https://github.com/ppl-ai/pplx-kernels.

DeepEP

DeepSeek-V3 tehniskais ziņojums

NVIDIA NVSHMEM

BibTeX

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

Ievads

Sadalīšana (Sharding) un paralelisms

NVSHMEM

Portatīvie kodoli

Nosūtīšana (Dispatch)

Kombinēšana (Combine)

Veiktspējas testi (Benchmarks)

Turpmāki uzlabojumi

Secinājums

Atsauces

Citēšana