Ефикасна и преносива комуникација са мјешавином стручњака
Преглед преносиве комуникације засноване на мјешавини стручњака (MoE), са фокусом на оптимизацију паралелизма графичких процесора и смањење кашњења у великим моделима вјештачке интелигенције
Представљамо високих перформанси, преносиву, отворену библиотеку за комуникацију са мјешавином стручњака (MoE) која остварује 10 пута брже перформансе у поређењу са стандардним комуникацијским примитивима „све-са-свега“ (All-to-All).
Наша имплементација садржи неколико кључних техничких иновација које пружају врхунску ефикасност MoE комуникације:
Комуникација коју покреће графички процесор (IBGDA): Подржава директну комуникацију између графичког процесора и мрежне картице (GPU-to-NIC), значајно смањујући кашњење заобилажењем учешћа централног процесора
Преклапање комуникације и рачунања: Архитектура подељеног језгра са одвојеним фазама слања и пријема омогућава да се рачунање одвија истовремено са мрежним преносима
Најбрже перформансе на једном чвору: 2,5 пута мање кашњење него код претходно најбрже имплементације на конфигурацијама са једним чвором
Ефикасне и преносиве перформансе на више чворова: Наша имплементација постиже брзине до 10 пута веће од стандардне комуникације „све-са-свега“. Иако је отприлике 2 пута спорија од веома специјализованих имплементација, наш приступ нуди бољу преносивост на различитим верзијама NVSHMEM-а и мрежним окружењима (NVLink, CX-7 и EFA)
Библиотека је потпуно отвореног кода и доступна је на адреси https://github.com/ppl-ai/pplx-kernels.
У овом чланку истражујемо изазове паралелизма стручњака у великим MoE моделима и описујемо наш приступ ефикасном прослеђивању и комбиновању токена у дистрибуираним окружењима са графичким процесорима.
Увод
Модели засновани на мјешавини стручњака (MoE), као што су DeepSeek R1 и Mixtral 8x7B, унапређују густе моделе тако што ограничавају број тежина које се активирају за сваки токен. Умјесто постизања великог броја параметара повећањем броја слојева или повећањем величине линеарних слојева у вишеслојном перцептрону (MLP) сваког слоја декодера, MoE модели замјењују традиционални MLP са више стручњака и рутером. На примјер, од 671 милијарде параметара модела DeepSeek R1, само 37 милијарди се множи са датим токеном током инференције. Ово у великој мјери смањује количину прорачуна потребних за декодирање токена и побољшава кашњење у поређењу са густим моделама.
MoE модели представљају додатне изазове за инференцију у поређењу са густим моделима. Док су сами стручњаци мали MLP слојеви, сваки слој декодера садржи рутер који одређује на које стручњаке се токен прослеђује, при чему се сваки токен прослеђује више стручњака. Рутер је обично мали линеарни слој који производи расподјелу вјероватноће. Обично се бирају стручњаци са најбољих K резултата, а коначна активација се израчунава као пондерисани просјек, сабирањем излаза стручњака помножених тежином добијеном из расподјеле вјероватноће рутера.
Да би се минимизирало кашњење у дистрибуираним MoE системима, паралелизам се може искористити на више уређаја, али то уводи комуникацијске изазове. Модели попут Mixtral 8x7B или Llama-70B стају у 8 уређаја на једном чвору, користећи предности брзих NVLink међувеза са малим кашњењем (до 900 Gbps). Међутим, већи модели захтијевају дијељење (шардирање) стручњака на више чворова користећи InfiniBand (са вршном вриједношћу од 400 Gbps), што уводи додатне изазове у погледу кашњења.
У овом чланку истражујемо проблем паралелизма стручњака, описујући нашу имплементацију језгара графичког процесора која ефикасно прослеђују токене уређајима који држе додељене стручњаке и прикупљају их назад израчунавањем пондерисаног збира излаза стручњака. Док постоје агретивно оптимизоване имплементације усмјерене на специфични мрежни хардвер, ми представљамо високо преносиву имплементацију која се ослања на минимални скуп NVSHMEM примитива, а ипак испоручује перформансе 10 пута брже од стандардне комуникације „све-са-свега“. Наша имплементација постиже врхунске перформансе на конфигурацијама са једним чвором уз одржавање одличне преносивости у различитим мрежним окружењима.
Шардирање и паралелизам
За ефикасну инференцију, тежине модела се морају чувати у меморији уређаја, остављајући притом довољно простора за активације, KV кешове и друге бафере потребне за пролаз кроз модел напријед (forward pass). Најспособнији модели премашују капацитет чак и најспособнијих графичких процесора, па се инференција мора расподијелити на више уређаја који заједно могу да складиште тежине модела. На основу шема шардирања тежина, морају се користити различите комуникацијске и рачунске шеме како би се растеретио прорачун и синхронизовали уређаји.
Паралелизам стручњака (EP), приказан на првој слици, паралелизује само прорачун стручњака. Различити стручњаци су додијељени различитим уређајима који држе њихове тежине. Након рутирања, токени се шаљу одговарајућем уређају, а резултати се након тога прикупљају и акумулирају. Сложеност рутирања зависи од степена паралелизма у другим дијеловима модела: реплицирање других слојева могло би у потпуности елиминисати потребу за рутирањем, пошто сваки ранк може одабрати токене из локално реплициране табеле рутирања. Међутим, ако само један ранк извршава рутирање, потребно је емитовање (broadcast) да би се токени, индекси и тежине прослиједили одговарајућим стручњацима. Коначно, операција прикупљања свих података (all-gather) или емитовања свима (all-to-all) синхронизује излазне токене са оним ранком који наставља извршавање модела. Таква имплементација је релативно једноставна пошто PyTorch већ излаже потребне примитиве, мада се одређене предности могу остварити спајањем акумулације са операцијом прикупљања која прикупља токене од стручњака.
Паралелизам који укључује само стручњаке не скалира се идеално, јер чворови у кластеру могу бити неактивни док модел извршава слојеве који нису повезани са стручњацима, као што су пажња (attention), нормализација и узорковање. Међутим, прорачун ових слојева, примарно пажње, такође може имати користи од тензорског паралелизма (TP). Већина модела се ослања на пажњу са више глава (multi-head attention), што значи да главе пажње и њихове одговарајуће Q, K и V пројекције такође могу бити шардиране на уређајима, при чему се ијесеци између различитих слојева реплицирају или прикупљају. Ако је пажња распоређена на све уређаје, операција прикупљања свих података може синхронизовати активације, омогућавајући да се рутирање реплицира, што за ефикасну имплементацију захтијева примитиве синхронизације сличне онима код паралелизма стручњака. Међутим, постоје границе паралелизма на овом нивоу, пошто ће смањење броја глава пажње испод одређеног прага дати све мање повратне резултате.
Да би се на најбољи начин искористили сви уређаји и подржао веома висок степен паралелизма стручњака, до 128 или 256 графичких процесора, потребан је паралелизам података (DP). Према овој шеми, уређаји на које је модел подељен груписани су да би истовремено обрађивали захтјеве, израчунавајући пажњу и одржавајући KV кешове шардиране у оквиру своје локалне групе. Више инстанци ових паралелних група сарађује у евалуацији стручњака, при чему свака домаћини различити подскуп тежина стручњака. На основу броја глава пажње, група може обично да се простире до величине цијелог чвора, пошто је комуникација у оквиру чвора бржа. На примјер, на слици изнад, један DP ранк независно опслужује два захтјева, A и B, рукујући пажњом, нормализацијом и свим другим евидентирањем за захтјеве. Други DP ранк обрађује посебан захтјев C. Међутим, први чвор садржи половину стручњака, док други чвор садржи другу половину, тако да након рутирања токени из A и B могу бити послати на други чвор и обрнуто. Ово води до проблема ријетке комуникације: сваки уређај може послати различит број токена било којем другом одредишном ранку. Постојећи примитиви из библиотеке PyTorch, примарно all_to_all, нису посебно погодни, пошто могу захтијевати неки облик додавања подложака (padding) или синхронизацију између графичког и централног процесора и емитовање метаподатака. Да би се комуникација ефикасно имплементирала, потребна су прилагођена језгра за обраду информација о рутирању и покретање комуникације са уређаја без интервенције централног процесора. Након рутирања, језгро за прослеђивање мора послати токене на ранкове на које су рутирани, док се на страни комбиновања морају прикупити активације које припадају захтјевима у тренутној DP групи. Поред тога, рад се мора уравнотежити у оквиру DP групе која може имати више уређаја, како би се минимизирало емитовање и правилно реплицирали токени на свим ранковима за наредне слојеве.
NVSHMEM
NVSHMEM је NVIDIA-специфична имплементација стандарда OpenSHMEM, која пружа преносиве објекте за комуникацију између уређаја који апстрахују сложеност основног хардвера. АПИ може изразити читања и уписавања између уређаја, која се пресликавају на примитиве појединачних транспортних слојева. Подржани су и NVLink и RDMA, што омогућава одређени степен преносивости и дозвољава NVSHMEM језгрима да раде у оквиру једног чвора или широм кластера са више чворова. У нашим језгрима користимо удаљене меморијске уписе и операције слања сигназа како бисмо пренијели податке и синхронизовали уређаје који учествују.
NVSHMEM операције су изграђене око концепта симетричне меморије: оне раде на баферима који су алоцирани на свим уређајима који учествују у комуникацији између уређаја. Локални уређај може користити адресу изведену у оквиру сопственог локалног бафера заједно са индексом циљног ранка како би одредио одредиште операције. Слика испод илуструје овај концепт: оба графичка процесора алоцирају симетричне бафере исте величине, задржавајући src показиваче на њих. Први ранк жели да пошаље 3 цијела броја другом ранку, смјештајући их на почетак бафера: nvshmem_int_put_nbi изводи почетну адресу из локалног бафера, одређујући циљни уређај. Други ранк изводи помјерај из сопственог бафера, шаљући један елемент првом уређају, са помјерајем од један. Док одредишне адресе увијек морају бити симетрични бафери алоцирани помоћу nvshmem_alloc, изворни бафери могу бити произвољни региони меморије уређаја, под условом да су претходно регистровани у NVSHMEM-у.
Иако NVSHMEM пружа широк спектар примитива, наша језгра се ослањају на само 3 функције, градећи сву синхронизацију и ограђивање (fencing) на њима.
nvshmemx_putmem_signal_nbi_warp: Преноси блок података са једног уређаја на други, а такође подешава и флаг на удаљеном уређају. Операција или поставља (NVSHMEM_SIGNAL_SET) или увећава (NVSHMEM_SIGNAL_ADD) 64-битну локацију. Флаг се ажурира након што се цијели блок меморије пренесе. Ако удаљени уређај примијети промјену у флагу, може безбједно да приступи баферу у сопственој меморији. Ова функција је корисна за спајање преноса података и синхронизације.
nvshmemx_signal_op ради на једној меморијској локацији, обично 64-битном флагу, атомски га постављајући или увећавајући. Користан је за слање метаподатака и синхронизацију уређаја.
nvshmem_uint64_wait_until се користи на пријемној страни сигнала, за испитивање (poll) флага све док га удаљени уређај не ажурира.
Осигуравање било каквог облика поретка између операција кроз nvshmem_quiet или nvshmem_fence је скупо, стога наша имплементација језгара избјегава употребу било којих других NVSHMEM примитива и минимизира све зависности осим појединачних парова слања и пријема. Када шаљемо податке, увијек користимо неблокирајућу верзију функција, без чекања да подаци уопште буду послати ван локалног ранка. Друга, имплицитна синхронизација широм групе језгара осигурава да ови бафери не буду деструктивно преписани док су мрежним картицама потребни за приступ.
NVSHMEM користи GPUDirect RDMA тако да мрежна картица може директно да чита и уписује меморију графичког процесора ради размјене података без укључивања централних процесора. Штавише, на ConnectX мрежним картицама, NVSHMEM подржава комуникацију коју покреће графички процесор (такође позната као InfiniBand GPUDirect Async или IBGDA), што омогућава графичком процесору да поставља RDMA операције директно на мрежну картицу и да директно испитује ред чекања за завршетак мрежне картице. На платформама које не подржавају комуникацију коју покреће графички процесор, NVSHMEM покреће „прокси“ тред на хост централном процесору да иницира комуникацију и испитује завршетак у име графичког процесора. NVSHMEM програм је преносив без обзира на то да ли је комуникација коју покреће графички процесор подржана или не. Међутим, комуникација коју покреће графички процесор значајно смањује кашњење јер у потпуности заобилази заобилазни пут до централног процесора.
Преносива језгра
Имплементирамо MoE комуникацију кроз пар језгара за прослеђивање и комбиновање. Језгра за прослеђивање су одговорна за читање токена и информација о рутирању на сваком ранку, прослеђујући их одговарајућим стручњацима. Језгра за комбиновање прикупљају активације које су произвели стручњаци и шаљу их назад на њихове изворне ранкове, истовремено израчунавајући пондерисани просјек од одабраних стручњака на основу тежина које је израчунао рутер. Језгра су додатно подељена на компоненту слања и пријема, како би се омогућило да се преноси података преклапају са прорачунима. Језгра за слање су неблокирајућа и не синхронизују се: она једноставно прослеђују сва уписавања на удаљене уређаје. Са друге стране, језгра за комбиновање само читају из меморије, чекајући док се сви потребни подаци не пренесу. Након прослеђивања посла мрежним картицама, док се подаци преносе асинхроно преко мреже, графички процесори могу локално да обављају други корисни посао, као што је примјена дијељених стручњака или израчунавање пажње.
Сваки пар језгара има сопствене бафере симетричне меморије, при чему сваки ранк алоцира приватно складиште за пријем података од свих осталих ранкова. Максимални број токена који се може прослиједити из DP групе је фиксан на свим ранковима, што такође поставља горњу границу коју сваки ранк може примити за сваког локалног стручњака из сваког DP ранка. Ово омогућава ранк-пошиљаоцима да изведу јединствену адресу на одредишном ранку на коју ће писати, без потребе за било каквом синхронизацијом међу њима. Након што се подаци приме, језгра мијешају токене како би их распоредила у меморији у оптималном формату који захтијевају језгра за пријем. Иако бафери имају значајну димензију, они се поново користе кроз све секвенцијалне слојеве модела.
Једини облик глобалне синхронизације, као што је илустратовано на слици изнад, имплементиран је у језгру за комбиновање и пријем (combine-receive). Када се подаци проследе из слања, ранк се враћа и наставља извршавање када прими све токене који су му потребни од других ранкова и улази у језгра за комбиновање и слање чим се они обраде. Баријера у комбиновању и пријему осигурава да ниједан ранк не може да појури напријед и започне прослеђивање и слање док било који други ранк још увијек чека да прими податке, што би потенцијално изазвало деструктивно преклапање. Синхронизација се ради у језгру за комбиновање ради једноставности; као оптимизација, иста баријера би се могла имплементирати између језгара за пријем при прослеђивању и језгара за слање при комбиновању, додатно смањујући кашњење омогућавајући мрежним картицама више времена да заврше преносе док графички процесор извршава друга језгра.
Оба језгра су подјељена на све доступне мултипроцесоре (SM) уређаја: док језгра за прослеђивање и слање и комбиновање и пријем морају да се паралелизују према максималном броју токена по ранку (max_m), језгра за прослеђивање и пријем и комбиновање и слање морају да обраде толики број токена са сваког ранка и локалног стручњака, до максимално max_m * (num_experts // EP) * (EP // TP).
Прослеђивање
Језгро за прослеђивање је одговорно за прослеђивање токена и фактора скалирања блока вршњацима који су домаћини стручњацима којима су токени рутирани. Страна пошиљаоца се ослања на специјализацију снопа нити (warp specialization) да паралелизује два задатка: агрегирање информација о рутирању како би се утврдило колико се токена шаље са сваког ранка и паковање токена у поруке које треба прослиједити удаљеном уређају. Страна пријема прво чека да се прими број свих токена, затим чита и распакује корисне терете из заједничке меморије у тензоре који ће бити прослијеђени доњим груписаним GEMM језгрима. Она такође чува информације у баферима који се дијеле са језгрима за комбиновање, како би се назначило гдје сваки токен треба да се пошаље назад. Ова информација је потребна пошто језгро за пријем мијеша токене у контигуалним баферима у недетерминистичком редослиједу.
У дијелу за пошиљаоца, један сноп од 32 нити (warp) на сваком блоку је одговоран за читање информација о рутирању и бројање броја токена за које се очекује да сваки одредишни стручњак прими од овог ранка, шаљући број увећан за један помоћу nvshmem_signal_op. Бројач се увећава за један, пошто прелаз са нуле на ненулу на удаљеном крају означава пријем бројева из изворног ранка. Паралелно, преостали снопови сарађују да копирају токене у симетричну меморију паралелно преко свих блокова, пакујући активације, факторе скалирања и њихов индекс на локалном ранку у контигуални комад. Индекс је потребан пошиљаоцу за комбиновање да одреди адресу на коју ће токен бити уписан. Затим, након што се осигура да су сви подаци копирани кроз баријеру широм група снопова, снопови поново раде независно, при чему сваки шаље исти бафер другом стручњаку паралелно. Токени се шаљу помоћу nvshmemx_putmem_signal_nbi_warp, што такође атомски повећава број послатих токена са локалног ранка на удаљеном уређају. У оквиру DP групе, пошто сваки ранк посједује реплику токена који се шаље, прослеђивање је уравнотежено равномерно, при чему сваки уређај шаље подскуп токена.
На страни пријема, сва језгра се прво синхронизују са ранковима пошиљаоца чекајући да се укупан број токена пошаље. Након тога, сви они чекају да се атомски увећани бројеви послатих токена стабилизују на укупне бројеве, што указује да су сви корисни терети са изворних ранкова такође послати, захваљујући семантици позива putmem. Језгра врше испитивање бројача користећи nvshmem_uint64_wait_until, паралелизујући операцију на свим блоковима и нитима. Након тога, баријера кроз блокове осигурава да ниједан блок не чита бафере осим ако су сви подаци исправно примљени. Распоређени по блоковима и синхронизовани преко атомски увећаног бројача, токени се копирају из симетричне меморије у редовне контигуалне бафере који постају улази за доња језгра која имплементирају стручњаке. Изворни ранк, индекс стручњака и индекс токена се чувају одвојено, тачно лоцирајући мјесто гдје језгро за комбиновање треба да пошаље активације. Иако се токени из DP групе шаљу са различитих уређаја, они су сви груписани заједно да би се прослиједили одговарајућем стручњаку.
Комбиновање
Језгра за комбиновање су поново подељена на половну слања и пријема: пошиљаоци копирају неквантизоване 16-битне активације на удаљене уређаје, при чему језгра за пријем чекају да се подаци пошаљу и локално израчунавају пондерисани просјек свих доприноса стручњака. Поред тога, они такође дјелују као баријера за синхронизацију секвенце прослеђивања и комбиновања: сваки ранк поставља флаг на сваком вршњаку при уласку у језгра за слање, а језгрима за пријем није дозвољено да се врате осим ако не примијете да је флаг постављен. Кашњење синхронизације је минимално, јер се преклапа са стварном комуникацијом и прорачуном.
На страни пошиљаоца, језгра пролазе кроз листу токена додијељених свим локалним стручњацима паралелно, уписујући их у бафер на одредишном ранку. Циљни ранк, индекс стручњака и индекс токена читају се из бафера по токену које су попунила језгра за расипање (scatter) након пријема токена. Сваки пошиљалац има сопствени приватни меморијски регион по стручњаку у који пише, као што је назначено на слици изнад, избегавајући потребу за синхронизацијом. Слично прослеђивању, комбиновање атомски увећава бројаче по токену на одредишном ранку како би указало на пријем података: када бројач одговара броју стручњака којима је токен прослијеђен, садржају токена се може приступити.
У језгру за пријем, листа токена се пролази паралелно преко више блокова, чекајући да њихов садржај стигне испитивањем флага који је поставила операција сигнализирања. Након доласка, корисни терети се читају из приватних бафера, при чему табела рутирања показује из којег бафера треба читати и коју тежину додијелити сваком стручњаку. Резултати се затим уписују у споља алоциране тензоре, а језгро завршава извршавање када сви уређаји прођу баријеру. Широм DP групе, сви ранкови примају копију сваке активације стручњака како би израчунали сопствене реплике.
Мјерења перформанси
Евалуирамо наша језгра на кластеру од 128 графичких процесора распоређених на 16 чворова повезаних преко InfiniBand и CX-7 мрежних картица. Испробавамо и GPUDirect Async (IBGDA) и поуздану везу (RC) са проксијем за централни процесор. Упоређујемо их са високо оптимизованом имплементацијом модела DeepSeek, као и са густим примитивима које обезбеђује PyTorch (преко NCCL-а) или NVSHMEM.
Иако је код густе NVSHMEM операције „све-са-свега“ перформанса IBRC-a и IBGDA-е слична (6378 µs наспрам 6180 µs), IBGDA је значајно бржа са ријетким језгрима. Када се саберу прослеђивање и комбиновање, IBGDA користи 902 µs док IBRC захтијева 3223 µs – што је побољшање кашњења од 3,6 пута. Док је операција „све-са-свега“ ограничена пропусным опсегом, ријетка језгра емитују мање података за неколико редова величине, будући да су ограничена кашњењем. Покретањем мрежних преноса директно из графичког процесора, без потребе за проксијем централног процесора који би координисао графички процесор и мрежну картицу, крајње кашњење је значајно смањено.
Иако су наша преносива језгра око 2 пута спорија од високо оптимизованих језгара модела DeepSeek, она побољшавају кашњење 10 пута у поређењу са густим језгрима (902 µs наспрам 9944 µs).
Поред тога, подјела на компоненте пошиљаоца и пријема такође омогућава да се дио кашњења сакрије, за разлику од примитива библиотеке.
На једном чвору (EP8), NVSHMEM користи NVLINK за транспорт, испоручујући мање кашњење и већи проток од умрежавања између чворова. Наша преносива језгра су око 2,5 пута бржа од библиотеке DeepEP на једном чвору (186 µs наспрам 481 µs).
Даља побољшања
Језгра описана овдје надмашују уграђене примитиве оквира за вјештачку интелигенцију и нуде пристојне перформансе без претјеране специјализације за одређене транспорте између уређаја, као што су NVLink или InfiniBand. Поред већ поменутих могућности, даљи добици у перформансама могу се постићи замјеном комуникацијских примитива специјализованијим верзијама. На примјер, преко NVLink-а, употреба симетричне меморије могла би се замјенити баферима који се дијеле међу уређајима, елиминишући одређено копирање и омогућавајући финије шеме синхронизације, преко појединачних токена умјесто серија токена. Преко InfiniBand-а, имплементација би могла директно да приступа основним паровима редова чекања, елиминишући потребу за преносом и очекиваног и послатог броја токена и пражњењем токена из редова чекања у излазне бафере како стижу. Међутим, такве имплементације долазе по цијену преносивости, док наша имплементација омогућава да се кластери са специфичним конфигурацијама евалуирају прије улагања развојног напора у оптимизацију за одређени хардвер.
Закључак
Представили смо високих перформанси, преносиву библиотеку за MoE комуникацију која постиже 10 пута брже перформансе у поређењу са стандардном комуникацијом „све-са-свега“ уз одржавање компатибилности са различитим хардверским конфигурацијама. На распоређивањима на једном чвору са NVLink-ом, наше рјешење показује 2,5 пута мање кашњење од претходних имплементација.
Наш приступ балансира перформансе са преносивошћу кроз кључне иновације, укључујући подршку за комуникацију коју покреће графички процесор, архитектуру са раздвојеним језгрима која омогућава преклапање прорачуна и комуникације, и ефикасно прослеђивање токена користећи минималне NVSHMEM примитиве. Иако је отприлике 2 пута спорија од високо специјализованих имплементација на поставкама са више чворова, наша библиотека нуди врхунску флексибилност у различитим мрежним окружењима (NVLink, CX-7 и EFA).
Како MoE модели настављају да се скалирају, ефикасне комуникацијске стратегије попут наше постати све важније за практично постављање. Наша имплементација потпуно отвореног кода доступна је на адреси https://github.com/ppl-ai/pplx-kernels.
Референце
Технички извештај за DeepSeek-V3
Цитирање
@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}
}