The Pulse
AWS mostra un recupero in 10 secondi dai guasti durante l’addestramento su GPU
AWS e NVIDIA dimostrano una configurazione di addestramento PyTorch tollerante ai guasti su Amazon EKS, usando NVIDIA Resiliency Extension. L’approccio riduce il tempo di recupero misurato da 270 secondi con i riavvii standard di Kubernetes
AI.info Team ·
AWS afferma che i processi di addestramento distribuito su GPU in Amazon Elastic Kubernetes Service possono riprendere in pochi secondi anziché in minuti dopo il guasto di un worker, ma il miglioramento dipende dalla sostituzione dei riavvii a livello di Kubernetes con meccanismi di recupero all’interno del processo di addestramento.
In un post pubblicato il 16 settembre, Aravind Neelakantan, Specialist Solutions Architect di AWS, e Shreya Gupta, AI Architect di NVIDIA, descrivono l’integrazione di NVIDIA Resiliency Extension, o NVRx, con l’addestramento PyTorch Fully Sharded Data Parallel su Amazon EKS. Il loro test ha confrontato il recupero con NVRx con un riavvio convenzionale di Kubernetes, rilevando una netta differenza: circa 10 secondi per il recupero di NVRx all’interno del processo, 17 secondi per il launcher a livello di job di NVRx e 270 secondi per la configurazione di riferimento con Kubernetes.
I risultati provengono da iniezioni controllate di guasti, non da un incidente in produzione. Gli autori hanno inserito cinque guasti deterministici in ogni sessione di addestramento di 2.000 passaggi, usando due nodi p5.48xlarge con 16 GPU H100 e salvando checkpoint ogni 500 passaggi.
Perché il guasto di un singolo rank può fermare l’intero job
I grandi job di addestramento distribuito funzionano in modo sincrono, quindi il guasto di un worker può costringere gli altri, ancora operativi, ad aspettare. AWS descrive una sequenza di guasti in cui un rank si arresta, quelli ancora attivi raggiungono il timeout di 60 secondi della NVIDIA Collective Communication Library e Kubernetes riavvia poi i pod in momenti non sincronizzati. Il conseguente CrashLoopBackOff e gli ulteriori cicli di timeout prolungano l’interruzione ben oltre il guasto iniziale.
Questa architettura crea una costosa discrepanza tra il guasto e la risposta. Un’eccezione transitoria o un blocco della comunicazione possono richiedere soltanto il ripristino del processo, mentre un worker terminato per esaurimento della memoria o per un guasto del sistema operativo ha bisogno di essere sostituito. Trattare entrambi i casi come un riavvio completo del container scarta più stato e lascia inattive più GPU.
NVRx distingue i guasti temporanei dagli arresti anomali gravi
NVRx offre tre meccanismi indipendenti. Il suo livello di checkpoint asincrono sposta la scrittura dei checkpoint in un processo in background, consentendo al ciclo di addestramento di proseguire mentre ogni rank FSDP scrive la propria partizione. Il wrapper all’interno del processo intercetta le eccezioni transitorie e i blocchi NCCL rilevati dal watchdog, interrompe il gruppo di processi interessato, verifica lo stato di GPU, NVLink e rete e reintegra i worker superstiti ripartendo dall’ultimo checkpoint.
Il secondo livello di recupero usa ft_launcher di NVRx. Un monitor dei segnali di attività controlla ogni rank e può terminare i processi superstiti bloccati, liberare la memoria GPU, stabilire un nuovo punto di rendezvous e avviare worker nuovi dopo eventi come SIGKILL, una terminazione per esaurimento della memoria o un blocco a livello di sistema operativo. I worker ricaricano quindi l’ultimo checkpoint.
AWS presenta i livelli come risposte distinte a diversi tipi di guasto: il wrapper Python gestisce quelli che restano all’interno del processo, ft_launcher gestisce i guasti che terminano o bloccano il processo e l’orchestratore del cluster resta responsabile della perdita di nodi.
Il benchmark premia il recupero all’interno del job
Nel test di recupero dai guasti, il riavvio di NVRx all’interno del processo ha raggiunto un goodput di addestramento del 31% e un goodput dell’infrastruttura dell’87%. La configurazione con ft_launcher ha raggiunto un goodput di addestramento del 25,5% e un goodput dell’infrastruttura dell’85,9%, mentre la configurazione di riferimento con riavvio di Kubernetes ha ottenuto un goodput di addestramento dell’11,5% e un goodput dell’infrastruttura del 35,8%.
Il goodput misura i progressi utili dell’addestramento al netto delle interruzioni. I dati di NVRx non significano che il sistema abbia continuato ad addestrare normalmente durante i guasti: mostrano quanto lavoro utile in più è stato prodotto nella sessione una volta conteggiato il costo del recupero.
Il test usa Llama 3.1 8B con FSDP e una sequenza prestabilita di cinque guasti. AWS non presenta i risultati come una garanzia valida per ogni modello, sistema di archiviazione o configurazione di cluster. Il tempo di recupero dipende anche dal caricamento dei checkpoint, che secondo gli autori, su larga scala, costituisce il principale fattore di ritardo, più del meccanismo di riavvio.
I checkpoint asincroni eliminano un altro collo di bottiglia
Il post misura anche il costo del salvataggio dei checkpoint. AWS ha testato Llama 3.1 8B da due a otto nodi, ovvero da 16 a 64 GPU H100, con checkpoint ogni 1.000 passaggi. In questo intervallo, il checkpointing asincrono ha mantenuto l’efficienza di addestramento tra il 99,2% e il 99,8%, mentre quello sincrono è rimasto tra il 57% e il 61%.
La differenza dipende dai tempi di archiviazione. AWS ha misurato circa 275 secondi per il percorso di scrittura su FSx for Lustre, un ritardo rimasto sostanzialmente costante mentre il cluster passava da 16 a 64 GPU. I salvataggi sincroni costringono tutti i rank ad aspettare, mentre NVRx sovrappone la scrittura al segmento di addestramento successivo.
Salvataggi più frequenti mettono in evidenza il compromesso. Con otto nodi e un checkpoint ogni 100 passaggi, l’efficienza del checkpointing sincrono è scesa al 14,7%, contro il 29,6% di quello asincrono. Con un intervallo di 1.000 passaggi, i valori corrispondenti erano il 60,3% e il 99,8%.
Cosa AWS chiede ai team di implementare
La configurazione di riferimento usa nodi p5.48xlarge autogestiti, ciascuno con otto GPU NVIDIA H100 da 80 GB e 32 interfacce di rete Elastic Fabric Adapter. I pod di addestramento vengono eseguiti come Kubernetes Job, mentre Amazon FSx for Lustre fornisce l’archiviazione condivisa dei checkpoint e Amazon ECR archivia l’immagine di addestramento.
L’esempio richiede un cluster EKS 1.28 o successivo, i plugin dei dispositivi NVIDIA ed EFA, PyTorch 2.9 o successivo e NVRx. Il post afferma che il benchmark usa NVRx 0.4.1, mentre per un’implementazione attuale è necessario usare la versione 0.6.0 con una configurazione aggiornata del launcher.
AWS pubblica i moduli Terraform, la configurazione del container e i manifest Kubernetes nella sua implementazione di riferimento. Il cambiamento pratico è più circoscritto di una riprogettazione del sistema di addestramento: i team possono aggiungere indipendentemente il checkpointing asincrono, il recupero all’interno del processo o il riavvio a livello di job, a seconda dei guasti che devono contenere.