Efektivní a přenositelná komunikace mixture-of-experts

Představujeme vysoce výkonnou, přenositelnou a open-source knihovnu pro komunikaci typu Mixture-of-Experts (MoE), která dosahuje 10krát vyššího výkonu ve srovnání se standardními komunikačními primitivy All-to-All.

Naše implementace přináší několik klíčových technických inovací, které zajišťují vynikající efektivitu MoE komunikace:

Komunikace iniciovaná GPU (IBGDA): Podporuje přímou komunikaci mezi GPU a síťovou kartou (NIC), což výrazně snižuje latenci tím, že zcela obchází zapojení CPU.

Překryv komunikace a výpočtů: Architektura rozděleného jádra s oddělenými fázemi odesílání a příjmu umožňuje, aby výpočty probíhaly současně se síťovými přenosy.

Nejvyšší výkon na jednom uzlu: 2,5krát nižší latence než u dosud nejrychlejší implementace v konfiguracích s jedním uzlem.

Efektivní a přenositelný výkon na více uzlech: Naše implementace dosahuje rychlosti až 10krát vyšší než standardní komunikace all-to-all. Ačkoliv je přibližně 2krát pomalejší než vysoce specializované implementace, náš přístup nabízí lepší přenositelnost napříč verzemi NVSHMEM a síťovými prostředími (NVLink, CX-7 a EFA).

Knihovna je plně open-source a je k dispozici na adrese https://github.com/ppl-ai/pplx-kernels.

V tomto článku se zabýváme výzvami paralelismu expertů ve velkých modelech MoE a popisujeme náš přístup k efektivnímu odesílání a kombinování tokenů v distribuovaných prostředích GPU.

Úvod

Modely Mixture-of-Experts (MoE), jako jsou DeepSeek R1 a Mixtral 8x7B, vylepšují husté modely tím, že omezují počet vah, které se aktivují pro každý token. Namísto dosahování vysokého počtu parametrů zvyšováním počtu vrstev nebo zvětšováním lineárních vrstev v multilayer perceptronu (MLP) každé vrstvy dekodéru nahrazují modely MoE tradiční MLP několika experty a routerem. Například z 671 miliard parametrů modelu DeepSeek R1 se během inference násobí s daným tokenem pouze 37 miliard. To výrazně snižuje množství výpočtů potřebných k dekódování tokenu a zlepšuje latenci ve srovnání s hustým modelem.

Modely MoE přinášejí pro inferenci určité dodatečné výzvy ve srovnání s hustými modely. Zatímco samotní experti jsou malé vrstvy MLP, každá vrstva dekodéru obsahuje router, který rozhoduje, ke kterým expertům je token odeslán, přičemž každý token je odeslán k několika expertům. Router je obvykle malá lineární vrstva vytvářející rozdělení pravděpodobnosti. Obvykle se vyberou experti s nejvyšším skóre (top-K) a konečná aktivace se vypočítá jako vážený průměr, přičemž se sečtou výstupy expertů vynásobené váhou odvozenou z rozdělení pravděpodobnosti routeru.

Pro minimalizaci latence v distribuovaných systémech MoE lze využít paralelismus napříč více zařízeními, což však přináší komunikační výzvy. Modely jako Mixtral 8x7B nebo Llama-70B se vejdou do 8 zařízení v rámci jednoho uzlu a těží z rychlých a nízkolatenčních propojení NVLink (až 900 Gb/s). Větší modely však vyžadují shardování expertů napříč více uzly pomocí technologie InfiniBand (s vrcholem na 400 Gb/s), což přináší další problémy s latencí.

Zvětšit

V tomto článku se zabýváme problémem paralelismu expertů a popisujeme naši implementaci jader GPU, která efektivně odesílají tokeny na zařízení obsahující přiřazené experty a sbírají je zpět výpočtem váženého součtu výstupů expertů. Přestože existují agresivně optimalizované implementace cílené na specifický síťový hardware, představujeme vysoce přenositelnou implementaci, která spoléhá na minimální sadu primitiv NVSHMEM a přitom poskytuje výkon 10krát rychlejší než standardní komunikace all-to-all. Naše implementace dosahuje špičkového výkonu v konfiguracích s jedním uzlem a zároveň si zachovává vynikající přenositelnost napříč různými síťovými prostředími.

Shardování a paralelismus

Pro efektivní inferenci musí být váhy modelu uchovávány v paměti zařízení a zároveň musí být ponechán dostatečný prostor pro aktivace, mezipaměti KV a další vyrovnávací paměti vyžadované pro průchod dopředným směrem (forward pass) modelem. Schopnější modely překračují kapacitu i těch nejvýkonnějších GPU, takže inferenci je nutné rozprostřít mezi více zařízení, která mohou společně ukládat váhy modelu. Na základě schémat shardování vah je nutné použít různá komunikační a výpočetní schémata pro odlehčení výpočtů a synchronizaci zařízení.

Zvětšit

Paralelismus expertů (Expert Parallelism – EP), znázorněný na prvním obrázku, paralelezuie pouze výpočty expertů. Různí experti jsou přiřazeni k různým zařízením, která uchovávají jejich váhy. Po routování jsou tokeny odeslány na příslušné zařízení a výsledky jsou následně sesbírány a akumulovány. Složitost routování závisí na stupni paralelismu v ostatních částech modelu: replikace ostatních vrstev by mohla zcela eliminovat potřebu routování, protože každá úroveň (rank) může vybírat tokeny z lokálně replikované tabulky routování. Pokud však routování provádí pouze jedna úroveň, je k odeslání tokenů, indexů a vah příslušným expertům vyžadováno všesměrové vysílání (broadcast). Nakonec operace all-gather nebo all-to-all synchronizuje výstupní tokeny s tou úrovní, která pokračuje v provádění modelu. Taková implementace je poměrně jednoduchá, protože knihovna torch již potřebná primitiva zpřístupňuje, ačkoliv určité výhody by mohlo přinést sloučení (fusing) akumulace s operací sběru, která shromažďuje tokeny od expertů.

Zvětšit

Paralelismus zaměřený pouze na experty se neškáluje ideálně, protože uzly v clusteru mohou být nečinné, zatímco model provádí vrstvy, které nejsou spojeny s experty, jako je pozornost (attention), normalizace a vzorkování. Výpočty těchto vrstev, především pozornosti, však mohou těžit také z tenzorového paralelismu (Tensor Parallelism – TP). Většina modelů spoléhá na vícenásobnou pozornost (multi-head attention), což znamená, že hlavy pozornosti a jejich odpovídající projekce Q, K a V mohou být rovněž rozshardovány mezi zařízení, přičemž výřezy mezi různými vrstvami se replikují nebo shromažďují. Pokud je pozornost rozprostřena na všechna zařízení, operace all-gather může synchronizovat aktivace, což umožňuje replikovat routování a pro efektivní implementaci vyžaduje synchronizační primitiva podobná případu paralelismu expertů. Na této úrovni však existují limity paralelismu, protože snížení počtu hlav pozornosti pod určitou hranici přinese klesající návratnost.

Zvětšit

Pro optimální využití všech zařízení a podporu velmi vysokého stupně paralelismu expertů (až 128 nebo 256 GPU) je vyžadován datový paralelismus (Data Parallelism – DP). V tomto schématu jsou zařízení, mezi která je model rozdělen, seskupena tak, aby současně zpracovávala požadavky, počítala pozornost a spravovala mezipaměti KV rozshardované v rámci jejich lokální skupiny. Více instancí těchto paralelních skupin spolupracuje na vyhodnocování expertů, přičemž každá hostí jinou podmnožinu vah expertů. V závislosti na počtu hlav pozornosti může skupina obvykle pokrývat až velikost celého uzlu, protože komunikace uvnitř uzlu je rychlejší. Například na obrázku výše obsluhuje jedna úroveň DP nezávisle dva požadavky, A a B, a zajišťuje pozornost, normalizaci a veškerou další administrativu pro tyto požadavky. Druhá úroveň DP zpracovává odlišný požadavek C. První uzel však hostí polovinu expertů a druhý uzel druhou polovinu, takže po routování mohou být tokeny z A a B odeslány do druhého uzlu a naopak. To vede k problému řídké komunikace: každé zařízení může odeslat jiný počet tokenů do libovolného jiného cílového umístění. Existující primitiva z knihovny torch, především all_to_all, nejsou pro tento účel příliš vhodná, protože mohou vyžadovat určitou formu zarovnání (padding), synchronizaci GPU s CPU nebo všesměrové vysílání metadat. Pro efektivní implementaci komunikace jsou zapotřebí vlastní jádra, která zpracují informace o routování a iniciují komunikaci ze zařízení bez zásahu CPU. Po routování musí odesílací jádro (dispatch kernel) odeslat tokeny na ty úrovně, kam byly routovány, zatímco na straně kombinování je nutné sesbírat aktivace patřící k požadavkům v aktuální skupině DP. Práce musí být navíc vyvážena v rámci skupiny DP, která může mít více zařízení, aby se minimalizovalo vysílání a správně se replikovaly tokeny na všech úrovních pro následující vrstvy.

NVSHMEM

NVSHMEM je implementace OpenSHMEM specifická pro společnost NVIDIA, která poskytuje přenositelné prostředky mezizařízení komunikace abstrahující složitost podkladového hardwaru. Rozhraní API dokáže vyjádřit čtení a zápisy mezi zařízeními, které se mapují na primitiva jednotlivých přenosových vrstev. Podporováno je rozhraní NVLink i RDMA, což zajišťuje určitý stupeň přenositelnosti a umožňuje jádrům NVSHMEM fungovat v rámci jednoho uzlu nebo napříč clusterem tvořeným více uzly. V našich jádrech používáme operace vzdáleného zápisu do paměti a signalizace k přenosu dat a synchronizaci zúčastněných zařízení.

Operace NVSHMEM jsou postaveny na konceptu symetrické paměti: pracují s vyrovnávacími paměťmi, které byly alokovány na všech zařízeních účastnících se komunikace mezi zařízeními. Místní zařízení může použít adresu odvozenou v rámci vlastní lokální vyrovnávací paměti spolu s indexem cílové úrovně k určení cíle operace. Následující obrázek tento koncept ilustruje: obě GPU alokují symetrické vyrovnávací paměti stejné velikosti a uchovávají k nim ukazatele src. První úroveň chce odeslat 3 celá čísla do druhé úrovně a umístit je na začátek vyrovnávací paměti: funkce nvshmem_int_put_nbi odvozuje počáteční adresu z lokální vyrovnávací paměti a specifikuje cílové zařízení. Druhá úroveň odvozuje posun ze své vlastní vyrovnávací paměti a odesílá jeden prvek na první zařízení s posunem o jedna. Zatímco cílové adresy musí být vždy symetrické vyrovnávací paměti alokované pomocí nvshmem_alloc, zdrojové vyrovnávací paměti mohou být libovolnými oblastmi paměti zařízení za předpokladu, že jsou předem registrovány u NVSHMEM.

Přestože NVSHMEM poskytuje širokou škálu primitiv, naše jádra spoléhají na pouhé 3 funkce, na nichž stavějí veškerou synchronizaci a zajištění pořadí (fencing).

Zvětšit

nvshmemx_putmem_signal_nbi_warp: Přenáší blok dat z jednoho zařízení na druhé a zároveň nastavuje příznak na vzdáleném zařízení. Operace buď nastaví (NVSHMEM_SIGNAL_SET), nebo inkrementuje (NVSHMEM_SIGNAL_ADD) 64bitové umístění. Příznak se aktualizuje poté, co je přenesen celý blok paměti. Pokud vzdálené zařízení zaznamená změnu příznaku, může bezpečně přistupovat k vyrovnávací paměti ve vlastní paměti. Tato funkce je užitečná pro slučování přenosu dat a synchronizace.

nvshmemx_signal_op pracuje s jedním umístěním paměti, obvykle 64bitovým příznakem, který atomicky nastavuje nebo inkrementuje. Je užitečný při odesílání metadat a synchronizaci zařízení.

nvshmem_uint64_wait_until se používá na přijímací straně signálu k dotazování (pollování) příznaku, dokud jej vzdálené zařízení neaktualizuje.

Zajištění jakékoli formy uspořádání mezi operacemi prostřednictvím funkcí nvshmem_quiet nebo nvshmem_fence je nákladné, a proto naše implementace jader vyhýbá použití jakýchkoli jiných primitiv NVSHMEM a minimalizuje závislosti kromě jednotlivých párů odeslání-příjem. Při odesílání dat vždy používáme neblokující verzi funkcí, aniž bychom čekali, až budou data vůbec odeslána z lokální úrovně. Další implicitní synchronizace napříč skupinou jader zajišťuje, že tyto vyrovnávací paměti nejsou destruktivně přepsány, dokud síťové karty potřebují přístup k nim.

NVSHMEM využívá technologii GPUDirect RDMA, takže síťová karta (NIC) může přímo číst a zapisovat do paměti GPU za účelem výměny dat bez zapojení CPU. Na síťových kartách ConnectX navíc NVSHEMEM podporuje komunikaci iniciovanou GPU (známou též jako Infiniband GPUDirect Async neboli IBGDA), která umožňuje GPU odesílat operace RDMA přímo na síťovou kartu a přímo se dotazovat její fronty dokončení (completion queue). Na platformách, které nepodporují komunikaci iniciovanou GPU, vytváří NVSHMEM na hostitelském CPU „proxy“ vlákno, které iniciuje komunikaci a provádí dotazování na dokončení jménem GPU. Program NVSHMEM je přenositelný bez ohledu na to, zda je komunikace iniciovaná GPU podporována, či nikoliv. Komunikace iniciovaná GPU však výrazně snižuje latenci, protože zcela obchází obchůzku přes CPU.

Přenositelná jádra

Komunikaci MoE implementujeme prostřednictvím páru odesílacích (dispatch) a kombinujících (combine) jader. Odesílací jádra mají na starosti čtení tokenů a informací o routování na každé úrovni a jejich odesílání příslušným expertům. Kombinující jádra shromažďují aktivace produkované experty a odesílají je zpět na jejich zdrojové úrovně a zároveň vypočítávají vážený průměr z vybraných expertů na základě vah spočítaných routerem. Jádra jsou dále rozdělena na komponentu odesílání a příjmu, což umožňuje překrývat přenosy dat s výpočty. Odesílací jádra jsou zcela neblokující a nesynchronizační: jednoduše odešlou všechny zápisy na vzdálená zařízení. Na druhé straně kombinující jádra pouze čtou z paměti a čekají, dokud nejsou přenesena všechna potřebná data. Po odeslání práce na síťové karty mohou GPU – zatímco jsou data asynchronně přenášena po drátě – provádět jinou užitečnou práci lokálně, například aplikovat sdílené experty nebo počítat pozornost.

Každý pár jader má své vlastní symetrické paměťové vyrovnávací paměti, přičemž každá úroveň alokuje soukromé úložiště pro příjem dat ze všech ostatních úrovní. Maximální počet tokenů, které lze odeslat ze skupiny DP, je napříč úrovněmi fixní, což rovněž stanovuje horní mez, kterou může každá úroveň přijmout pro každého lokálního experta z každé úrovně DP. To odesílajícím úrovním umožňuje odvodit jedinečnou adresu na cílové úrovni, kam se má zapisovat, aniž by byla nutná jakákoliv synchronizace mezi nimi. Po přijetí dat jádra přemístí tokeny tak, aby je uspořádala v paměti do optimálního formátu vyžadovaného přijímajícími jádry. Ačkoliv mají vyrovnávací paměti značné rozměry, jsou opakovaně používány ve všech sekvenčních vrstvách modelu.

Zvětšit

Jediná forma globální synchronizace, jak je znázorněno na obrázku výše, je implementována v jádru combine-receive. Jakmile jsou data odeslána z odesílací části, úroveň se vrátí a pokračuje v provádění, jakmile obdrží všechny tokeny, které potřebuje od ostatních úrovní, a vstoupí do kombinujících odesílacích jader (combine send) ihned po jejich zpracování. Bariéra v combine-receive zajišťuje, že žádná úroveň nemůže předběhnout ostatní a spustit dispatch-send, zatímco jiná úroveň stále čeká na příjem dat, což by mohlo potenciálně způsobit destruktivní překrývání. Synchronizace se kvůli jednoduchosti provádí v kombinujícím jádru; jako optimalizace by stejná bariéra mohla být implementována mezi jádry dispatch-receive a combine-send, čímž by se dále snížila latence tím, že by se síťovým kartám poskytlo více času na usazení přenosů, zatímco GPU spouští jiná jádra.

Obě jádra jsou rozdělena mezi všechny dostupné streaming multiprocesory (SM) zařízení: zatímco jádra pro odesílání a kombinování (dispatch send a combine receive) musí paralelně pracovat s maximálním počtem tokenů na jednu úroveň (max_m), jádra pro příjem a odesílání kombinování (dispatch receive a combine send) musí zpracovat tento počet tokenů z každé úrovně a lokálního experta, a to až do maxima max_m * (num_experts // EP) * (EP // TP).

Odesílání (Dispatch)

Odesílací jádro zodpovídá za odesílání tokenů a faktorů škálování bloků partnerům (peers), kteří hostí experty, ke kterým byly tokeny routovány. Odesílací strana spoléhá na specializaci warpů (warp specialization) za účelem paralelezuace dvou úkolů: agregace informací o routování pro určení počtu tokenů odesílaných z každé úrovně a zabalení tokenů do zpráv určených k odeslání na vzdálené zařízení. Přijímací strana nejprve čeká na přijetí všech počtů tokenů, poté přečte a rozbalí užitečná data (payloads) ze sdílené paměti do tenzorů, které budou předány následným jádrům Grouped GEMM. Rovněž ukládá informace do vyrovnávacích pamětí sdílených s kombinujícími jádry, čímž indikuje, kam má být každý token odeslán zpět. Tyto informace jsou vyžadovány, protože přijímací jádro přemísťuje tokeny v souvislých vyrovnávacích pamětích v nedeterministickém pořadí.

Zvětšit

V odesílací části je jediný warp čítající 32 vláken v každém bloku odpovědný za čtení informací o routování a počítání počtu tokenů, které má každý cílový expert obdržet z této úrovně, přičemž odesílá tento počet zvětšený o jedna pomocí funkce nvshmem_signal_op. Počet se inkrementuje o jedna, protože přechod z nuly na nenulovou hodnotu na vzdálené straně signalizuje příjem počtů ze zdrojové úrovně. Souběžně zbývající warpy spolupracují na paralelním kopírování tokenů do symetrické paměti napříč všemi bloky, přičemž balí aktivace, faktory škálování a svůj index na lokální úrovni do souvislého bloku. Index je vyžadován kombinujícím odesílatelem k určení adresy, kam bude token zapsán. Dále, po zajištění toho, že všechna data byla zkopírována prostřednictvím bariéry napříč skupinami warpů, pracují warpy opět nezávisle a každý paralelně odesílá stejnou vyrovnávací paměť jinému expertovi. Tokeny se odesílají pomocí nvshmemx_putmem_signal_nbi_warp, což rovněž atomicky zvyšuje počet odeslaných tokenů z lokální úrovně na vzdáleném zařízení. V rámci skupiny DP je odesílání rovnoměrně vyvážené, protože každá úroveň vlastní repliku tokenu, který se má odeslat, přičemž každé zařízení odesílá podmnožinu tokenů.

Na přijímací straně se všechna jádra nejprve synchronizují se zdrojovými úrovněmi tím, že čekají na odeslání celkových počtů tokenů. Následně všechna čekají, až se atomicky inkrementované počty odeslaných tokenů ustálí na celkových počtech, což signalizuje, že díky sémantice volání putmem byla odeslána také všechna užitečná data ze zdrojových úrovní. Jádra provádějí dotazování na počty pomocí funkce nvshmem_uint64_wait_until, čímž operaci paralelezuijí napříč všemi bloky a vlákny. Následně mezibloková bariéra zajišťuje, že žádný blok nečte vyrovnávací paměti, pokud nebyla všechna data správně přijata. Tokeny rozprostřené mezi bloky a synchronizované prostřednictvím atomicky inkrementovaného čítače jsou zkopírovány ze symetrické paměti do běžných souvislých vyrovnávacích pamětí, které se stávají vstupy pro následná jádra implementující experty. Zdrojová úroveň, index experta a index tokenu jsou uloženy zvlášť, což přesně určuje umístění, kam musí kombinující jádro odeslat aktivace. Ačkoli jsou tokeny ze skupiny DP odesílány z různých zařízení, jsou všechny seskupeny dohromady, aby mohly být předány příslušnému expertovi.

Kombinování (Combine)

Kombinující jádra jsou opět rozdělena na poloviční komponenty odesílání a příjmu: odesílatelé kopírují nekvantované 16bitové aktivace na vzdálená zařízení, přičemž přijímací jádra čekají na odeslání dat a lokálně počítají vážený průměr všech příspěvků expertů. Navíc fungují také jako bariéra k synchronizaci sekvence dispatch-combine: každá úroveň nastaví příznak na každém partnerovi při vstupu do odesílacích jader, přičemž přijímacím jádrům není dovoleno vrátit se, dokud nezaznamenají, že byl příznak nastaven. Latence synchronizace je minimální, protože se překrývá s vlastní komunikací a výpočty.

Zvětšit

Na odesílací straně jádra paralelně procházejí seznam tokenů přiřazených všem lokálním expertům a zapisují je do vyrovnávací paměti na cílové úrovni. Cílová úroveň, index experta a index tokenu se čtou z vyrovnávacích pamětí pro jednotlivé tokeny, které byly naplněny jádry rozptylu (scatter) po přijetí tokenů. Každý odesílatel má svou vlastní vyhrazenou paměťovou oblast pro každého experta, do které zapisuje (jak je znázorněno na obrázku výše), což eliminuje potřebu synchronizace. Podobně jako při odesílání, kombinování atomicky inkrementuje čítače pro jednotlivé tokeny na cílovém zařízení, což značí příjem dat: jakmile čítač odpovídá počtu expertů, ke kterým byl token odeslán, lze k obsahu tokenu přistoupit.

V přijímacím jádru je seznam tokenů procházen paralelně napříč několika bloky a čeká se na příchod jejich obsahu dotazováním příznaku nastaveného signalizační operací. Po příjezdu jsou užitečná data přečtena ze soukromých vyrovnávacích pamětí, přičemž tabulka routování určuje, ze které vyrovnávací paměti se má číst a jaká váha se má přiřadit každému expertovi. Výsledky se pak zapisují do externě alokovaných tenzorů, přičemž jádro dokončí provádění, jakmile všechna zařízení projdou bariérou. V rámci skupiny DP obdrží všechny úrovně kopii každé aktivace experta pro výpočet vlastních replik.

Srovnávací testy (Benchmarks)

Naše jádra vyhodnocujeme na clusteru 128 GPU rozmístěných v 16 uzlech propojených pomocí rozhraní InfiniBand a síťových karet CX-7. Testujeme jak GPUDirect Async (IBGDA), tak spolehlivé připojení (Reliable Connection – RC) s proxy serverem CPU. Porovnáváme je s vysoce optimalizovanou implementací DeepSeek a také s hustými primitivy poskytovanými knihovnou PyTorch (prostřednictvím NCCL) nebo NVSHMEM.

Zvětšit

Ačkoli je u husté operace all-to-all v NVSHMEM výkon IBRC a IBGDA podobný (6378 µs vs 6180 µs), u řídkých jader je IBGDA výrazně rychlejší. Po sečtení fází Dispatch a Combine využívá IBGDA 902 µs, zatímco IBRC 3223 µs – což představuje 3,6krát lepší latenci. Zatímco operace all-to-all je limitována šířkou pásma, řídká jádra vysílají o řád méně dat a jsou limitována latencí. Přímým spuštěním síťových přenosů z GPU, aniž by bylo nutné k corespondenci mezi GPU a síťovou kartou využívat proxy CPU, se celková latence (end-to-end) výrazně snižuje.

Zatímco naše přenositelná jádra jsou přibližně 2krát pomalejší než vysoce optimalizovaná jádra DeepSeek, zlepšují latenci 10krát ve srovnání s hustými jádry (902 µs vs 9944 µs).

Navíc rozdělení na komponenty odesílatele a příjemce umožňuje skrýt část latence, na rozdíl od primitiv knihovny.

Zvětšit

V konfiguraci s jedním uzlem (EP8) využívá NVSHMEM pro přenos technologii NVLINK, což přináší nižší latenci a vyšší propustnost než síťování mezi uzly. Naše přenositelná jádra jsou na jednom uzlu přibližně 2,5krát rychlejší než DeepEP (186 µs vs 481 µs).

Další vylepšení

Zde popsaná jádra překonávají vestavěná primitiva rámců strojového učení a nabízejí slušný výkon bez nadměrné specializace na konkrétní přenosové mechanismy mezi zařízeními, jako jsou NVLink nebo InfiniBand. Kromě již zmíněných příležitostí lze dalšího zisku ve výkonu dosáhnout nahrazením komunikačních primitiv specifičtějšími verzemi. Například v rámci NVLink lze použití symetrické paměti nahradit vyrovnávacími paměťmi sdílenými napříč zařízeními, čímž se eliminují některé operace kopírování a umožní se jemnozrnnější synchronizační schémata pro jednotlivé tokeny namísto dávek tokenů. V rámci sítě InfiniBand by implementace mohla přistupovat přímo k podkladovým párům front, což by eliminovalo potřebu přenášet očekávané i odeslané počty tokenů a vyprázdňovat tokeny z front do výstupních vyrovnávacích pamětí tak, jak přicházejí. Takové implementace však jdou na úkor přenositelnosti, zatímco naše implementace umožňuje vyhodnotit clustery se specifickými konfiguracemi předtím, než se investuje vývojové úsilí do optimalizace pro konkrétní hardware.

Závěr

Představili jsme vysoce výkonnou a přenositelnou knihovnu pro komunikaci MoE, která dosahuje 10krát vyššího výkonu ve srovnání se standardní komunikací all-to-all při zachování kompatibility napříč různými hardwarovými konfiguracemi. V nasazeních na jednom uzlu s NVLink naše řešení demonstruje 2,5krát nižší latenci než předchozí implementace.

Náš přístup vyvažuje výkon a přenositelnost prostřednictvím klíčových inovací, mezi které patří podpora komunikace iniciované GPU, architektura rozděleného jádra umožňující překryv výpočtů a komunikace a efektivní odesílání tokenů s použitím minimálních primitiv NVSHMEM. Přestože je naše knihovna na víceuuzlových sestavách přibližně 2krát pomalejší než vysoce specializované implementace, nabízí vynikající flexibilitu napříč různými síťovými prostředími (NVLink, CX-7 a EFA).

S tím, jak se modely MoE nadále škálují, budou efektivní komunikační strategie, jako je tato, stále důležitější pro praktické nasazení. Naše plně open-source implementace je k dispozici na adrese https://github.com/ppl-ai/pplx-kernels.

Odkazy

DeepEP

Technická zpráva DeepSeek-V3

NVIDIA NVSHMEM

Citace

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