Addestramento PyTorch rapido e tollerante ai guasti su AI Runtime

Su larga scala, l'efficienza dell'addestramento è determinata da un'unica metrica: il "goodput", ovvero la percentuale di tempo che le GPU dedicano al calcolo produttivo anziché all'attesa o al ripristino dopo un guasto. Poiché i guasti delle GPU sono un evento previsto su larga scala, la capacità di ripristinare il sistema in modo rapido e automatico dopo un guasto è l'unico modo per mantenere un buon livello di goodput e gestire la spesa totale per le GPU.

Due sottosistemi determinano il successo o il fallimento di questo ripristino, eppure entrambi vengono regolarmente considerati come secondari: la pipeline di dati che alimenta gli acceleratori e il meccanismo di checkpointing che salva lo stato in modo che un job possa riprendere. Se si sbaglia anche solo uno di questi aspetti, ogni guasto costerà molto più tempo di inattività delle GPU del dovuto. Anche al di fuori degli scenari di guasto, una pipeline di dati che non riesce a tenere il passo con gli acceleratori lascerà le GPU inattive senza dare segnali evidenti, riducendo il goodput proprio come farebbe un crash. Esamineremo i meccanismi e i compromessi di entrambi, e come ciascuno di essi influisca sul goodput e sulla spesa totale per le GPU. Consulta la Guida alle prestazioni e alla resilienza dell'addestramento per indicazioni sul codice ed esempi.

Per l'aspetto legato all'infrastruttura dello stesso problema, ovvero come una flotta rileva e isola le GPU non integre prima che interrompano un job, consulta il post correlato, Come garantiamo l'affidabilità delle GPU in Databricks AI.

Perché i guasti sono un evento previsto su larga scala

All'aumentare del numero di GPU in un job, la probabilità che questo completi l'intera durata senza interruzioni diminuisce rapidamente. Un utile modello approssimativo descritto nel post correlato di Databricks ipotizza che ogni GPU abbia un tasso di guasto annualizzato di circa l'1%. In base a questa ipotesi, il post rileva che "un job con 256 GPU in esecuzione per 30 giorni ha circa il 19% di probabilità di subire un guasto. Con 1.024 GPU, questa percentuale sale al 57%." e questi sono solo problemi a livello di infrastruttura.

Per calare questa stima nella realtà, il supercomputer delta con 608 GPU H100 ha registrato guasti ogni 1,9 ore; ciò significa che per un job con 32 GPU, il tempo medio prima di un guasto sarebbe di 36 ore. La conclusione principale è che il job di addestramento probabilmente subirà un guasto a un certo punto, e prendere le decisioni corrette può rendere il modello resiliente e ridurre il tempo totale perso quando ciò accade.

Impatto 1: Il formato del checkpoint determina la frequenza di salvataggio sostenibile

La resilienza si decide tutta nel checkpointing, e il meccanismo scelto ha un effetto primario sulla frequenza con cui è possibile salvare. Questa è la leva principale per il tuo goodput: se esegui un checkpoint una volta al giorno, un guasto richiederà di ripetere in media 12 ore di lavoro duplicato per ripristinare lo stato in cui si trovava il sistema al momento del guasto.

Il collo di bottiglia monolitico di torch.save

Il primo checkpoint che la maggior parte dei team scrive è un semplice torch.save sul rank 0. A seconda di come viene addestrato il modello, possono presentarsi due problemi:

1. Per l'addestramento distribuito, raccoglie tutti gli stati sul rank 0 e scrive un singolo file.
2. Un singolo processo scrive l'intero checkpoint in modo sincrono. Questo processo può essere bloccato da fattori come i trasferimenti di rete quando si salva su storage di oggetti remoti come Unity Catalog (UC).

Questo comportamento di blocco lascia le GPU inattive, riducendo il goodput. Esiste tuttavia un modo per ridurre il tempo dedicato al checkpointing dalle GPU: l'API di checkpoint distribuito di Torch.

Distributed checkpoint (DCP): ogni rank scrive il proprio shard

Il checkpoint distribuito di PyTorch inverte questo approccio. Ogni rank scrive il proprio shard distinto in parallelo, insieme a un piccolo file .metadata che descrive come gli shard si compongono nei tensori completi.

Il tempo di salvataggio diminuisce all'incirca come 1/N con il numero di rank e, poiché il file .metadata registra il layout globale, lo stesso checkpoint può essere ricaricato su un numero diverso di GPU. DCP pianifica nuovamente i byte necessari per ogni nuovo rank, quindi il ripristino su un cluster a capacità ridotta dopo la perdita di nodi funziona senza problemi.

DCP conviene anche per i normali job con parallelismo dei dati

Spesso si pensa che DCP sia utile solo per i modelli con partizionamento (sharding), e che un job con parallelismo dei dati (DDP), in cui ogni rank contiene una replica identica dei pesi, non tragga alcun vantaggio. Non è così: DCP suddivide in shard lo stato del modello e lo scrive in parallelo su ciascun worker, anche per i task di addestramento DDP.

Inoltre, è la stessa API di cui avrai bisogno quando passerai a FSDP o al parallelismo dei tensor; adottarla fin da subito significa non dover mai riscrivere il codice di resilienza nel momento meno opportuno.

I salvataggi asincroni rendono la frequenza quasi a costo zero

Anche con le scritture parallele, un salvataggio sincrono blocca l'addestramento finché i byte non sono memorizzati in modo permanente nello storage; per un checkpoint di grandi dimensioni su un volume remoto, questo comporta decine di secondi di inattività dell'acceleratore. async_save suddivide l'operazione: una copia rapida in un buffer di staging, seguita da un caricamento in background che si sovrappone alla continuazione dell'addestramento.

Il ciclo di addestramento paga solo per la copia di staging, non per il caricamento. Un checkpoint che prima costava decine di secondi di inattività ora non costa quasi nulla, ed è proprio questo che rende sostenibile il checkpointing frequente descritto nella sezione successiva.

Su AI Runtime, UCVolumeWriter e UCVolumeReader implementano DCP sui volumi UC, eseguendo lo staging dell'I/O tramite NVMe locale e contrassegnando un checkpoint come completato solo quando i relativi dati sono stati interamente scritti. Consulta la guida alle prestazioni e alla resilienza per tutti i dettagli e gli esempi di codice.

I dati sopra riportati escludono il tempo di storage di rete per torch.save.

Impatto 2: La frequenza dei checkpoint determina il costo di ripristino

È qui che i vari elementi si sommano. Quando un job fallisce, perde tutto ciò che è stato fatto dall'ultimo checkpoint valido e deve ricalcolarlo. Pertanto, il lavoro sprecato previsto per guasto è pari a circa la metà dell'intervallo di checkpoint, e i salvataggi asincroni economici consentono di ridurre tale intervallo.

Ridurre l'intervallo di un fattore pari a 10 riduce il tempo di ripristino previsto dello stesso fattore. Ricorda il dato di Llama 3 di circa 8,6 interruzioni al giorno: con questo tasso di guasto, eseguire un checkpoint ogni 2 ore significa prevedere di sprecare 8,6 ore al giorno in riaddestramento, con un goodput del 64%. Eseguendo il checkpoint ogni 30 minuti, si perdono solo 2,15 ore, con un goodput del 91%.

Inoltre, il ripristino deve essere automatico. Al riavvio, il job deve trovare il checkpoint più recente completato con successo, saltando quelli rimasti scritti a metà a causa del crash, e riprendere da lì senza alcun intervento umano. DCP rende questo processo affidabile: il file .metadata viene scritto solo dopo che tutti gli shard sono stati salvati, quindi la sua presenza è un indicatore affidabile del tipo "questo salvataggio è completo" su cui basare la selezione.

Impatto 3: Il caricamento dei dati determina se le GPU rimangono inattive

Un job di addestramento procede alla velocità del suo input più lento. Quando gli acceleratori attendono il batch successivo, il goodput si riduce poiché le GPU rimangono semplicemente inattive. L'unico modo per risolvere questo problema è garantire che la pipeline di input sovrapponga la preparazione dei dati per lo step successivo al calcolo di quello corrente, come mostrato nella figura seguente:

Spesso notiamo che i clienti che passano alla sovrapposizione del caricamento dei dati con il calcolo registrano una riduzione del 20-50% del tempo effettivo.

Il costo della lettura diretta dallo storage remoto

Su una piattaforma gestita, i dati di addestramento risiedono in uno storage di oggetti remoto. Su AI Runtime, i volumi di Unity Catalog (UC) vengono presentati come montaggi di rete.

Leggere i file direttamente da quel montaggio a ogni accesso vincola il tempo di step alla latenza di rete e scarica nuovamente gli stessi file a ogni epoca. La soluzione è un dataloader che copia ciascun file in uno storage locale rapido al primo accesso, gestisce le letture successive da quella cache locale e recupera i file successivi in parallelo mentre la GPU esegue i calcoli.

Con AI Runtime, UCVolumeDataset e DataLoader fanno esattamente questo (vedi la guida per esempi di codice). UCVolumeDataset distribuisce in streaming i file da un volume UC, memorizzando ciascuno di essi nella cache NVMe locale al primo accesso, e partiziona i file tra rank e worker in modo che ogni acceleratore riceva una porzione disgiunta e non sovrapposta. Il nostro DataLoader è una sottoclasse drop-in del DataLoader di PyTorch i cui valori predefiniti sono ottimizzati per questo percorso, in modo che i file vengano recuperati e memorizzati nella cache contemporaneamente mentre la GPU elabora, anziché uno alla volta sul thread di addestramento.

Esempio: addestramento di un modello di immagini da file UC

Consideriamo un semplice carico di lavoro di classificazione delle immagini: decodificare i file JPEG da un volume UC, applicare l'augmentation e addestrare un modello di computer vision. Vediamo due modi per farlo sulla stessa GPU, con lo stesso modello e la stessa dimensione del batch: il Dataset standard di PyTorch che legge da un volume UC rispetto a UCVolumeDataset combinato con i valori predefiniti del DataLoader di Databricks.

Non devi indovinare dove viene speso il tempo

Durante lo sviluppo di DataLoader, ci siamo assicurati che registri le sue metriche su MLflow, rendendo facile capire a colpo d'occhio se la pipeline di dati sta bloccando l'addestramento.

La metrica fetch_seconds misura esplicitamente quanto tempo impiega il dataloader a produrre un batch; durante questo tempo la GPU rimane inattiva.

Impatto 4: Dimenticare la pipeline di dati corrompe silenziosamente il modello

Esiste un ultimo bug di resilienza che non produce alcun messaggio di errore, nessun crash e nessun job non riuscito, ma solo un modello leggermente peggiore di quanto dovrebbe essere. Si verifica quando salvi un checkpoint del modello, dell'ottimizzatore e dello step, ma non della posizione della pipeline di dati all'interno del dataset.

Consideriamo un job interrotto a metà di un'epoca. Ripristina correttamente il modello e riprende il loop di addestramento, ma il dataloader ricomincia dall'inizio del dataset.

Il job ripreso esegue nuovamente l'addestramento su esempi già visti in questa epoca e potenzialmente salta quelli che non aveva ancora raggiunto. Tra i molti riavvii che diventano routine su larga scala, questo introduce silenziosamente un bias nella distribuzione dei dati. Il modello continua ad addestrarsi, ma lo fa sul campionamento errato dei dati: proprio il tipo di errore silenzioso più costoso, perché il job viene completato e nessuno nota il problema finché le metriche non si rivelano deludenti.

La soluzione consiste nel trattare la posizione dei dati come parte del checkpoint. A seconda della pipeline, ciò significa tracciare l'offset di un campione o di uno shard e saltare in avanti al riavvio, fare in modo che un dataset personalizzato serializzi la propria posizione o salvare i checkpoint ai limiti dell'epoca. Tutto questo si basa su un unico prerequisito: il determinismo. Lo shuffling e l'augmentation attingono da generatori di numeri casuali, quindi anche i seed e gli stati RNG devono far parte del checkpoint, altrimenti l'ordine dei dati dopo un riavvio non corrisponderà a quello precedente e una posizione salvata punterà ai campioni errati.

Seed, ordine riproducibile e pipeline di dati ripristinabile sono tre espressioni di un'unica idea. La guida illustra ciascuna strategia con esempi di codice.

Riepilogo

Un addestramento rapido e tollerante ai guasti deriva da una serie di decisioni che si sommano:

1. Usa Distributed Checkpoint invece di torch.save, anche per DDP, in modo che i salvataggi siano paralleli ed economici anziché un collo di bottiglia seriale.
2. Salva in modo asincrono così i checkpoint sono quasi gratuiti, il che ti consente di salvare spesso.
3. Ripristina automaticamente l'ultimo checkpoint valido, in modo che un guasto costi minuti di ricalcolo, non ore.
4. Sovrapponi il caricamento dei dati al calcolo tramite caching e prefetching dallo storage remoto, in modo che gli acceleratori non rimangano mai inattivi in attesa dell'input. Si tratta di ore-GPU risparmiate a ogni step.
5. Salva il checkpoint della pipeline di dati e dello stato RNG, in modo che un job ripreso continui sui dati corretti anziché corrompere silenziosamente il modello.

Il principio unificatore: checkpoint frequenti, economici e completi trasformano un guasto hardware da un evento che interrompe il job a un semplice errore di arrotondamento, e una pipeline di input sovrapposta mantiene gli acceleratori occupati nel frattempo. I salvataggi economici (asincroni) rendono sostenibile la frequenza; i salvataggi completi (modello, dati e RNG) garantiscono la correttezza del ripristino. Con entrambi questi elementi attivi e una flotta che rileva e isola l'hardware guasto, il tempo di addestramento effettivo si avvicina al limite massimo consentito dall'hardware, indipendentemente da quanto sia instabile il cluster sottostante.

Riferimenti

• Meta, The Llama 3 Herd of Models (2024): §3.3.1, affidabilità dell'addestramento e analisi delle interruzioni su un massimo di 16.384 GPU H100.
• Characterizing the Resilience of Hopper H100 and Ampere A100 GPUs (2025): studio sul campo di 2,5 anni del sistema Delta; MTBE e probabilità di guasto in base alla dimensione del job.
• Databricks, How we keep GPUs reliable across Databricks AI: controllo dello stato delle GPU a livello di flotta e modello di guasto probabilistico.

Vuoi provarlo? Consulta la Guida alle prestazioni e alla resilienza dell'addestramento nella documentazione di Databricks AI Runtime per il codice completo, e leggi How we keep GPUs reliable across Databricks AI per scoprire l'aspetto legato all'infrastruttura.

(Questo post sul blog è stato tradotto utilizzando strumenti basati sull'intelligenza artificiale) Post originale