Zprovoznění modelů s biliony parametrů na AWS EFA

V Perplexity používáme ty nejlepší modely pro náš produkt, naše API a naše výzkumné týmy. Velké open-source modely typu Mixture-of-Experts, jako je Kimi-K2, představují specifické výzvy, protože největší inferenční uzly s 8x GPU NVIDIA H200 je nedokážou efektivně pojmout, což si vyžaduje nasazení ve více uzlech. Představujeme sadu jader pro paralelismus expertů, která dosahují špičkových latencí na ConnectX-7 a překonávají výkon DeepEP. Stejná jádra také jako první dosahují životaschopných latencí na AWS Elastic Fabric Adapter (EFA), což umožňuje nasazení modelů s biliony parametrů.

Vyzkoušejte naše jádra na GitHub a přečtěte si kompletní výzkumnou práci na arXiv.

Úvod

Mixture-of-Experts (MoE) se objevila jako standardní architektura pro škálování na stovky miliard nebo biliony parametrů a zároveň nabízí rozumnou latenci v době inference. Modely MoE nahrazují hustou vrstvu transformátoru sadou expertů a malou lineární vrstvou směrování, která určuje, s kterým expertem se má token vynásobit. Seskupením tokenů spojených se stejným expertem se problém stává řídkým a triviálně paralelizovatelným, protože různé experty lze nyní hostovat na různých GPU rozložených v rámci několika uzlů.

Na rozdíl od jiných forem paralelismu, jako je tenzorový paralelismus (TP) a datový paralelismus (DP), které lze snadno implementovat pomocí knihoven kolektivní komunikace jako torch.distributed a NCCL, zahrnuje směrování MoE řídkou peer-to-peer komunikaci. Jádro pro odesílání (dispatch kernel) musí rozdělit sadu příchozích tokenů a odeslat je do řad hostujících experty, zatímco jeho duální problém, jádro pro kombinování (combine kernel), musí přijmout jejich zpracované protějšky a vypočítat jejich vážený průměr. Tyto přenosy se nemapují na kolektivní operace a vyžadují přizpůsobená jádra k dosažení nízké latence.

V Perplexity jsme vyvíjeli jak přenositelná mezimandantní (inter-node) jádra, tak specializovaná vnitřní (intra-node) jádra, a naši práci podrobněji představujeme v předchozím článku o komunikaci MoE. V dalším příspěvku na blogu jsme také ukázali, že nasazení ve více uzlech může přinést nižší latenci a vyšší propustnost pro velké modely MoE.

V tomto článku se hlouběji ponoříme do implementace našich nových mezimandantních jader, která výrazně vylepšují výkon naší předchozí práce založené na NVSHMEM. Kromě předvedení špičkových výsledků na ConnectX-7 také ukazujeme, že v rámci našeho inferenčního běhového prostředí lze jádra efektivně využít k obsluze největších open-source modelů.

Inference přes InfiniBand

Tenzorový paralelismus a paralelismus expertů se mohou efektivně škálovat v rámci jednoho uzlu, protože NVLink nabízí vysokou propustnost (900 GB/s) s latencí v řádu mikrosekund. Zatímco systémy NVIDIA GB200 a GB300 jsou vybaveny dostatečnou pamětí pro uložení vah a KV mezipamětí modelů s biliony parametrů, nasazení na instanci AWS p5en s až 8 GPU H200 je omezeno na 1120 GB HBM, která musí být rozdělena mezi váhy modelu a KV mezipaměti. Některé modely si proto vyžadují mezimandantní nasazení přes InfiniBand.

Ačkoli je InfiniBand stále neuvěřitelně rychlý, nabízí propustnost pouze 400 Gb/s (50 GB/s), což přidává desítky až stovky mikrosekund latence pro směrování MoE. Naštěstí lze přidané latence skrýt pomocí sdílených expertů, mikro-dávkování a překrývání výpočtů a komunikace, což zajišťuje životaschopnou propustnost za konkurenceschopnou cenu. Společnost DeepSeek byla v těchto technikách pionýrem prostřednictvím modelu DeepSeek-R1 nasazeného na clusterech GPU H800 a adaptérů NVIDIA ConnectX-7 InfiniBand.

Na clusterech AWS je situace jiná kvůli používání jejich vlastních adaptérů Elastic Fabric Adapters (EFA). Přestože tyto adaptéry poskytují špičkovou propustnost 400 Gb/s u kolektivních operací běžně používaných při tréninkových úlohách, mírně zaostávají za ConnectX-7 u velikostí zpráv vyměňovaných během odesílání a kombinování MoE. EFA navíc nepodporuje GPUDirect Async, což vyžaduje, aby proxy vlákno procesoru propojovalo GPU s NIC za účelem zahájení přenosů. Kromě zvýšené složitosti přidávají další transakce PCIe režijní náklady v řádu mikrosekund ke všem přenosům. Rámce jako NVSHMEM umožňují tvůrcům jader poskytovat implementace nezávislé na zařízení, obecná API však nesou režijní náklady, které činí nasazení neefektivními z hlediska nákladů i latence.

pplx-kernels založené na NVSHMEM

Naše předchozí jádra abstrahovala podkladovou implementaci přenosu pomocí NVSHMEM. Spoléhala na nvshmemx_putmem_signal_nbi_warp k přenosu tokenů jeden po druhém a k synchronizaci peerů využívala atomické čítače. Zatímco to nabízelo přiměřený výkon na síťových kartách ConnectX-7 s IBGDA, implementace ConnectX-7 IBRC založená na proxy byla výrazně pomalejší a implementace EFA nedokázala zajistit latenci pod jednu milisekundu.

Jádra jsme použili jako testovací prostředí pro základní hardware a zároveň jsme určili nejvhodnější rozhraní pro vystavení našim modelům. Vyslovili jsme hypotézu, že časté přenosy a interakce mezi jádry a síťovými kartami nebo proxy vlákny brzdí výkon. Naše implementace se navíc nespoléhala na žádné jiné záruky ze strany hardwaru než na spolehlivost, což naznačovalo, že bychom mohli věci zrychlit snížením režie synchronizačních primitiv.

Na základě těchto pozorování jsme navrhli naše jádra založená na EFA, která odhalují stejné rozhraní a optimalizují interakce mezi jádrem zařízení a hostitelským proxy pod kapotou pro špičkový výkon.

MoE all-to-all přes EFA

Zvětšit

Nová jádra se řídí hybridní architekturou CPU-GPU, kde jádra GPU interagují s modelem na zařízení a proxy vlákno spuštěné na hostitelském CPU zajišťuje interakce se síťovou kartou. Nová jádra znovu využívají stejný TransferEngine, který jsme sestavili pro přenosy KV mezipaměti.

Směrování MoE provádí odesílající jádro (dispatch kernel), které od modelu přijímá tokeny a experty pro jednotlivé tokeny, aby je mohlo směrovat, a odesílá je přes síť do hodností (ranks), které hostují experty. Poté, co seskupená jádra GEMM vynásobí tokeny s vahami expertů, kombinující jádro přenese tokeny zpět do jejich zdrojových řad a spočítá jejich vážený průměr, než je vrátí modelu.

Odesílací i kombinující jádra jsou rozdělena na odesílací a přijímací polovinu, což usnadňuje překrývání a mikro-dávkování. Odesílatel přeuspořádá tokeny do odesílacích vyrovnávacích pamětí, takže je lze zapsat k peerům prostřednictvím jediné operace zápisu RDMA. Po spuštění odesílatele může GPU provádět jinou práci (například násobení sdílených expertů), zatímco přenosy po síti čekají na pozadí. Přijímací jádra se zablokují, dokud nebudou přijaty všechny přenosy. Na pozadí zpracovává přenosy pracovní vlákno na straně hostitele: GPU ho upozorní, když jsou odesílací vyrovnávací paměti zaplněny,čehož využije k zahájení zápisů RDMA. Jakmile síťová karta signalizuje dokončení, odblokuje přijímač, který vrátí tokeny modelu.

Jádra jsou úzce provázena s hostitelem a komunikují spolu prostřednictvím sjednocené paměti (unified memory) i GDRCopy. Sjednocená paměť CUDA umožňuje mapovat paměť zařízení nebo hostitele do adresních prostorů obou, přičemž transparentně zvládá komunikaci přes PCIe.

Zatímco sjednocená paměť je vhodná pro hromadné přenosy, dotazování s nízkou latencí (low-latency polling) se provádí prostřednictvím GDRCopy. GDRCopy využívá GPUDirect RDMA pro čtení a zápis do paměti GPU s mikrosekundovou latencí.

Kromě RDMA využívají jádra v rámci jednoho uzlu také NVLink k výměně tokenů. To může odlehčit značnou část provozu, až čtvrtinu pro EP32 a až osminu pro EP64, čímž se dosáhne celkově nižších latencí.

Odesílací a přijímací vyrovnávací paměti

Návrh odesílacích vyrovnávacích pamětí pro odesílací jádro je klíčový pro dosažení životaschopných latencí zahrnujících hostitelský proxy. Pokud existuje N řad (ranks), z nichž každá odesílá T tokenů, přičemž každý je směrován na R expertů z E, může řada v nejhorším možném případě přijmout až N * T * max(R, E / N) tokenů, pokud jsou všechny směrovány na ni. Navíc lze této minimální hranice dosáhnout pouze tehdy, když odesílatelé rozvrhnou své zápisy souvisle, protože soukromé přijímací vyrovnávací paměti by spotřebovaly příliš mnoho paměti. Aby odesílatelé mohli rozvrhnout své zápisy po sobě bez konfliktů, musí být úplné znalosti o informacích o směrování distribuovány do každé řady. Naše odesílací jádra proto nejprve vymění počty tokenů na experta, na jejichž základě pak každé z nich může určit, kam směřují jejich zápisy a z jakého offsetu přijímají tokeny od jednotlivých řad.

Pouhá výměna informací o směrování by vedla k zbytečně vysokým latencím, protože každá řada by musela čekat na dorazení zápisů před odesláním jakýchkoli vlastních zápisů. Místo toho si na každé přijímací řadě vyhrazujeme malé množství soukromého prostoru pro odesílatele a odešleme některé tokeny spolu se zápisy informací o směrování. To zajišťuje plné využití šířky pásma po většinu procesu odesílání. Po výměně tras jsou zbývající tokeny přeneseny prostřednictvím jediného zápisu na vzdáleného peera. Výměna informací o směrování a vyrovnávacích pamětí pevné velikosti spolu s přenosy, které je plní, jsou znázorněny na obrázku níže.

Zvětšit

Kombinující jádra jsou podstatně jednodušší, protože informace o směrování jsou jim k dispozici již z fáze odesílání. Mohou zkopírovat tokeny v jediném zápisu na peera, načež čtou z přijímací vyrovnávací paměti, aby vypočítaly vážený průměr tokenů do výstupních tenzorů.

V případě NVLink jsou odesílací a přijímací vyrovnávací paměti v rámci uzlu mapovány do peerů, aby mohly přímo tlačit nebo tahat tokeny ze sousedních řad.

Odesílání a kombinování

Následující obrázky podrobně popisují časování a interakce mezi GPU, CPU a síťovou kartou (NIC) a ukazují jednu iteraci odesílání a kombinování proloženou s jádry modelu.

Práce začíná v odesílání (dispatch send), přičemž jádro agreguje informace o směrování a počítá počty tokenů na experta prostřednictvím atomické akumulace ve sdílené paměti v rámci jednoho bloku. Jakmile jsou počty tokenů k dispozici, jsou zkopírovány do hostitelské paměti a proxy je okamžitě upozorněn, aby zahájil odesílání informací o směrování všem ostatním řadám peerů. Paralelně odesílání pokračuje výpočtem kumulativního součtu počtu tokenů, přičemž identifikuje offsety potřebné k zabalení tokenů do souvislých odesílacích vyrovnávacích pamětí pro jednotlivé peery. Jakmile jsou offsety k dispozici, všechny tokeny, které mají být odeslány přes EFA, jsou zabaleny do odesílacích vyrovnávacích pamětí. Proxy je upozorněn, aby zveřejnil zápisy do odesílacích vyrovnávacích pamětí mezimandantních peerů. V rámci jednoho uzlu jádra také okamžitě tlačí tokeny do soukromých vyrovnávacích pamětí ve stejném uzlu prostřednictvím NVLink.

Zvětšit

Zvětšit

Jakmile řada obdrží informace o směrování od všech peerů, zcentralizuje je a vypočítá různé offsety nezbytné k identifikaci tokenů pro odeslání a kombinování. Tato práce se provádí na CPU, zatímco přenosy již čekají, takže není kritická na latenci. Jakmile jsou známé všechny offsety, jsou zveřejněny zápisy pro zbývající část tokenů, které přesahují velikost soukromé vyrovnávací paměti, a odesílací fáze odesílání je dokončena.

Na přijímací straně jádro nejprve vytáhne ze stejného uzlu tokeny, které jsou nad rámec kapacity soukromých vyrovnávacích pamětí, prostřednictvím NVLink. Dále se zablokuje, dokud hostitelský proxy nesignalizuje dokončení všech přenosů. Jakmile je signál přijat, tokeny se přesunou z přijímacích vyrovnávacích pamětí do výstupního tenzoru předaného skupinovému GEMM. Vzhledem k tomu, že vyrovnávací paměti se znovu používají mezi odesíláním a kombinováním, bariéry NVLink i RDMA zajišťují, že přijímací vyrovnávací paměti byly přečteny, než bude umožněno dalšímu jádru je přepsat.

Kombinující jádra se řídí zhruba podobným vzorcem: přeuspořádávají tokeny, signalizují proxy a čekají na dokončení. Metadata pro zveřejnění jejich zápisů jsou však okamžitě k dispozici a všechny vzdálené offsety jsou známé, takže všechny tokeny přenesené v rámci uzlu jsou tlačeny přes NVLink k přijímačům a mezimandantní tokeny jsou přenášeny prostřednictvím jediného zápisu. Před vrácením výsledků jádro akumuluje tokeny z přijímacích vyrovnávacích pamětí a zapisuje je do výstupních tenzorů.

TransferEngine pro směrování MoE

K vytvoření prototypu naší implementace jsme zpočátku znovu použili Transfer Engine vytvořený pro disagregované předběžné plnění (prefill) a dekódování za účelem vydávání a sledování zápisů RDMA. Na rozdíl od disagregovaného dekódování, kde každý uzel publikuje velký počet zápisů ke stejnému peeru, však odesílání MoE publikuje malé přenosy ke všem peerům v jedné operaci. Výkon nebyl optimální hned po vybalení z krabice a naše implementace zaostávala za DeepEP na ConnectX-7 přibližně o 20 us. Proto jsme TransferEngine specializovali tak, aby podporoval dvě operace: scatter (rozptyl) a barrier (bariéra), abychom optimalizovali interakce s velkým počtem peerů.

Operace scatter i barrier odhalují skupiny peerů, které lze předem zaregistrovat k předběžnému naplnění datových struktur používaných podkladovou implementací přenosu (šablonování pracovních požadavků - WR templating). Scatter zkopíruje různé výřezy ze stejného zdrojového vyrovnávacího paměti do různých offsetů na všech členech skupiny peerů. Bariéra signalizuje všechny peery prostřednictvím immediate hodnoty. Zařazení a zpracování těchto operací ve velkém šetří cenné mikrosekundy na kritické cestě. Na EFA, kde se dvě 200 Gb/s síťové karty agregují a poskytují šířku pásma 400 Gb/s, rozdělujeme peery skupiny mezi dvě síťové karty namísto rozdělování bajtů každého přenosu.

Ačkoli jsme TransferEngine zpočátku vyvíjeli výhradně na EFA pomocí libfabric, jádra MoE fungují na ConnectX-7 bez změny jediného řádku kódu poté, co jsme přidali podporu ConnectX-7 do TransferEngine pomocí libibverbs. Ve srovnání s EFA založenou na Scalable Reliable Datagram (SRD) vyžaduje ConnectX-7 další práci na nastavení připojení a správě peerů. K vytvoření spolehlivého připojení (RC) používáme frontu nespolehlivých datagramů (UD queue pair). Pro každého vzdáleného peera vytvoříme dvě párové fronty: jednu pro obousměrný SEND/RECV a jednu pro jednosměrný WRITE_IMM, protože jak RECV, tak WRITE_IMM spotřebovávají operaci RECV postupně, zatímco obě operace vyžadují různé velikosti vyrovnávací paměti. Kromě toho uvolňujeme řazení transakcí PCIe mezi ConnectX-7 a pamětí GPU. Podobně jako u EFA implementujeme optimalizaci šablonování WR. Navíc při odesílání více malých zpráv stejnému peeru používáme zřetězení WR ke zlepšení frekvence zpráv.

Jelikož je práce na straně GPU identická u obou adaptérů, posunuli jsme podkladovou implementaci přenosu na její limity, abychom překonali výkon DeepEP a zdůraznili, že mezera mezi směrováním EFA a ConnectX-7 MoE je určena schopnostmi podkladového hardwaru.

Vyhodnocení

Latence odesílání/kombinování

Vyhodnotili jsme latence odesílání a kombinování na systémech H200 vybavených adaptéry ConnectX-7 a EFA na standardní konfiguraci DeepSeek-V3 a Kimi-K2 s隱ovaným (hidden) rozměrem 7168, blokově škálovaným fp8 odesíláním a bf16 kombinováním. Uvažujeme velikost dávky 128 pro odesílání a 4096 pro kombinování. Na ConnectX-7 porovnáváme s DeepEP a našimi předchozími pplx-kernels založenými na NVSHMEM za použití jak GPUDirect Async (IBGDA), tak hostitelského proxy (IBRC). Dosahujeme kombinovaných latencí 459 us, 582 us a 692 us na EP16, EP32 a EP64. Překonáváme ostatní jádra založená na EFA, která spoléhají na obecné proxy, jako je UCCL-EP, jež dosahují 519 us, 966 us a 1159 us.

Zvětšit

Při dekódování, přestože jsou naše jádra o několik mikrosekund pomalejší než DeepEP, dosahujeme celkově nižších latencí, protože naše kombinující jádra jsou podstatně rychlejší. Ačkoli používáme hostitelský proxy, díky našim hromadným přenosům překonáváme výkon všech implementací bez proxy založených na IBGDA. Také ukazujeme, že na implementaci proxy záleží, protože naše vysoce specializovaná proxy je o řád rychlejší než obecná implementace NVSHMEM. Nehledě na to, že EFA je dvakrát pomalejší než ConnectX-7 při přenosu 256KB paketů, které se blízko shodují s typickou velikostí paketů vyměňovaných při směrování MoE, je mezera mezi oběma adaptéry celkově méně podstatná.

Zvětšit

I když se naše jádra optimalizovaná pro dekódování mohou škálovat na velikosti dávek předběžného plnění (prefill) bez jakýchkoli dalších úprav, výkon zaostává za DeepEP. To je způsobeno především tím, že akumulujeme tokeny pouze na přijímací straně kombinování, zatímco DeepEP akumuluje v odesílateli, aby se snížilo množství přenášených dat. Tato optimalizace však může mít dopad na přesnost. Navíc, jelikož jádra přenášejí data v jednom zápisu, využíváme více paměti pro odesílací a přijímací vyrovnávací paměti.

Propustnost Kimi-K2 / DeepSeek-V3

Abychom dokázali životaschopnost našich jader, provádíme benchmark end-to-end výkonu našeho vlastního inferenčního enginu obsluhujícího DeepSeek-V3 (671B parametrů) a Kimi-K2 (1T parametrů) na více instancích p5en vybavených adaptéry Elastic Fabric Adapters (EFA) a GPU H200. Spouštíme model v režimu ověřování konceptu (draft verification mode) a uvádíme latenci za předpokladu 80% míry přijetí za použití jedné vrstvy MTP.

U DeepSeek-V3 porovnáváme s vnitrouzlovým (intranode) srovnávacím základem EP=8, který spouští efektivní odesílací a kombinující jádra pouze s NVLink. Mezimandantní nasazení buď odpovídá, nebo překonává výkon jednouzlového nasazení, zejména u středně velkých dávek, což ukazuje, že škálování napříč uzly přináší vylepšení celkové propustnosti.

U Kimi-K2 tato jádra umožňují obsluhu. Vzhledem ke své velikosti se model nevejde do stejného uzlu za použití datového paralelismu (DP) na uzlu H200, takže k obsluze modelu je nutné mezimandantní nasazení. Jelikož má model méně hlaviček pozornosti (attention heads), dosahujeme lepších latencí než u DeepSeek-V3 a můžeme model obsluhovat na EFA při životaschopné latenci.

Zvětšit

Zvětšit

Zvažovali jsme také mikro-dávkování, ačkoliv zlepšilo latence pouze při velmi vysokých velikostech dávek nebo při použití více než jednoho tokenu návrhu. Aby bylo možné těžit z rozdělení problémů se seskupeným GEMM na dva, musí být velikosti dávek tokenů dostatečně husté. Většina času se stále tráví odesíláním a kombinováním MoE, což naznačuje, že jakákoli další vylepšení jader následně zlepší latence a usnadní mikro-dávkování.

Budoucí práce

Úzce spolupracujeme s inženýry AWS na zlepšení výkonu na EFA. Z naší strany budeme sledovat aktualizace libfabric, které sníží režii datové roviny. Budeme také experimentovat s efa-direct, abychom dále snížili režii v uživatelském prostoru a komunikovali přímo s hardwarem.

Odkazy

Technická zpráva DeepSeek R1

Technická zpráva Kimi K2

DeepEP

UCCL-EP

Zlepšení síťového výkonu systémů HPC pomocí NVIDIA Magnum IO NVSHMEM a GPUDirect Async

Citace

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