Aktivácia modelov s biliónmi parametrov na AWS EFA

V spoločnosti Perplexity používame tie najlepšie modely pre náš produkt, naše rozhrania API a naše výskumné tímy. Veľké open-source modely typu Mixture-of-Experts, ako je napríklad Kimi-K2, predstavujú jedinečné výzvy, pretože najväčšie uzly inferencie s 8× GPU NVIDIA H200 ich nedokážu efektívne pojmout, čo si vyžaduje nasadenie vo viacerých uzloch. Predstavujeme súbor jadier pre paralelizm odborníkov (expert parallelism), ktoré dosahujú špičkovú latenciu na rozhraní ConnectX-7 a prekonávajú výkon DeepEP. Tieto isté jadrá zároveň ako prvé dosahujú životaschopnú latenciu na rozhraní AWS Elastic Fabric Adapter (EFA), čím umožňujú nasadenie modelov s biliónmi parametrov.

Vyskúšajte naše jadrá na GitHub a prečítajte si kompletný výskumný článok na arXiv.

Úvod

Architektúra Mixture-of-Experts (MoE) sa stala štandardnou architektúrou na škálovanie na stovky miliárd až bilióny parametrov a zároveň ponúka primeranú latenciu v čase inferencie. Modely MoE nahrádzajú hustú vrstvu transformátora súborom expertov a malou lineárnou smerovacou vrstvou, ktorá rozhoduje, s ktorým expertom sa má vynásobiť token. Zoskupením tokenov asociovaných s rovnakým expertom sa problém stáva riedkym a triviálne paralelizovateľným, pretože rôzni experti teraz môžu byť umiestnení na rôznych GPU rozložených vo viacerých uzloch.

Na rozdiel od iných foriem paralelizmu, akými sú tenzorový paralelizmus (Tensor Parallelism, TP) a dátový paralelizmus (Data Parallelism, DP), ktoré možno ľahko implementovať pomocou knižníc kolektívnej komunikácie ako torch.distributed a NCCL, smerovanie MoE zahŕňa riedku peer-to-peer komunikáciu. Smerovacie jadro (dispatch kernel) musí rozdeliť množinu prichádzajúcich tokenov a odoslať ich do hodností (ranks) hostujúcich expertov, zatiaľ čo jeho duálny problém, kombinovacie jadro (combine kernel), musí prijať ich spracované náprotivky a vypočítať ich vážený priemer. Tieto prenosy sa nemapujú na kolektívne operácie a vyžadujú prispôsobené jadrá na dosiahnutie nízkej latencie.

V spoločnosti Perplexity sme vyvíjali prenosné medzizuzlové jadrá aj špecializované vnútrozuzlové jadrá a našu prácu podrobnejšie prezentujeme v predchádzajúcom článku o komunikácii MoE. V ďalšom blogovom príspevku sme tiež ukázali, že viaczuzlové nasadenie môže priniesť nižšiu latenciu a vyššiu priepustnosť pre veľké modely MoE.

V tomto článku sa hlboko ponoríme do implementácie našich nových medzizuzlových jadier, ktoré výrazne vylepšujú výkon našej predchádzajúcej práce založenej na NVSHMEM. Okrem predstavenia špičkových výsledkov na rozhraní ConnectX-7 tiež ukazujeme, že v rámci nášho inferenčného runtime možno jadrá efektívne použiť na obsluhu najväčších open-source modelov.

Inferencia cez InfiniBand

Tenzorový paralelizmus a paralelizmus expertov sa môžu efektívne škálovať v rámci jedného uzla, pretože NVLink ponúka vysokú priepustnosť (900 GB/s) s latenciou v ráde mikrosekúnd. Zatiaľ čo systémy NVIDIA GB200 a GB300 sú vybavené dostatočnou pamäťou na uchovanie váh a vyrovnávacích pamätí KV (KV caches) modelov s biliónmi parametrov, nasadenie na inštancii AWS p5en s až 8 GPU H200 je obmedzené na 1120 GB HBM, ktoré sa musí rozdeliť medzi váhy modelu a vyrovnávacie pamäte KV. V dôsledku toho si určité modely vyžadujú medzizuzlové nasadenie cez InfiniBand.

Hoci je InfiniBand stále extrémne rýchly, ponúka priepustnosť iba 400 Gbps (50 GB/s), čím pridáva desiatky až stovky mikrosekúnd latencie pre smerovanie MoE. Našťastie možno pridané latencie skryť použitím zdieľaných expertov, mikro-dávkovania (micro-batching) a prekrývania výpočtov a komunikácie, aby sa zabezpečila životaschopná priepustnosť za konkurencieschopnú cenu. DeepSeek bol pionierom v týchto technikách prostredníctvom DeepSeek-R1, nasadeného na clusteroch GPU H800 a adaptérov NVIDIA ConnectX-7 InfiniBand.

Na clusteroch AWS je situácia iná kvôli používaniu ich vlastných adaptérov Elastic Fabric Adapters (EFA). Hoci tieto poskytujú špičkovú priepustnosť 400 Gbps pri kolektívnych operáciách typicky používaných pri tréningových záťažiach, v porovnaní s ConnectX-7 mierne zaostávajú pri veľkostiach správ vymieňaných počas smerovania a kombinovania MoE. EFA navyše nepodporuje GPUDirect Async, čo si vyžaduje proxy vlákno CPU premostené medzi GPU a sieťovou kartou (NIC) na iniciovanie prenosov. Okrem pridanej zložitosť pridávajú dodatočné transakcie PCIe réžiu v hodnote mikrosekúnd ku všetkým prenosom. Rámce ako NVSHMEM umožňujú tvorcom jadier poskytovať implementácie nezávislé od zariadenia, avšak generické rozhrania API spôsobujú režijné náklady, v dôsledku ktorých je nasadenie neefektívne z hľadiska nákladov aj latencie.

pplx-kernels založené na NVSHMEM

Naše predchádzajúce jadrá abstrahovali základnú implementáciu prenosu prostredníctvom NVSHMEM. Spoliehali sa na nvshmemx_putmem_signal_nbi_warp na prenos tokenov po jednom, pričom na synchronizáciu rovesníkov (peers) používali atomické počítadlá. Zatiaľ čo to ponúkalo primeraný výkon na sieťových kartách ConnectX-7 s IBGDA, implementácia ConnectX-7 IBRC založená na proxy bola výrazne pomalšia a implementácia EFA nedokázala dodať latencie pod jednu milisekundu.

Jadrá sme použili ako testovacie prostredie pre základný hardvér a zároveň sme určili najvhodnejšie rozhranie na sprístupnenie našim modelom. Predpokladali sme, že časté prenosy a interakcie medzi jadrami a sieťovými kartami alebo proxy vláknami znižujú výkon. Naša implementácia sa navyše spoliehala iba na spoľahlivosť a žiadne iné záruky od hardvéru, čo naznačuje, že by sme veci mohli urýchliť znížením réžie synchronizačných primitív.

Na základe týchto pozorovaní sme vytvorili naše jadrá založené na EFA, ktoré vystavujú rovnaké rozhranie a pod kapotou optimalizujú interakcie medzi jadrom zariadenia a hostiteľským proxy pre dosiahnutie špičkového výkonu.

MoE all-to-all cez EFA

Zväčšiť

Nové jadrá sa riadia hybridnou architektúrou CPU-GPU, kde jadrá GPU interagujú s modelom na zariadení a proxy vlákno bežiace na hostiteľskom CPU spracováva interakcie so sieťovou kartou. Nové jadrá opätovne využívajú rovnaký TransferEngine, ktorý sme vytvorili pre prenosy vyrovnávacej pamäte KV.

Smerovanie MoE vykonáva smerovacie jadro (dispatch kernel), ktoré od modelu prijíma tokeny a expertov pre každý token na ich smerovanie, pričom ich posiela cez sieť do hodností, ktoré hostia expertov. Po tom, čo zoskupené jadrá GEMM vynásobia tokeny váhami expertov, kombinovacie jadro prenesie tokeny do ich zdrojových hodností a vypočíta ich vážený priemer pred ich vrátením modelu.

Smerovacie aj kombinovacie jadro sú rozdelené na polovicu odosielateľa a prijímača s cieľom uľahčiť prekrývanie a mikro-dávkovanie. Odosielateľ usporiada tokeny do odosielacích vyrovnávacích pamätí, aby sa dali zapísať do rovesníckych uzlov prostredníctvom jedinej operácie zápisu RDMA. Po spustená odosielateľa môže GPU vykonávať inú prácu (napríklad násobenie zdieľaných expertov), zatiaľ čo prenosy cez sieť čakajú na pozadí. Jadrá prijímača sú zablokované, kým sa neprijmú všetky prenosy. Na pozadí spracováva prenosy pracovné vlákno na strane hostiteľa: GPU ho upozorní, keď sú odosielacie vyrovnávacie pamäte naplnené, načo odošle zápisy RDMA. Akonáhle sieťová karta signalizuje dokončenie, odblokuje prijímač, ktorý vráti tokeny modelu.

Jadrá sú úzko spojené s hostiteľom a komunikujú medzi sebou prostredníctvom zjednotenej pamäte (unified memory) aj GDRCopy. Zjednotená pamäť CUDA umožňuje, aby bola pamäť zariadenia alebo hostiteľská pamäť mapovaná do adresných priestorov oboch, pričom transparentne riadi komunikáciu cez PCIe.

Zatiaľ čo zjednotená pamäť je vhodná pre hromadné prenosy, dotazovanie s nízkou latenciou sa vykonáva prostredníctvom GDRCopy. GDRCopy využíva GPUDirect RDMA na čítanie a zápis do pamäte GPU s mikrosekundovou latenciou.

Okrem RDMA jadrá v rámci uzla využívajú aj NVLink na výmenu tokenov. To môže odľahčiť významnú časť prevádzky, až štvrtinu pre EP32 a až osminu pre EP64, čím sa dosiahnu celkovo nižšie latencie.

Odosielacie a prijímacie vyrovnávacie pamäte

Návrh odosielacích vyrovnávacích pamätí pre smerovacie jadro je kľúčový pre dosiahnutie životaschopných latencií zahŕňajúcich hostiteľský proxy. Ak existuje N hodností, z ktorých každá odosiela T tokenov smerovaných na R expertov z E, hodnosť môže v najhoršom možnom prípade prijať až N * T * max(R, E / N) tokenov, ak sú všetky smerované na ňu. Okrem toho túto minimálnu hranicu možno dosiahnuť iba vtedy, keď odosielatelia rozložia svoje zápisy súvisle, pretože súkromné prijímacie vyrovnávacie pamäte by spotrebovali príliš veľa pamäte. Aby odosielatelia mohli usporiadať svoje zápisy za sebou bez konfliktov, úplné znalosti o smerovacích informáciách musia byť distribuované do každej hodnosti. V dôsledku toho naše smerovacie jadrá najprv vymenia počty tokenov pre jednotlivých expertov, po čom môžu každá určiť, kam smerujú ich zápisy a z ktorej hodnosti prijímajú tokeny s akým posunom.

Jednoduchá výmena smerovacích informácií by mala za následok zbytočne vysokú latenciu, pretože každá hodnosť by musela čakať na príchod zápisov pred odoslaním akýchkoľvek vlastných zápisov. Namiesto toho si na každej hodnosti prijímača vyhradíme malé množstvo súkromného priestoru pre každého odosielateľa a odošleme niektoré tokeny spolu so zápismi smerovacích informácií. To zaisťuje úplné využitie šírky pásma počas väčšiny procesu smerovania. Po výmene trás sa zostávajúce tokeny prenesú prostredníctvom jediného zápisu na vzdialeného rovesníka. Výmena smerovacích informácií a vyrovnávacích pamätí s pevnou veľkosťou spolu s prenosmi, ktoré ich napĺňajú, sú znázornené na obrázku nižšie.

Zväčšiť

Kombinovacie jadrá sú podstatne jednoduchšie, pretože smerovacie informácie sú im už k dispozícii z fázy smerovania. Tokeny môžu skopírovať v jedinom zápise na rovesníka, po čom čítajú z prijímacej vyrovnávacej pamäte, aby vypočítali vážený priemer tokenov do výstupných tenzorov.

V prípade NVLink sú odosielacie a prijímacie vyrovnávacie pamäte v rámci uzla mapované do rovesníkov, aby mohli priamo presúvať alebo vyťahovať tokeny susedných hodností.

Smerovanie a kombinovanie (Dispatch and Combine)

Obrázky nižšie podrobne popisujú načasovanie a interakcie medzi GPU, CPU a sieťovou kartou a ukazujú jednu iteráciu smerovania a kombinovania preloženú s jadrami modelu.

Práca sa začína v odosielaní smerovania (dispatch send), pričom jadro agreguje smerovacie informácie a počíta počty tokenov pre jednotlivých expertov prostredníctvom atomickej kumulácie v zdieľanej pamäti v rámci jedného bloku. Akonáhle sú počty tokenov k dispozícii, skopírujú sa do hostiteľskej pamäte a proxy je okamžite upozornený, aby začal odosielať smerovacie informácie všetkým ostatným rovesníckym hodnostiam. Paralelne smerovanie pokračuje v počítaní kumulatívneho súčtu počtov tokenov, pričom identifikuje posuny potrebné na zabalenie tokenov do súvislých odosielacích vyrovnávacích pamätí pre jednotlivé hodnosti. Akonáhle sú posuny k dispozícií, všetky tokeny, ktoré sa majú odoslať cez EFA, sú zabalené do odosielacích vyrovnávacích pamätí. Proxy je upozornený, aby odoslal zápisy do odosielacích vyrovnávacích pamätí medzizuzlových rovesníkov. V rámci uzla jadrá taktiež okamžite tlačia tokeny do súkromných vyrovnávacích pamätí v rámci toho istého uzla cez NVLink.

Zväčšiť

Zväčšiť

Akonáhle hodnosť prijme smerovacie informácie od všetkých rovesníkov, centralizuje ich a vypočíta rôzne posuny potrebné na identifikáciu tokenov na smerovanie a kombinovanie. Táto práca sa vykonáva na CPU, zatiaľ čo prenosy už čakajú na vybavenie, takže nie je kritická na latenciu. Následne, akonáhle sú známe všetky posuny, zápisy pre zvyšok tokenov, ktoré presahujú veľkosť súkromnej vyrovnávacej pamäte, sa odošlú a odosielacia fáza smerovania je dokončená.

Na strane prijímača jadro najskôr vyťahuje tokeny z toho istého uzla, ktoré sú za kapacitou súkromných vyrovnávacích pamätí, cez NVLink. Ďalej sa zablokuje, kým hostiteľský proxy nesignalizuje dokončenie všetkých prenosov. Po prijatí signálu sa tokeny preskupia z prijímacích vyrovnávacích pamätí do výstupného tenzora odovzdaného do skupinového GEMM. Keďže sa vyrovnávacie pamäte opätovne používajú medzi smerovaním a kombinovaním, bariéry NVLink aj RDMA zaisťujú, že prijímacie vyrovnávacie pamäte boli prečítané predtým, ako sa umožní ďalšiemu jadru ich prepísať.

Kombinovacie jadrá sa riadia zhruba podobným vzorom: preskupujú tokeny, signalizujú proxy a čakajú na dokončenie. Metadáta na odosielanie ich zápisov sú však hneď k dispozícii a všetky vzdialené posuny sú známe, takže všetky tokeny prenesené v rámci uzla sú zatlačené cez NVLink k prijímačom a medzizuzlové tokeny sú prenesené prostredníctvom jediného zápisu. Pred vrátením výsledkov jadro kumuluje tokeny z prijímacích vyrovnávacích pamätí a zapisuje ich do výstupných tenzorov.

TransferEngine pre smerovanie MoE

Na vytvorenie prototypu našej implementácie sme pôvodne opätovne využili Transfer Engine postavený na deみなcovaných (disaggregate) predbežných vyplneniach (prefill) a dekódovaní s cieľom vydávať a sledovať zápisy RDMA. Avšak na rozdiel od disagregovaného dekódovania, kde každý uzol odošle veľký počet zápisov rovnakému rovesníkovi, smerovanie MoE odošle malé prenosy všetkým rovesníkom v jednej operácii. Výkon nebol optimálny a naša implementácia zaostávala za DeepEP na ConnectX-7 o približne 20 us. V dôsledku toho sme špecializovali TransferEngine tak, aby podporoval dve operácie, rozptyl (scatter) a bariéru (barrier), s cieľom optimalizovať interakcie s veľkým počtom rovesníkov.

Operácie rozptylu aj bariéry sprístupňujú skupiny rovesníkov, ktoré je možné vopred zaregistrovať na predbežné naplnenie štruktúr údajov používaných základnou implementáciou prenosu (šablónovanie pracovných požiadaviek Work Request - WR templating). Rozptyl skopíruje rôzne časti z toho istého zdrojového vyrovnávacieho priestoru do rôznych posunov na všetkých členoch rovesníckej skupiny. Bariéra signalizuje všetkým rovesníkom prostredníctvom bezprostrednej správy. Zaradenie do frontu a spracovanie týchto operácií hromadne šetrí cenné mikrosekundy na kritickej ceste. Na EFA, kde sa dve 200Gbps sieťové karty agregujú, aby poskytli šírku pásma 400Gbps, rozdeľujeme (shard) rovesníkov skupiny naprieč oboma sieťovými kartami namiesto rozdeľovania bajtov každého prenosu.

Hoci sme pôvodne vyvinuli TransferEngine výhradne na EFA s použitím libfabric, jadrá MoE fungujú na ConnectX-7 bez zmeny jediného riadku kódu po pridaní podpory ConnectX-7 v TransferEngine pomocou libibverbs. V porovnaní s EFA založenou na Scalable Reliable Datagram (SRD) vyžaduje ConnectX-7 prácu navyše na nastavenie pripojenia a správu rovesníkov. Na spustenie spoľahlivého pripojenia (Reliable Connection - RC) používame pár frontov nespätného datagramu (Unreliable Datagram - UD). Vytvárame dva páry frontov pre každého vzdialeného rovesníka, jeden pre obojsmerný SEND/RECV a jeden pre jednosmerný WRITE_IMM, pretože RECV aj WRITE_IMM spotrebovávajú operáciu RECV v poradí, zatiaľ čo tieto dve operácie vyžadujú rôzne veľkosti vyrovnávacej pamäte. Okrem toho uvoľňujeme usporiadanie transakcií PCIe medzi ConnectX-7 a pamäťou GPU. Podobne ako pri EFA implementujeme optimalizáciu šablónovania WR. Navyše pri odosielaní viacerých malých správ rovnakému rovesníkovi používame reťazenie WR na zlepšenie rýchlosti správ.

Keďže práca na strane GPU je na oboch adaptéroch identická, posunuli sme implementáciu základného prenosu na jej limity, abychom prekonali výkon DeepEP a zdôraznili, že medzera medzi smerovaním MoE EFA a ConnectX-7 je určená schopnosťami základného hardvéru.

Vyhodnotenie

Latencia smerovania / kombinovania (Dispatch/Combine Latency)

Vyhodnotili sme latencie smerovania a kombinovania na systémoch H200 vybavených adaptérami ConnectX-7 a EFA na štandardnej konfigurácii DeepSeek-V3 a Kimi-K2 so skrytou dimenziou 7168, blokovo škálovaným fp8 smerovaním a bf16 kombinovaním. Zvažujeme veľkosť dávky 128 pre smerovanie a 4096 pre kombinovanie. Na ConnectX-7 porovnávame s DeepEP a našimi predchádzajúcimi pplx-kernels založenými na NVSHMEM pomocou GPUDirect Async (IBGDA) aj hostiteľského proxy (IBRC). Dosahujeme kombinované latencie 459 us, 582 us a 692 us na EP16, EP32 a EP64. Prekonávamos iné jadrá založené na EFA, ktoré sa spoliehajú na generické proxy, ako je UCCL-EP, ktoré dosahujú 519 us, 966 us a 1159 us.

Zväčšiť

Pri dekódovaní, napriek tomu, že naše jadrá sú o pár mikrosekúnd pomalšie ako DeepEP, dosahujeme celkovo nižšiu latenciu, pretože naše kombinovacie jadrá sú podstatne rýchlejšie. Hoci používame hostiteľský proxy, vďaka našim hromadným prenosom prekonávame výkon všetkých implementácií bez proxy založených na IBGDA. Tiež ukazujeme, že na implementácii proxy záleží, pretože naše vysoko špecializované proxy je o rádovo rýchlejšie ako generická implementácia NVSHMEM. Napriek skutočnosti, že EFA je dvakrát pomalšia ako ConnectX-7 pri prenose 256 kB paketov, ktoré sa blížia k typickej veľkosti paketov vymieňaných smerovaním MoE, medzera medzi týmito dvoma adaptérami je celkovo menej podstatná.

Zväčšiť

Aj keď sa naše jadrá optimalizované na dekódovanie môžu škálovať na veľkosti predbežných dávok (prefill batch sizes) bez akýchkoľvek dodatočných úprav, výkon zaostáva za DeepEP. Je to primárne spôsobené tým, že tokeny kumulujeme iba na strane prijímača kombinovania, zatiaľ čo DeepEP kumuluje v odosielateľovi, čím sa znižuje množstvo prenesených údajov. Táto optimalizácia však môže mať vplyv na presnosť. Navyše, keďže jadrá prenášajú údaje v jednom zápise, využívame viac pamäte pre odosielacie a prijímacie vyrovnávacie pamäte.

Priepustnosť Kimi-K2 / DeepSeek-V3

Na preukázanie životaschopnosti našich jadier testujeme komplexný (end-to-end) výkon nášho vlastného inferenčného enginu obsluhujúceho DeepSeek-V3 (671B parametrov) a Kimi-K2 (1T parametrov) na viacerých inštanciách p5en vybavených adaptérmi Elastic Fabric Adapters (EFA) a GPU H200. Spúšťame model v režime overovania návrhu (draft verification mode) a uvádzame latenciam za predpokladu 80 % miery akceptácie s použitím jednej vrstvy MTP.

V prípade DeepSeek-V3 porovnávame s vnútrozuzlovou základnou líniou EP=8, ktorá spúšťa efektívne jadrá smerovania a kombinovania iba s NVLink. Viaczuzlové nasadenie celkovo buď zodpovedá, alebo prekračuje výkon jednóuzlového nasadenia, obzvlášť pri stredne veľkých dávkach, čo dokazuje, že škálovanie naprieč uzlami prináša zlepšenie celkovej priepustnosti.

Pri Kimi-K2 jadrá umožňujú obsluhu. Kvôli svojej veľkosti sa model nezmestí do rovnakého uzla pomocou dátového paralelizmu (DP) na uzle H200, takže na obsluhu modelu sa vyžaduje medzizuzlové nasadenie. Keďže model má menej hláv pozornosti (attention heads), dosahujeme lepšiu latenciu ako DeepSeek-V3 a môžeme model obsluhovať na EFA s životaschopnou latenciou.

Zväčšiť

Zväčšiť

Zvážili sme aj mikro-dávkovanie, hoci zlepšilo latenciu iba pri veľmi veľkých dávkach alebo pri použití viac ako jedného tokenu návrhu. Aby sa využilo rozdelenie problémov zoskupeného GEMM na dva, dávky tokenov musia byť dostatočne husté. Väčšina času sa stále míňa na smerovanie a kombinovanie MoE, čo naznačuje, že akékoľvek ďalšie vylepšenia jadier následne zlepšia latenciu a uľahčia mikro-dávkovanie.

Budúca práca

Úzko spolupracujeme s inžiniermi AWS na zlepšení výkonu na EFA. Z našej strany budeme sledovať aktualizácie libfabric, ktoré znížia réžiu dátovej roviny (data plane overhead). Okrem toho budeme experimentovať s efa-direct s cieľom ďalej znížiť réžiu používateľského priestoru (userspace overheads) a interagovať priamo s hardvérom.

Referencie

Technická správa DeepSeek R1

Technická správa Kimi K2

DeepEP

UCCL-EP

Zlepšenie sieťového výkonu systémov HPC pomocou NVIDIA Magnum IO NVSHMEM a GPUDirect Async

Citácia

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