Омогућавање модела са трилионима параметара на AWS EFA мрежи

Учинимо моделе са трилионима параметара доступним уз преносивост клауд платформи

АуториPerplexity Research

У компанији Perplexity користимо најбоље моделе за наш производ, наше АПИ-је и наше истраживачке тимове. Велики отворени Mixture-of-Experts модели, као што је Kimi-K2, постављају посебне изазове, пошто највећи инференсни чворови са 8x NVIDIA H200 ГПУ-ова не могу ефикасно да их сместе, што захтева постављање на више чворова. Предручујемо скуп кернела за паралелизам експерата који постижу кашњења на нивоу водећих на тржишту на мрежи ConnectX-7, превазилазећи перформансе библиотеке DeepEP. Исти кернели су такође први који постижу одржива кашњења на AWS Elastic Fabric Adapter (EFA) мрежи, омогућавајући имплементацију модела са трилионима параметара.

Испробајте наше кернеле на платформи GitHub и прочитајте комплетан истраживачки рад на платформи arXiv.

Увод

Mixture-of-Experts (MoE) се појавио као стандардна архитектура за скалирање на стотине милијарди или трилионе параметара, истовремено нудећи разумно кашњење у време инференце. Модели MoE замењују густи слој transformer модела скупом експерата и малим линеарним слојем за рутирање који одлучује са којим експерима треба помножити токен. Груписањем токена повезаних са истим експертом, проблем постаје редак и једноставно паралелизабилан, пошто различити експерти сада могу бити смештени на различитим ГПУ-овима распоређеним на више чворова.

За разлику од других облика паралелизма, као што су тензорски паралелизам (TP) и паралелизам података (DP), који се могу лако имплементирати помоћу библиотеке колективне комуникације као што су torch.distributed и NCCL, MoE рутирање укључује ретку peer-to-peer комуникацију. Кернел за слање мора да подели скуп долазних токена и да их пошаље ранговима који хостују експецте, док његов дуални проблем, кернел за комбиновање, мора да прими њихове обрађене пандане и израчуна њихов пондерисани просек. Ови преноси се не мапирају на колективне операције и захтевају прилагођене кернеле како би се постигло ниско кашњење.

У компанији Perplexity развијамо и преносиве међучворне кернеле и специјализоване унутарчворне кернеле, представљајући наш рад детаљније у претходном чланку о MoE комуникацији. Такође смо показали у другом блог посту да постављање на више чворова може донети мање кашњење и већи проток за велике MoE моделе.

У овом чланку дубље се бавимо имплементацијом наших нових међучворних кернела који значајно побољшавају перформансе нашег претходног рада заснованог на технологији NVSHMEM. Поред представљања врхунских резултата на мрежи ConnectX-7, такође демонстрирамо да се у оквиру нашег инференсног окружења кернели могу користити за ефикасно опслуживање највећих отворених модела.

Инференција преко InfiniBand мреже

Тензорски паралелизам и паралелизам експерата могу се ефикасно скалирати у оквиру једног чвора, пошто NVLink нуди висок проток (900GB/s) уз кашњење од неколико микросекунди. Док су NVIDIA GB200 и GB300 системи опремљени са довољно меморије да садрже тежине и KV кеш меморије модела са трилионима параметара, постављања на AWS p5en инстанцу са до 8 H200 ГПУ-ова ограничена су на 1120GB HBM меморије, која се мора поделити између тежина модела и KV кеш меморије. Сходно томе, одређени модели захтевају постављање на више чворова преко InfiniBand мреже.

Иако је InfiniBand и даље изузетно брз, он нуди проток од само 400Gbps (50GB/s), додајући десетине до стотине микросекунди кашњења за MoE рутирање. Срећом, додата кашњења се могу сакрити употребом дељених експерата, микро-бачевања и преклапања израчунавања и комуникације како би се испоручио одржив проток уз конкурентну цену. DeepSeek је био пионир ових техника путем DeepSeek-R1 модела, постављеног на кластере H800 ГПУ-ова и NVIDIA ConnectX-7 InfiniBand адаптера.

На AWS кластерима ситуација је другачија због употребе њихових прилагођених Elastic Fabric Adapters (EFA) адаптера. Иако они остварују вршни проток од 400Gbps на колективним операцијама које се обично користе у радним оптерећењима за обуку, они мало заостају за мрежом ConnectX-7 када су у питању величине порука које се размењују током MoE слања и комбиновања. Осим тога, EFA не подржава GPUDirect Async, што захтева CPU прокси тред који повезује ГПУ са мрежном картицом (NIC) како би се иницирали преноси. Поред додате сложености, додатне PCIe трансакције додају режијске трошкове вредне микросекундама свим преносима. Фрејворкови као што је NVSHMEM омогућавају ауторима кернела да обезбеде имплементације независне од уређаја, али генерички АПИ-ји стварају режијске трошкове који чине постављања неефикасним у погледу и цене и кашњења.

pplx-кернели засновани на NVSHMEM-у

Наши претходни кернели су апстраговали основну имплементацију преноса кроз NVSHMEM. Ослоњали су се на nvshmemx_putmem_signal_nbi_warp за пренос токена један по један, користећи атомске бројаче за синхронизацију парњака. Иако је ово нудило разумне перформансе на ConnectX-7 NIC картицама са IBGDA технологијом, IBRC имплементација на бази проксија за ConnectX-7 била је знатно спорија, док EFA имплементација није успела да испоручи кашњења испод једне милисекунде.

Користили смо кернеле као тест окружење за оперативни хардвер, истовремено утврђујући најприкладнији интерфејс за излагање нашим моделима. Претпоставили смо да чести преноси и интеракције између кернела и NIC картица или прокси тредова ометају перформансе. Поред тога, наша имплементација се није ослањала ни на какве гаранције хардвера осим поузданости, што сугерише да бисмо могли да убрзамо ствари смањењем режијских трошкова примитива синхронизације.

Након ових запажања, осмислили смо наше кернеле засноване на EFA мрежи који излажу исти интерфејс, оптимизујући интеракције између кернела уређаја и хост проксија "иза кулиса" ради врхунских перформанси.

MoE all-to-all преко EFA мреже

Хибридна комуникација RDMA и NVLink са proxy-јем домаћина, која приказује иницирајуће и примајуће rank-ове, proxy worker-е и NIC-ове.

Нови кернели прате хибридну CPU-GPU архитектуру, где ГПУ кернели интерагују са модела на уређају, док прокси тред који се извршава на хост процесору (CPU) управља интеракцијама са мрежном картицом. Нови кернели поново користе исти TransferEngine који смо изградили за пренос KV кеш меморије.

MoE рутирање се врши помоћу кернела за слање, који прима токене и експерте по токену за рутирање од модела, шаљући их преко мреже до рангова који хостују експецте. Након што груписани GEMM кернели помноже токене са тежинама експерата, кернел за комбиновање преноси токена на њихове изворне рангове и израчунава њихов пондерисани просек пре него што их врати моделу.

И кернел за слање и кернел за комбиновање подељени су на пошиљалачку и пријемну половину како би се олакшало преклапање и микро-бачевање. Пошиљалац меша токене у бафере за слање тако да се они могу уписати у парњаке помоћу једне RDMA операције уписа. Након извршавања пошиљаоца, ГПУ може да обавља друге задатке (као што је множење дељених експерата) док су преноси преко мреже на чекању у позадини. Пријемни кернели се блокирају све док се не приме сви преноси. У позадини, радни тред на страни хоста управља преносима: ГПУ му шаље сигнал када су бафери за слање попуњени, након чега он отпрема RDMA уписе. Када мрежна картица сигналом потврди завршетак, она деблокира пријемник, који враћа токене моделу.

Кернели су чврсто повезани са хостом, комуницирајући међусобно путем и унификоване меморије и GDRCopy технологије. CUDA унификована меморија омогућава да се меморија уређаја или хост меморија мапира у меморијске просторе оба ентитета, транспарентно управљајући комуникацијом путем PCIe везе.

Док је унификована меморија погодна за масовне преносе, анкетирање са ниским кашњењем се врши преко GDRCopy технологије. GDRCopy користи GPUDirect RDMA за читање из ГПУ меморије и уписивање у њу уз кашњење од микросекунди.

Поред RDMA, у оквиру чвора кернели такође користе NVLink за размену токена. Ово може да оффлоудује значајан део саобраћаја, до четвртине за EP32 и до осмине за EP64, постижући укупно мања кашњења.

Бафери за слање и пријем

Дизајн бафера за слање за кернел за слање је кључан за постизање одрживих кашњења која укључују хост прокси. Ако постоји N рангова од којих сваки шаље T токена, при чему је сваки рутиран до R експерата од укупно E, ранг може примити до N * T * max(R, E / N) токена у најгорем могућем случају, ако су сви рутирани ка њему. Штавише, ова минимална граница се може постићи само када пошиљаоци распореде своје уписе контигуирано, пошто би приватни бафери за пријем потрошили превише меморије. Да би пошиљаоци могли узастопно да распореде своје уписе без конфликата, потпуно знање о информацијама о рутирању мора бити дистрибуирано сваком рангу. Сходно томе, наши кернели за слање прво размењују број токена по експерту, након чега сваки од њих може да утврди где њихови уписи иду и са којег офсета примају токене од сваког ранга.

Једноставна размена информација о рутирању резултирала би непотребно високим кашњењима, пошто би сваки ранг морао да чека да уписи стигну пре него што пошаље сопствене уписе. Уместо тога, резервишемо мању количину приватног простора по пошиљаоцу на сваком пријемном рангу и шаљемо неке од токена заједно са уписима информација о рутирању. Ово обезбеђује пуну искоришћеност пропусног опсега током већег дела процеса слања. Након што се руте размене, преостали токени се преносе путем једног уписа по удаљеном парњаку. Размена информација о рутирању и бафера фиксне величине, заједно са преносима који их попуњавају, приказани су на слици испод.

Распоред бафера за MoE dispatch, који приказује како су подаци токена и routing метаподаци распоређени и размијењени између два rank-а.

Кернели за комбиновање су знатно једноставнији, пошто су им информације о рутирању већ доступне из фазе слања. Они могу да копирају токена једним уписом по парњаку, након чега читају из пријемног бафера како би израчунали пондерисани просек токена у излазним тензорима.

За NVLink, бафери за слање и пријем унутар чвора се мапирају у парњаке како би они могли директно да гурају или повлаче токена из суседних рангова.

Слање и комбиновање

Слике испод детаљно приказују време и интеракције између ГПУ-а, процесора и NIC картице, показујући једну итерацију слања и комбиновања у преплитању са кернелима модела.

Рад почиње у слању слања (dispatch send), при чему кернел агрегира информације о рутирању и израчунава бројеве токена по експерту путем атомске акумулације у дељеној меморији у оквиру једног блока. Када су бројеви токена доступни, они се копирају у хост меморију и прокси се одмах обавештава да почне са слањем информација о рутирању свим осталим ранговима парњака. Паралелно, слање наставља да израчунава кумулативни збир бројева токена, идентификујући офсете потребне за паковање токена у контигуиране бафере за слање по рангу. Када су офсети доступни, сви токени који се шаљу преко EFA мреже пакују се у бафере за слање. Прокси је обавештен да објави уписе у бафере за слање међучворних парњака. Унутар чвора, кернели такође одмах гурају токена у приватне бафере у оквиру истог чвора путем NVLink везе.

Тајминг и координација за MoE dispatch, који приказују фазне GPU, CPU и NIC операције и њихову синхронизацију.
Тајминг и координација за MoE combine, који приказују GPU, CPU и NIC операције и њихову синхронизацију.

Када ранг прими информације о рутирању од свих парњака, он их централизује и израчунава различите офсете неопходне за идентификацију токена за слање и комбиновање. Овај посао се обавља на процесору док су преноси већ на чекању, тако да није критичан за кашњење. Након тога, када су сви офсети познати, објављују се уписи за остатак токена, који премашују величину приватног бафера, и фаза слања је завршена.

На страни пријемника, кернел прво повлачи токена из истог чвора који премашују капацитет приватних бафера путем NVLink везе. Затим се блокира све док хост прокси не сигнализира завршетак свих преноса. Када се сигнал прими, токени се мешају из бафера за пријем у излазни тензор прослеђен групи GEMM. Пошто се бафери поново користе између слања и комбиновања, и NVLink и RDMA баријере обезбеђују да су бафери пријемника прочитани пре него што се дозволи следећем кернелу да их препише.

Кернели за комбиновање прате отприлике сличан образац: мешају токена, сигнализирају проксију и чекају завршетак. Међутим, метаподаци за објављивање њихових уписа су лако доступни и сви удаљени офсети су познати, па се тако сви токени пренети унутар чвора гурају преко NVLink везе до пријемника, док се међучворни токени преносе једним уписом. Пре враћања резултата, кернел акумулира токена из бафера пријемника, уписујући их у излазне тензоре.

TransferEngine за MoE рутирање

Да бисмо направили прототип наше имплементације, прво смо поново искористили Transfer Engine изграђен за дезагрегирани prefill и декодирање како бисмо издали и пратили RDMA уписе. Међутим, за разлику од дезагрегираног декодирања где сваки чвор објављује велики број уписа истом парњаку, MoE слање објављује мала слања свим парњацима у једној операцији. Иза кулиса, перформансе нису биле оптималне и наша имплементација је заостајала за DeepEP-ом на ConnectX-7 мрежи за око 20us. Сходно томе, специјализовали смо TransferEngine да подржава две операције, расипање (scatter) и баријеру (barrier), како бисмо оптимизовали интеракције са великим бројем парњака.

И операције расипања и баријере излажу групе парњака које се могу унапред регистровати да би се унапред попуниле структуре података које користи основна имплементација преноса (моделирање захтева за рад - WR templating). Расипање копира различите исечке из истог изворног бафера у различите офсете на свим члановима групе парњака. Баријера сигнализира свим парњацима путем непосредне поруке. Стављање у ред чекања и руковање овим операцијама у маси штеди вредне микросекунде на критичној путањи. На EFA мрежи, где се две 200Gbps NIC картице агрегирају како би обезбедиле пропусни опсег од 400Gbps, ми делимо групе парњака на две NIC картице уместо да делим бајтове сваког преноса.

Иако смо у почетку развијали TransferEngine искључиво на EFA мрежи користећи libfabric, MoE кернели раде на ConnectX-7 без промене ниједне линије кода након додавања подршке за ConnectX-7 у TransferEngine користећи libibverbs. У поређењу са EFA мрежом заснованом на Scalable Reliable Datagram (SRD) протоколу, ConnectX-7 захтева додатни рад за успостављање везе и управљање парњацима. Користимо пар редова чекања са непоузданим датаграмом (UD) да бисмо покренули поуздану везу (RC). Креирамо два пара редова чекања за сваког удаљеног парњака, један за двострани SEND/RECV и један за једнострани WRITE_IMM, зато што и RECV и WRITE_IMM троше RECV операцију редоследом, док две операције захтевају различите величине бафера. Осим тога, ублажавамо редослед PCIe трансакција између ConnectX-7 и ГПУ меморије. Слично као код EFA мреже, имплементирамо оптимизацију моделирања WR захтева. Поред тога, при слању више малих порука истом парњаку, користимо WR уланчавање да бисмо побољшали брзину порука.

Пошто је рад на страни ГПУ-а идентичан на оба адаптера, померили смо основну имплементацију преноса до њених граница како бисмо превазишли перформансе библиотеке DeepEP и истакли да је јаз између EFA и ConnectX-7 MoE рутирања одређен могућностима основног хардвера.

Евалуација

Кашњење слања/комбиновања

Извршили смо евалуацију кашњења слања и комбиновања на H200 системима опремљеним са ConnectX-7 и EFA адаптерима на стандардној DeepSeek-V3 и Kimi-K2 конфигурацији са скривеном димензијом од 7168, блок-скалираним fp8 слањем и bf16 комбиновањем. Разматрамо величину бача од 128 за слање и 4096 за комбиновање. На ConnectX-7 вршимо поређење са DeepEP-ом и нашим претходним pplx-кернелима заснованим на NVSHMEM-у, користећи и GPUDirect Async (IBGDA) и хост прокси (IBRC). Постижемо комбинована кашњења од 459us, 582us и 692us на EP16, EP32 и EP64 редом. Надигравамо друге кернеле засноване на EFA мрежи који се ослањају на генеричке проксије, као што је UCCL-EP, који постижу 519us, 966us и 1159us.

Dispatch и combine латенција за EP64 MoE decode, уз поређење имплементација и комуникационих путања приказаних овдје.

При декодирању, упркос томе што су наши кернели спорији за неколико микросекунди од DeepEP-а, постижемо укупно мања кашњења пошто су наши кернели за комбиновање знатно бржи. Иако користимо хост прокси, због наших масовних преноса превазилазимо перформансе свих имплементација без проксија заснованих на IBGDA. Такође показујемо да је имплементација проксија важна, пошто је наш високо специјализовани прокси ред величина бржи од генеричке NVSHMEM имплементације. Упркос чињеници да је EFA двоструко спорија од ConnectX-7 мреже при преносу пакета од 256KB, који су близу типичне величине пакета која се размењује MoE рутирањем, јаз између ова два адаптера је у целини мање значајан.

Dispatch и combine латенција за EP64 MoE prefill, уз поређење pplx и DeepEP комуникационих путања.

Иако наши кернели оптимизовани за декодирање могу да се скалирају на величине бача за prefill без икаквих додатних подешавања, перформансе заостају за DeepEP-ом. Ово је првенствено због чињенице да акумулирамо токена само на страни пријемника комбиновања, док DeepEP акумулира код пошиљаоца како би смањио количину пренетих података. Међутим, ова оптимизација може имати утицај на тачност. Осим тога, пошто кернели преносе податке у једном упису, користимо више меморије за бафере слања и пријема.

Проток за Kimi-K2 / DeepSeek-V3

Да бисмо докаали одрживост наших кернела, вршимо бенчмарковање перформанси перформанси од краја до краја нашег интерног инференсног енџина који опслужује DeepSeek-V3 (671 милијарда параметара) и Kimi-K2 (1 трилион параметара) на више p5en инстанци опремљених са Elastic Fabric Adapters (EFA) и H200 ГПУ-овима. Покрећемо модел у режиму верификације нацрта и пријављујемо кашњење претпостављајући стопу прихватања од 80% користећи један MTP слој.

На модела DeepSeek-V3, вршимо поређење са EP=8 унутарчворном основицом која покреће ефикасне кернеле за слање и комбиновање искључиво преко NVLink везе. Постављања на више чворова у целини или поклапају перформансе једночворног постављања или их превазилазе, посебно на бачевима средње величине, што показује да скалирање преко чворова доноси побољшања укупног протока.

На модела Kimi-K2, кернели омогућавају опслуживање. Због своје величине, модел не може да стане у исти чвор користећи паралелизам података (DP) на H200 чвору, тако да је за опслуживање модела потребно постављање на више чворова. Како модел има мање глава пажње (attention heads), постижемо боља кашњења него код модела DeepSeek-V3 и можемо да опслужујемо модел на EFA мрежи уз одрживо кашњење.

Проток по GPU-у за DeepSeek-V3 у односу на брзину излаза за конфигурације EP32, EP16 и EP8.
Проток по GPU-у за Kimi-K2 у односу на брзину излаза за конфигурације EP32 и EP16.

Такође смо разматрали микро-бачевање, мада је оно побољшало кашњења само при веома великим величинама бача или када се користи више од једног токена нацрта. Да би се имало користи од дељења проблема са груписаним GEMM-ом на два дела, величине бачева токена морају бити довољно густе. Главни део времена се и даље троши на MoE слање и комбиновање, што указује на то да ће свака даља побољшања кернела последично побољшати кашњења и олакшати микро-бачевање.

Будући рад

Блиско сарађујемо са AWS инжењерима како бисмо побољшали перформансе на EFA мрежи. Са наше стране, пратићемо ажурирања за libfabric која ће смањити режијске трошкове равни података. Поред тога, експериментисаћемо са efa-direct технологијом како бисмо додатно смањили режијске трошкове у корисничком простору и директно интегрисали са хардвером.

Референце

Технички извештај за DeepSeek R1

Технички извештај за Kimi K2

DeepEP

UCCL-EP

Побољшање мрежних перформанси HPC система коришћењем NVIDIA Magnum IO NVSHMEM и GPUDirect Async технологија

Цитирање

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