Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures

Il paper presenta Mozart, un framework di co-progettazione algoritmo-hardware che ottimizza l'addestramento efficiente dei modelli LLM basati su Mixture-of-Experts su architetture di chiplet su scala wafer 3.5D, risolvendo le sfide di località della memoria e sovraccarico di comunicazione attraverso strategie di allocazione degli esperti e meccanismi di scheduling a grana fine.

L'essenziale

🎻 Mozart: L'Orchestra Perfetta per i Cervelli Artificiali

Immagina di dover costruire un cervello artificiale gigante (una Large Language Model o LLM) capace di scrivere poesie, risolvere equazioni e conversare come un umano. Fino a poco tempo fa, questi cervelli erano come giganti monolitici: un unico blocco di cemento enorme, difficile da spostare e molto lento a pensare.

Oggi, però, gli scienziati usano un'architettura chiamata MoE (Mixture of Experts). Immagina invece di costruire il cervello non come un blocco unico, ma come un enorme team di specialisti.

• C'è lo specialista in matematica.
• C'è quello in letteratura.
• C'è quello in cucina.

Quando il cervello riceve una domanda ("Come si cucina la pasta?"), non sveglia tutti gli specialisti. Chiama solo lo specialista in cucina. Questo rende il sistema molto più veloce ed efficiente.

Il problema?
Gestire questo team di migliaia di specialisti è un incubo logistico.

1. Il traffico: Se ogni volta che serve uno specialista devi chiamarlo da un altro edificio (un altro chip), perdi tempo prezioso in telefonate (comunicazione).
2. La memoria: Gli specialisti hanno bisogno di libri di testo (i dati) che sono troppo grandi per stare sulla scrivania di ognuno. Devono essere presi da un archivio centrale (la memoria DRAM) ogni volta, creando code.
3. Il caos: A volte tutti chiedono lo stesso specialista, creando un ingorgo.

🚀 La Soluzione: Mozart

Gli autori del paper propongono Mozart, un sistema che unisce due mondi:

1. L'Algoritmo (Il Direttore d'Orchestra): Come organizzare gli specialisti.
2. L'Hardware (La Sala da Concerto): Il tipo di edificio fisico dove vivono gli specialisti.

1. La Sala da Concerto: Il Chiplet 3.5D

Invece di costruire un unico chip gigante (che è costoso e fragile), Mozart usa una torta a più piani (architettura 3.5D).

• I Chiplet: Immagina di avere 16 piccoli appartamenti (chiplet) collegati tra loro su un unico grande piano (il wafer).
• La Torre: Ogni appartamento ha la sua cucina (logica) e un piccolo frigo a portata di mano (memoria SRAM).
• L'Ascensore: Tra i piani c'è un ascensore super veloce (connessioni 3D) che porta i dati direttamente alla cucina, senza dover scendere al piano terra (il disco rigido esterno).

2. Il Direttore d'Orchestra: L'Algoritmo Mozart

Qui entra in gioco l'intelligenza del sistema. Mozart non mette gli specialisti a caso.

• La Mappa degli Amici (Expert Clustering):
  Mozart osserva come lavora il cervello. Scopre che certi specialisti lavorano spesso insieme (es. lo specialista in grammatica e quello in sintassi).

  • Metafora: Invece di mettere lo specialista in grammatica a New York e quello in sintassi a Tokyo, Mozart li mette nella stessa stanza. Se devono lavorare insieme, si passano il foglio di carta senza dover chiamare un corriere. Questo riduce drasticamente il traffico.

• Il Sistema a Nastro Trasportatore (Streaming & Scheduling):
  Il problema principale è che gli specialisti devono aspettare che i loro "libri di testo" arrivino dall'archivio centrale.

  • Metafora: Immagina una catena di montaggio. Mentre un operaio sta dipingendo un'auto (calcolo), il prossimo operaio sta già ricevendo i pezzi per la prossima auto (memoria). Mozart fa la stessa cosa: mentre un chip sta calcolando, il prossimo chip sta già scaricando i dati necessari. Nessuno aspetta mai. Si sovrappongono i tempi di attesa con i tempi di lavoro.

📊 I Risultati: Quanto è veloce?

Gli scienziati hanno testato Mozart su tre famosi modelli di intelligenza artificiale (Qwen, OLMoE, DeepSeek).
I risultati sono stati impressionanti:

• Il sistema è diventato da 2 a 2,4 volte più veloce rispetto ai metodi tradizionali.
• Ha ridotto gli ingorghi di dati e ha fatto lavorare le risorse al 100%, senza tempi morti.

🎯 In Sintesi

Mozart è come se avessimo preso un'orchestra caotica dove ogni musicista era in un edificio diverso e costretto a correre per prendere lo spartito.
Con Mozart:

1. Abbiamo costruito un palazzo unico (chiplet) dove i musicisti sono vicini.
2. Abbiamo riorganizzato la partitura (algoritmo) per mettere insieme i musicisti che suonano spesso insieme.
3. Abbiamo creato un sistema di consegna (streaming) che porta gli spartiti mentre i musicisti stanno già suonando.

Il risultato? Un'orchestra che suona in modo fluido, veloce ed efficiente, capace di gestire composizioni (modelli di intelligenza) sempre più grandi senza andare in tilt.

Sommario tecnico

Ecco un riassunto tecnico dettagliato del paper "Mozart: Modularized and Efficient MoE Training on 3.5D Wafer-Scale Chiplet Architectures", presentato in italiano.

1. Il Problema

Le architetture Mixture-of-Experts (MoE) hanno rivoluzionato l'efficienza dei Large Language Models (LLM) permettendo un'espansione massiccia dei parametri senza un aumento proporzionale del costo computazionale, attivando dinamicamente solo un sottoinsieme di "esperti" per ogni token. Tuttavia, l'addestramento di questi modelli su hardware convenzionale (come GPU tradizionali) incontra ostacoli significativi:

• Scarsa località della memoria: I dati necessari per gli esperti attivati sono dispersi, causando accessi frequenti alla memoria esterna (DRAM).
• Overhead di comunicazione: La necessità di sincronizzazione "all-to-all" tra le unità di calcolo per ridistribuire i token crea colli di bottiglia nella larghezza di banda.
• Utilizzo inefficiente delle risorse: Il carico di lavoro dinamico e disomogeneo degli esperti porta a sottoutilizzo delle risorse computazionali e a tempi di attesa (stall) dovuti alla memoria.
• Limiti di scalabilità: Le architetture monolitiche tradizionali raggiungono i limiti fisici della litografia e della scalabilità dei transistor.

2. Metodologia: Co-Design Algoritmo-Hardware

Mozart propone un framework di co-design algoritmo-hardware specifico per l'addestramento efficiente di MoE-LLM su architetture 3.5D Wafer-Scale Chiplet. L'approccio si divide in due dimensioni principali:

A. Ottimizzazioni Algoritmiche

1. Strategia di Allocazione degli Esperti (Expert Allocation):
   • Analizzando i pattern di attivazione pre-addestrati, Mozart raggruppa gli esperti che vengono attivati frequentemente insieme (co-attivazione) nello stesso chiplet o su chiplet adiacenti.
   • Questo riduce drasticamente il volume dei dati da trasferire durante la comunicazione all-to-all, poiché molti token non necessitano di uscire dal chiplet locale per raggiungere l'esperto richiesto.
   • Viene formulato come un problema di programmazione intera binaria per bilanciare il carico di lavoro tra i gruppi di chiplet.
2. Scheduling Fine-Grained con Streaming:
   • Per sovrapporre la comunicazione (DRAM) al calcolo, Mozart introduce lo streaming di token ed esperti.
   • Gli esperti con carico di lavoro più elevato vengono caricati per primi.
   • I token vengono suddivisi in micro-batch, permettendo di calcolare un gruppo di token mentre si caricano i dati per il successivo, massimizzando l'overlap tra I/O e computazione.

B. Architettura Hardware (3.5D Wafer-Scale)

Mozart introduce un'architettura fisica innovativa che combina:

• Stack 3D Logic-on-Memory: Ogni chiplet di calcolo integra verticalmente un die logico e un die SRAM tramite hybrid bonding. Questo riduce la latenza di accesso ai dati intermedi (attivazioni) e aumenta la banda interna.
• Topologia 2.5D NoP-Tree (Network-on-Package): Una struttura ad albero che disaccoppia le operazioni di Attention (nodi centrali, vicino alla DRAM per alta banda) dagli esperti (nodi foglia).
• Switch con Computazione In-Network: Gli switch della rete non si limitano a instradare i dati, ma aggregano localmente le uscite degli esperti, riducendo ulteriormente il traffico sulla rete.
• Gerarchia di Memoria: Pesi del modello in DRAM distribuita (accesso statico) e attivazioni in SRAM locale (accesso dinamico e frequente).

3. Contributi Chiave

1. Framework Mozart: Un sistema completo che mappa modelli MoE modulari su piattaforme chiplet wafer-scale, risolvendo i problemi di località e comunicazione.
2. Strategia di Clustering degli Esperti: Un metodo basato su dati empirici per posizionare gli esperti in modo da minimizzare la comunicazione inter-chiplet e bilanciare il carico.
3. Architettura 3.5D Ibrida: La progettazione di un'interconnessione 2.5D NoP-Tree combinata con stack 3D, ottimizzata specificamente per il pattern di accesso "sparsamente attivato" degli MoE.
4. Scheduling Dinamico: Un meccanismo di pipeline che nasconde la latenza di memoria sovrapponendo il caricamento degli esperti e il calcolo dei token.

4. Risultati Sperimentali

Il framework è stato valutato su tre modelli MoE open-source di diverse scale: Qwen3-30B-A3B, OLMoE-1B-7B e DeepSeek-MoE-16B.

• Accelerazione: Mozart ha dimostrato un'accelerazione superiore a 1.9x rispetto ai metodi baseline (senza ottimizzazioni specifiche) su tutti i modelli testati. In particolare:
  • 1.92x per Qwen3-30B-A3B.
  • 2.37x per OLMoE-1B-7B.
  • 2.17x per DeepSeek-MoE-16B.
• Efficienza della Comunicazione: La strategia di allocazione degli esperti ha ridotto significativamente il volume dei dati nella comunicazione all-to-all (misurato tramite il parametro $C_T$), abbassando la latenza end-to-end.
• Robustezza: Le prestazioni rimangono superiori anche con sequenze più lunghe (fino a 512 token) e in scenari con larghezza di banda DRAM variabile (HBM2 vs SSD), dimostrando che l'overlap calcolo-comunicazione è il fattore critico di successo.
• Analisi dei Colli di Bottiglia: Lo studio ha rivelato che il sistema è principalmente limitato dalla memoria (memory-bound) a causa del caricamento sequenziale dei pesi, ma Mozart mitiga efficacemente questo problema rispetto alle architetture tradizionali.

5. Significato e Impatto

Il lavoro di Mozart è significativo perché:

• Colma il divario tra Algoritmo e Hardware: Dimostra che le architetture MoE, spesso considerate difficili da implementare su hardware standard, possono essere ottimizzate radicalmente attraverso un design hardware dedicato che ne sfrutta la modularità intrinseca.
• Abilita la Scalabilità: Offre una via praticabile per l'addestramento di modelli LLM su scala wafer-scale, superando i limiti fisici dei chip monolitici.
• Efficienza Energetica e di Risorse: Migliorando la località dei dati e riducendo la comunicazione non necessaria, Mozart promette di ridurre il consumo energetico e i costi operativi per l'addestramento di modelli di grandi dimensioni.
• Futuro dei Sistemi Modulari: Fornisce un blueprint per l'integrazione di sistemi eterogenei (memoria, logica, switch) in architetture 3.5D, aprendo la strada a nuove generazioni di acceleratori AI specializzati.

In sintesi, Mozart rappresenta un passo avanti fondamentale verso l'addestramento efficiente di modelli di intelligenza artificiale di prossima generazione, trasformando le sfide della sparsità degli MoE in opportunità per architetture hardware modulari e scalabili.