Внедряване на модели с трилиони параметри на AWS EFA

В Perplexity използваме най-добрите модели за нашия продукт, нашите API и нашите изследователски екипи. Големите модели с отворен код от тип смесица от експерти (Mixture-of-Experts), като Kimi-K2, поставят специфични предизвикателства, тъй като най-големите инферентни нодове с 8x NVIDIA H200 GPU не могат ефективно да ги погодят, което налага внедряване в много нодове (multi-node). Представяме набор от ядра за експертен паралелизъм, които постигат върхови латентности на ConnectX-7, надминавайки производителността на DeepEP. Същите ядра също така са първите, които постигат жизнеспособни латентности на AWS Elastic Fabric Adapter (EFA), позволявайки внедряването на модели с трилиони параметри.

Изпробвайте нашите ядра в GitHub и прочетете пълния изследователски доклад на arXiv.

Въведение

Архитектурата Mixture-of-Experts (MoE) се наложи като стандартна архитектура за мащабиране до стотици милиарди или трилиони параметри, като същевременно предлага разумна латентност по време на инференция. Моделите MoE заменят плътния слой на трансформер с набор от експерти и малък линеен слой за маршрутизиране, който решава с кой експерт да се умножи токенът. Чрез групиране на токените, свързани със същия експерт, проблемът става разреден (sparse) и тривиално паралелизируем, тъй като различните експерти вече могат да бъдат разположени на различни GPU-та, разпределени в множество нодове.

За разлика от други форми на паралелизъм, като тензорен паралелизъм (Tensor Parallelism - TP) и блоков паралелизъм (Data Parallelism - DP), които могат лесно да бъдат реализирани с помощта на библиотеките за колективна комуникация като torch.distributed и NCCL, маршрутизирането при MoE включва разредена (sparse) peer-to-peer комуникация. Ядрото за изпращане (dispatch kernel) трябва да разделя набора от входящи токени и да ги насочва към ранговете, хостващи експертите, докато неговият дуален проблем, ядрото за обединяване (combine kernel), трябва да получи техните обработени еквиваленти и да изчисли тяхната претеглена средна стойност. Тези трансфери не се мапират към колективни операции и изискват персонализирани ядра за постигане на ниска латентност.

В Perplexity разработваме както преносими междувъзлови ядра, така и специализирани вътревъзлови ядра, като представихме работата си по-подробно в предишна статия за MoE комуникацията. В друга публикация в блога също показахме, че разгръщането на няколко възела може да осигури по-ниска латентност и по-висока пропускателна способност за големи MoE модели.

В тази статия разглеждаме в дълбочина реализацията на нашите нови междунодови ядра, които значително подобряват производителността на предишната ни работа, базирана на NVSHMEM. Освен че показваме върхови резултати на ConnectX-7, ние също така демонстрираме, че в рамките на нашата среда за изпълнение на инференция (inference runtime) ядрата могат да се използват за ефективно обслужване на най-големите модели с отворен код.

Инференция през InfiniBand

Тензорният паралелизъм и експертният паралелизъм могат ефективно да се мащабират в рамките на един нод, тъй като NVLink предлага висока пропусквателна способност (900 GB/s) при микросекунди латентност. Въпреки че системите NVIDIA GB200 и GB300 са оборудвани с достатъчно памет за съхранение на теглата и KV кешовете на модели с трилиони параметри, внедряванията на AWS p5en инстанция с до 8 H200 GPU са ограничени до 1120 GB HBM памет, която трябва да бъде разделена между теглата на модела и KV кешовете. Следователно определени модели изискват междунодови внедрявания през InfiniBand.

Въпреки че InfiniBand все още е изключително бърз, той предлага пропусквателна способност от едва 400 Gbps (50 GB/s), добавяйки десетки до стотици микросекунди латентност за MoE маршрутизирането. За щастие, добавените латентности могат да бъдат скрити чрез използването на споделени експерти, микропартидиране (micro-batching) и припокриване на изчисленията с комуникацията, за да се осигури жизнеспособна пропусквателна способност на конкурентна цена. DeepSeek е пионер в тези техники чрез DeepSeek-R1, внедрен върху клъстери от H800 GPU-та и адаптери NVIDIA ConnectX-7 InfiniBand.

В клъстерите на AWS ситуацията е различна поради използването на техните собствени Elastic Fabric Adapters (EFA). Въпреки че те осигуряват пикова пропусквателна способност от 400 Gbps при колективните операции, които обикновено се използват в тренировъчните работни натоварвания, те леко отстъпват на ConnectX-7 по отношение на размерите на съобщенията, обменяни по време на MoE разпределението (dispatch) и обединяването (combine). Освен това EFA не поддържа GPUDirect Async, което изисква нишка прокси на централния процесор (CPU proxy thread), свързваща GPU-то към мрежовия адаптер (NIC), за да се инициират трансферите. Освен добавената сложност, допълнителните PCIe транзакции добавят овърхед на стойност микросекунди към всички трансфери. Рамки като NVSHMEM позволяват на авторите на ядра да предоставят независими от устройството реализации, но общите API внасят овърхед, който прави внедряванията неефективни както по отношение на разходите, така и на латентността.

pplx-kernels, базирани на NVSHMEM

Нашите предишни ядра абстрахираха основната транспортна реализация чрез NVSHMEM. Те разчитаха на nvshmemx_putmem_signal_nbi_warp за прехвърляне на токени един по един, използвайки атомни броячи за синхронизиране на партньорите (peers). Въпреки че това предлагаше разумна производителност на мрежови адаптери ConnectX-7 с IBGDA, базираната на прокси ConnectX-7 IBRC реализация беше значително по-бавна, докато реализацията с EFA не успя да осигури латентности под една милисекунда.

Използвахме ядрата като тестова среда за хардуера, като същевременно определихме и най-подходящия интерфейс за излагане към нашите модели. Предположихме, че честите трансфери и взаимодействия между ядрата и мрежовите адаптери или прокси темите (threads) възпрепятстват производителността. Освен това нашата реализация не разчиташе на никакви гаранции от хардуера, освен надеждността, което предполагаше, че бихме могли да ускорим нещата чрез намаляване на овърхеда от примитивите за синхронизация.

След тези наблюдения разработихме нашите ядра, базирани на EFA, които излагат същия интерфейс, оптимизирайки взаимодействията между устройството (device kernel) и хост проксито под капака за постигане на пикова производителност.

MoE all-to-all през EFA

Уголеми

Новите ядра следват хибридна CPU-GPU архитектура, при която GPU ядрата взаимодействат с модела на устройството, а прокси тема (thread), работеща на хост CPU-то, управлява взаимодействията с мрежовия адаптер (NIC). Новите ядра преизползват същия TransferEngine, който изградихме за трансфери на KV кеш.

Маршрутизирането при MoE се изпълнява от ядро за изпращане (dispatch kernel), което получава токените и експертите за съответния токен, към които да ги насочи от модела, изпращайки ги през мрежата до ранговете, които хостват експертите. След като груповите GEMM ядра умножат токените с теглата на експертите, ядрото за обединяване (combine kernel) прехвърля токените обратно към техните изходни рангове и изчислява тяхната претеглена средна стойност, преди да ги върне на модела.

Както ядрото за изпращане, така и ядрото за обединяване са разделени на половин изпращач (sender) и половин получател (receiver), за да се улесни припокриването и микропартидирането. Изпращачът разбърква токените в буфери за изпращане, така че те да могат да бъдат записани в партньорите (peers) чрез една единствена операция за RDMA запис. След изпълнението на изпращача, GPU-то може да изпълнява друга работа (като умножение със споделен експерт), докато трансферите през мрежата чакат във фонов режим. Ядрата получатели блокират, докато всички трансфери бъдат получени. На фонов режим работна тема (worker thread) от страната на хоста управлява трансферите: тя получава сигнал от GPU-то, когато буферите за изпращане са били запълнени, след което изпраща RDMA записите. Щом мрежовият адаптер сигнализира за завършване, тя деблокира получателя, който връща токените на модела.

Ядрата са тясно свързани с хоста, комуникирайки помежду си както чрез обединена памет (unified memory), така и чрез GDRCopy. CUDA unified memory позволява паметта на устройството или паметта на хоста да бъдат мапирани в адресното пространство на двете, като прозрачно управлява комуникацията чрез PCIe.

Въпреки че обединената памет е подходяща за масови трансфери, полълването с ниска латентност се извършва чрез GDRCopy. GDRCopy използва GPUDirect RDMA за четене и запис в GPU паметта при микросекундна латентност.

В допълнение към RDMA, в рамките на един нод ядрата използват и NVLink за обмен на токени. Това може да разтовари значителна част от трафика — до една четвърт за EP32 и до една осма за EP64 — постигане на общо по-ниски латентности.

Буфери за изпращане и получаване (Send and Receive Buffers)

Дизайнът на буферите за изпращане за ядрото за изпращане е от решаващо значение за постигането на жизнеспособни латентности, включващи хост проксито. Ако има N ранга, всеки от които изпраща T токена, като всеки е насочен към R експерта от общо E, в най-лошия възможен случай даден ранг може да получи до N * T * max(R, E / N) токена, ако всички те са насочени към него. Освен това тази минимална граница може да бъде постигната само когато изпращачите подредят своите записи последователно, тъй като частните буфери за получаване биха изразходвали твърде много памет. За да могат изпращачите да подреждат своите записи последователно без конфликти, пълната информация за маршрутизирането трябва да бъде разпределена на всеки ранг. Следователно нашите ядра за изпращане първо обменят броя на токените за всеки експерт, след което всяко от тях може да определи къде отиват неговите записи и от кой ранг получава токени при съответното отместване (offset).

Простият обмен на информация за маршрутизиране би довел до ненужно високи латентности, тъй като всеки ранг би трябвал да чака записите да пристигнат, преди да изпрати каквито и да било собствени записи. Вместо това ние запазваме малко количество частно пространство на всеки ранг получател и изпращаме част от токените заедно със записите на информацията за маршрутизиране. Това гарантира пълно използване на пропускателната способност през по-голямата част от процеса на изпращане. След като маршрутите бъдат обменени, останалите токени се прехвърлят чрез един единствен запис на отдалечен партньор (peer). Обменът на информация за маршрутизиране и буферите с фиксиран размер, заедно с трансферите, които ги запълват, са илюстрирани на фигурата по-долу.

Уголеми

Ядрата за обединяване (combine kernels) са съществено по-опростени, тъй като информацията за маршрутизиране вече им е достъпна от етапа на изпращане. Те могат да копират токените с един единствен запис на партньор, след което четат от буфера за получаване, за да изчислят претеглената средна стойност на токените в изходните тензори.

За NVLink, буферите за изпращане и получаване в рамките на един нод са мапирани в партньорите (peers), за да могат те директно да избутват (push) или изтеглят (pull) токени от съседните рангове.

Изпращане и обединяване (Dispatch and Combine)

Фигурите по-долу описват подробно времето и взаимодействията между GPU, CPU и мрежовия адаптер (NIC), показвайки една итерация на изпращане и обединяване, интерлеирани с ядрата на модела.

Работата започва в изпращането на изпращача (dispatch send), като ядрото агрегира информацията за маршрутизиране и изчислява броя на токените за експерт чрез атомно натрупване (atomic accumulation) в споделената памет в рамките на един блок. След като броят на токените стане достъпен, те се копират в паметта на хоста и проксито незабавно се уведомява да започне изпращането на информацията за маршрутизиране до всички останали рангове партньори. Паралелно с това, изпращането продължава да изчислява кумулативната сума на броя токени, идентифицирайки отместванията (offsets), необходими за пакетиране на токените в непрекъснати буфери за изпращане за всеки ранг. След като отместванията са налични, всички токени, които трябва да бъдат изпратени през EFA, се пакетират в буферите за изпращане. Проксито се уведомява да пусне записите в буферите за изпращане на междунодовите партньори. Вътре в нода ядрата също незабавно избутват токени в частни буфери в рамките на същия нод чрез NVLink.

Уголеми

Уголеми

След като даден ранг получи информация за маршрутизиране от всички партньори, той я централизира и изчислява различни отмествания (offsets), необходими за идентифициране на токени за изпращане и обединяване. Тази работа се извършва на CPU-то, докато трансферите вече са чакащи (pending), поради което тя не е критична за латентността. Впоследствие, след като всички отмествания са известни, се пускат записи за остатъка от токените, които надвишават размера на частния буфер, и фазата на изпращане е завършена.

В края на получателя, ядрото първо изтегля токени от същия нод, които са извън капацитета на частните буфери чрез NVLink. След това то блокира, докато хост проксито сигнализира за завършването на всички трансфери. След като сигналът бъде получен, токените се разбъркват от буферите за получаване в изходния тензор, подаден към груповия GEMM. Тъй като буферите се преизползват между изпращането и обединяването, както NVLink, така и RDMA бариерите гарантират, че буферите на получателя са били прочетени, преди да се позволи на следващото ядро да ги презапише.

Ядрата за обединяване следват приблизително подобен модел — разбъркване на токени, сигнализиране на проксито и изчакване за завършване. Метадантите за пускане на техните записи обаче са лесно достъпни и всички отдалечени отмествания са известни, следователно всички токени, прехвърлени вътре в нода, се избутват чрез NVLink към получателите, а междунодовите токени се прехвърлят чрез един единствен запис. Преди да върне резултатите, ядрото натрупва токените от буферите за получаване, записвайки ги в изходните тензори.

TransferEngine за MoE маршрутизиране

За да прототипираме нашата реализация, първоначално преизползвахме Transfer Engine, изграден за дезагрегирани операции на предварително попълване (prefill) и декодиране, за да инициираме и проследяваме RDMA записи. За разлика от дезагрегираното декодиране обаче, където всеки нод публикува голям брой записи към същия партньор, MoE разпределението (dispatch) изпраща малки трансфери към всички партньори в една операция. Първоначално производителността не беше оптимална и нашата реализация изоставаше спрямо DeepEP на ConnectX-7 с около 20us. В резултат на това специализирахме TransferEngine да поддържа две операции — разпръскване (scatter) и бариера (barrier) — за оптимизиране на взаимодействията с голям брой партньори.

Операциите за разпръскване (scatter) и бариера (barrier) излагат групи от партньори, които могат да бъдат предварително регистрирани за предварително попълване на структурите от данни, използвани от основната транспортна реализация (създаване на шаблони за работни заявки - Work Request (WR) templating). Scatter копира различни редове от един и същ изходен буфер в различни отмествания (offsets) за всички членове на групата партньори. Бариерата (Barrier) сигнализира на всички партньори чрез immediate стойност. Опаковането (enqueuing) и обработката на тези операции в масови обеми спестява ценни микросекунди по критичния път. На EFA, където два мрежови адаптера по 200 Gbps се агрегират, за да осигурят 400 Gbps пропусквателна способност, ние разделяме партньорите в групата на шард-ове между двата адаптера, вместо да разделяме байтовете на всеки трансфер.

Въпреки че първоначално разработихме TransferEngine изцяло на EFA, използвайки libfabric, MoE ядрата работят на ConnectX-7 без промяна дори на един ред код, след като добавихме поддръжка за ConnectX-7 в TransferEngine с помощта на libibverbs. В сравнение с EFA, базиран на Scalable Reliable Datagram (SRD), ConnectX-7 изисква допълнителна работа за настройка на връзката и управление на партньорите. Използваме Unreliable Datagram (UD) двойка опашки (queue pair) за бутстрапване на Reliable Connection (RC). Създаваме две двойки опашки за всеки отдалечен партньор — една за двупосочни SEND/RECV и една за еднопосочни WRITE_IMM — тъй като и RECV, и WRITE_IMM консумират RECV операции последователно, докато двете операции изискват различни размери на буфера. Освен това отпускаме подредбата на PCIe транзакциите между ConnectX-7 и GPU паметта. Подобно на EFA, реализираме оптимизацията с шаблони за работни заявки (WR templating). Допълнително, когато изпращаме множество малки съобщения до същия партньор, използваме WR chaining за подобряване на скоростта на съобщенията.

Тъй като работата от страната на GPU е идентична за двата адаптера, избутахме основната транспортна реализация до нейните лимити, за да надминем производителността на DeepEP и да подчертаем, че разликата между EFA и ConnectX-7 MoE маршрутизирането се определя от възможностите на базовия хардуер.

Оценка (Evaluation)

Латентност на разпределение и обединяване (Dispatch/Combine Latency)

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

Уголеми

При декодиране, въпреки че нашите ядра са с няколко микросекунди по-бавни от DeepEP, постигаме общо по-ниски латентности, тъй като нашите ядра за обединяване са съществено по-бързи. Въпреки че използваме хост прокси, благодарение на масовите ни трансфери надминаваме производителността на всички реализации без прокси, базирани на IBGDA. Също така показваме, че изборът на прокси реализация има значение, тъй като нашето силно специализирано прокси е с порядък по-бързо от общата NVSHMEM реализация. Въпреки факта, че EFA е два пъти по-бавна от ConnectX-7 при прехвърляне на пакети от 256 KB, които са близки до типичния размер на пакетите, обменяни при MoE маршрутизиране, разликата между двата адаптера като цяло е по-малко съществена.

Уголеми

Въпреки че нашите оптимизирани за декодиране ядра могат да се мащабират до размери на партидите за предварително попълване (prefill) без никакви допълнителни настройки, производителността изостава спрямо DeepEP. Това се дължи главно на факта, че ние натрупваме токени само от страната на получателя при обединяването, докато DeepEP натрупва в изпращача, за да намали количеството прехвърлени данни. Тази оптимизация обаче може да има отражение върху точността. Освен това, тъй като ядрата прехвърлят данни с един единствен запис, използваме повече памет за буферите за изпращане и получаване.

Пропусквателна способност на Kimi-K2 / DeepSeek-V3

За да докажем жизнеспособността на нашите ядра, правим бенчмарк на цялостната производителност (end-to-end) на нашия собствен инференсен двигател, обслужващ DeepSeek-V3 (671B параметри) и Kimi-K2 (1T параметри) на множество p5en инстанции, оборудвани с Elastic Fabric Adapters (EFA) и H200 GPU-та. Изпълняваме модела в режим на верификация на чернова (draft verification mode) и докладваме латентността, приемайки процент на приемане от 80% с използване на един MTP слой.

При DeepSeek-V3 правим сравнение с базова линия EP=8 в рамките на един нод, изпълняваща ефективни ядра за разпределение и обединяване само през NVLink. Междунодовите внедрявания като цяло или съответстват на, или надминават производителността на внедряването в рамките на един нод, особено при партиди със среден размер, което демонстрира, че мащабирането между нодовете води до подобрения на общата пропусквателна способност.

При Kimi-K2 ядрата дават възможност за обслужване. Поради своя размер моделът не може да се побере в рамките на същия нод с използване на блоков паралелизъм (Data Parallelism - DP) на H200 нод, поради което се изисква междунодово внедряване за обслужване на модела. Тъй като моделът има по-малко глави на внимание (attention heads), постигаме по-добри латентности от DeepSeek-V3 и можем да обслужваме модела на EFA при жизнеспособна латентност.

Уголеми

Уголеми

Обмислили сме и микропартидирането (micro-batching), макар че то подобри латентностите само при много големи размери на партидите или при използване на повече от един токен за чернова. За да се извлече полза от разделянето на проблемите с групов GEMM на две, размерите на партидите с токени трябва да бъдат достатъчно плътни. Лъвският пай от времето все още се отива за MoE разпределение и обединяване, което показва, че всякакви по-нататъшни подобрения на ядрата съответно ще подобрят латентностите и ще улеснят микропартидирането.

Бъдеща работа

Работим в тясно сътрудничество с инженерите на AWS за подобряване на производителността на EFA. От своя страна ще следим актуализациите на libfabric, които ще намалят овърхеда на плоскостта за данни (data plane overhead). Допълнително ще експериментираме с efa-direct, за да намалим допълнително овърхеда в потребителското пространство (userspace overheads) и да взаимодействаме директно с хардуера.

Референции

Технически доклад за DeepSeek R1

Технически доклад за Kimi K2

DeepEP

UCCL-EP

Подобряване на мрежовата производителност на HPC системи с помощта на NVIDIA Magnum IO NVSHMEM и GPUDirect Async

Цитат

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