Habilitació de models d'un bilió de paràmetres a l'AWS EFA

Poseu a disposició models d'un bilió de paràmetres amb la portabilitat de la plataforma al núvol

AutorsPerplexity Research

A Perplexity, fem servir els millors models per al nostre producte, les nostres API i els nostres equips de recerca. Els grans models de barreja d'experts (Mixture-of-Experts) de codi obert, com ara Kimi-K2, plantegen reptes particulars, ja que els nodes d'inferència més grans amb GPUs 8x NVIDIA H200 no els poden allotjar de manera eficient, cosa que fa necessaris desplegaments multinode. Presentem un conjunt de nuclis per al paral·lelisme d'experts que aconsegueixen latències límit en ConnectX-7, superant el rendiment de DeepEP. Els mateixos nuclis també són els primers a aconseguir latències viables a l'Elastic Fabric Adapter (EFA) d'AWS, permetent el desplegament de models d'un bilió de paràmetres.

Proveu els nostres nuclis a GitHub i llegiu l'article de recerca complet a arXiv.

Introducció

La barreja d'experts (Mixture-of-Experts o MoE) s'ha convertit en l'arquitectura estàndard per escalar a centenars de milers de milions o bilions de paràmetres, oferint alhora una latència raonable en el moment de la inferència. Els models MoE substitueixen la capa densa d'un transformador per un conjunt d'experts i una petita capa d'encaminament lineal que decideix amb quin expert s'ha de multiplicar un testimoni. En agrupar els testimonis associats al mateix expert, el problema es torna dispers i trivialment paral·lelitzable, ja que ara es poden allotjar diferents experts en diferents GPU repartides per diversos nodes.

A diferència d'altres formes de paral·lelisme, com ara el paral·lelisme de tensors (TP) i el paral·lelisme de dades (DP), que es poden implementar fàcilment utilitzant biblioteques de comunicació col·lectiva com torch.distributed i NCCL, l'encaminament MoE implica comunicació dispersa peer-to-peer. Un nucli d'enviament ha de dividir el conjunt de testimonis entrants i enviar-los als rangs que allotgen els experts, mentre que el seu problema dual, el nucli de combinació, ha de rebre les seves contrapartides processades i calcular la seva mitjana ponderada. Aquestes transferències no es mapegen a col·lectius i requereixen nuclis personalitzats per aconseguir una baixa latència.

A Perplexity, hem estat desenvolupant tant nuclis portàtils entre nodes com nuclis especialitzats dins del node, presentant el nostre treball amb més profunditat en un article anterior sobre la comunicació MoE. També vam mostrar en una altra entrada de blog que el desplegament multinode pot aportar una latència més baixa i un rendiment més elevat per a models MoE grans.

En aquest article, fem una immersió profunda en la implementació dels nostres nous nuclis entre nodes, que milloren considerablement el rendiment del nostre treball anterior basat en NVSHMEM. A més de mostrar resultats capdavanters en ConnectX-7, també demostrem que, dins del nostre entorn d'execució d'inferència, els nuclis es poden utilitzar per donar servei de manera eficaç als models de codi obert més grans.

Inferència sobre InfiniBand

El paral·lelisme de tensors i el paral·lelisme d'experts poden escalar eficaçment dins d'un únic node, ja que NVLink ofereix un alt rendiment (900 GB/s) amb microsegons de latència. Tot i que els sistemes NVIDIA GB200 i GB300 estan equipats amb prou memòria per contenir els pesos i les catèsis KV de models d'un bilió de paràmetres, els desplegaments a la instància p5en d'AWS amb fins a 8 GPU H200 es limiten a 1120 GB de HBM, que s'han de repartir entre els pesos del model i les catèsis KV. En conseqüència, determinats models necessiten desplegaments entre nodes a través d'InfiniBand.

Tot i que InfiniBand continua sent extremadament ràpid, només ofereix un rendiment de 400 Gbps (50 GB/s), afegint desenes o centenars de microsegons de latència per a l'encaminament MoE. Afortunadament, les latències afegides es poden amagar mitjançant l'ús d'experts compartits, microlots i solapament de còmput i comunicació per oferir un rendiment viable a un cost competitiu. DeepSeek va ser pioner en aquestes tècniques a través de DeepSeek-R1, desplegat en clústers de GPU H800 i adaptadors NVIDIA ConnectX-7 InfiniBand.

En els clústers d'AWS, la situació és diferent a causa de l'ús dels seus adaptadors Elastic Fabric Adapters (EFA) personalitzats. Tot i que aquests ofereixen un pic de rendiment de 400 Gbps en els col·lectius utilitzats habitualment en càrregues de treball d'entrenament, es queden lleugerament per sota de ConnectX-7 pel que fa a les mides de missatge intercanviades durant l'enviament i la combinació de MoE. A més, EFA no admet GPUDirect Async, la qual cosa requereix un fil intermediari de CPU que connecti la GPU amb la targeta de xarxa per iniciar les transferències. A part de la complexitat afegida, les transaccions PCIe addicionals afegeixen un cost de microsegons a totes les transferències. Els marcs com NVSHMEM permeten als escriptors de nuclis proporcionar implementacions independents del dispositiu, però les API genèriques comporten costos generals que fan que els desplegaments siguin ineficaços tant en termes de cost com de latència.

pplx-kernels basats en NVSHMEM

Els nostres nuclis anteriors abstreien la implementació de transport subjacent mitjançant NVSHMEM. Es basaven en nvshmemx_putmem_signal_nbi_warp per transferir testimonis un per un, utilitzant comptadors atòmics per sincronitzar els iguals. Tot i que això ofereix un rendiment raonable en targetes de xarxa ConnectX-7 amb IBGDA, la implementació IBRC de ConnectX-7 basada en intermediaris era significativament més lenta, mentre que la implementació d'EFA no va aconseguir latències inferiors a un mil·lisecond.

Vam utilitzar els nuclis com a banc de proves per al maquinari subjacent, alhora que vam determinar la interfície més adequada per exposar als nostres models. Vam formular la hipòtesi que les freqüents transferències i interaccions entre els nuclis i les targetes de xarxa o els fils intermediaris obstaculitzaven el rendiment. A més, la nostra implementació no depenia de cap garantia del maquinari a part de la fiabilitat, cosa que suggeria que podríem accelerar les coses reduint el cost general de les primitives de sincronització.

Seguint aquestes observacions, vam dissenyar els nostres nuclis basats en EFA que exposen la mateixa interfície, optimitzant les interaccions entre el nucli del dispositiu i l'intermediari de l'amfitrió sota el capó per obtenir un rendiment màxim.

MoE all-to-all sobre EFA

Comunicació híbrida RDMA i NVLink amb un proxy d'amfitrió, que mostra els rangs d'inici i de recepció, els treballadors del proxy i les NIC.

Els nous nuclis segueixen una arquitectura híbrida de CPU-GPU, on els nuclis de GPU interactuen amb el model al dispositiu i un fil intermediari que s'executa a la CPU de l'amfitrió s'encarrega de les interaccions amb la targeta de xarxa. Els nous nuclis reutilitzen el mateix TransferEngine que vam construir per a les transferències de catèsis KV.

L'encaminament MoE el duu a terme un nucli d'enviament, que rep del model els testimonis i els experts per testimoni per encaminar-los, enviant-los a través de la xarxa als rangs que allotgen els experts. Després que els nuclis GEMM agrupats multipliquin els testimonis amb els pesos dels experts, el nucli de combinació transfereix els testimonis als seus rangs d'origen i calcula la seva mitjana ponderada abans de retornar-los al model.

Tant el nucli d'enviament com el de combinació estan dividits en una meitat emissora i una altra de receptora per facilitar el solapament i els microlots. L'emissor reorganitza els testimonis en memòria intermèdia d'enviament perquè es puguin escriure als iguals mitjançant una sola operació d'escriptura RDMA. Després d'executar l'emissor, la GPU pot realitzar altres tasques (com la multiplicació d'experts compartits) mentre les transferències a través de la xarxa estan pendents en segon pla. Els nuclis receptors es bloquegen fins que s'han rebut totes les transferències. En segon pla, un fil de treball del costat de l'amfitrió gestiona les transferències: la GPU l'avisa quan s'han omplert les memòries intermèdies d'enviament, moment en què envia les escriptures RDMA. Un cop la targeta de xarxa senyalitza la finalització, desbloqueja el receptor, el qual retorna els testimonis al model.

Els nuclis estan estretament acoblats amb l'amfitrió i es comuniquen entre si mitjançant la memòria unificada i GDRCopy. La memòria unificada de CUDA permet assignar la memòria del dispositiu o de l'amfitrió als espais d'adreça de tots dos, gestionant de manera transparent la comunicació a través de PCIe.

Tot i que la memòria unificada és adequada per a transferències massives, el sondeig de baixa latència es fa mitjançant GDRCopy. GDRCopy utilitza GPUDirect RDMA per llegir i escriure a la memòria de la GPU amb una latència de microsegons.

A més de RDMA, dins d'un node els nuclis també aprofiten NVLink per intercanviar testimonis. Això pot descarregar una fracció significativa del trànsit, fins a una quarta part per a EP32 i fins a una vuitena part per a EP64, aconseguint globalment latències més baixes.

Memòries intermèdies d'enviament i recepció

El disseny de les memòries intermèdies d'enviament per al nucli d'enviament és crucial per aconseguir latències viables que impliquin l'intermediari de l'amfitrió. Si hi ha N rangs que envien T testimonis cadascun, cadascun encaminat a R experts d'entre E, un rang pot rebre fins a N * T * max(R, E / N) testimonis en el pitjor cas possible, si tots s'encaminen cap a ell. A més, aquest límit mínim només es pot assolir quan els emissors disposen les seves escriptures de manera contigua, ja que les memòries intermèdies de recepció privades utilitzarien massa memòria. Perquè els emissors puguin disposar les seves escriptures consecutivament sense conflictes, s'ha de distribuir a cada rang el coneixement complet de la informació d'encaminament. En conseqüència, els nostres nuclis d'enviament primer intercanvien els recomptes de testimonis per expert, després dels quals cadascun pot determinar on van les seves escriptures i de quin desplaçament reben els testimonis de cada rang.

El simple intercanvi d'informació d'encaminament donaria lloc a latències innecessàriament altes, ja que cada rang hauria d'esperar que arribessin les escriptures abans d'enviar cap escriptura pròpia. En canvi, reservem una petita quantitat d'espai privat per emissor a cada rang receptor i enviem alguns dels testimonis al costat de les escriptures de la informació d'encaminament. Això garanteix la utilització completa de l'amplada de banda durant la major part del procés d'enviament. Un cop intercanuades les rutes, els testimonis restants es transfereixen mitjançant una sola escriptura per peer remot. L'intercanvi d'informació d'encaminament i les memòries intermèdies de mida fixa, juntament amb les transferències que les omplen, es il·lustren a la figura següent.

Disposició del buffer per a la distribució MoE, que mostra com s'organitzen i s'intercanvien les dades dels tokens i les metadades d'encaminament entre dos rangs.

Els nuclis de combinació són substancialment més senzills, ja que la informació d'encaminament ja està disponible per a ells des de l'etapa d'enviament. Poden copiar els testimonis en una sola escriptura per peer, després de la qual llegeixen de la memòria intermèdies de recepció per calcular la mitjana ponderada dels testimonis als tensors de sortida.

Per a NVLink, les memòries intermèdies d'enviament i recepció dins d'un node es mapegen als peer perquè puguin empènyer o extreure directament testimonis dels rangs veïns.

Enviament i combinació

Les figures següents detallen la temporització i les interaccions entre la GPU, la CPU i la targeta de xarxa, mostrant una iteració d'enviament i combinació intercalada amb els nuclis del model.

El treball comença a l'enviament (dispatch send), amb el nucli agrupant la informació d'encaminament i calculant els recomptes de testimonis per expert mitjançant acumulació atòmica a la memòria compartida dins d'un sol bloc. Un d'aquests recomptes de testimonis està disponible, es copien a la memòria de l'amfitrió i s'avisa immediatament l'intermediari perquè comenci a enviar la informació d'encaminament a tots els altres rangs peer. En paral·lel, l'enviament continua calculant la suma acumulada dels recomptes de testimonis, identificant els desplaçaments necessaris per empaquetar els testimonis en memòries intermèdies d'enviament contigües per peer. Un cop disponibles els desplaçaments, tots els testimonis que s'han d'enviar a través d'EFA s'empaqueten a les memòries intermèdies d'enviament. S'avisa l'intermediari perquè publiqui les escriptures a les memòries intermèdies d'enviament dels peer entre nodes. Dins del mateix node, els nuclis també empenyen immediatament els testimonis a les memòries intermèdies privades del mateix node mitjançant NVLink.

Temporització i coordinació per a la distribució MoE, que mostra operacions esglaonades de GPU, CPU i NIC i la seva sincronització.
Temporització i coordinació per al combine MoE, que mostra operacions de GPU, CPU i NIC i la seva sincronització.

Un cop un rang rep la informació d'encaminament de tots els peer, la centralitza i calcula diversos desplaçaments necessaris per identificar els testimonis per a l'enviament i la combinació. Aquest treball es fa a la CPU mentre les transferències ja estan pendents, per la qual cosa no és crític per a la latència. Posteriorment, un cop es coneixen tots els desplaçaments, es publiquen les escriptures per a la resta de testimonis, que superen la mida de la memòria intermèdia privada, i la fase d'enviament es completa.

Al costat del receptor, el nucli primer extreu del mateix node els testimonis que superen la capacitat de les memòries intermèdies privades a través de NVLink. A continuació, es bloca fins que l'intermediari de l'amfitrió senyalitza la finalització de totes les transferències. Un cop rebut el senyal, els testimonis es reorganitzen des de les memòries intermèdies de recepció al tensor de sortida passat a GEMM de grup. Com que les memòries intermèdies es reutilitzen entre l'enviament i la combinació, tant les barreres NVLink com RDMA garanteixen que s'hagin llegit les memòries intermèdies del receptor abans de permetre que el nucli següent les sobrescrigui.

Els nuclis de combinació segueixen un patró similar, reorganitzant testimonis, senyalitzant l'intermediari i esperant la finalització. Tanmateix, les metadades per publicar les seves escriptures estan immediatament disponibles i es coneixen tots els desplaçaments remots, de manera que tots els testimonis transferits dins del node s'empenyen mitjançant NVLink als receptors i els testimonis entre nodes es transfereixen mitjançant una sola escriptura. Abans de retornar els resultats, el nucli acumula els testimonis de les memòries intermèdies del receptor, escrivint-los als tensors de sortida.

TransferEngine per a l'encaminament MoE

Per crear el prototip de la nostra implementació, inicialment vam reutilitzar el Transfer Engine construït per a la pre-ompliment desagregada i la descodificació per tal d'emetre i fer el seguiment de les escriptures RDMA. No obstant això, a diferència de la descodificació desagregada on cada node publica un gran nombre d'escriptures al mateix peer, l'enviament MoE publica petites transferències a tots els peer en una sola operació. De sèrie, el rendiment no era òptim i la nostra implementació es quedava endarrerida respecte a DeepEP en ConnectX-7 per uns 20 us. En conseqüència, vam especialitzar el TransferEngine perquè admeti dues operacions, dispersió (scatter) i barrera, per optimitzar les interaccions amb un gran nombre de peer.

Tant les operacions de dispersió com de barrera exposen grups de peer que es poden preregistrar per preomplir les estructures de dades utilitzades per la implementació de transport subjacent (plantilles de sol·licitud de treball o WR templating). La dispersió copia diferents talls de la mateixa memòria intermèdia d'origen en diferents desplaçaments de tots els membres d'un grup de peer. La barrera senyalitza a tots els peer mitjançant un missatge immediat. Encua i gestionar aquestes operacions de manera massiva estalvia microsegons valuosos en el camí crític. A l'EFA, on dues targetes de xarxa de 200 Gbps s'agreguen per proporcionar una amplada de banda de 400 Gbps, fragmentem els peer de grup entre les dues targetes de xarxa en lloc de fragmentar els bytes de cada transferència.

Tot i que inicialment vam desenvolupar el TransferEngine només a l'EFA utilitzant libfabric, els nuclis MoE funcionen a ConnectX-7 sense canviar ni una sola línia de codi després d'afegir el suport de ConnectX-7 al TransferEngine utilitzant libibverbs. En comparació amb l'EFA basat en datagrames fiable escalable (SRD), ConnectX-7 requereix treball addicional per a la configuració de la connexió i la gestió de peer. Utilitzem un parell de cues de datagrames no fiables (UD) per iniciar la connexió fiable (RC). Creem dos parells de cues per a cada peer remot, un per a SEND/RECV bidireccional i un altre per a WRITE_IMM unidireccional, perquè tant RECV com WRITE_IMM consumeixen l'operació RECV en ordre mentre les dues operacions requereixen mides de memòria intermèdia diferents. A més, relaxem l'ordenació de transaccions PCIe entre ConnectX-7 i la memòria de la GPU. De manera similar a l'EFA, implementem l'optimització de plantilles WR. A més, quan s'envien múltiples missatges petits al mateix peer, fem servir el encadenament de WR per millorar la taxa de missatges.

Com que el treball del costat de la GPU és idèntic en ambdós adaptadors, vam portar la implementació de transport subjacent al límit per superar el rendiment de DeepEP i destacar que la bretxa entre l'encaminament MoE d'EFA i ConnectX-7 està determinada per les capacitats del maquinari subjacent.

Avaluació

Latència d'enviament/combinació

Vam avaluar les latències d'enviament i combinació en sistemes H200 equipats amb adaptadors ConnectX-7 i EFA a la configuració estàndard de DeepSeek-V3 i Kimi-K2 amb una dimensió oculta de 7168, enviament fp8 escalat per blocs i combinació bf16. Considerem una mida de lot de 128 per a l'enviament i 4096 per a la combinació. A ConnectX-7, comparem amb DeepEP i els nostres pplx-kernels anteriors basats en NVSHMEM, utilitzant tant GPUDirect Async (IBGDA) com un intermediari d'amfitrió (IBRC). Aconseguim latències combinades de 459 us, 582 us i 692 us en EP16, EP32 i EP64, respectivament. Superem altres nuclis basats en EFA que depenen d'intermediaris genèrics, com UCCL-EP, que aconsegueixen 519 us, 966 us i 1159 us.

Latència de distribució i combine per a la descodificació MoE EP64, comparant les implementacions i els camins de comunicació mostrats.

En la descodificació, tot i que els nostres nuclis són uns quants microsegons més lents que DeepEP, aconseguim latències generals més baixes gràcies al fet que els nostres nuclis de combinació són substancialment més ràpids. Tot i que fem servir un intermediari d'amfitrió, a causa de les nostres transferències massives superem el rendiment de totes les implementacions sense intermediari basades en IBGDA. També mostrem que la implementació de l'intermediari importa, ja que el nostre intermediari altament especialitzat és un ordre de magnitud més ràpid que la implementació genèrica de NVSHMEM. Malgrat el fet que l'EFA és dues vegades més lent que ConnectX-7 en transferir paquets de 256 KB, que s'acosten a la mida típica de paquet intercanviada per l'encaminament MoE, la bretxa entre els dos adaptadors és globalment menys substancial.

Latència de distribució i combine per a l'emplenament previ MoE EP64, comparant els camins de comunicació de pplx i DeepEP.

Tot i que els nostres nuclis optimitzats per a la descodificació poden escalar a mides de lot de pre-ompliment sense cap ajust addicional, el rendiment es queda endarrerit respecte a DeepEP. Això es deu principalment al fet que acumulem testimonis només al costat receptor de la combinació, mentre que DeepEP acumula a l'emissor per reduir la quantitat de dades transferides. Tanmateix, aquesta optimització pot tenir implicacions en la precisió. A més, com que els nuclis transfereixen dades en una sola escriptura, utilitzem més memòria per a les memòries intermèdies d'enviament i recepció.

Rendiment de Kimi-K2 / DeepSeek-V3

Per demostrar la viabilitat dels nostres nuclis, fem proves de rendiment de punta a punta del nostre motor d'inferència intern que dona servei a DeepSeek-V3 (671 mil milions de paràmetres) i Kimi-K2 (1 bilió de paràmetres) en múltiples instàncies p5en equipades amb Elastic Fabric Adapters (EFA) i GPU H200. Executem el model en mode de verificació d'esborrany i informem de la latència suposant una taxa d'acceptació del 80% utilitzant una sola capa MTP.

A DeepSeek-V3, comparem amb una línia de base intra-node EP=8 que executa nuclis d'enviament i combinació eficients només amb NVLink. Els desplegaments entre nodes generalment igualen o superen el rendiment del desplegament d'un sol node, particularment en lots de mida mitjana, demostrant que l'escalat entre nodes aporta millores al rendiment general.

A Kimi-K2, els nuclis permeten donar servei. A causa de la seva mida, el model no pot cabre dins del mateix node utilitzant el paral·lelisme de dades (DP) en un node H200, per la qual cosa es requereix un desplegament entre nodes per donar servei al model. Com que el model té menys caps d'atenció, aconseguim millors latències que DeepSeek-V3 i podem donar servei al model a l'EFA amb una latència viable.

Rendiment per GPU de DeepSeek-V3 en relació amb la velocitat de sortida per a les configuracions EP32, EP16 i EP8.
Rendiment per GPU de Kimi-K2 en relació amb la velocitat de sortida per a les configuracions EP32 i EP16.

També hem considerat els microlots, tot i que només van millorar les latències a mides de lot molt altes o quan s'utilitza més d'un testimoni d'esborrany. Per beneficiar-se de dividir els problemes GEMM agrupats en dos, les mides dels lots de testimonis han de ser prou denses. El gruix del temps encara es inverteix en l'enviament i la combinació MoE, cosa que indica que qualsevol altra millora dels nuclis millorarà consegüentment les latències i facilitarà els microlots.

Treballs futurs

Estem treballant estretament amb els enginyers d'AWS per millorar el rendiment a l'EFA. Per la nostra banda, seguirem les actualitzacions de libfabric que reduiran el cost general del pla de dades. A més, experimentarem amb efa-direct per reduir encara més els costos generals de l'espai d'usuari i interactuar directament amb el maquinari.

Referències

Informe tècnic de DeepSeek R1

Informe tècnic de Kimi K2

DeepEP

UCCL-EP

Millora del rendiment de xarxa de sistemes d'alt rendiment (HPC) utilitzant NVIDIA Magnum IO NVSHMEM i GPUDirect Async

Citació

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