LazyTrain: Limited-resource Allocation toward Zero-waste Yield Optimization in Large Language Model Training

Riepilogo Tecnico: LazyTrain

Problematica
L'addestramento di Large Language Models (LLM) su hardware limitato si è evoluto in un complesso problema di pianificazione (scheduling) che coinvolge il calcolo GPU, la memoria host (CPU DRAM), il trasferimento PCIe e la larghezza di banda di archiviazione (NVMe). Sebbene i sistemi esistenti come ZeRO e MegaTrain riducano con successo la residenza in GPU tramite l'offloading degli stati del modello e lo streaming dei layer, essi spesso si affidano a euristiche fisse per il checkpointing e il posizionamento delle attivazioni. Queste strategie fisse possono inavvertitamente esporre i trasferimenti di comunicazione (CPU-GPU o NVMe) sul percorso critico, creando colli di bottiglia che degradano il throughput anche quando i vincoli di memoria sono soddisfatti. La sfida centrale non è semplicemente far rientrare i tensori nella memoria, ma ottimizzare lo scheduling per nascondere il movimento dei dati all'interno delle finestre di calcolo.

Metodologia
LazyTrain affronta questo problema formulando l'addestramento di LLM a risorse limitate come un problema di scheduling di programmazione lineare intera mista (MIP) vincolato. Esso opera come uno strato di ottimizzazione sopra un esecutore di streaming dei layer (simile a MegaTrain).

1. MIP Scheduler:

   • Variabili Decisionali: Lo scheduler utilizza variabili binarie per selezionare i percorsi di checkpoint delle attivazioni, determinare le "case" di archiviazione (GPU HMB, CPU DRAM o NVMe) e definire i blocchi di ricalcolo. Le variabili continue assegnano il traffico delle attivazioni a specifiche finestre di calcolo.
   • Funzione Obiettivo: L'obiettivo è minimizzare la somma dei costi di ricalcolo e l'esposizione incrementale della comunicazione indotta dal posizionamento delle attivazioni. A differenza degli approcci precedenti che massimizzano il volume di offload, LazyTrain penalizza specificamente il traffico che non può essere nascosto sotto il calcolo della GPU.
   • Vincoli: Il modello impone i limiti di capacità di memoria per GPU, CPU e NVMe, nonché i vincoli di larghezza di banda per gli endpoint PCIe e NVMe. Tratta l'NVMe sia come consumatore PCIe che come endpoint separato, assicurando che il traffico SSD sia utilizzato solo quando può essere nascosto o quando è necessario per la fattibilità.
   • Esecuzione: Il solver (SCIP) viene eseguito offline prima dell'addestramento per generare una policy. Questa policy viene poi eseguita dal runtime, che esegue lo streaming dei parametri dalla CPU, calcola i layer e ricarica i checkpoint dai loro tier assegnati secondo lo schedule ottimizzato.

2. Operatore Ibrido a 8 bit:

   • Per ottimizzare ulteriormente il percorso di aggiornamento lato CPU, LazyTrain accoppia gli stati dell'ottimizzatore a 8 bit con il clipping rapido dei gradienti in un singolo operatore.
   • Meccanismo: Gli stati a 8 bit riducono l'impronta di memoria degli stati dell'ottimizzatore sulla CPU. Tuttavia, la quantizzazione e la dequantizzazione aggiungono overhead alla CPU. Per contrastare ciò, il sistema impiega il fast gradient clipping utilizzando kernel foreach della CPU per calcolare le norme e applicare la scalatura in modo efficiente.
   • Integrazione: Questo operatore è trattato come una componente formale del sistema piuttosto che come un'ottimizzazione esterna, essendo specificamente progettato per bilanciare il risparmio di memoria con i costi di calcolo della CPU.

Contributi Chiave

• Formulazione: Il paper formula l'addestramento di LLM a risorse limitate come un modello MILP che ottimizza congiuntamente la selezione dei checkpoint, il posizionamento delle attivazioni, il ricalcolo e la sovrapposizione della comunicazione.
• Implementazione dello Scheduler: Instanzia questo scheduler MIP su un esecutore di streaming dei layer, dimostrando che il baseline MegaTrain può essere visto come un'euristica fissa all'interno dello stesso spazio di scheduling.
• Componente Ibrida a 8 bit: Introduce un operatore ibrido a 8 bit accoppiato che riduce la memoria degli stati dell'ottimizzatore utilizzando il fast gradient clipping per neutralizzare l'overhead associato alla CPU.

Risultati Sperimentali
Gli esperimenti sono stati condotti su singole GPU H800 (80GB) e RTX 3090 (24GB) su scale di modelli da 3B a 27B (utilizzando architetture Qwen2.5 e Qwen3.6) sul dataset MetaMathQA.

• Throughput: Su H800 con un modello da 27B, LazyTrain ha raggiunto 219,95 TFLOPS e 1361 token/s con un batch size di 72, rispetto ai 176,90 TFLOPS e 1075,8 token/s del baseline MegaTrain. Ciò rappresenta un miglioramento del throughput sostenuto di 1,24×.
• Fattibilità: Su RTX 3090 con risorse limitate, LazyTrain ha aumentato il batch size massimo fattibile di uno rispetto a MegaTrain ad ogni scala di modello. Ad esempio, ha completato con successo l'addestramento del modello da 27B con batch size 2, mentre ZeRO-3 Offload falliva con batch size 1.
• Qualità: L'esecuzione di LazyTrain su 27B ha ottenuto un'accuratezza exact-match del 95,42% sull'intero split di valutazione, comparabile o leggermente superiore ai baseline.
• Studi di Ablazione:
  • Rimuovendo lo scheduler MILP (mantenendo gli altri miglioramenti del runtime), il throughput è sceso a 193,17 TFLOPS (una riduzione del 12,2%), identificando lo scheduler come la fonte primaria dei guadagni di performance.
  • Rimuovendo l'operatore Ibrido a 8 bit, si è verificata una diminuzione trascurabile dello 0,3% (219,29 TFLOPS), confermando il suo ruolo nel bilanciare l'overhead della CPU.

Significatività e Rivendicazioni
Il paper sostiene che l'addestramento a risorse limitate debba essere trattato come un problema di scheduling vincolato piuttosto che come un semplice compito di offloading della memoria. Ottimizzando congiuntamente dove e quando i dati si muovono, LazyTrain riesce a nascondere con successo la comunicazione sul percorso critico, ottenendo miglioramenti significativi del throughput senza sacrificare la qualità del modello. Gli autori sottolineano che lo scheduler MILP è l'innovazione centrale, consentendo al sistema di adattare dinamicamente le strategie di posizionamento e ricalcolo alle specifiche restrizioni hardware (GPU, CPU, NVMe) e ai profili di larghezza di banda. Il lavoro dimostra che anche su configurazioni a singola GPU con budget di memoria ristretti, l'addestramento ad alto throughput di grandi modelli è realizzabile attraverso un'ottimizzazione rigorosa della gerarchia della memoria e del movimento dei dati.