veScale-FSDP: Flexible and High-Performance FSDP at Scale

Il paper presenta veScale-FSDP, un sistema FSDP ridisegnato che combina il formato di sharding flessibile RaggedShard con un algoritmo di pianificazione consapevole della struttura per supportare nativamente metodi di addestramento avanzati come la quantizzazione a blocchi e ottimizzatori non elementari, ottenendo al contempo un aumento del throughput e una riduzione dell'uso di memoria rispetto alle implementazioni esistenti.

L'essenziale

Immagina di dover costruire un grattacielo enorme (un'intelligenza artificiale gigante) con un esercito di migliaia di operai (i processori GPU). Il problema è che il cantiere è piccolo e gli operai devono lavorare tutti insieme, ma non possono stare tutti nello stesso posto.

Il sistema attuale, chiamato FSDP, è come un capocantiere molto intelligente che divide i mattoni tra gli operai. Tuttavia, questo capocantiere ha due grossi limiti:

1. È troppo rigido: Divide i mattoni uno per uno (o riga per riga). Se i mattoni devono essere usati in blocchi specifici (come mattoni quadrati per un muro speciale), il capocantiere li taglia a metà, creando confusione e sprechi.
2. È lento e disordinato: Quando gli operai devono scambiarsi i mattoni, lo fanno in modo disorganizzato, perdendo tempo a correre avanti e indietro e occupando troppo spazio nel magazzino (memoria).

veScale-FSDP è il nuovo, rivoluzionario capocantiere che ha risolto questi problemi. Ecco come funziona, spiegato con analogie semplici:

1. Il "Taglio Personalizzato" (RaggedShard)

Immagina che i mattoni non siano tutti uguali. Alcuni sono piccoli, altri sono enormi blocchi di cemento.

• Il vecchio sistema: Tagliava tutto in fette sottili e uguali, ignorando la forma originale dei blocchi. Se dovevi usare un blocco intero, dovevi ricomporlo con mille pezzi, perdendo tempo.
• veScale-FSDP: Usa un metodo chiamato RaggedShard (che potremmo chiamare "Taglio Irregolare Intelligente"). Invece di tagliare tutto in modo uniforme, guarda ogni blocco e dice: "Tu sei un blocco di 100 mattoni? Ti do a te. Tu sei un blocco di 50? Ti do a te".
• Il risultato: Non c'è bisogno di ricomporre nulla. Gli operai ricevono i pezzi esattamente come sono fatti, pronti per essere usati subito. Questo permette di usare tecniche avanzate (come la "quantizzazione a blocchi") che prima erano impossibili.

2. L'Organizzatore di Magazzini (Planning Algorithm)

Ora immagina che tutti questi pezzi irregolari debbano essere messi in un camion per essere spediti tra gli operai.

• Il vecchio sistema: Metteva i pezzi nel camion uno dopo l'altro, spesso lasciando buchi vuoti o sovrapponendoli male. Per farli entrare, aggiungeva tanta "paglia" (padding) per riempire gli spazi, sprecando spazio nel camion.
• veScale-FSDP: Usa un algoritmo di pianificazione che è come un super-organizzatore di Tetris. Prima di caricare il camion, calcola esattamente come impilare i pezzi irregolari per occupare il minimo spazio possibile e non lasciare buchi.
• Il risultato: Il camion è pieno al 100%, non c'è paglia sprecata, e gli operai ricevono tutto molto più velocemente.

3. Il Nastro Trasportatore Senza Attrito (Distributed Buffer)

Infine, quando gli operai devono passare i mattoni l'uno all'altro:

• Il vecchio sistema: Dovevano prendere il mattone, metterlo in una scatola, passarlo, toglierlo dalla scatola e rimetterlo sul muro. Ogni passaggio richiedeva tempo e fatica.
• veScale-FSDP: Usa un nastro trasportatore magico (Distributed Buffer). I mattoni scorrono direttamente da un operario all'altro senza mai essere messi in una scatola o toccati due volte.
• Il risultato: Il lavoro è fluido, veloce e non si consuma energia inutile.

Perché è importante?

Grazie a queste innovazioni, veScale-FSDP permette di:

• Costruire più velocemente: Le macchine lavorano dal 5% al 66% più velocemente rispetto ai sistemi attuali.
• Risparmiare spazio: Usano fino al 30% di memoria in meno, permettendo di costruire modelli più grandi senza esplosioni di costi.
• Scalare all'infinito: Funziona perfettamente anche con 10.000 operai (GPU) che lavorano insieme, cosa che i sistemi precedenti facevano fatica a gestire.

In sintesi, veScale-FSDP è come passare da un cantiere caotico dove si perde tempo a tagliare mattoni e riorganizzare scatole, a un cantiere di alta tecnologia dove ogni pezzo arriva esattamente dove serve, nel modo giusto, e il lavoro procede a velocità della luce.

Sommario tecnico

1. Il Problema

I modelli linguistici su larga scala (LLM) moderni richiedono tecniche di parallelizzazione avanzate per essere addestrati su migliaia di GPU. La Fully Sharded Data Parallel (FSDP), nota anche come ZeRO, è lo standard di fatto per la sua flessibilità e la bassa intrusività sul codice del modello. Tuttavia, i sistemi FSDP attuali (come DeepSpeed ZeRO, PyTorch FSDP1 e FSDP2) presentano limitazioni critiche che ne impediscono l'efficienza nell'era dei modelli di nuova generazione:

• Incompatibilità con metodi di addestramento "Structure-Aware": I modelli all'avanguardia (es. Gemini, DeepSeek-V3) utilizzano ottimizzatori non element-wise (come Shampoo e Muon) e tecniche di quantizzazione a blocchi (Block-wise Quantization). Questi metodi richiedono che i tensori siano manipolati come blocchi atomici interi.
• Conflitto di allineamento: Le implementazioni FSDP esistenti frammentano i parametri, i gradienti e gli stati degli ottimizzatori in modo element-wise o riga-per-riga (row-wise). Questo crea confini di sharding che spesso non si allineano con i blocchi strutturati richiesti dagli ottimizzatori o dalla quantizzazione. Di conseguenza, gli sviluppatori devono modificare intrusivamente i modelli o gestire logiche complesse di padding e comunicazione, introducendo overhead significativi.
• Inefficienza di memoria e comunicazione: Le implementazioni attuali soffrono di frammentazione della memoria, copie di dati intercalate (interleaved copy) durante le operazioni di AllGather e ReduceScatter, e buffer di comunicazione non allineati. Questo limita la scalabilità a poche decine di migliaia di GPU e porta a un uso inefficiente delle risorse hardware.

2. Metodologia e Architettura

Per risolvere questi problemi, gli autori hanno riprogettato il sistema FSDP introducendo veScale-FSDP, che combina un nuovo formato di sharding, un algoritmo di pianificazione e un buffer distribuito.

A. RaggedShard (Nuovo Formato di Sharding)

Il cuore della flessibilità è RaggedShard, un nuovo formato di placement per i DTensor di PyTorch.

• Granularità Arbitraria: A differenza dello sharding rigido (elemento o riga), RaggedShard permette di definire blocchi atomici di dimensioni personalizzate (es. blocchi 2D per la quantizzazione).
• Distribuzione Non Uniforme: Consente di assegnare un numero diverso di blocchi a ciascun dispositivo, supportando nativamente lo sharding irregolare necessario per ottimizzatori come Muon.
• Componibilità: RaggedShard è ortogonale alle strategie di parallelismo esistenti (Tensor Parallelism, Expert Parallelism) e si integra con le API native di PyTorch (fully_shard) senza richiedere modifiche intrusive al codice del modello.

B. Algoritmo di Pianificazione (Planning Algorithm)

Per massimizzare l'efficienza della comunicazione con tensori RaggedShard, è stato sviluppato un algoritmo di pianificazione che risolve un problema di ottimizzazione NP-hard.

• Obiettivo: Riorganizzare i tensori RaggedShard in un buffer di comunicazione globale per minimizzare il padding e bilanciare il carico tra i dispositivi.
• Strategia: L'algoritmo permuta i tensori e applica il padding tra i tensori invece che dentro di essi, garantendo che i blocchi atomici rimangano contigui e non vengano spezzati dai confini di comunicazione.
• Euristiche: Utilizza un approccio euristico basato sulla programmazione dinamica (DP) che sfrutta la regolarità dei modelli Transformer (dove i pesi lineari dominano) per trovare soluzioni quasi ottimali in tempo polinomiale, evitando i tempi di calcolo proibitivi dei solver ILP.

C. Distributed Buffer (DBuffer)

Per supportare RaggedShard, è stato introdotto un nuovo primitivo chiamato Distributed Buffer.

• Zero-Copy: Mappa i dati dei tensori direttamente in un buffer globale, eliminando le costose operazioni di copia (copy-in/copy-out) tipiche di FSDP2.
• Fusione dei Kernel: Esegue operazioni a livello di gruppo (es. add, scale, zero) prima della comunicazione, riducendo il tempo di blocco e frammentazione computazionale.
• Gestione della Memoria: Utilizza allocazioni batched e dipendenze di stream deterministiche per ridurre la frammentazione della memoria e prevenire l'Out-Of-Memory (OOM).

3. Contributi Chiave

1. Supporto Nativo per Ottimizzatori e Quantizzazione Avanzata: veScale-FSDP è il primo sistema FSDP a supportare nativamente ottimizzatori non element-wise (Muon, Shampoo) e quantizzazione a blocchi (es. 8-bit Adam) senza modifiche intrusive al codice.
2. Flessibilità Estrema: L'introduzione di RaggedShard permette di adattarsi a qualsiasi granularità di sharding richiesta dal modello, superando i limiti dei formati rigidi attuali.
3. Efficienza Scalabile: L'architettura è progettata per scalare efficientemente fino a 10.000+ GPU e modelli con trilioni di parametri, risolvendo i colli di bottiglia di comunicazione e memoria.
4. Integrazione con PyTorch: Il sistema è implementato come modulo Python plug-and-play che sostituisce il backend di FSDP2, mantenendo la stessa API nativa (fully_shard).

4. Risultati Sperimentali

Le valutazioni sono state condotte su cluster GPU con modelli come LLaMA-3-70B, GPT-OSS-120B e modelli interni MoE (Mixture of Experts).

• Throughput: veScale-FSDP supera tutti i sistemi baselines (DeepSpeed, FSDP1, FSDP2, Megatron-FSDP) con un aumento del throughput del 5% al 66%. I miglioramenti sono particolarmente evidenti nei modelli MoE e con ottimizzatori avanzati.
• Utilizzo della Memoria: Riduce l'uso di memoria di picco (peak memory) del 16% al 30% grazie alla gestione deterministica del DBuffer e alla riduzione del padding.
• Scalabilità:
  • Weak Scaling: Dimostra una scalabilità quasi lineare fino a 8.000 GPU.
  • Strong Scaling: Scalabilità lineare fino a 10.000 GPU con batch globali di 128M token.
  • Model Scaling: Capacità di addestrare modelli da 2.4 trilioni di parametri su 1.000 GPU senza degradazione delle prestazioni.
• Casi d'Uso Specifici:
  • 8-bit Adam: Implementato con poche righe di codice, mantenendo la convergenza e riducendo drasticamente la memoria degli stati dell'ottimizzatore.
  • Muon: Supporto nativo per l'aggiornamento Newton-Schulz su tensori 2D completi, raggiungendo un MFU (Model FLOPS Utilization) del 47.3% su 256 GPU.
• Overhead di Pianificazione: L'algoritmo di pianificazione richiede meno di 0.3 secondi per essere eseguito, rendendo il suo costo trascurabile rispetto al tempo di addestramento.

5. Significato e Impatto

veScale-FSDP rappresenta un passo fondamentale nell'evoluzione dell'addestramento distribuito di LLM.

• Democratizzazione delle Tecniche Avanzate: Rimuove le barriere tecniche che costringevano i ricercatori a scegliere tra flessibilità architetturale e efficienza di sistema. Ora è possibile utilizzare ottimizzatori complessi e quantizzazione senza sacrificare le prestazioni.
• Efficienza Industriale: Permette di sfruttare appieno cluster di migliaia di GPU, riducendo i costi di addestramento e i tempi di iterazione per modelli su scala triliardaria.
• Ecosistema Open Source: Essendo basato su PyTorch e integrato con DTensor, veScale-FSDP favorisce l'interoperabilità con l'ecosistema esistente (checkpointing distribuito, TP, EP) ed è stato open-sourced, contribuendo al progresso della comunità.

In sintesi, veScale-FSDP risolve il conflitto storico tra la necessità di sharding flessibile per modelli moderni e la necessità di comunicazione efficiente, fornendo un sistema che è sia flessibile (adattabile a qualsiasi architettura) che ad alte prestazioni (ottimizzato per hardware su larga scala).