Abilitazione di modelli da trilioni di parametri su AWS EFA

In Perplexity utilizziamo i modelli migliori per il nostro prodotto, le nostre API e i nostri team di ricerca. I grandi modelli open-source Mixture-of-Experts, come Kimi-K2, pongono sfide particolari poiché i nodi di inferenza più grandi con 8 GPU NVIDIA H200 non possono ospitarli in modo efficiente, rendendo necessarie implementazioni multi-nodo. Presentiamo una serie di kernel per il parallelismo degli esperti che raggiungono latenze all'avanguardia su ConnectX-7, superando le prestazioni di DeepEP. Gli stessi kernel sono anche i primi a raggiungere latenze praticabili su AWS Elastic Fabric Adapter (EFA), consentendo implementazioni di modelli con trilioni di parametri.

Prova i nostri kernel su GitHub e leggi l'articolo di ricerca completo su arXiv.

Mixture-of-Experts (MoE) è emerso come l'architettura standard per il ridimensionamento a centinaia di miliardi o trilioni di parametri, offrendo al contempo una latenza ragionevole al momento dell'inferenza. I modelli MoE sostituiscono lo strato denso di un transformer con una serie di esperti e un piccolo strato di routing lineare che decide con quale esperto moltiplicare un token. Raggruppando i token associati allo stesso esperto, il problema diventa sparso e banalmente parallelizzabile, poiché diversi esperti possono ora essere ospitati su GPU diverse distribuite su più nodi.

A differenza di altre forme di parallelismo, come il Parallelismo dei Tensori (TP) e il Parallelismo dei Dati (DP), che possono essere facilmente implementati utilizzando librerie di comunicazione collettiva come torch.distributed e NCCL, il routing MoE comporta una comunicazione sparsa peer-to-peer. Un kernel di invio deve dividere l'insieme dei token in arrivo e inviarli ai ranghi che ospitano gli esperti, mentre il suo problema duale, il kernel di combinazione, deve ricevere le controparti elaborate e calcolarne la media ponderata. Questi trasferimenti non mappano su collettivi e richiedono kernel personalizzati per ottenere una bassa latenza.

In Perplexity, abbiamo sviluppato sia kernel inter-nodo portatili che kernel intra-nodo specializzati, presentando il nostro lavoro in modo più approfondito in un precedente articolo sulla comunicazione MoE. Abbiamo anche mostrato in un altro post del blog che l' implementazione multi-nodo può ridurre la latenza e aumentare il throughput per i modelli MoE di grandi dimensioni.

In questo articolo, approfondiamo l'implementazione dei nostri nuovi kernel inter-nodo che migliorano notevolmente le prestazioni del nostro precedente lavoro basato su NVSHMEM. Oltre a mostrare risultati all'avanguardia su ConnectX-7, dimostriamo anche che all'interno del nostro runtime di inferenza, i kernel possono essere utilizzati per servire in modo efficace i più grandi modelli open-source.

Il parallelismo dei tensori e il parallelismo degli esperti possono scalare efficacemente all'interno di un singolo nodo, poiché NVLink offre un throughput elevato (900 GB/s) con microsecondi di latenza. Sebbene i sistemi NVIDIA GB200 e GB300 siano dotati di memoria sufficiente per contenere i pesi e le cache KV dei modelli da trilioni di parametri, le implementazioni su istanza AWS p5en con un massimo di 8 GPU H200 sono limitate a 1120 GB di HBM, che devono essere suddivisi tra pesi del modello e cache KV. Di conseguenza, alcuni modelli richiedono implementazioni inter-nodo su InfiniBand.

Sebbene InfiniBand sia ancora estremamente veloce, offre solo 400 Gbps (50 GB/s) di throughput, aggiungendo da decine a centinaia di microsecondi di latenza per il routing MoE. Fortunatamente, le latenze aggiunte possono essere nascoste attraverso l'uso di esperti condivisi, micro-batching e sovrapposizione tra calcolo e comunicazione per fornire un throughput praticabile a un costo competitivo. DeepSeek ha aperto la strada a queste tecniche tramite DeepSeek-R1, implementato su cluster di GPU H800 e adattatori NVIDIA ConnectX-7 InfiniBand.

Sui cluster AWS, la situazione è diversa a causa dell'uso dei loro adattatori personalizzati Elastic Fabric Adapters (EFA). Anche se questi forniscono un throughput di picco di 400 Gbps sui collettivi tipicamente utilizzati nei carichi di lavoro di addestramento, non raggiungono le prestazioni di ConnectX-7 sulle dimensioni dei messaggi scambiati durante l'invio e la combinazione MoE. Inoltre, EFA non supporta GPUDirect Async, richiedendo un thread proxy della CPU che colleghi la GPU alla scheda di rete per avviare i trasferimenti. Oltre alla maggiore complessità, le transazioni PCIe aggiuntive aggiungono microsecondi di overhead a tutti i trasferimenti. Framework come NVSHMEM consentono agli autori di kernel di fornire implementazioni indipendenti dal dispositivo; tuttavia, le API generiche comportano overhead che rendono le implementazioni inefficaci sia in termini di costi che di latenza.

I nostri kernel precedenti astraggono l'implementazione del trasporto sottostante tramite NVSHMEM. Si basavano su nvshmemx_putmem_signal_nbi_warp per trasferire i token uno alla volta, utilizzando contatori atomici per sincronizzare i peer. Sebbene ciò offrisse prestazioni ragionevoli su schede di rete ConnectX-7 con IBGDA, l'implementazione IBRC ConnectX-7 basata su proxy era significativamente più lenta, mentre l'implementazione EFA non è riuscita a fornire latenze inferiori al millisecondo.

Abbiamo utilizzato i kernel come banco di prova per l'hardware sottostante, definendo al contempo l'interfaccia più adatta da esporre ai nostri modelli. Abbiamo ipotizzato che i frequenti trasferimenti e le interazioni tra i kernel e le schede di rete o i thread proxy ostacolassero le prestazioni. Inoltre, la nostra implementazione non si basava su alcuna garanzia da parte dell'hardware oltre all'affidabilità, suggerendo che avremmo potuto velocizzare le cose riducendo l'overhead delle primitive di sincronizzazione.

In seguito a queste osservazioni, abbiamo ideato i nostri kernel basati su EFA che espongono la stessa interfaccia, ottimizzando le interazioni kernel del dispositivo - proxy host per prestazioni ottimali.

Comunicazione ibrida RDMA e NVLink con un proxy host, che mostra i rank di inizializzazione e ricezione, i worker del proxy e le NIC.

I nuovi kernel seguono un'architettura ibrida CPU-GPU, dove i kernel GPU interagiscono con il modello sul dispositivo e un thread proxy in esecuzione sulla CPU host gestisce le interazioni con la scheda di rete. I nuovi kernel riutilizzano lo stesso TransferEngine che abbiamo creato per i trasferimenti di cache KV.

Il routing MoE viene eseguito da un kernel di invio, che riceve i token e gli esperti per token a cui instradarli dal modello, spedendoli attraverso la rete ai ranghi che ospitano gli esperti. Dopo che i kernel GEMM raggruppati moltiplicano i token con i pesi degli esperti, il kernel di combinazione trasferisce i token ai loro ranghi di origine e calcola la loro media ponderata prima di restituirli al modello.

Sia il kernel di invio che quello di combinazione sono suddivisi in una metà mittente e una ricevente per facilitare la sovrapposizione e il micro-batching. Il mittente mescola i token nei buffer di invio in modo che possano essere scritti sui peer tramite una singola operazione di scrittura RDMA. Dopo aver eseguito il mittente, la GPU può eseguire altri lavori (come la moltiplicazione degli esperti condivisi) mentre i trasferimenti sulla rete sono in sospeso in background. I kernel riceventi si bloccano finché tutti i trasferimenti non sono stati ricevuti. In background, un thread di lavoro lato host gestisce i trasferimenti: viene segnalato dalla GPU quando i buffer di invio sono stati popolati, dopodiché invia le scritture RDMA. Una volta che la scheda di rete segnala il completamento, sblocca il ricevitore, che restituisce i token al modello.

I kernel sono strettamente accoppiati con l'host, comunicando tra loro attraverso sia la memoria unificata che GDRCopy. La memoria unificata CUDA consente alla memoria del dispositivo o alla memoria dell'host di essere mappata negli spazi di indirizzamento di entrambi, gestendo in modo trasparente la comunicazione tramite PCIe.

Sebbene la memoria unificata sia adatta per trasferimenti di massa, il polling a bassa latenza viene eseguito tramite GDRCopy. GDRCopy utilizza GPUDirect RDMA per leggere e scrivere nella memoria GPU con latenza di microsecondi.

Oltre all'RDMA, all'interno di un nodo i kernel sfruttano anche NVLink per scambiare token. Ciò può scaricare una frazione significativa del traffico, fino a un quarto per EP32 e fino a un ottavo per EP64, ottenendo latenze complessivamente inferiori.

La progettazione dei buffer di invio per il kernel di invio è cruciale per ottenere latenze praticabili che coinvolgono il proxy host. Se ci sono N ranghi che inviano ciascuno T token, ciascuno instradato a R esperti su E, un rango può ricevere fino a N * T * max(R, E / N) token nel caso peggiore, se tutti sono instradati verso di esso. Inoltre, questo limite minimo può essere raggiunto solo quando i mittenti dispongono le loro scritture in modo contiguo, poiché i buffer di ricezione privati consumerebbero troppa memoria. Affinché i mittenti siano in grado di disporre le loro scritture consecutivamente senza conflitti, la conoscenza completa delle informazioni di routing deve essere distribuita a ciascun rango. Di conseguenza, i nostri kernel di invio scambiano innanzitutto i conteggi dei token per esperto, dopodiché possono determinare ciascuno dove vanno le loro scritture e da quale offset ricevono i token da ciascun rango.

Il semplice scambio di informazioni di routing comporterebbe latenze inutilmente elevate, poiché ogni rango dovrebbe attendere l'arrivo delle scritture prima di inviarne di proprie. Invece, riserviamo una piccola quantità di spazio privato per mittente su ciascun rango ricevente e inviamo alcuni dei token insieme alle scritture delle informazioni di routing. Ciò garantisce la piena utilizzazione della larghezza di banda durante la maggior parte del processo di invio. Dopo lo scambio dei percorsi, i token rimanenti vengono trasferiti tramite una singola scrittura per peer remoto. Lo scambio di informazioni di routing e i buffer a dimensione fissa, insieme ai trasferimenti che li popolano, sono illustrati nella figura seguente.

Layout del buffer per la dispatch MoE, che mostra come i dati dei token e i metadati di routing vengono organizzati e scambiati tra due rank.

I kernel di combinazione sono sostanzialmente più semplici, poiché le informazioni di routing sono già a loro disposizione dalla fase di invio. Possono copiare i token in una singola scrittura per peer, dopodiché leggono dal buffer di ricezione per calcolare la media ponderata dei token nei tensori di output.

Per NVLink, i buffer di invio e ricezione all'interno di un nodo sono mappati nei peer affinché possano spingere o estrarre direttamente i token dai ranghi vicini.

Le figure seguenti descrivono i tempi e le interazioni tra GPU, CPU e scheda di rete, mostrando un'iterazione di invio e combinazione intercalata con i kernel del modello.

Il lavoro inizia nell'invio di dispatch, con il kernel che aggrega le informazioni di routing e calcola i conteggi dei token per esperto tramite accumulo atomico nella memoria condivisa all'interno di un singolo blocco. Una volta che i conteggi dei token sono disponibili, vengono copiati nella memoria host e il proxy viene immediatamente avvisato di iniziare a inviare le informazioni di routing a tutti gli altri ranghi peer. In parallelo, il dispatch continua a calcolare la somma cumulativa dei conteggi dei token, identificando gli offset necessari per impacchettare i token in buffer di invio contigui per peer. Una volta disponibili gli offset, tutti i token da inviare tramite EFA vengono impacchettati nei buffer di invio. Il proxy viene avvisato di pubblicare le scritture nei buffer di invio dei peer inter-nodo. Intra-nodo, i kernel spingono anche immediatamente i token in buffer privati all'interno dello stesso nodo tramite NVLink.

Temporizzazione e coordinamento per la dispatch MoE, che mostra operazioni GPU, CPU e NIC in fasi successive e la loro sincronizzazione.

Temporizzazione e coordinamento per la combine MoE, che mostra operazioni GPU, CPU e NIC e la loro sincronizzazione.

Una volta che un rango riceve le informazioni di routing da tutti i peer, le centralizza e calcola vari offset necessari per identificare i token per l'invio e la combinazione. Questo lavoro viene svolto sulla CPU mentre i trasferimenti sono già in sospeso, quindi non è critico per la latenza. Successivamente, una volta noti tutti gli offset, le scritture per il resto dei token, che superano la dimensione del buffer privato, vengono pubblicate e la fase di invio del dispatch è completa.

Sul lato ricevente, il kernel estrae innanzitutto i token dallo stesso nodo che superano la capacità dei buffer privati tramite NVLink. Successivamente, si blocca finché il proxy host non segnala il completamento di tutti i trasferimenti. Una volta ricevuto il segnale, i token vengono spostati dai buffer di ricezione al tensore di output passato al gruppo GEMM. Poiché i buffer vengono riutilizzati tra invio e combinazione, le barriere sia NVLink che RDMA assicurano che i buffer del ricevitore siano stati letti prima di consentire al kernel successivo di sovrascriverli.

I kernel di combinazione seguono un modello simile, spostando i token, segnalando il proxy e attendendo il completamento. Tuttavia, i metadati per la pubblicazione delle loro scritture sono prontamente disponibili e tutti gli offset remoti sono noti, quindi tutti i token trasferiti intra-nodo vengono spinti tramite NVLink ai ricevitori e i token inter-nodo vengono trasferiti tramite una singola scrittura. Prima di restituire i risultati, il kernel accumula i token dai buffer di ricezione, scrivendoli nei tensori di output.

Per prototipare la nostra implementazione, abbiamo inizialmente riutilizzato il Transfer Engine costruito per il prefill e la decodifica disaggregati al fine di emettere e tracciare le scritture RDMA. Tuttavia, a differenza della decodifica disaggregata in cui ogni nodo pubblica un gran numero di scritture sullo stesso peer, l'invio MoE pubblica piccoli trasferimenti a tutti i peer in un'unica operazione. Non è stato ottimale e la nostra implementazione era in ritardo rispetto a DeepEP su ConnectX-7 di circa 20us. Di conseguenza, abbiamo specializzato il TransferEngine per supportare due operazioni, scatter e barrier, per ottimizzare le interazioni con un gran numero di peer.

Entrambe le operazioni scatter e barrier espongono gruppi di peer che possono essere pre-registrati per pre-popolare le strutture dati utilizzate dall'implementazione del trasporto sottostante (templating delle richieste di lavoro (WR)). Scatter copia diverse sezioni dallo stesso buffer di origine in diversi offset su tutti i membri di un gruppo di peer. Barrier segnala tutti i peer attraverso un immediato. Accodare e gestire queste operazioni in blocco fa risparmiare preziosi microsecondi sul percorso critico. Su EFA, dove due schede di rete da 200 Gbps si aggregano per fornire 400 Gbps di larghezza di banda, suddividiamo i peer del gruppo tra le due schede di rete invece di suddividere i byte di ogni trasferimento.

Sebbene inizialmente avessimo sviluppato il TransferEngine esclusivamente su EFA utilizzando libfabric, i kernel MoE funzionano su ConnectX-7 senza modificare una singola riga di codice dopo aver aggiunto il supporto ConnectX-7 nel TransferEngine utilizzando libibverbs. Rispetto all'EFA basato su Scalable Reliable Datagram (SRD), ConnectX-7 richiede lavoro aggiuntivo per la configurazione della connessione e la gestione dei peer. Utilizziamo una coppia di code Unreliable Datagram (UD) per avviare la connessione affidabile (RC). Creiamo due coppie di code per ogni peer remoto, una per SEND/RECV a due vie e una per WRITE_IMM a una via, perché sia RECV che WRITE_IMM consumano l'operazione RECV in ordine, mentre le due operazioni richiedono dimensioni di buffer diverse. Inoltre, rilassiamo l'ordinamento delle transazioni PCIe tra ConnectX-7 e la memoria GPU. Analogamente all'EFA, implementiamo l'ottimizzazione del templating WR. Inoltre, quando inviamo più messaggi piccoli allo stesso peer, utilizziamo il concatenamento WR per migliorare la velocità dei messaggi.

Poiché il lavoro lato GPU è identico su entrambi gli adattatori, abbiamo spinto l'implementazione del trasporto sottostante ai suoi limiti per superare le prestazioni di DeepEP, per evidenziare che il divario tra il routing MoE EFA e ConnectX-7 è determinato dalle capacità dell'hardware sottostante.

Abbiamo valutato le latenze di invio e combinazione su sistemi H200 dotati di adattatori ConnectX-7 ed EFA sulla configurazione standard DeepSeek-V3 e Kimi-K2 con una dimensione nascosta di 7168, invio fp8 a blocchi scalati e combinazione bf16. Consideriamo una dimensione del batch di 128 per l'invio e 4096 per la combinazione. Su ConnectX-7, confrontiamo con DeepEP e i nostri precedenti pplx-kernels basati su NVSHMEM, utilizzando sia GPUDirect Async (IBGDA) che un proxy host (IBRC). Otteniamo latenze combinate di 459us, 582us e 692us rispettivamente su EP16, EP32 e EP64. Superiamo altri kernel basati su EFA che si basano su proxy generici, come UCCL-EP, che ottengono 519us, 966us e 1159us.

Latenza di dispatch e combine per il decode MoE EP64, a confronto tra le implementazioni e i percorsi di comunicazione mostrati.

Sulla decodifica, nonostante i nostri kernel siano di pochi microsecondi più lenti di DeepEP, otteniamo latenze complessivamente inferiori poiché i nostri kernel di combinazione sono sostanzialmente più veloci. Sebbene utilizziamo un proxy host, a causa dei nostri trasferimenti in blocco superiamo le prestazioni di tutte le implementazioni senza proxy basate su IBGDA. Mostriamo anche che l'implementazione del proxy è importante, poiché il nostro proxy altamente specializzato è un ordine di grandezza più veloce dell'implementazione NVSHMEM generica. Nonostante il fatto che EFA sia due volte più lenta di ConnectX-7 durante il trasferimento di pacchetti da 256 KB, che sono vicini alla dimensione tipica del pacchetto scambiato dal routing MoE, il divario tra i due adattatori è complessivamente meno sostanziale.

Latenza di dispatch e combine per il prefill MoE EP64, a confronto tra i percorsi di comunicazione pplx e DeepEP.

Anche se i nostri kernel ottimizzati per la decodifica possono scalare alle dimensioni del batch di prefill senza alcuna modifica aggiuntiva, le prestazioni sono in ritardo rispetto a DeepEP. Ciò è dovuto principalmente al fatto che accumuliamo token solo sul lato ricevente della combinazione, mentre DeepEP accumula nel mittente per ridurre la quantità di dati trasferiti. Tuttavia, questa ottimizzazione può avere implicazioni sull'accuratezza. Inoltre, poiché i kernel trasferiscono i dati in una sola scrittura, utilizziamo più memoria per i buffer di invio e ricezione.

Per dimostrare la praticabilità dei nostri kernel, effettuiamo il benchmark delle prestazioni end-to-end del nostro motore di inferenza interno che serve DeepSeek-V3 (671B parametri) e Kimi-K2 (1T parametri) su più istanze p5en dotate di Elastic Fabric Adapters (EFA) e GPU H200. Eseguiamo il modello in modalità di verifica della bozza e riportiamo la latenza assumendo un tasso di accettazione dell'80% utilizzando un singolo strato MTP.

Su DeepSeek-V3, confrontiamo con una baseline intra-nodo EP=8 che esegue kernel di invio e combinazione efficienti solo NVLink. Le implementazioni inter-nodo nel complesso eguagliano o superano le prestazioni dell'implementazione a nodo singolo, in particolare su batch di medie dimensioni, dimostrando che la scalabilità tra i nodi porta miglioramenti al throughput complessivo.

Su Kimi-K2, i kernel consentono il servizio. A causa delle sue dimensioni, il modello non può entrare nello stesso nodo utilizzando il parallelismo dei dati (DP) su un nodo H200, pertanto è necessaria un'implementazione inter-nodo per servire il modello. Poiché il modello ha meno teste di attenzione, otteniamo latenze migliori rispetto a DeepSeek-V3 e possiamo servire il modello su EFA con una latenza praticabile.

Throughput per GPU di DeepSeek-V3 rispetto alla velocità di output per le configurazioni EP32, EP16 ed EP8.

Throughput per GPU di Kimi-K2 rispetto alla velocità di output per le configurazioni EP32 ed EP16.

Abbiamo anche preso in considerazione il micro-batching, sebbene abbia migliorato le latenze solo a dimensioni di batch molto elevate o quando si utilizzano più di un token di bozza. Per trarre vantaggio dalla suddivisione dei problemi GEMM raggruppati in due, le dimensioni dei batch di token devono essere sufficientemente dense. La maggior parte del tempo viene ancora spesa per l'invio e la combinazione MoE, indicando che eventuali ulteriori miglioramenti ai kernel miglioreranno di conseguenza le latenze e faciliteranno il micro-batching.

Stiamo lavorando a stretto contatto con gli ingegneri AWS per migliorare le prestazioni su EFA. Da parte nostra, seguiremo gli aggiornamenti a libfabric che ridurranno l'overhead del piano dati. Inoltre, sperimenteremo efa-direct per ridurre ulteriormente l'overhead dello spazio utente e interfacciarci direttamente con l'hardware.

Rapporto tecnico DeepSeek R1

Rapporto tecnico Kimi K2

DeepEP

UCCL-EP

Miglioramento delle prestazioni di rete dei sistemi HPC utilizzando NVIDIA Magnum IO NVSHMEM e 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}
}

Introduzione

Inferenza su InfiniBand

pplx-kernels basati su NVSHMEM

MoE all-to-all su EFA

Buffer di invio e ricezione

Invio e combinazione

TransferEngine per il routing MoE

Valutazione

Latenza di invio/combinazione

Throughput di Kimi-K2 / DeepSeek-V3

Lavori futuri

Riferimenti

Citazione