Riassunto Tecnico: MatrixFSDP

Problematica

L'addestramento di modelli linguistici su larga scala si affida sempre più a ottimizzatori a struttura matriciale come Muon, Shampoo e SOAP. A differenza degli ottimizzatori per singola coordinata (es. AdamW), questi metodi richiedono la matrice di gradiente 2D completa come input per eseguire operazioni quali l'ortogonalizzazione di Newton–Schulz o il precondizionamento della matrice. Ciò crea un fondamentale disallineamento sistemico con FSDP (Fully Sharded Data Parallel) e ZeRO-3, che frammentano (shard) parametri, gradienti e stati dell'ottimizzatore tra i rank per massimizzare l'efficienza della memoria.

Sotto lo standard ZeRO-3, ogni rank possiede solo una porzione del modello. Per applicare un ottimizzatore matriciale, il sistema deve ricostruire l'intera matrice di gradiente 2D ad ogni step dell'ottimizzatore. Questo richiede un all-gather delle dimensioni del peso ad ogni step del gradiente dopo il passaggio backward. Questa comunicazione è un collo di bottiglia critico sul percorso critico che non può essere sovrapposto al calcolo dei layer, portando a una latenza significativa che scala negativamente con il numero di nodi e la dimensione del modello.

Le soluzioni esistenti sono incomplete:

1. Ricostruzione (Stock FSDP2-Muon): Mantiene i risparmi di memoria di ZeRO-3 ma comporta un pesante overhead di comunicazione ad ogni step dell'ottimizzatore.
2. Posizionamento ZeRO-1 Owner: Assegna intere matrici a specifici rank (owner) per rendere locale lo step dell'ottimizzatore, ma richiede che i parametri completi siano residenti su ogni rank, abbandonando l'efficienza di memoria di ZeRO-3.

L'obiettivo è ottenere input dell'ottimizzatore a matrice intera senza sacrificare la memoria dei parametri su scala ZeRO-3.

Metodologia: MatrixFSDP

MatrixFSDP introduce una terza via: shard ZeRO-3 a forma di proprietario (Owner-shaped). Invece di cambiare l'algoritmo dell'ottimizzatore o abbandonare la frammentazione, esso ridefinisce dove risiedono gli shard all'interno del gruppo di data-parallelism.

1. Posizionamento a forma di Proprietario (Owner-Shaped Placement)

Per ogni matrice di peso 2D nel modello, MatrixFSDP designa un singolo rank di data-parallel come proprietario (owner).

• Rank Proprietario (Owner Rank): Detiene la matrice 2D completa.
• Rank Non-Proprietari: Detengono uno shard vuoto per quel parametro specifico.
• Tensori Non-Matriciali: I vettori 1D e altri tensori non 2D vengono impacchettati in una "coda" (tail) e assegnati a un proprietario separato, dove continuano a utilizzare lo standard AdamW.

Ciò garantisce che, dopo il passaggio backward e la riduzione del gradiente, l'input completo di Muon (la matrice di gradiente 2D completa) risieda naturalmente sul rank proprietario. Di conseguenza, l'aggiornamento di Newton–Schulz viene eseguito localmente sul proprietario con zero comunicazione di matrice durante lo step dell'ottimizzatore.

2. Sfide e Soluzioni Runtime

Questa disposizione disomogenea (alcuni rank che detengono matrici complete, altri che detengono shard vuoti) viola le assunzioni di dimensioni uguali dei collettivi standard di FSDP2. MatrixFSDP affronta ciò attraverso quattro meccanismi runtime chiave:

• Metadati MatrixShard & Global Planner: Un pianificatore globale assegna i proprietari per bilanciare i byte residenti e il carico di lavoro dell'ottimizzatore attraverso il cluster. Questi metadati tracciano quale rank possiede quale ruolo di tensore, garantendo che il sistema rimanga un layout ZeRO-3 coerente invece di una collezione di casi speciali.
• Collettivi P2P Owner-Segment: Gli all-gather standard sono sostituiti da collettivi deterministici owner-segment. Durante la materializzazione forward/backward, solo i segmenti non vuoti vengono trasferiti dal proprietario ai rank che ne hanno bisogno. Durante la riduzione del gradiente, i risultati vengono scritti solo nei segmenti non vuoti del proprietario. Ciò evita il padding dei rank vuoti ed elimina gli all-gather di matrice completa.
• Blocco del Buffer del Proprietario (Owner-Buffer Pinning): Per gestire il ciclo di vita di FSDP2 (unshard $\to$ compute $\to$ reshard), MatrixFSDP introduce uno stato "pinned". Quando un buffer viene rimpicciolito dopo il passaggio forward, lo storage non viene restituito al pool condiviso ma viene bloccato (pinned). Ciò consente al passaggio backward di "rivivere" la matrice completa nella stessa posizione di memoria, prevenendo errori di visualizzazione (view errors) di autograd e garantendo che il proprietario possa eseguire l'aggiornamento prima di rilasciare il buffer.
• Checkpointing e Routing: Il checkpointing (DCP) e il routing dello stato dell'ottimizzatore utilizzano i metadati MatrixShard. Gli stati dell'ottimizzatore sono instradati ai corretti proprietari (Muon per le matrici, AdamW per le code), e i checkpoint salvano la specifica disposizione del layout per supportare il resharding attraverso diverse dimensioni del mondo (world sizes).

Contributi Chiave

1. ZeRO-3 a forma di Proprietario (Owner-Shaped): Una nuova strategia di posizionamento in cui le matrici 2D sono intere su un singolo proprietario e vuote altrove, consentendo l'ottimizzazione matriciale locale all'interno di un regime di memoria frammentata.
2. Step dell'Ottimizzatore Senza Comunicazione: Allineando la destinazione della riduzione del gradiente con il requisito di input dell'ottimizzatore, MatrixFSDP elimina il collo di bottiglia dell'all-gather della dimensione del peso precedentemente inerente agli ottimizzatori matriciali in ZeRO-3.
3. Infrastruttura Runtime: Una suite di componenti runtime (planner, collettivi personalizzati, buffer pinning) che rendono efficienti e corretti i layout disomogenei dei proprietari all'interno del framework FSDP2.
4. Generalità: L'approccio è agnostico rispetto all'ottimizzatore, dimostrando di funzionare con Muon, Shampoo e SOAP senza richiedere varianti distribuite dei precondizionatori.

Risultati Sperimentali

Valutato su 64 GPU NVIDIA A100 (8 nodi) con varie dimensioni di modello (fino a 32B parametri):

• Latenza dello Step dell'Ottimizzatore: MatrixFSDP riduce la latenza dello step dell'ottimizzatore di 4,2× su un singolo nodo e di 54,6× su otto nodi rispetto allo stock FSDP2-Muon. Il miglioramento cresce con la scala perché la comunicazione di base attraversa il fabric inter-nodo, mentre lo step di MatrixFSDP rimane locale.
• Speedup End-to-End: Raggiunge fino a un 2,15× di speedup di addestramento end-to-end.
• Efficienza della Memoria: MatrixFSDP mantiene i livelli di memoria di ZeRO-3. Per un modello da 32B, utilizza 10 GB/rank, mentre un approccio di posizionamento owner ZeRO-1 supera il limite di 80 GB per GPU (proiezione >54 GB/rank per 14B). È 3–7× più efficiente in termini di memoria rispetto agli approcci "gather-once" di FSDP2-Muon.
• Correttezza: La traiettoria di addestramento corrisponde a un riferimento DDP Muona a matrice completa con alta precisione (differenza di perdita stampata $\leq 0$), confermando che il proprietario riceve esattamente la matrice di gradiente necessaria per l'aggiornamento.

Significato e Rivendicazioni

Il documento afferma che MatrixFSDP risolve il "disallineamento sistemico" tra ottimizzatori matriciali e frammentazione dei parametri. La sua importanza primaria risiede nel consentire l'addestramento su larga scala con ottimizzatori matriciali (che offrono una migliore efficienza dei token e convergenza) senza incorrere nei costi di comunicazione proibitivi o negli overhead di memoria delle precedenti approcci.

Gli autori sottolineano che MatrixFSDP è particolarmente efficace per le corse ZeRO-3 multi-nodo con ampi span di shard, dove il costo di ricostruzione delle matrici ad ogni step altrimenti dominerebbe il tempo di addestramento. Il documento nota modestamente che i guadagni sono minori nelle corse su singolo nodo o sotto lo scaling forte del modello fisso, dove il numero di ruoli di proprietario diventa rado, rendendo potenzialmente il fanout del proprietario un nuovo collo di bottiglia. Il lavoro non pretende di sostituire gli ottimizzatori per singola coordinata per i parametri non matriciali, ma piuttosto fornisce un percorso unificato ed efficiente per gestire i componenti di matrice 2D dei modelli moderni.