Comunicazione Mixture-of-Experts efficiente e portatile

Presentiamo una libreria open-source portatile e ad alte prestazioni per la comunicazione Mixture-of-Experts (MoE) che raggiunge prestazioni 10 volte superiori rispetto alle primitive di comunicazione standard All-to-All.

La nostra implementazione presenta diverse innovazioni tecniche chiave che offrono un'efficienza di comunicazione MoE superiore:

Comunicazione avviata da GPU (IBGDA): supporta la comunicazione diretta GPU-to-NIC, riducendo significativamente la latenza grazie all'esclusione della CPU

Sovrapposizione di comunicazione e calcolo: un'architettura a kernel suddiviso con fasi di invio e ricezione separate consente al calcolo di procedere contemporaneamente ai trasferimenti di rete

Prestazioni a nodo singolo più veloci: latenza 2,5 volte inferiore rispetto alla precedente implementazione più veloce su configurazioni a nodo singolo

Prestazioni multi-nodo efficienti e portatili: la nostra implementazione raggiunge velocità fino a 10 volte superiori rispetto alla comunicazione standard all-to-all. Sebbene sia circa 2 volte più lenta delle implementazioni altamente specializzate, il nostro approccio offre una migliore portabilità tra le versioni di NVSHMEM e gli ambienti di rete (NVLink, CX-7 e EFA)

La libreria è completamente open-source e disponibile all'indirizzo https://github.com/ppl-ai/pplx-kernels.

In questo articolo esploriamo le sfide del parallelismo degli esperti nei modelli MoE su larga scala e descriviamo il nostro approccio all'invio e alla combinazione efficiente dei token in ambienti GPU distribuiti.

I modelli Mixture-of-Experts (MoE), come DeepSeek R1 e Mixtral 8x7B, migliorano i modelli densi limitando il numero di pesi attivati per ogni token. Invece di ottenere un numero elevato di parametri aumentando il numero di strati o le dimensioni degli strati lineari nel Multi-Layer Perceptron (MLP) di ogni strato del decoder, i modelli MoE sostituiscono il tradizionale MLP con esperti multipli e un router. Ad esempio, dei 671B parametri di DeepSeek R1, solo 37B vengono moltiplicati con un determinato token durante l'inferenza. Ciò riduce notevolmente la quantità di calcolo necessaria per decodificare un token e migliora la latenza rispetto a un modello denso.

I modelli MoE presentano ulteriori sfide per l'inferenza rispetto ai modelli densi. Sebbene gli esperti stessi siano piccoli strati MLP, ogni strato del decoder include un router che decide a quali esperti inviare un token, con ogni token inviato a esperti multipli. Il router è solitamente un piccolo strato lineare che produce una distribuzione di probabilità. Di solito vengono selezionati gli esperti con i punteggi top-K e l'attivazione finale viene calcolata come media ponderata, sommando gli output degli esperti moltiplicati per un peso derivato dalla distribuzione di probabilità del router.

Per ridurre al minimo la latenza nei sistemi MoE distribuiti, è possibile sfruttare il parallelismo su più dispositivi, ma ciò introduce sfide di comunicazione. Modelli come Mixtral 8x7B o Llama-70B si adattano a 8 dispositivi su un singolo nodo, beneficiando di interconnessioni NVLink veloci e a bassa latenza (fino a 900 Gbps). Tuttavia, i modelli più grandi richiedono lo sharding degli esperti su più nodi utilizzando InfiniBand (con un picco di 400 Gbps), il che introduce ulteriori sfide di latenza.

Routing dei token, elaborazione esperta e raccolta degli output in un livello mixture-of-experts.

In questo articolo esploriamo il problema del parallelismo degli esperti, descrivendo la nostra implementazione di kernel GPU che inviano in modo efficiente i token ai dispositivi che ospitano gli esperti assegnati e li raccolgono calcolando la somma ponderata degli output degli esperti. Sebbene esistano implementazioni ottimizzate rivolte a specifici hardware di rete, presentiamo un'implementazione altamente portatile che si basa su un set minimo di primitive NVSHMEM, offrendo prestazioni 10 volte più veloci della comunicazione standard all-to-all. La nostra implementazione raggiunge prestazioni all'avanguardia su configurazioni a nodo singolo pur mantenendo un'eccellente portabilità in vari ambienti di rete.

Per un'inferenza efficiente, i pesi di un modello devono essere conservati nella memoria del dispositivo, lasciando spazio sufficiente per attivazioni, cache KV e altri buffer richiesti dal forward pass attraverso il modello. I modelli più capaci superano la capacità anche delle GPU più performanti, pertanto l'inferenza deve essere distribuita su più dispositivi, che possono archiviare collettivamente i pesi del modello. In base agli schemi di sharding dei pesi, è necessario utilizzare diversi schemi di comunicazione e calcolo per scaricare il calcolo e sincronizzare i dispositivi.

Il parallelismo degli esperti (EP), illustrato nella prima figura, parallelizza solo il calcolo degli esperti. Diversi esperti vengono assegnati a diversi dispositivi, che ne detengono i pesi. Dopo il routing, i token vengono inviati al dispositivo corrispondente, con i risultati raccolti e accumulati in seguito. La complessità del routing dipende dal grado di parallelismo nelle altre parti del modello: la replica di altri strati potrebbe eliminare del tutto la necessità di routing, poiché ogni rango può selezionare i token da una tabella di routing replicata localmente. Tuttavia, se solo uno dei ranghi esegue il routing, è richiesto un broadcast per inviare token, indici e pesi ai rispettivi esperti. Infine, un all-gather o un broadcast all-to-all sincronizza i token di output con qualsiasi rango prosegua l'esecuzione del modello. Tale implementazione è relativamente semplice poiché torch espone già le primitive richieste, sebbene si potrebbero ottenere alcuni vantaggi fondendo l'accumulo con l'operazione di gather che raccoglie i token dagli esperti.

Combinazione del parallelismo degli esperti con il parallelismo tensoriale su due GPU.

Il parallelismo basato solo sugli esperti non scala idealmente, poiché i nodi in un cluster potrebbero essere inattivi mentre il modello esegue strati non basati su esperti, come attenzione, norm e campionamento. Tuttavia, il calcolo di questi strati, principalmente l'attenzione, può beneficiare anche del parallelismo dei tensori (TP). La maggior parte dei modelli si affida alla multi-head attention, il che significa che le teste di attenzione e le relative proiezioni Q, K e V possono anche essere suddivise tra i dispositivi, replicando o raccogliendo le porzioni tra vari strati. Se l'attenzione è distribuita su tutti i dispositivi, un all-gather può sincronizzare le attivazioni, consentendo di replicare il routing, richiedendo primitive di sincronizzazione simili al caso del parallelismo degli esperti per un'implementazione efficiente. Tuttavia, ci sono limiti al parallelismo a questo livello, poiché ridurre il numero di teste di attenzione al di sotto di una certa soglia produrrà rendimenti decrescenti.

Combinazione del parallelismo degli esperti con il parallelismo dei dati per richieste separate su due GPU.

Per utilizzare al meglio tutti i dispositivi e supportare un grado molto elevato di parallelismo degli esperti, fino a 128 o 256 GPU, è necessario il parallelismo dei dati (DP). Con questo schema, i dispositivi su cui è suddiviso un modello vengono raggruppati per gestire le richieste simultaneamente, calcolando l'attenzione e mantenendo le cache KV suddivise nel loro gruppo locale. Istanze multiple di questi gruppi paralleli collaborano alla valutazione degli esperti, ognuna ospitando un sottoinsieme diverso dei pesi degli esperti. In base al numero di teste di attenzione, un gruppo può solitamente estendersi fino alla dimensione di un intero nodo, poiché la comunicazione intra-nodo è più veloce. Ad esempio, nella figura sopra, un rango DP gestisce autonomamente due richieste, A e B, gestendo attenzione, norm e qualsiasi altra operazione per le richieste. L'altro rango DP elabora una richiesta distinta C. Tuttavia, il primo nodo ospita metà degli esperti, mentre l'altro nodo l'altra metà, quindi dopo il routing i token da A e B potrebbero essere inviati al secondo nodo e viceversa. Ciò porta a un problema di comunicazione sparsa: ogni dispositivo potrebbe inviare un numero diverso di token a qualsiasi altro rango di destinazione. Le primitive esistenti da torch, principalmente all_to_all, non sono particolarmente adatte, poiché potrebbero richiedere una forma di padding o sincronizzazione GPU-to-CPU e broadcast di metadati. Per implementare la comunicazione in modo efficace, sono necessari kernel personalizzati per elaborare le informazioni di routing e avviare la comunicazione dai dispositivi senza l'intervento della CPU. Dopo il routing, un kernel di invio deve inviare i token ai ranghi a cui sono stati instradati, mentre sul lato di combinazione devono essere raccolte le attivazioni appartenenti alle richieste nel gruppo DP corrente. Inoltre, il carico di lavoro deve essere bilanciato all'interno di un gruppo DP che può avere più dispositivi, al fine di ridurre al minimo il broadcast e replicare correttamente i token su tutti i ranghi per gli strati successivi.

NVSHMEM è un'implementazione OpenSHMEM specifica per NVIDIA, che fornisce strutture di comunicazione inter-dispositivo portatili che astraggono la complessità dell'hardware sottostante. L'API può esprimere letture e scritture da dispositivo a dispositivo, che vengono mappate alle primitive dei singoli livelli di trasporto. Sono supportati sia NVLink che RDMA, garantendo un grado di portabilità e consentendo ai kernel NVSHMEM di operare all'interno di un nodo o attraverso un cluster di più nodi. Nei nostri kernel, utilizziamo scritture di memoria remota e operazioni di segnalazione per trasferire dati e sincronizzare i dispositivi partecipanti.

Le operazioni NVSHMEM sono costruite attorno al concetto di memoria simmetrica: operano su buffer che sono stati allocati su tutti i dispositivi che partecipano alla comunicazione inter-dispositivo. Un dispositivo locale può utilizzare un indirizzo derivato all'interno del proprio buffer locale insieme all'indice del rango di destinazione per specificare la destinazione di un'operazione. La figura seguente illustra questo concetto: entrambe le GPU allocano buffer simmetrici della stessa dimensione, mantenendo i puntatori src ad essi. Il primo rango vuole inviare 3 interi al secondo, posizionandoli all'inizio del buffer: nvshmem_int_put_nbi deriva l'indirizzo di inizio dal buffer locale, specificando il dispositivo di destinazione. Il secondo rango deriva un offset dal proprio buffer, inviando un elemento al primo dispositivo, spostato di uno. Mentre gli indirizzi di destinazione devono essere sempre buffer simmetrici allocati utilizzando nvshmem_alloc, i buffer di origine possono essere regioni arbitrarie della memoria del dispositivo, a condizione che siano pre-registrati con NVSHMEM.

Sebbene NVSHMEM fornisca una vasta gamma di primitive, i nostri kernel si basano solo su 3 funzioni, costruendo su di esse tutta la sincronizzazione e il fencing.

Memoria simmetrica e scritture remote tra due GPU usando NVSHMEM.

nvshmemx_putmem_signal_nbi_warp: trasferisce un blocco di dati da un dispositivo a un altro, impostando anche un flag sul dispositivo remoto. L'operazione imposta (NVSHMEM_SIGNAL_SET) o incrementa (NVSHMEM_SIGNAL_ADD) una posizione a 64 bit. Il flag viene aggiornato dopo che l'intero blocco di memoria è stato trasferito. Se il dispositivo remoto osserva una modifica nel flag, può accedere in sicurezza al buffer nella propria memoria. Questa funzione è utile per coalescere il trasferimento dei dati e la sincronizzazione.

nvshmemx_signal_op opera su una singola posizione di memoria, solitamente un flag a 64 bit, impostandolo o incrementandolo in modo atomico. È utile per inviare metadati e sincronizzare i dispositivi.

nvshmem_uint64_wait_until viene utilizzato sul lato ricevente di un segnale, per eseguire il polling di un flag finché il remoto non lo aggiorna.

Garantire qualsiasi forma di ordinamento tra le operazioni tramite nvshmem_quiet o nvshmem_fence è costoso, pertanto la nostra implementazione dei kernel evita l'uso di qualsiasi altra primitiva NVSHMEM e riduce al minimo le dipendenze diverse dalle singole coppie invio-ricezione. Quando inviamo dati, utilizziamo sempre la versione non bloccante delle funzioni, senza attendere che i dati vengano effettivamente inviati dal rango locale. Altra sincronizzazione implicita tra il gruppo di kernel garantisce che questi buffer non vengano sovrascritti in modo distruttivo mentre le NIC necessitano dell'accesso ad essi.

NVSHMEM utilizza GPUDirect RDMA in modo che la NIC possa leggere e scrivere direttamente nella memoria della GPU per lo scambio di dati senza coinvolgere le CPU. Inoltre, sulle NIC ConnectX, NVSHMEM supporta la comunicazione avviata da GPU (nota anche come Infiniband GPUDirect Async o IBGDA), che consente alla GPU di pubblicare operazioni RDMA direttamente sulla NIC e di eseguire il polling direttamente sulla coda di completamento della NIC. Sulle piattaforme che non supportano la comunicazione avviata da GPU, NVSHMEM genera un thread proxy sulla CPU host per avviare la comunicazione ed eseguire il polling del completamento per conto della GPU. Il programma NVSHMEM è portatile indipendentemente dal fatto che la comunicazione avviata da GPU sia supportata o meno. Tuttavia, la comunicazione avviata da GPU riduce significativamente la latenza perché aggira completamente il passaggio attraverso la CPU.

Implementiamo la comunicazione MoE tramite una coppia di kernel di invio (dispatch) e combinazione (combine). I kernel di dispatch sono responsabili della lettura dei token e delle informazioni di routing su ogni rango, inviandoli agli esperti appropriati. I kernel di combinazione raccolgono le attivazioni prodotte dagli esperti e le inviano ai ranghi di origine, calcolando anche la media ponderata dagli esperti selezionati in base ai pesi calcolati dal router. I kernel sono ulteriormente suddivisi in un componente di invio e uno di ricezione, per consentire la sovrapposizione dei trasferimenti di dati con il calcolo. I kernel di invio sono tutti non bloccanti e non sincronizzanti: inviano semplicemente tutte le scritture ai remoti. Dall'altro lato, i kernel di combinazione leggono solo dalla memoria, attendendo che tutti i dati richiesti siano stati trasferiti. Dopo aver inviato il lavoro alle NIC, mentre i dati vengono trasferiti in modo asincrono sul cavo, le GPU possono eseguire altri lavori utili localmente, come l'applicazione di esperti condivisi o il calcolo dell'attenzione.

Ogni coppia di kernel ha i propri buffer di memoria simmetrici, con ogni rango che alloca spazio privato per ricevere dati da tutti gli altri ranghi. Il numero massimo di token che possono essere inviati da un gruppo DP è fisso tra i ranghi, il che imposta anche il limite superiore che ogni rango può ricevere per ogni esperto locale da ciascun rango DP. Ciò consente ai ranghi mittenti di derivare un indirizzo univoco sul rango di destinazione a cui scrivere, senza richiedere alcuna sincronizzazione tra loro. Dopo che i dati sono stati ricevuti, i kernel mescolano i token per disporli in memoria nel formato ottimale richiesto dai kernel riceventi. Sebbene i buffer abbiano dimensioni considerevoli, vengono riutilizzati in tutti gli strati sequenziali di un modello.

Kernel separati di invio e ricezione per le operazioni di dispatch e combine.

L'unica forma di sincronizzazione globale, come illustrato nella figura sopra, è implementata nel kernel combine-receive. Una volta che i dati vengono inviati dal dispatch, un rango restituisce e continua l'esecuzione quando riceve tutti i token richiesti dagli altri ranghi ed entra nei kernel combine-send non appena vengono elaborati. La barriera in combine-receive garantisce che nessun rango possa procedere in anticipo e avviare dispatch-send mentre qualsiasi altro rango è ancora in attesa di ricevere dati, causando potenzialmente una sovrapposizione distruttiva. La sincronizzazione viene eseguita nel kernel di combinazione per semplicità; come ottimizzazione, la stessa barriera potrebbe essere implementata tra i kernel dispatch-receive e combine-send, riducendo ulteriormente la latenza consentendo alle NIC più tempo per completare i trasferimenti mentre la GPU esegue altri kernel.

Entrambi i kernel sono suddivisi tra tutti gli SM disponibili dei dispositivi: mentre i kernel dispatch-send e combine-receive devono parallelizzarsi su un conteggio massimo di token per rango (max_m), i kernel dispatch-receive e combine-send devono gestire quel numero di token da ogni rango ed esperto locale, per un massimo di max_m * (num_experts // EP) * (EP // TP).

Il kernel di dispatch è responsabile dell'invio di token e fattori di scala a blocchi ai peer che ospitano gli esperti a cui i token sono stati instradati. Il lato mittente si affida alla specializzazione warp per parallelizzare due attività: aggregare le informazioni di routing per determinare quanti token vengono inviati da ogni rango e impacchettare i token in messaggi da inviare al remoto. Il lato ricevente attende prima che tutti i conteggi dei token vengano ricevuti, quindi legge e disimballa i payload dalla memoria condivisa nei tensori che verranno passati ai kernel Grouped GEMM a valle. Memorizza inoltre informazioni in buffer condivisi con i kernel di combinazione, per indicare dove ogni token debba essere rimandato indietro. Queste informazioni sono necessarie poiché il kernel di ricezione mescola i token in buffer contigui in un ordine non deterministico.

Dispatch dei token e segnalazione del loro arrivo tra i rank sorgente e destinazione.

Nella parte del mittente, un singolo warp di 32 thread su ogni blocco è responsabile della lettura delle informazioni di routing e del conteggio del numero di token che ogni esperto di destinazione dovrebbe ricevere da questo rango, inviando il conteggio più uno utilizzando nvshmem_signal_op. Il conteggio viene incrementato di uno, poiché la transizione da zero a non zero sull'estremità remota indica la ricezione dei conteggi da un rango di origine. In parallelo, i restanti warp cooperano per copiare i token nella memoria simmetrica tra tutti i blocchi in parallelo, impacchettando attivazioni, fattori di scala e il loro indice sul rango locale in un blocco contiguo. L'indice è richiesto dal mittente combine per determinare l'indirizzo in cui verrà scritto il token. Successivamente, dopo aver garantito che tutti i dati siano stati copiati tramite una barriera tra i gruppi warp, i warp operano ancora una volta in modo indipendente, inviando ciascuno lo stesso buffer a un esperto diverso in parallelo. I token vengono inviati utilizzando nvshmemx_putmem_signal_nbi_warp, che aumenta anche in modo atomico il conteggio dei token inviati dal rango locale sul dispositivo remoto. All'interno di un gruppo DP, poiché ogni rango possiede una replica del token da inviare, il dispatch è bilanciato in modo uniforme, con ogni dispositivo che invia un sottoinsieme dei token.

Sul lato di ricezione, tutti i kernel si sincronizzano prima con i ranghi mittenti attendendo che i conteggi totali dei token vengano inviati. Successivamente, attendono tutti che i conteggi dei token inviati incrementati in modo atomico si stabilizzino sui conteggi totali, indicando che anche tutti i payload dai ranghi di origine sono stati inviati, grazie alla semantica della chiamata putmem. I kernel eseguono il polling sui conteggi utilizzando nvshmem_uint64_wait_until, parallelizzando l'operazione tra tutti i blocchi e i thread. Successivamente, una barriera inter-blocco garantisce che nessun blocco legga i buffer a meno che tutti i dati non siano stati ricevuti correttamente. Distribuiti tra i blocchi e sincronizzati tramite un contatore incrementato in modo atomico, i token vengono copiati dalla memoria simmetrica in normali buffer contigui che diventano gli input per i kernel a valle che implementano gli esperti. Il rango di origine, l'indice dell'esperto e l'indice del token vengono archiviati separatamente, individuando esattamente la posizione in cui il kernel di combinazione deve inviare le attivazioni. Anche se i token all'interno di un gruppo DP vengono inviati da dispositivi diversi, sono tutti raggruppati insieme per essere passati all'esperto corrispondente.

I kernel di combinazione sono ancora una volta suddivisi in metà di invio e ricezione: i mittenti copiano le attivazioni a 16 bit non quantizzate sui dispositivi remoti, con i kernel di ricezione che attendono l'invio dei dati e calcolano localmente la media ponderata di tutti i contributi degli esperti. Inoltre, fungono anche da barriera per sincronizzare la sequenza dispatch-combine: ogni rango imposta un flag su ogni peer all'ingresso dei kernel di invio, con i kernel di ricezione a cui non è consentito tornare a meno che non osservino il flag impostato. La latenza di sincronizzazione è minima, poiché si sovrappone alla comunicazione e al calcolo effettivi.

Combinazione degli output degli esperti nei token su ciascun rank di parallelismo dei dati.

Sul lato del mittente, i kernel attraversano l'elenco di token assegnati a tutti gli esperti locali in parallelo, scrivendoli in un buffer sul rango di destinazione. Il rango di destinazione, l'indice dell'esperto e l'indice del token vengono letti dai buffer per-token popolati dai kernel scatter al ricevimento dei token. Ogni mittente ha la propria regione di memoria privata per esperto in cui scrivere, come indicato nella figura sopra, evitando la necessità di sincronizzazione. Analogamente al dispatch, combine incrementa in modo atomico i contatori per token sul rango di destinazione per indicare la ricezione dei dati: quando il contatore corrisponde al numero di esperti a cui è stato inviato un token, è possibile accedere ai contenuti del token.

Nel kernel di ricezione, l'elenco di token viene attraversato in parallelo tra blocchi multipli, in attesa che arrivino i loro contenuti tramite il polling del flag impostato dall'operazione di segnalazione. All'arrivo, i payload vengono letti dai buffer privati, con la tabella di routing che indica da quale buffer leggere e quale peso assegnare a ciascun esperto. I risultati vengono quindi scritti su tensori allocati esternamente, con il kernel che termina l'esecuzione una volta che tutti i dispositivi hanno superato la barriera. In un gruppo DP, tutti i ranghi ricevono una copia di ogni attivazione dell'esperto per calcolare le proprie repliche.

Valutiamo i nostri kernel su un cluster di 128 GPU distribuite su 16 nodi collegati tramite InfiniBand e NIC CX-7. Proviamo sia GPUDirect Async (IBGDA) che Reliable Connection (RC) con un proxy CPU. Li confrontiamo con l'implementazione altamente ottimizzata di DeepSeek, così come con le primitive dense fornite da PyTorch (tramite NCCL) o NVSHMEM.

Latenza di comunicazione per 128 token per GPU in una configurazione EP128 multi-node.

Sebbene nell'operazione all-to-all densa di NVSHMEM le prestazioni di IBRC e IBGDA siano simili (6378 µs contro 6180 µs), IBGDA è significativamente più veloce con i kernel sparsi. Sommando Dispatch e Combine, IBGDA utilizza 902 µs mentre IBRC ne richiede 3223 µs: un miglioramento della latenza di 3,6 volte. Mentre all-to-all è limitato dalla larghezza di banda, i kernel sparsi trasmettono ordini di grandezza inferiori di dati, essendo limitati dalla latenza. Attivando i trasferimenti di rete direttamente da una GPU, senza richiedere un proxy CPU per coordinare la GPU e la NIC, la latenza end-to-end viene ridotta significativamente.

Mentre i nostri kernel portatili sono circa 2 volte più lenti dei kernel DeepSeek altamente ottimizzati, migliorano la latenza di 10 volte rispetto ai kernel densi (902 µs contro 9944 µs).

Inoltre, la suddivisione in componenti mittente-ricevente consente anche di nascondere parte della latenza, a differenza delle primitive della libreria.

Latenza di comunicazione per 128 token per GPU in una configurazione EP8 single-node.

Sul nodo singolo (EP8), NVSHMEM utilizza NVLINK per il trasporto, offrendo una latenza inferiore e un throughput più elevato rispetto al networking inter-nodo. I nostri kernel portatili sono circa 2,5 volte più veloci di DeepEP su nodo singolo (186 µs contro 481 µs).

I kernel qui descritti superano le primitive integrate dei framework ML e offrono prestazioni decenti senza sovra-specializzarsi per particolari trasporti inter-dispositivo, come NVLink o InfiniBand. Oltre alle opportunità già menzionate, sono ottenibili ulteriori guadagni in termini di prestazioni sostituendo le primitive di comunicazione con versioni più specializzate. Ad esempio, su NVLink, l'uso della memoria simmetrica potrebbe essere sostituito con buffer condivisi tra i dispositivi, eliminando alcune copie e consentendo schemi di sincronizzazione a grana più fine, tra i singoli token anziché batch di token. Su InfiniBand, un'implementazione potrebbe accedere direttamente alle coppie di code sottostanti, eliminando la necessità di trasferire sia i conteggi dei token attesi che quelli inviati e svuotando i token dalle code nei buffer di output man mano che arrivano. Tuttavia, tali implementazioni comportano il costo della portabilità, mentre le nostre implementazioni consentono di valutare cluster con configurazioni specifiche prima di impegnare sforzi di sviluppo per l'ottimizzazione per un particolare hardware.

Abbiamo presentato una libreria portatile ad alte prestazioni per la comunicazione MoE che raggiunge prestazioni 10 volte superiori rispetto alla comunicazione standard all-to-all, mantenendo la compatibilità tra diverse configurazioni hardware. Nelle implementazioni su nodo singolo con NVLink, la nostra soluzione dimostra una latenza 2,5 volte inferiore rispetto alle precedenti implementazioni.

Il nostro approccio bilancia le prestazioni con la portabilità attraverso innovazioni chiave, tra cui il supporto alla comunicazione avviata da GPU, un'architettura a kernel suddiviso che abilita la sovrapposizione calcolo-comunicazione e l'invio efficiente dei token utilizzando primitive NVSHMEM minime. Sebbene sia circa 2 volte più lenta delle implementazioni altamente specializzate su configurazioni multi-nodo, la nostra libreria offre una flessibilità superiore in vari ambienti di rete (NVLink, CX-7 e EFA).

Poiché i modelli MoE continuano a scalare, strategie di comunicazione efficienti come la nostra diventeranno sempre più importanti per l'implementazione pratica. La nostra implementazione completamente open-source è disponibile all'indirizzo https://github.com/ppl-ai/pplx-kernels.

DeepEP

Rapporto tecnico DeepSeek-V3

NVIDIA NVSHMEM

BibTeX

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

Introduzione

Sharding e parallelismo

NVSHMEM

Kernel portatili

Dispatch

Combine

Benchmark

Ulteriori miglioramenti

Conclusione

Riferimenti

Citazione