Spotlight: Synergizing Seed Exploration and Spot GPUs for DiT RL Post-Training

Riepilogo Tecnico: SPOTLIGHT

Problema

L'addestramento post-training tramite apprendimento per rinforzo (RL) per i Diffusion Transformer (DiT) è proibitivamente costoso, richiedendo spesso migliaia di GPU ad alte prestazioni. Il tipico flusso di lavoro in stile GRPO prevede due fasi: una fase di rollout (generazione di campioni con seed casuali) e una fase di training (calcolo dei gradienti e aggiornamento dei pesi).

Le strategie esistenti di riduzione dei costi affrontano due colli di bottiglia primari quando applicate ai DiT:

1. Overhead di Esplorazione dei Seed: Per accelerare la convergenza, metodi come DanceGRPO e TreeGRPO esplorano molteplici seed casuali per trovare campioni ad alto contrasto (reward elevati e bassi). Tuttavia, questa fase di esplorazione si trova sul percorso critico, aggiungendo tempo di calcolo a ogni iterazione e annullando il potenziale risparmio sui costi.
2. Idleness e Frammentazione delle GPU Spot: Sebbene le GPU spot offrano costi inferiori del 69–77%, esse sono preemttibili e hanno una larghezza di banda inter-nodo limitata.
   • Idleness (Inattività): A differenza dei Large Language Models (LLM) dove le latenze di rollout variano (permettendo il pipelining con il training), i rollout dei DiT completano in un tempo quasi uniforme. Ciò crea una barriera di sincronizzazione in cui le GPU spot rimangono inattive durante l'intera fase di training.
   • Frammentazione: Le frequenti preemption delle GPU spot interrompono i gruppi di Sequence Parallelism (SP), essenziali per eseguire grandi modelli DiT. La riconfigurazione di questi gruppi richiede tipicamente minuti per ricaricare i modelli e reinizializzare i motori, sprecando risorse significative.

Metodologia: SPOTLIGHT

Gli autori presentano SPOTLIGHT, il primo sistema progettato per sfruttare le GPU spot per il post-training RL dei DiT, sinergizzando l'esplorazione dei seed con le risorse spot. Il sistema si basa su due intuizioni chiave e tre tecniche principali.

Intuizioni Chiave

1. L'Esplorazione Tolera Pesi Obsoleti (Stale Weights): Gli autori dimostrano che il ranking relativo dei seed casuali (quali seed producono reward alti rispetto a bassi) rimane altamente coerente anche utilizzando i pesi del modello dell'iterazione precedente invece di quelli aggiornati correntemente. Ciò consente all'esplorazione dei seed di girare su GPU spot utilizzando pesi obsoleti mentre il cluster riservato esegue la fase di training, offrendo efficacementamente l'esplorazione dal percorso critico.
2. La Riconfigurazione SP Può Riutilizzare lo Stato: L'alto costo della riconfigurazione dei gruppi SP deriva dall'inizializzazione dello scheduler CPU e dal caricamento remoto del modello. Decoupling (disaccoppiando) lo scheduler e utilizzando il caricamento dei pesi intra-nodo (tramite NVLink) da peer co-localizzati, il tempo di riconfigurazione SP può essere ridotto da minuti a frazioni di secondo.

Architettura di Sistema e Tecniche

SPOTLIGHT è implementato sulla piattaforma RL open-source ROLL e impiega tre tecniche principali:

1. Planner di Esplorazione Basato su Bandit:

   • Il sistema formula la pianificazione dell'esplorazione come un problema multi-armed bandit.
   • Seleziona dinamicamente il numero di sequenze e i passi di denoising effettivi (utilizzando strategie di caching come Teacache) per massimizzare la varianza del reward entro il budget temporale fornito dalla fase di training.
   • Utilizza una strategia Upper Confidence Bound (UCB) per bilanciare l'esplorazione di nuove configurazioni e lo sfruttamento di impostazioni a alta varianza già note.

2. Sequence Parallelism (SP) Elastica:

   • Scheduler Persistente: Lo scheduler CPU è disaccoppiato dai worker GPU e mantenuto residente sul nodo attraverso i cambiamenti di SP, eliminando i costi ripetuti di inizializzazione.
   • Caricamento Pesi Intra-Nodo: Quando un nuovo worker GPU si unisce a un gruppo SP sullo stesso nodo, copia i pesi da un peer co-localizzato tramite NVLink invece di richiederli da un nodo remoto, riducendo drasticamente la latenza di riconfigurazione.

3. Scheduler di Richiesta Pull-Based Consapevole della Preemption:

   • I worker di rollout prelevano (pull) le richieste da una coda centralizzata invece di avere compiti assegnati staticamente, adattandosi alle velocità eterogenee dei worker.
   • Impegno dello Stato (State Commitment): Gli stati intermedi di rollout (latenti e metadati) sono salvati in un Tensor Store persistente sui nodi riservati. Se una GPU spot viene preemptionata, il sistema può recuperare lo stato in volo e riprendere la richiesta su un altro worker, minimizzando la perdita di progresso.

Risultati

Il sistema è stato valutato sul post-training di Qwen-Image (20B) utilizzando i dataset DeepSeek-OCR e Geneval a risoluzioni di 512×512 e 1280×1280.

• Riduzione dei Costi: SPOTLIGHT riduce il costo totale di training di 1,4× fino a 6,4× rispetto ai baseline allo stato dell'arte (inclusi i sistemi che utilizzano GPU spot senza queste ottimizzazioni e i sistemi basati solo su risorse riservate).
• Velocità di Convergenza: Raggiunge lo stesso punteggio di validazione target 4 volte più velocemente rispetto ai baseline. Ciò è ottenuto riducendo il numero di iterazioni necessarie (grazie all'efficace esplorazione dei seed) e sovrapponendo l'esplorazione al training.
• Qualità: SPOTLIGHT raggiunge punteggi di validazione superiori su tutti i dataset rispetto ai baseline, grazie alla selezione dinamica di seed ad alto contrasto.
• Efficienza:
  • L'overhead di esplorazione è mantenuto basso (2–3% del tempo di iterazione) inserendo l'esplorazione all'interno della finestra di training.
  • La SP elastica permette al sistema di seguire i cambiamenti di capacità spot quasi in tempo reale, mentre gli approcci ingenui soffrono ritardi di diversi minuti durante la riconfigurazione.
  • Il sistema dimostra una scalabilità quasi lineare del throughput all'aumentare del numero di GPU spot da 8 a 64.

Significato e Rivendicazioni

Il paper afferma che SPOTLIGHT è il primo sistema che riesce a sfruttare con successo le GPU spot per il post-training RL dei DiT. La sua importanza risiede nel fatto che:

1. Rompe la Dipendenza dal Percorso Critico: Dimostrando che l'esplorazione dei seed può tollerare pesi obsoleti, il sistema elimina il compromesso tra una convergenza più rapida (tramite esplorazione) e la latenza di iterazione.
2. Risolve l'Idleness delle GPU Spot: Risolve il problema specifico dell'inattività dei rollout dei DiT utilizzando la capacità idle del training per l'esplorazione, una sinergia trascurata dai sistemi precedenti focalizzati sugli LLM.
3. Robustezza alla Preemption: Attraverso la SP elastica e la persistenza dello stato, rende l'addestramento dei DiT resiliente alle frequenti preemption intrinseche nei mercati spot, trasformando un pool di risorse frammentato in un'infrastruttura di addestramento efficiente e conveniente.

Gli autori concludono che questi contributi consentono un sistema di post-training RL per DiT efficiente in termini di costi che non sacrifica la qualità delle immagini, rendendo l'addestramento di modelli di diffusione di alta qualità più accessibile.