Riepilogo Tecnico: ShardTensor – Parallelismo di Dominio per l'Apprendimento Automatico Scientifico

1. Enunciato del Problema

L'Apprendimento Automatico Scientifico (SciML) affronta un collo di bottiglia critico: l'impossibilità di addestrare ed eseguire inferenze su dati scientifici ad alta risoluzione a causa dei vincoli di memoria GPU. A differenza dei Modelli Linguistici di Grandi Dimensioni (LLM), dove la memoria è spesso dominata dal numero di parametri, i carichi di lavoro SciML sono frequentemente limitati dalle attivazioni intermedie.

Il documento identifica una "maledizione computazionale della dimensionalità" nei dati scientifici. Mentre raddoppiare la lunghezza del testo in un LLM raddoppia approssimativamente i token di input, raddoppiare la risoluzione di dati scientifici N-dimensionali aumenta la dimensione dei dati di un fattore $2^N$. Per dati 3D o di dimensioni superiori (ad esempio, fluidodinamica, imaging medico, modellazione climatica), la memoria richiesta per memorizzare le attivazioni intermedie durante il passaggio all'indietro dell'addestramento cresce esponenzialmente.

Le soluzioni esistenti spesso non scalano o degradano l'accuratezza:

• Precisione Ridotta: Offre solo risparmi di memoria modesti (2x) e può soffrire di problemi di intervallo dinamico nei surrogati scientifici.
• Sottocampionamento Spaziale: Riduce la memoria ma sacrifica l'alta fedeltà della risoluzione richiesta per le scoperte scientifiche.
• Riduzione del Modello: Diminuisce l'accuratezza.
• Checkpointing delle Attivazioni/Offloading su CPU: Efficace ma introduce un significativo sovraccarico computazionale (ricalcolo o trasferimenti host-dispositivo).
• Parallelismo Standard: Il parallelismo dei dati (divisione del batch) raggiunge un limite invalicabile a batch size 1. Il parallelismo del modello (divisione dei pesi) è progettato per i parametri, non per i dati di input, e assume forme tensoriali statiche che non valgono per l'elaborazione di dati di input dinamici (ad esempio, convoluzioni che cambiano forme globali).

Di conseguenza, non esiste un framework generalizzato per parallelizzare i dati di input al di sotto di un batch size di uno per dispositivo, impedendo allo SciML di raggiungere le risoluzioni native richieste dagli scienziati.

2. Metodologia: ShardTensor

Gli autori introducono ShardTensor, un nuovo paradigma di parallelismo di dominio che disaccoppia la dimensionalità spaziale dei dati di input dai vincoli hardware. È un'estensione dell'astrazione DTensor (Distributed Tensor) di PyTorch, progettata specificamente per gestire i dati di input anziché i pesi del modello.

Principi Progettuali Fondamentali

• Parallelismo di Dominio: Definito come parallelismo attraverso i dispositivi per i dati di input, che va oltre la dimensione del batch. Quando il batch size raggiunge 1, ShardTensor suddivide ulteriormente le dimensioni spaziali (ad esempio, altezza, larghezza, profondità) su più GPU.
• Gestione di Forme Dinamiche: A differenza di DTensor, che assume forme di distribuzione statiche adatte ai pesi del modello, ShardTensor supporta forme dinamiche. Introduce un quarto componente di metadati: le forme di sharding, che descrivono la forma del chunk locale di ogni tensore lungo l'asse della mesh shardata. Ciò consente la suddivisione arbitraria di dati non strutturati o non uniformi.
• Dispatch Imperativo: Il framework è imperativo piuttosto che statico. Utilizza i meccanismi __torch_dispatch__ e __torch_function__ di PyTorch per instradare le operazioni dinamicamente. Ciò permette al sistema di gestire il parallelismo di dominio strato per strato senza richiedere grafi computazionali precompilati o patch specifiche per ogni strato.
• Estensibilità: Gli utenti possono estendere il sistema tramite interfacce funzionali di alto livello (@custom_op) o logica di dispatch di basso livello (operazioni aten). Ciò garantisce la compatibilità con architetture di modelli in evoluzione e kernel personalizzati.

Dettagli di Implementazione

ShardTensor è costruito come un'estensione di DTensor di PyTorch.

• Architettura di Dispatch: Intercetta le operazioni a livello Python. Se un'operazione specifica ha un gestore per il parallelismo di dominio, viene eseguita; altrimenti, fa ricorso alla semantica di DTensor o al comportamento standard di PyTorch.
• Comunicazione: Il framework gestisce operazioni collettive (ad esempio, all-reduce, all-gather) e scambi "halo" (recupero di pixel adiacenti da dispositivi vicini) necessari per operazioni come convoluzione o normalizzazione per mantenere la coerenza numerica.
• Integrazione: È progettato per essere componibile con strategie di parallelismo esistenti (Parallelismo dei Dati e Parallelismo del Modello), consentendo strategie di parallelizzazione 2D o di dimensioni superiori (ad esempio, dividendo il batch su un asse e le dimensioni spaziali su un altro).

3. Contributi Chiave

1. Framework ShardTensor: Un framework generalizzato e open-source (disponibile tramite NVIDIA PhysicsNeMo) che abilita il parallelismo di dominio per i dati di input al di sotto di batch size uno.
2. Astrazione Innovativa: L'introduzione delle "forme di sharding" per gestire distribuzioni dinamiche di tensori di input, distinguendosi dagli strumenti di parallelismo focalizzati sui pesi come DTensor.
3. Modello di Programmazione Non Invasivo: Un design che permette agli utenti di avvolgere gli input del modello con minime modifiche al codice, evitando la necessità di riscrivere strati o modelli per il parallelismo.
4. Componibilità: Capacità dimostrata di combinare il parallelismo di dominio con il parallelismo dei dati e del modello su assi di mesh ortogonali.

4. Risultati e Benchmark

Il documento convalida ShardTensor attraverso benchmark sintetici e applicazioni industriali reali su GPU NVIDIA Blackwell (GB200) e H100.

• Ring Attention: Su dimensioni di sequenza molto piccole, il sovraccarico di parallelizzazione domina. Tuttavia, su dimensioni di sequenza molto grandi, il sistema raggiunge una scalabilità forte quasi lineare con il numero di GPU sia nell'addestramento che nell'inferenza.
• Vision Transformer (ViT):
  • Scalabilità Forte: Ad alte risoluzioni (ad esempio, $4096^2$), l'inferenza è 15 volte più veloce su 16 GPU rispetto a una singola GPU. L'addestramento su $2048^2$ è 5 volte più veloce con 8 GPU.
  • Efficienza di Memoria: L'uso di memoria scala quadraticamente per dati 2D e cubicamente per dati 3D, confermando che le attivazioni intermedie dominano. La scalabilità forte con ShardTensor riduce proporzionalmente la memoria per dispositivo, consentendo l'addestramento di modelli 3D con oltre 1 miliardo di punti di input su hardware standard.
• Transolver (Aerodinamica Automobilistica):
  • Addestrato sul dataset DrivaerML con fino a 1,2 milioni di punti per dominio.
  • Il sistema ha mantenuto stabilità numerica e accuratezza competitiva (errore L2 per pressione e velocità) rispetto alle baseline a singola GPU, anche quando la dimensione del dominio è aumentata di fattori 2, 4 e 8.
  • L'intera pipeline di pre-elaborazione è stata parallelizzata, garantendo la scalabilità end-to-end.
• StormScope (Previsioni Meteorologiche):
  • Un modello su scala continentale (CONUS) a risoluzione di 3 km (griglia $1024 \times 1792$).
  • L'uso di memoria di picco (114 GB) ha superato la capacità di una singola GPU H100 (80 GB).
  • ShardTensor ha consentito un addestramento stabile dividendo le attivazioni su 2 GPU all'interno di 16 gruppi di parallelismo dei dati. Il modello ha convergito stabilmente, tracciando la traiettoria della perdita di una baseline a singola GPU a risoluzione inferiore (6 km), consentendo la risoluzione di singoli temporali convettivi precedentemente inaccessibili senza sacrificare l'estensione del dominio.

5. Significato e Affermazioni

Il documento afferma che ShardTensor affronta un collo di bottiglia fondamentale nello SciML: l'impossibilità di addestrare ed eseguire inferenze su dati alla risoluzione nativa richiesta dagli scienziati.

• Sblocco dei Carichi di Lavoro Limitati dalla Risoluzione: Converte calcoli precedentemente impossibili (a causa di soffitti di memoria) in calcoli gestibili, abilitando l'addestramento ad alta fedeltà in campi come l'imaging medico volumetrico, la fluidodinamica computazionale e la modellazione climatica.
• Componibilità: Rimuove il compromesso tra volume di dati, dimensione del modello e risoluzione. Gli scienziati possono ora scalare lungo più assi di parallelismo simultaneamente.
• Accessibilità: Fornendo un modello di programmazione non invasivo, riduce la barriera all'ingresso per i professionisti che non sono esperti di sistemi distribuiti, in contrasto con i precedenti sforzi di parallelizzazione specifici per architettura e su misura.

Limitazioni Riconosciute

Gli autori notano modestamente che il modello di dispatch imperativo, strato per strato, introduce un sovraccarico a livello Python, rendendolo meno efficiente per operazioni piccole o dati a bassa risoluzione. Le prestazioni dipendono anche dalla larghezza di banda dell'interconnessione. Inoltre, non tutte le operazioni PyTorch hanno attualmente percorsi di dispatch per il parallelismo di dominio, e l'integrazione completa con torch.compile per ottimizzazioni di grafi statici è in corso.

In conclusione, ShardTensor rappresenta un passo verso la normalizzazione del parallelismo di dominio ad alto livello nell'apprendimento automatico scientifico, così come il parallelismo dei dati lo è oggi, facilitando l'era dei modelli fondazionali scientifici.