Beyond Task-Agnostic: Task-Aware Grouping for Communication-Efficient Multi-Task MoE Inference

Riepilogo Tecnico: Oltre l'Agnosticismo rispetto ai Task: Raggruppamento Consapevole dei Task per l'Inferenza Efficiente nella Comunicazione di MoE Multi-Task

1. Definizione del Problema

I modelli Mixture-of-Experts (MoE) a attivazione sparsa scalano efficientemente la capacità tramite il calcolo condizionale, ma affrontano sfide sistemiche significative durante l'inferenza distribuita, specificamente un eccessivo overhead di comunicazione cross-GPU e un grave squilibrio del carico. Le strategie esistenti di posizionamento degli esperti, che tentano di mitigare questi costi co-localizzando gli esperti co-attivati frequentemente, derivano tipicamente un piano di distribuzione globale unico da tracce di routing task-agnostic (aggregati globalmente).

Gli autori identificano un difetto critico in questo approccio: la co-attivazione degli esperti è fortemente condizionata dal task. Come dimostrato dall'analisi del paper (Figura 1), coppie di esperti strettamente accoppiate in una famiglia di task (es. generazione di codice) sono spesso non correlate o addirittura in competizione in un'altra (es. ragionamento matematico). Mediare questi pattern eterogenei in un singolo grafo di affinità globale fa collassare i gruppi di esperti specifici per task, lasciando inesploitata la località di comunicazione e fallendo nell'adattarsi ai cambiamenti di distribuzione nei carichi di lavoro online.

2. Metodologia

Il framework proposto, Task-Aware Coactivation Grouping (TACG) abbinato a Generic Expert Shared Replication (GESR), affronta queste problematiche attraverso un processo di distribuzione in due fasi:

A. Task-Aware Coactivation Grouping (TACG)
TACG converte le statistiche di co-attivazione condizionate dal task in un piano di posizionamento fattibile rispetto alla capacità:

1. Scoring della Preferenza della Famiglia di Task: Utilizzando tracce di calibrazione, il sistema calcola i vettori di utilizzo e le matrici di co-attivazione per ogni famiglia di task. Calcola un "punteggio di preferenza soft" ($s^{(f)}_{l,e}$) per ogni esperto rispetto a specifiche famiglie di task, normalizzando questi punteggi per creare una distribuzione di probabilità ($p^{(f)}_{l,e}$) che rappresenti l'affinità di un esperto per una famiglia.
2. Grafo di Co-attivazione Modulato dal Task: Invece di utilizzare un grafo di co-attivazione aggregato, TACG costruisce un grafo modulato dal task ($\tilde{A}_l$). Questo grafo ripesa gli archi di co-attivazione standard tramite il prodotto delle preferenze familiari degli esperti ($p^{(f)}_{l,e} p^{(f)}_{l,e'}$). Ciò assicura che gli archi tra esperti che si co-attivano e condividono una forte preferenza familiare siano amplificati, mentre gli archi cross-family siano attenuati.
3. Raggruppamento Fattibile rispetto alla Capacità: Il framework partiziona gli esperti in $M$ gruppi (uno per GPU) massimizzando l'affinità intra-gruppo sul grafo modulato dal task sotto rigorosi vincoli di capacità. Questo viene risolto tramite un approccio di clustering spettrale in due fasi seguito da un passaggio di "riparazione della capacità" locale per garantire che i limiti di carico esatti della GPU siano rispettati.

B. Generic Expert Shared Replication (GESR)
Per gestire lo skew del carico causato dagli esperti "generici" (quelli attivi in tutte le famiglie e inclini a diventare hotspot), GESR introduce un meccanismo di accompagnamento leggero:

1. Identificazione: Identifica gli esperti generici utilizzando un punteggio che premia la centralità globale e la coerenza cross-family, penalizzando al contempo la forte deviazione specifica del task.
2. Replicazione: Questi esperti vengono replicati su un piccolo set di GPU secondarie scelte in base alla compatibilità strutturale (affinità di co-attivazione).
3. Guardia Online: Al momento del serving, una regola di selezione consapevole di località e carico sceglie la migliore GPU candidata per un esperto generico, utilizzando segnali di carico con decadimento esponenziale per prevenire hotspot senza alterare il router MoE.

3. Contributi Chiave

• Insight: Il paper stabilisce che la co-attivazione degli esperti nel serving MoE multi-task è fortemente condizionata dal task, sostenendo che un posizionamento efficace deve essere guidato da una co-attivazione consapevole del task piuttosto che da un singolo grafo di affinità mediato globalmente.
• TACG: Una strategia consapevole del task che ripesa il grafo di co-attivazione utilizzando le preferenze esperto-famiglia, producendo un posizionamento che generalizza strettamente il raggruppamento task-agnostic.
• GESR: Un modulo di replicazione leggero che replica selettivamente gli esperti generici cross-family e esegue una selezione online consapevole del carico per mantenere l'equilibrio sotto i cambiamenti del carico di lavoro.
• Validazione Empirica: Valutazione estesa su tre modelli MoE open-source rappresentativi (DeepSeek-MoE-16B, Qwen1.5-MoE-A2.7B, Moonlight-16B-A3B) che dimostrano miglioramenti costanti rispetto a forti baseline.

4. Risultati Sperimentali

Il framework è stato valutato su una configurazione expert-parallel a 16 GPU su tre modelli con quattro famiglie di task (Code, Query, Math, Complex Reasoning).

• Riduzione della Comunicazione: Il framework completo (TACG + GESR) ha ridotto il costo medio di comunicazione per token del 31,39% rispetto alla baseline, superando la più forte strategia di posizionamento precedente (ExFlow) del 12,58%.
• Bilanciamento del Carico: Il metodo ha mantenuto un alto indice di equità di Jain medio di 0,9975 e una bassa violazione media del carico massimo (MaxVio) di 0,0736, eguagliando o superando il bilanciamento del carico delle baseline pur riducendo significativamente la comunicazione.
• Robustezza: Sotto severi spostamenti della distribuzione dei task (carichi di lavoro single-task Out-of-Distribution), l'approccio consapevole del task ha mantenuto riduzioni della comunicazione stabili, mentre le varianti task-agnostic hanno subito ampi sbalzi di performance e formazione di hotspot (con MaxVio che aumentava di circa 7 volte sui task di Code).
• Ablazione: La rimozione del segnale di co-attivazione modulato dal task ha ridotto i guadagni di comunicazione dal 31,39% al 16,51%, confermando che il condizionamento del task è il fattore dominante. È stato dimostrato che GESR è essenziale per controllare le violazioni del carico (MaxVio) senza sacrificare l'efficienza della comunicazione.

5. Significato e Rivendicazioni

Il paper afferma che la sua importanza risiede nello spostamento del paradigma del posizionamento degli esperti MoE da una visione statica e task-agnostic a una prospettiva dinamica e consapevole del task. Modellando esplicitamente l'eterogeneità della co-attivazione degli esperti tra diverse famiglie di task, il framework sblocca una località di comunicazione che è invisibile ai metodi di aggregazione globale.

Gli autori asseriscono che questo approccio fornisce un miglioramento rigoroso sulla frontiera del bilanciamento comunicazione-carico. A differenza di metodi precedenti che spesso sacrificano il bilanciamento del carico per i guadagni di comunicazione (o viceversa), TACG + GESR ottiene una sostanziale riduzione della comunicazione preservando al contempo un bilanciamento del carico quasi perfetto. Il lavoro evidenzia che il serving multi-task efficace richiede strategie di deployment che rispettino la struttura intrinseca della collaborazione tra esperti specifica per task, anziché mediarla.

Il paper rimane modesto riguardo al proprio ambito, notando che si concentra sul posizionamento offline statico e lascia come lavoro futuro il raggruppamento dinamico all'interno delle finestre di richiesta, il co-design della memoria (prefetching/caching) e gli esperti specializzati in modelli partizionati per dominio.