Fast MoE Inference via Predictive Prefetching and Expert Replication

Riepilogo Tecnico: Inferenza MoE Veloce tramite Prefetching Predittivo e Replicazione degli Esperti

Enunciato del Problema
L'architettura Mixture of Experts (MoE) è diventata una pietra angolare per il ridimensionamento dei Large Language Models (LLM), consentendo conteggi di parametri massicci senza aumenti proporzionali del sovraccarico computazionale. Tuttavia, l'inferenza MoE presenta inefficienze critiche che ostacolano il dispiegamento pratico:

1. Utilizzo Subottimale della GPU: A causa della sparsità intrinseca dell'attivazione degli esperti, solo un piccolo sottoinsieme di esperti è attivo in qualsiasi momento, lasciando la maggior parte delle risorse GPU inattive.
2. Squilibrio del Carico e Latenza: Il routing dinamico spesso comporta una distribuzione disuguale del carico, dove più token fanno coda per lo stesso esperto fisico, causando arresti dei token e aumento della latenza.
3. Sovraccarico del Routing: I meccanismi di routing tradizionali basati su softmax introducono complessità computazionale e sovraccarico di comunicazione (routing all-to-all) che non si scalano favorevolmente con l'aumento del numero di esperti.

Le tecniche di ottimizzazione esistenti, come la distillazione o i miglioramenti a livello di kernel, offrono spesso miglioramenti limitati o compromettono l'accuratezza del modello.

Metodologia: MoE-MPMC
Gli autori propongono MoE-MPMC (Mixture of Experts con Multiple Producers Multiple Consumers), un framework che affronta queste inefficienze tramite prefetching predittivo e replicazione dinamica degli esperti. Il sistema opera utilizzando due thread concorrenti:

1. Thread di Inferenza:

   • Recupera una tabella hash pre-calcolata contenente assegnazioni degli esperti e conteggi di replicazione per il batch corrente.
   • Carica dinamicamente nella memoria GPU solo gli esperti predetti attivi, scaricando gli esperti inutilizzati nella memoria CPU.
   • Esegue l'inferenza utilizzando un layer MoE modificato in cui gli esperti attivi sono replicati fisicamente sulla GPU. Ciò consente a più token che richiedono lo stesso esperto di essere elaborati concorrentemente da distinte repliche fisiche, eliminando i ritardi di coda.

2. Thread di Costruzione dell'Hash:

   • Esegue in parallelo al thread di inferenza, elaborando il prossimo batch di token.
   • Utilizza una Simple Recurrent Unit (SRU) con attenzione sparsa per prevedere quali esperti verranno attivati per il batch imminente.
   • Sostituisce il predittore basato su LSTM utilizzato nei lavori precedenti (SiDA-MoE) con SRU per ottenere velocità di addestramento e inferenza più elevate grazie a ridotte interdipendenze e meccanismi di gating più semplici.
   • Calcola i conteggi di replicazione per ogni esperto previsto in base alla domanda dei token e aggiorna la tabella hash, che viene quindi messa in coda per il thread di inferenza.

Contributi Chiave
Il documento delinea tre contributi principali:

1. Predizione Basata su SRU: L'introduzione della Simple Recurrent Unit (SRU) per prevedere rapidamente le selezioni degli esperti per ogni token. Ciò sostituisce i predittori LSTM più lenti, accelerando significativamente la fase di predizione mantenendo l'accuratezza.
2. Replicazione Fisica Profonda: Una strategia per replicare fisicamente gli esperti previsti sulla GPU in base alla domanda dei token. Ciò garantisce che quasi nessun token sperimenti ritardi di attesa durante l'elaborazione del batch, massimizzando così il parallelismo e l'utilizzo della GPU.
3. Limitazione della Capacità: Un meccanismo pratico per limitare i conteggi di replicazione degli esperti in base a vincoli di memoria GPU predefiniti. Ciò garantisce che la replicazione sia ottimizzata entro i limiti hardware per raggiungere un utilizzo delle risorse quasi ottimale senza causare overflow.

Risultati Sperimentali
Gli autori hanno valutato MoE-MPMC su modelli SwitchTransformers su larga scala (Switch-base-8, 64, 128 e 256) utilizzando dataset come SST2 e MRPC. Le scoperte chiave includono:

• Utilizzo della GPU: Il metodo raggiunge un utilizzo quasi completo della GPU (fino a ~100% in alcune configurazioni), un miglioramento significativo rispetto ai SwitchTransformers di base (che mostravano un utilizzo basso come 1,3% per 128 esperti) e all'approccio SiDA-MoE.
• Velocità di Inferenza: L'approccio produce un miglioramento fino a 3 volte nella velocità di inferenza rispetto alle architetture di base.
• Accuratezza: Il metodo preserva circa il 90-95% delle prestazioni delle architetture di base, con un'accuratezza mantenuta meglio rispetto a SiDA-MoE grazie alle capacità di predizione superiori di SRU rispetto a LSTM.
• Throughput: Il throughput aumenta significativamente all'aumentare del numero di esperti, guidato dalla capacità di elaborare più token in parallelo tramite replicazione degli esperti.

Significato e Affermazioni
Il documento afferma che MoE-MPMC rappresenta un passo significativo verso il dispiegamento pratico ed economicamente sostenibile di modelli linguistici su larga scala basati su MoE. Spostandosi dal bilanciamento del carico reattivo a una strategia proattiva di prefetching predittivo e replicazione, il metodo:

• Mitiga i colli di bottiglia nell'accesso alla memoria e gli arresti dei token.
• Ottimizza l'utilizzo delle risorse GPU senza alterare il grafo computazionale MoE sottostante.
• Dimostra che la previsione delle selezioni degli esperti è fattibile ed efficace quando il fine-tuning e il pre-addestramento avvengono in concomitanza con i compiti di predizione dei token.

Gli autori posizionano questo lavoro come un avanzamento rispetto ai metodi reattivi (come SiDA-MoE) anticipando le future esigenze e regolando dinamicamente l'allocazione delle risorse, affrontando così le tipiche insidie della memoria GPU inattiva e degli squilibri di carico nelle architetture sparse.