Riepilogo Tecnico: DisagMoE

1. Enunciato del Problema

Le architetture Mixture-of-Experts (MoE) su larga scala abilitano l'addestramento di Modelli Linguistici (LLM) con trilioni di parametri attivando in modo sparso un sottoinsieme di esperti. Tuttavia, l'addestramento di questi modelli incontra gravi colli di bottiglia a causa della comunicazione all-to-all richiesta per il parallelismo degli esperti (EP). Man mano che i modelli si scalano su più nodi GPU, la limitata larghezza di banda di rete inter-nodo (es. InfiniBand) esacerba la latenza di comunicazione, che può consumare fino al 50% del tempo di un passo di addestramento.

Le strategie di ottimizzazione esistenti si concentrano principalmente sull'intersezione tra comunicazione e calcolo a livello di operatore (es. frammentazione dei calcoli FFN) o sull'uso del pipelining di micro-lotti. Tuttavia, questi approcci soffrono di due limitazioni fondamentali:

1. Rallentamenti Residui: Gli squilibri intrinseci nei rapporti calcolo-comunicazione tra gli strati di attenzione e FFN impediscono un sovrapposizione perfetta. Il calcolo dell'attenzione scala quadraticamente con la lunghezza della sequenza, mentre FFN e comunicazione scalano linearmente. Ciò crea "code di comunicazione" che non possono essere nascoste, specialmente all'aumentare delle lunghezze delle sequenze.
2. Inefficienza delle Risorse: Gli approcci disaggregati precedenti (che separano attenzione e FFN su hardware diverso) richiedono spesso che ogni gruppo mantenga l'intero set di parametri del modello, portando a errori di Out-of-Memory (OOM) su larga scala, oppure non riescono ad adattare l'allocazione delle risorse alle specifiche caratteristiche calcolo-comunicazione di ciascun componente.

2. Metodologia: DisagMoE

DisagMoE è un sistema di addestramento progettato per mitigare i colli di bottiglia all-to-all ottimizzando congiuntamente il posizionamento del modello e la pianificazione attraverso un'architettura disaggregata.

2.1 Posizionamento Disaggregato dei Componenti

A differenza del Parallelismo degli Esperti tradizionale, dove gli strati di attenzione e FFN coesistono sullo stesso gruppo di GPU, DisagMoE partiziona il modello per tipo di componente:

• Gruppi di Attenzione (A-Workers): Questi gruppi contengono gli strati di auto-attenzione da più strati transformer. Replicano i componenti densi e distribuiscono i dati di addestramento tramite Parallelismo dei Dati (DP).
• Gruppi FFN (F-Workers): Questi gruppi contengono gli strati FFN (esperti) da più strati transformer. Gli esperti sono frammentati (sharded) tra le GPU all'interno del gruppo tramite Parallelismo degli Esperti (EP).
• Assegnazione degli Strati: Gli strati vengono assegnati ai gruppi in modo intercalato (es. il Gruppo 0 gestisce gli strati 0, $p$, $2p$; il Gruppo 1 gestisce gli strati 1, $p+1$, ecc.), massimizzando l'utilizzo dei dispositivi sotto i vincoli di memoria.

2.2 AF-Pipe: Pipeline Multi-Stadio con Sovrapposizione

DisagMoE introduce AF-Pipe, una pianificazione di pipeline che tratta la comunicazione come una fase di primo piano piuttosto che come un effetto collaterale.

• Comunicazione Many-to-Many (M2N): Invece degli scambi all-to-all standard all'interno di un singolo gruppo di nodi, AF-Pipe stabilisce scambi di token unidirezionali, many-to-many, tra i gruppi di Attenzione e FFN.
• Allineamento delle Fasi: La pipeline allinea i confini tra i gruppi, fondendo le tradizionali operazioni di "dispatch" e "combine" in un'unica fase di comunicazione unificata. Ciò riduce l'overhead di comunicazione di circa $1/k$ (dove $k$ è il numero di top-$k$ esperti) ed elimina i trasferimenti di dati ridondanti.
• Esecuzione Asincrona: Il sistema utilizza tre flussi coordinati (forward, backward e comunicazione) all'interno di ogni worker. Ciò permette al gruppo FFN di inviare gradienti (N2M) mentre calcola simultaneamente il passaggio forward sui dati disponibili, garantendo un'esecuzione non bloccante e un utilizzo quasi continuo della GPU.

2.3 Allocazione Adattiva dei Worker

Per affrontare lo squilibrio tra attenzione (limitata dal calcolo) e FFN (limitata dalla comunicazione), DisagMoE impiega un allocatore adattivo guidato da un Modello Roofline Calcolo-Comunicazione.

• Modellazione: Il sistema modella il punto di svolta ($\hat{I}$) in cui un componente passa dall'essere limitato dalla rete a essere limitato dal calcolo. Gli strati di attenzione raggiungono il tetto del calcolo più rapidamente rispetto agli strati FFN all'aumentare della lunghezza della sequenza.
• Ottimizzazione: Il sistema formula la partizione di GPU e Schede di Interfaccia di Rete (NIC) come un problema di Programmazione Lineare Intera Mista (MILP). Alloca dinamicamente più larghezza di banda di rete ai gruppi FFN (per innalzare il loro tetto di comunicazione) e più risorse di calcolo ai gruppi di Attenzione, bilanciando le fasi della pipeline per minimizzare le bolle.
• Strategia di Ricerca: Un algoritmo utilizza il modello roofline per generare una configurazione iniziale, che viene poi raffinata tramite ricerca locale utilizzando i tempi misurati per passo per tenere conto degli overhead non modellati.

3. Contributi Chiave

1. Posizionamento Disaggregato dei Componenti: Un'architettura innovativa che separa gli strati di attenzione e FFN in gruppi di worker disgiunti, consentendo scalabilità e allocazione delle risorse indipendenti.
2. Pianificazione AF-Pipe: Una pipeline multi-stadio che sostituisce l'all-to-all convenzionale con scambi M2N inter-gruppo, sovrapposendo sistematicamente la comunicazione sia ai calcoli di attenzione che a quelli FFN.
3. Allocazione Adattiva delle Risorse: Un meccanismo guidato dal roofline che dimensiona congiuntamente le risorse GPU e NIC per i gruppi di attenzione e FFN, bilanciando efficacemente il rapporto calcolo-comunicazione attraverso diverse lunghezze di sequenza e configurazioni di modello.

4. Risultati Sperimentali

Gli autori hanno valutato DisagMoE su cluster di 16 nodi (128× GPU H800) utilizzando tre modelli MoE: DeepSeek-MoE, GPT-OSS-120B e Qwen3, con lunghezze di sequenza comprese tra 4K e 32K.

• Miglioramento del Throughput: DisagMoE ha costantemente superato basi di confronto solide, inclusi Megatron-LM (1F1B intercalato), Tutel, Comet e DualPipe.
  • Ha raggiunto un speedup fino a 1,81× rispetto a Megatron-1F1B.
  • Ha raggiunto un speedup fino a 1,34× rispetto ai sistemi di sovrapposizione MoE più avanzati.
• Riduzione della Comunicazione: Rispetto alle basi di confronto, DisagMoE ha ridotto il tempo di comunicazione non sovrapposto fino all'88% (vs. Tutel), al 75% (vs. Comet) e al 45% (vs. DualPipe).
• Scalabilità: Il sistema ha dimostrato efficacia attraverso diversi valori di top-$k$ e dimensioni del Parallelismo degli Esperti (EP).
• Sensibilità alle Risorse: Gli esperimenti hanno confermato che i rapporti GPU ottimali tra i gruppi di attenzione e FFN variano con la lunghezza della sequenza (es. sequenze più lunghe favoriscono più risorse di attenzione), validando la necessità di un'allocazione adattiva.

5. Significato e Affermazioni

Il paper afferma che DisagMoE affronta una limitazione fondamentale nell'addestramento MoE attuale: l'incapacità di sovrapporre perfettamente la comunicazione a causa dello squilibrio strutturale tra gli strati di attenzione e FFN. Disaggregando questi componenti e allocaando adattivamente le risorse basandosi su un modello roofline, il sistema espande la finestra di sovrapposizione effettiva oltre quanto possibile con una fornitura simmetrica di larghezza di banda.

Gli autori sottolineano che il loro approccio è particolarmente significativo per carichi di lavoro di pre-addestramento su larga scala con lunghezze di sequenza fisse. Osservano che, sebbene il loro design attuale assuma forme stabili (lunghezza di sequenza e micro-lotto fissi), l'idea centrale — che l'hardware e il software futuri dovrebbero mirare a rapporti calcolo-comunicazione favorevoli tra i componenti piuttosto che semplicemente scalare la larghezza di banda grezza — offre una nuova direzione per scalare efficientemente l'addestramento MoE. Il sistema è implementato su Megatron-LM ed è dimostrato essere una soluzione praticabile per l'addestramento di modelli con centinaia di miliardi di parametri senza i vincoli di memoria che hanno afflitto i precedenti tentativi di addestramento disaggregato.