Riepilogo Tecnico: Kareus – Riduzione Congiunta dell'Energia Dinamica e Statica nell'Addestramento di Grandi Modelli

Problema

La rapida crescita della domanda di calcolo per l'IA ha superato l'offerta di energia, rendendo l'energia una risorsa costosa e contesa per l'addestramento di grandi modelli. Gli attuali sforzi di ottimizzazione affrontano tipicamente o l'energia dinamica (consumata dal calcolo attivo, riducibile tramite la scalabilità della frequenza) o l'energia statica (consumata dall'hardware indipendentemente dall'attività, riducibile accorciando il tempo di esecuzione).

Le soluzioni esistenti soffrono di un limite critico: ottimizzano questi fattori in isolamento o tramite combinazioni ingenue.

• Gli approcci di scalabilità della frequenza (es. Perseus) riducono l'energia dinamica abbassando le frequenze delle GPU per i percorsi non critici, ma spesso seguono l'esecuzione sequenziale dei kernel, perdendo le opportunità di sovrapporre comunicazione e calcolo.
• Gli approcci di scheduling dei kernel a grana fine (es. Nanobatching) riducono l'energia statica sovrapponendo comunicazione e calcolo, ma ignorano l'energia dinamica e la scalabilità della frequenza.

Il documento identifica che l'allocazione degli SM, il timing del lancio dei kernel e la frequenza della GPU sono interdipendenti. Cambiare un fattore sposta la configurazione ottimale degli altri. Ad esempio, abbassare la frequenza rende i kernel relativamente più vincolati dal calcolo (compute-bound), alterando le strategie ottimali di contesa delle risorse e di sovrapposizione. Di conseguenza, l'ottimizzazione di sottoinsiemi di questi fattori o la loro combinazione ingenua non riesce a spingere la frontiera del compromesso tempo-energia verso il suo limite teorico.

Metodologia

Gli autori propongono Kareus, un sistema di addestramento progettato per ottimizzare congiuntamente l'allocazione degli SM, il timing del lancio dei kernel e la frequenza della GPU per trovare la frontiera Pareto-ottimale del compromesso tempo-energia.

1. Modello di Esecuzione a Sovrapposizione Partizionata

Per rendere trattabile l'ottimizzazione globale, Kareus decompone il problema globale in sottoproblemi locali utilizzando un modello di esecuzione a sovrapposizione partizionata:

• Partizionamento: Il grafo di esecuzione dei passaggi forward e backward viene diviso in partizioni fisse. Una partizione consiste in un kernel di comunicazione (es. AllReduce) e nella sequenza contigua più lunga di kernel di calcolo da un microbatch differente che non ha dipendenze di dati con il kernel di comunicazione.
• Generalizzazione: Questo modello generalizza le tecniche di nanobatching, aggiungendo un controllo fine-grained sull'allocazione degli SM e sul timing del lancio per ogni partizione.
• Composizione Gerarchica: Il sistema ottimizza ogni partizione indipendentemente e poi compone gerarchicamente queste frontiere locali in una frontiera globale a livello di microbatch e di iterazione.

2. Ottimizzazione Bayesiana Multi-Obiettivo Multi-Passaggio (MBO)

Poiché la ricerca esaustiva dello spazio di configurazione (che coinvolge conteggi SM, timing di lancio e frequenze) è proibitivamente costosa (richiedendo migliaia di ore-GPU e causando interferenze termiche), Kareus impiega un algoritmo MBO specializzato:

• Modelli Surrogate: Utilizza due modelli separati di Gradient-Boosted Decision Tree (XGBoost) per prevedere rispettivamente il tempo di esecuzione e l'energia dinamica.
• Selezione dei Candidati Multi-Passaggio: Per esplorare efficacemente la frontiera di Pareto, l'algoritmo utilizza più funzioni di acquiszione in parallelo:
  • Miglioramento dell'Ipervolume (HVI): Tre distinte funzioni HVI guidano la ricerca verso un'energia totale inferiore, un'energia dinamica inferiore (favorendo frequenze più basse) e un'energia statica inferiore (favorendo un'esecuzione più veloce).
  • Esplorazione dell'Incertezza: Un ensemble bootstrap di modelli surrogate stima l'incertezza della previsione, dando priorità ai candidati nelle regioni dello spazio di ricerca meno esplorate.
• Profiling Termicamente Stabile: Per garantire misurazioni energetiche accurate, il sistema esegue le partizioni su una finestra di 5 secondi e inserisce un periodo di raffreddamento di 5 secondi tra le configurazioni per prevenire bias termici.

3. Implementazione del Sistema

Kareus è implementato integrandosi con Megatron-LM e Perseus.

• Utilizza MSCCL++ per il controllo fine-grained dell'allocazione degli SM.
• Impiega una torch.autograd.Function personalizzata per avvolgere le partizioni, consentendo schedule distinte per i passaggi forward e backward.
• A runtime, seleziona uno schema di esecuzione dalla frontiera pre-calcolata in base ai vincoli dell'utente (es. deadline temporale o budget energetico) e regola dinamicamente le frequenze della GPU tramite un controller adattato da Perseus.

Risultati Chiave

Kareus è stato valutato su 14 carichi di lavoro, inclusi l'addestramento reale su Llama 3.2 3B e Qwen 3 1.7B, e una grande emulazione di Llama 3.3 70B su 16 GPU NVIDIA A100.

• Compromesso Tempo-Energia: Rispetto al baseline allo stato dell'arte (Megatron-LM + Perseus), Kareus ottiene:
  • Fino al 28,3% di riduzione dell'energia di addestramento allo stesso tempo di addestramento.
  • Fino al 27,5% di riduzione del tempo di addestramento allo stesso consumo di energia.
• Performance di Throughput Massimo: In modalità di throughput massimo, Kareus riduce il tempo di iterazione fino al 14,9% e l'energia fino al 22,1% rispetto al baseline Megatron-LM.
• Studi di Ablazione: Rimuovendo o la scalabilità della frequenza o lo scheduling dei kernel si verifica una significativa degradazione delle prestazioni (es. rimuovere lo scheduling dei kernel porta a un aumento dell'energia del 10,8%), confermando la necessità dell'ottimizzazione congiunta.
• Scalabilità: I risultati di emulazione su larga scala per Llama 3.3 70B mostrano tendenze coerenti con i risultati del testbed fisico, dimostrando l'applicabilità del sistema a modelli massivi.

Significato e Rivendicazioni

Il documento afferma che Kareus spinge la frontiera del compromesso tempo-energia per l'addestramento di grandi modelli dimostrando che l'energia dinamica e quella statica non possono essere ottimizzate indipendentemente. Il contributo principale è l'intuizione che gli schemi di esecuzione (allocazione SM e timing) e la frequenza della GPU sono congiuntamente e interdipendentemente determinati.

Gli autori sottolineano che:

1. Il Controllo Congiunto è Superiore: Ottimizzare congiuntamente frequenza e scheduling è provabilmente più efficiente dal punto di vista energetico rispetto a lasciare il controllo della frequenza ai controller di potenza hardware o ottimizzare lo scheduling senza consapevolezza della frequenza.
2. L'Interdipendenza è la Chiave: L'allocazione ottimale degli SM e il timing del lancio sono specifici per la frequenza; uno schema ottimale a una certa frequenza può essere subottimale a un'altra.
3. Efficienza Pratica: Decomponendo il problema in partizioni e utilizzando l'MBO, Kareus ottiene questi guadagni senza il costo proibitivo di una ricerca esaustiva, rendendolo una soluzione pratica per gli attuali e futuri carichi di lavoro di addestramento IA su larga scala.