🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

Il documento presenta Asteria, un sistema runtime che abilita un'ottimizzazione del secondo ordine pratica e scalabile per i grandi modelli linguistici distribuendo dinamicamente gli stati dell'ottimizzatore attraverso gerarchie di memoria eterogenee e disaccoppiando i costosi calcoli del precondizionatore dal percorso critico di addestramento sulla GPU per ridurre l'overhead e accelerare la convergenza.

Autori originali: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Pubblicato 2026-05-18

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di dover insegnare a un robot gigante e super-intelligente (un Modello Linguistico di grandi dimensioni) a scrivere storie. Per farlo, hai bisogno di un "insegnante" (un ottimizzatore) che esamini gli errori del robot e gli dica come migliorare.

Attualmente, la maggior parte degli insegnanti utilizza un metodo semplice chiamato AdamW. È come uno studente che controlla i compiti, vede alcune risposte sbagliate e apporta piccole e rapide correzioni. È veloce e funziona bene, ma richiede molta pratica (milioni di esempi) per diventare davvero bravo.

Esiste un metodo di insegnamento più intelligente e avanzato chiamato Ottimizzazione del Secondo Ordine (come SOAP o Shampoo). Questo insegnante non guarda solo cosa era sbagliato; analizza la forma degli errori per comprendere il problema in profondità. Impara molto più velocemente e ha bisogno di meno esempi. Tuttavia, c'è un enorme ostacolo: questo insegnante intelligente è incredibilmente pesante. Richiede enormi quantità di memoria e potenza di calcolo per eseguire i suoi complessi calcoli matematici, causando spesso il blocco del computer o un rallentamento tale che l'insegnante semplice completa effettivamente il lavoro per primo.

Ecco Asteria: Il "Responsabile Logistico Intelligente"

I ricercatori di Oxford hanno costruito un nuovo sistema chiamato Asteria. Pensate ad Asteria non come a un nuovo insegnante, ma come a un brillante responsabile logistico che riorganizza l'intera aula in modo che l'insegnante intelligente possa finalmente svolgere il suo lavoro senza distruggere la scuola.

Ecco come Asteria risolve i tre problemi più grandi, utilizzando semplici analogie:

1. Il Problema "Arredamento in una Stanza Piccola" (Memoria)

Il Problema: L'insegnante intelligente deve mantenere enormi e complessi grafici (matrici) nella sua testa (la memoria GPU). Su un computer standard, non c'è abbastanza spazio. È come cercare di infilare un letto matrimoniale, un tavolo da pranzo e un armadio in un monolocale. La stanza si riempie e l'insegnante deve arrendersi.
La Soluzione di Asteria: Asteria agisce come un esperto di mobili pieghevoli. Si rende conto che l'insegnante non ha bisogno di tenere tutto in mano contemporaneamente.

Mantiene le parti più attive dei grafici sulla GPU (la scrivania).
Sposta le parti pesanti e meno frequentemente utilizzate sulla CPU (il corridoio) o persino su un hard disk (il garage).
Fondamentalmente, riporta le parti pesanti sulla scrivania esattamente solo quando sono necessarie. Questo permette all'insegnante intelligente di lavorare su un piccolo portatile tanto bene quanto su un supercomputer.

2. Il Problema "Ingorgo Stradale" (Velocità)

Il Problema: Ogni pochi passaggi, l'insegnante intelligente deve eseguire un calcolo massiccio e complesso (come risolvere un gigantesco puzzle) per aggiornare i suoi grafici. Questo richiede molto tempo e blocca l'intera classe dal lavorare. È come un camion delle consegne che blocca l'intera autostrada mentre scarica un singolo pacco. La GPU (il cervello del computer) rimane inattiva, in attesa.
La Soluzione di Asteria: Asteria introduce una catena di montaggio parallela.

Invece di fermare la classe principale per eseguire i calcoli pesanti, Asteria invia il "lavoro pesante" a un team di lavoratori nell'ufficio retrostante (la CPU).
Mentre la classe principale (la GPU) continua a insegnare al robot, i lavoratori dell'ufficio retrostante risolvono silenziosamente i complessi puzzle in background.
Quando la classe ha bisogno dei nuovi grafici, l'ufficio retrostante li ha già completati e li ha fatti scivolare fuori. La classe principale non deve mai fermarsi. L'"ingorgo stradale" scompare.

3. Il Problema "Chat di Gruppo" (Addestramento Distribuito)

Il Problema: Quando si addestra con molti computer contemporaneamente, tutti devono solitamente fermarsi e concordare esattamente le stesse informazioni prima di passare al passo successivo. È come una chat di gruppo in cui tutti devono aspettare che la persona più lenta risponda prima che chiunque possa scrivere di nuovo.
La Soluzione di Asteria: Asteria utilizza una politica "Buona Abbastanza".

Invece di aspettare che tutti siano perfettamente sincronizzati, permette ai computer di lavorare con informazioni leggermente "obsolette" (leggermente datate) per un breve periodo.
Invia aggiornamenti solo quando sono davvero necessari.
Questo mantiene il gruppo in movimento veloce, come una squadra che si fida l'una dell'altra per continuare a lavorare mentre recupera gli aggiornamenti in seguito, invece di fermare l'intero progetto ogni cinque minuti.

I Risultati: Cosa Hanno Scoperto?

I ricercatori hanno testato Asteria su hardware reale, inclusi un potente singolo computer (Nvidia DGX Spark) e un grande cluster di computer (Nvidia GH200).

Funziona su macchine piccole: Sono riusciti ad addestrare un modello linguistico di grandi dimensioni (1 miliardo di parametri) su una singola macchina che precedentemente non poteva gestire i requisiti di memoria di questo insegnante intelligente.
È più veloce in tempo reale: Poiché nasconde gli "ingorghi stradali", l'addestramento si conclude in meno tempo reale, anche se la matematica è più complessa.
Risparmia energia: Mantenendo il cervello del computer (GPU) occupato e impedendogli di rimanere inattivo in attesa di calcoli, Asteria utilizza in realtà meno energia totale per ottenere lo stesso risultato rispetto ai vecchi metodi goffi.
Non perde qualità: Il modello impara esattamente allo stesso modo in cui lo farebbe con l'insegnante intelligente "nativo" (non ottimizzato), ma ci arriva molto più velocemente e a costi inferiori.

In Sintesi:
Asteria non inventa una nuova formula matematica. Invece, ripensa come il computer esegue quella matematica. Separa il lavoro pesante dal lavoro principale, utilizza intelligentemente tutto lo spazio di archiviazione disponibile e permette ai computer di lavorare in modo asincrono. Questo trasforma un metodo "teoricamente ottimo ma praticamente impossibile" in uno strumento pratico per l'addestramento della prossima generazione di intelligenza artificiale.

Riepilogo Tecnico: Asteria – Ottimizzazione del Secondo Ordine Orchestrata a Runtime per l'Addestramento Scalabile di LLM

1. Enunciato del Problema

I metodi di ottimizzazione del secondo ordine (ad es. Shampoo, SOAP, KL-Shampoo) offrono una maggiore efficienza statistica e una convergenza più rapida per l'addestramento di Modelli Linguistici di Grande Dimensione (LLM) rispetto agli ottimizzatori del primo ordine come AdamW. Tuttavia, il loro dispiegamento pratico su larga scala è ostacolato da gravi barriere a livello di sistema:

Vincoli di Memoria: I metodi del secondo ordine richiedono l'archiviazione di grandi matrici di covarianza ( $L$ e $R$ ) e l'esecuzione di decomposizioni spettrali dense, portando a un footprint di memoria di $O(N^2)$ e a un costo computazionale di $O(N^3)$ . In ambienti con vincoli di memoria (ad es. workstation single-node con Architetture di Memoria Unificata), questi stati esauriscono la memoria fisica, causando fallimenti per Out-of-Memory (OOM) o costringendo a riduzioni delle dimensioni del modello e della dimensione del batch.
Picchi di Latenza: Il calcolo periodico delle radici inverse (ad es. $L^{-1/4}$ ) comporta operazioni matriciali a costo cubico. Nelle implementazioni esistenti, questi aggiornamenti sincroni bloccano la pipeline di addestramento GPU, causando gravi picchi periodici di latenza che interrompono la sovrapposizione fine tra calcolo e comunicazione essenziale per un addestramento distribuito efficiente (ad es. FSDP).
Sovraccarico di Sincronizzazione: I metodi distribuiti del secondo ordine attuali si affidano a aggiornamenti globali dello stato rigidi e sincroni (AllGather) su tutti i rank. Ciò accoppia la comunicazione intra-nodo veloce con la sincronizzazione inter-nodo lenta, impedendo al sistema di sfruttare limiti di obsolescenza asincroni e consapevoli della topologia.

Le soluzioni esistenti, come lo sharding basato su DTensor, non riescono a risolvere questi problemi in ambienti single-node con risorse limitate o quando il sovraccarico di comunicazione della sincronizzazione globale supera i benefici.

2. Metodologia: Il Runtime Asteria

Asteria è un sistema runtime co-progettato hardware-software costruito su PyTorch FSDP. Decoupla la logica di ottimizzazione del secondo ordine dal percorso critico di addestramento GPU attraverso tre meccanismi fondamentali:

A. Tiering Asimmetrico della Memoria Adattivo all'Architettura

A differenza degli ottimizzatori del primo ordine che trattano lo stato come vettori omogenei, Asteria riconosce il ciclo di vita asimmetrico degli stati del secondo ordine:

GPU: Archivia factor_matrices (fattori di Kronecker) che vengono aggiornati frequentemente durante l'addestramento.
Host (CPU/UVM): Archivia inv_factor_matrices (fattori inversi). Questi sono mantenuti in memoria accessibile alla CPU per evitare pressioni sulla memoria GPU.
NVMe (Opzionale): Fornisce un livello di staging per esecuzioni con capacità limitata, consentendo lo swapping degli stati tramite calcolo serializzato o paging JIT.
Meccanismo: Il runtime cattura snapshot lato host delle matrici dei fattori, invia worker CPU asincroni per calcolare le radici inverse e scrive i risultati nella memoria host. Quando necessario, la GPU accede a questi tramite visualizzazioni di Memoria Virtuale Unificata (UVM) o prefetching esplicito, evitando copie ridondanti host-dispositivo.

B. Pipeline di Stato Ombra Orchestrata da Hook

Per eliminare i picchi di latenza, Asteria rimuove il calcolo delle radici inverse di costo $O(N^3)$ dal percorso critico GPU:

Stream Ombra: Uno stream CUDA a bassa priorità gestisce lo staging degli stati lato GPU (prefetching e paging) senza bloccare lo stream principale di addestramento.
Esecuzione Asincrona: Gli aggiornamenti costosi delle radici inverse vengono eseguiti asincronamente su thread worker CPU.
Integrazione Hook: Hook leggeri di PyTorch (register_forward_hook, register_full_backward_pre_hook) agiscono come segnali di scheduling. Attivano lo svuotamento dei trasferimenti di stato pronti e il prefetching dei fattori inversi prossimi al consumo, garantendo che la GPU abbia precondizionatori freschi senza bloccare il ciclo principale di addestramento.

C. Coerenza a Obsolescenza Limitata Consapevole della Topologia

Per l'addestramento distribuito, Asteria sostituisce la sincronizzazione globale rigida con un protocollo di coerenza selettivo:

Tracciamento dell'Obsolescenza: Un CoherenceRegistry traccia la freschezza di ogni blocco precondizionatore. Vengono sincronizzati solo i blocchi che superano un budget di obsolescenza configurato.
Sincronizzazione Gerarchica: La sincronizzazione avviene in due fasi: prima all'interno di gruppi di processi locali al nodo, poi attraverso i rappresentanti dei nodi. Ciò minimizza il traffico inter-nodo ed evita inutili round trip CPU-GPU mantenendo la comunicazione su memoria accessibile all'host.
Barriera Limitata: Il runtime impone una barriera locale che assicura che, se una versione del precondizionatore diventa troppo obsoleta mentre un job CPU è in esecuzione, l'ottimizzatore attenda il completamento, bilanciando freschezza e throughput.

3. Contributi Chiave

Decoupling a Runtime: Asteria dimostra che l'ottimizzazione del secondo ordine può essere resa pratica non semplificando l'algoritmo, ma ripensando la gestione dello stato, il calcolo in background e la sincronizzazione a livello di runtime.
Scalabilità della Memoria: Abilita l'addestramento del secondo ordine su hardware single-node con vincoli di memoria (ad es. una singola GPU GB10 con 128GB di memoria unificata) dove lo sharding tradizionale è impossibile.
Nascondimento della Latenza: Spostando i calcoli delle radici inverse sull'host e utilizzando una pipeline ombra, Asteria appiattisce i picchi periodici di latenza, ottenendo un profilo di esecuzione hardware fluido paragonabile ai metodi del primo ordine.
Efficienza Energetica: Il sistema migliora il compromesso perdita-energia, ottenendo una maggiore riduzione della perdita per unità di energia riducendo i tempi di inattività della GPU e ottimizzando la sovrapposizione del calcolo lato host.

4. Risultati Sperimentali

Gli autori hanno valutato Asteria su modelli della famiglia OLMo (660M, 1B, 7B parametri) utilizzando il framework di addestramento OLMo su due piattaforme:

Single-Node (DGX Spark, GPU GB10):
- Throughput: Asteria-SOAP e Asteria-KL-Shampoo hanno eliminato i picchi periodici di tempo per step osservati nelle implementazioni native (riducendo il tempo di picco per step da ~96s a ~66.5s).
- Energia: Sebbene i metodi del secondo ordine nativi consumassero circa il 17-20% in più di energia totale SoC rispetto ad AdamW, Asteria ha ridotto significativamente questo divario (ad es. Asteria-KL-Shampoo ha consumato solo circa il 7.9% in più di energia rispetto ad AdamW).
- Efficienza: Asteria-KL-Shampoo ha raggiunto la massima efficienza di riduzione della perdita normalizzata (3.3010), superando sia i metodi del secondo ordine nativi che AdamW.
Multi-Node (Cluster GH200):
- Convergenza: Le varianti di Asteria hanno mantenuto traiettorie di convergenza vicine alle loro controparti native del secondo ordine, superando significativamente AdamW nel tempo reale per raggiungere i livelli di perdita target.
- Tolleranza all'Obsolescenza: Esperimenti che variavano il budget di obsolescenza ( $S$ ) hanno mostrato che aumentare $S$ da 1 a 5 ha ridotto sostanzialmente il tempo di addestramento con un impatto trascurabile sulla qualità finale del modello.
- Scalabilità Forte: Su modelli da 7B che scalano da 2 a 16 nodi, Asteria ha raggiunto un speedup più elevato e un tempo per step inferiore rispetto ai baseline nativi, dimostrando efficacia anche quando il costo relativo della manutenzione dello stato aumenta.

5. Significato e Affermazioni

Il paper sostiene che il principale ostacolo all'addestramento pratico di LLM del secondo ordine non è il costo matematico degli aggiornamenti consapevoli della curvatura, ma il mismatch di sistema tra la manutenzione degli stati del secondo ordine e gli stack di addestramento moderni.

Riformulazione del Problema: Gli autori affermano che l'ottimizzazione del secondo ordine dovrebbe essere vista come un problema di co-design che coinvolge algoritmi, scheduling del runtime e sistemi di memoria eterogenea.
Fattibilità: Dimostrano che la memoria aggregata GPU di grandi dimensioni non è un prerequisito fondamentale per l'addestramento del secondo ordine; con un posizionamento consapevole del ciclo di vita e un accesso a stadi, questi metodi possono essere eseguiti con budget di memoria ristretti.
Praticità: Trattando la freschezza del precondizionatore e la comunicazione come risorse gestite a runtime piuttosto che come stati globali rigidi, Asteria riduce il divario tra la promessa statistica dell'ottimizzazione del secondo ordine e la sua deployabilità nel mondo reale.

Il paper conclude che Asteria rende l'ottimizzazione del secondo ordine una scelta di sistema praticabile per l'addestramento di LLM, offrendo una migliore convergenza ed efficienza energetica senza richiedere semplificazioni algoritmiche.