🤖 machine learning

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training

FORGE è un metodo di addestramento per LLM efficiente in termini di memoria che fonde lo step dell'ottimizzatore nel passaggio di backward per eliminare i gradienti materializzati processandoli interamente nei registri, dimezzando così l'uso della memoria dell'ottimizzatore, accelerando l'addestramento e preservando la fedeltà a piena precisione senza alterare la logica di aggiornamento sottostante.

Autori originali: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

Pubblicato 2026-06-23

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Dikshant Kukreja, Kritarth Prasad, Avinash Anand, Zhengkui Wang, Erik Cambria, Timothy Liu, Aik Beng Ng, Simon See, Bapi Chatterjee

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

Il Grande Problema: Il "Laboratorio Disordinato"

Immagina di stare addestrando un modello di IA gigante (come un cervello robotico) su un computer. Per insegnare a questo robot, il computer deve eseguire una quantità massiccia di calcoli matematici.

Nel modo standard di farlo (chiamato "Differenziazione in modalità inversa"), il computer segue un processo rigoroso in due fasi:

Calcolare gli Errori: Osserva i dati, capisce dove il robot ha sbagliato e scrive una gigantesca lista di "note di correzione" (gradienti) per ogni singola parte del cervello del robot. Scrive queste note su una gigantesca lavagna (la memoria del computer).
Applicare le Correzioni: Solo dopo che l'intera lavagna è piena, il computer prende una gomma e un pennarello per aggiornare effettivamente il cervello del robot basandosi su quelle note.

Il Collo di Bottiglia: Il problema è che il computer deve tenere quella gigantesca lavagna di note nella sua memoria mentre tiene anche il cervello del robot e le note per il passaggio successivo. Questa "lavagna" occupa così tanto spazio che spesso riempie la memoria del computer prima ancora che l'addestramento possa iniziare. È come cercare di riparare un'auto in un garage, ma devi tenere l'intero progetto dell'auto, gli attrezzi e il manuale di riparazione tutti stesi sul pavimento contemporaneamente. Se il garage è troppo piccolo, non riesci a farci stare l'auto.

La Soluzione: FORGE (Il Metodo della "Correzione Istantanea")

Gli autori di questo paper, FORGE (Fused On-Register Gradient Elimination), dicono: "Perché scrivere le note sulla lavagna affatto?"

Sostengono che la gigantesca lista di note sia solo un artefatto del modo in cui facciamo i calcoli, non qualcosa di cui il processo di apprendimento abbia realmente bisogno.

Come funziona FORGE:
Inveve di scrivere le note per poi rileggerle, FORGE esegue il calcolo e la correzione contemporaneamente, nella zona di archiviazione più veloce e piccola del computer (chiamata "registri").

L'Analogia: Immagina uno chef esperto che cucina un piatto complesso.
- Vecchio Modo: Lo chef trita una cipolla, scrive "cipolla tritata" su un taccuino, mette la cipolla in una ciotola e poi passa all'ingrediente successivo. Una volta che tutti gli ingredienti sono stati tritati ed elencati, lo chef legge il taccuino e mescola tutto insieme. Il taccuino occupa spazio sul bancone.
- Modo FORGE: Lo chef trita una cipolla e la getta immediatamente nella pentola. Poi trita una carota e la getta immediatamente. Non scrive mai nulla sul taccuino. Non ha mai bisogno di un taccuino. Il bancone resta libero.

Perché questo è un grande passo avanti

Il paper sostiene che questo approccio "senza note" porti tre grandi benefici:

1. Risparmia Spazio Massiccio (Memoria)
Poiché il computer non scrive mai la gigantesca lista di note nella memoria principale, libera una enorme quantità di spazio.

Il Risultato: Su un chip standard (H200), sono stati in grado di addestrare un modello da 8 miliardi di parametri utilizzando il 53% in meno di memoria.
L'Analogia: È come essere in grado di far stare un tavolo da pranzo per 10 persone in un monolocale perché hai smesso di dover conservare le sedie extra nel corridoio.

2. È Effettivamente Più Veloce
Poiché il computer non deve fermarsi per scrivere note, aspettare e poi rileggerle, l'intero processo accelera.

Il Risultato: I passaggi di addestramento sono circa 1,5 volte più veloci.
L'Analogia: È la differenza tra un corriere che consegna un pacco, corre indietro al camion per prendere quello successivo e ripete l'operazione (Vecchio Modo), rispetto a un corriere che ha un nastro trasportatore che carica il pacco direttamente sul camion mentre lo afferra (FORGE).

3. Non Perde Accuratezza
Di solito, quando si cerca di risparmiare spazio saltando dei passaggi, si perde precisione (il robot impara leggermente meno bene). Gli autori dimostrano che, poiché eseguono i calcoli nei "registri" ad alta qualità del computer (a piena precisione) prima di scartare i dati, l'apprendimento è matematicamente identico al vecchio metodo lento.

Il Risultato: Il robot impara altrettanto bene, ma in modo molto più efficiente.

Il Trucco del "Tile" (La Mattonella)

Come fanno a farlo senza creare il caos? Dividono il gigantesco problema matematico in piccoli pezzi gestibili chiamati "tile" (come quadrati da 128x128).

Elaborano un tile: Calcolano l'errore, correggono il cervello e scartano immediatamente l'errore.
Poi passano al tile successivo.
Poiché lo fanno tile per tile, il computer non deve mai tenere l'intera lista di errori contemporaneamente.

Cosa ci permette di fare

Il paper mostra che con FORGE:

Si possono addestrare modelli più grandi sullo stesso computer.
Si possono usare "batch" più grandi (insegnare al robot più esempi alla volta) senza esaurire la memoria.
In un test specifico, sono stati in grado di addestrare un modello su quattro GPU che avrebbero richiesto otto GPU usando il vecchio metodo.

Riassunto

FORGE è un nuovo modo per addestrare l'IA che impedisce al computer di ingolfare la propria memoria con temporanee "note di correzione". Calcolando gli errori e correggendo il modello istantaneamente nella parte più veloce del chip, riduce la metà dell'uso della memoria e accelera l'addestramento, il tutto senza rendere l'IA meno intelligente. Trasforma un laboratorio affollato e lento in una linea di montaggio snella e ad alta velocità.

Riepilogo Tecnico: FORGE (Fused On-Register Gradient Elimination)

1. Il Problema: Il Gradiente di Base (Gradient Floor)

Le attuali metodologie di addestramento di Large Language Models (LLM) efficienti dal punto di vista della memoria mirano principalmente a tre dei quattro pool di memoria su una GPU: gli stati dell'ottimizzatore (tramite compressione), le attivazioni (tramite ricomputazione) e i pesi del modello. Tuttavia, lasciano intatto il quarto pool: il gradiente del peso materializzato.

Nella differenziazione standard in modalità inversa (reverse-mode differentiation), il passaggio backward calcola il tensore del gradiente completo ( $\nabla W$ ) per ogni livello, lo scrive nella memoria ad alta larghezza di banda (HBM) e solo allora l'ottimizzatore lo rilegge per aggiornare i pesi. Questo schema a due fasi crea un "gradiente di base" dove, al confine tra il passaggio backward e il passo dell'ottimizzatore, i gradienti per tutti i livelli devono coesistere in memoria. Per un modello da 8 miliardi di parametri, questo pool di gradiente transitorio consuma da solo circa 16 GiB, dettando spesso se una configurazione di addestramento possa risiedere su una singola GPU. I metodi esistenti che fondono i passaggi backward e dell'ottimizzatore (ad es., LOMO, FlashOptim) operano alla granularità di interi parametri o livelli, il che significa che il tensore del gradiente completo per un parametro viene comunque materializzato in memoria prima che avvenga l'aggiornamento.

2. Metodologia: FORGE

FORGE (Fused On-Register Gradient Elimination) elimina interamente il tensore del gradiente materializzato fondendo il passo dell'ottimizzatore direttamente nel calcolo del gradiente del peso a una granularità di tile (specificamente tile $128 \times 128$ ).

Meccanismo Core

Invece di calcolare la matrice del gradiente completa e memorizzarla, FORGE elabora la matrice dei pesi un tile alla volta:

Accumulazione del Gradiente per Tile: Per un determinato tile di peso, il kernel accumula i prodotti parziali del gradiente a monte ( $\Delta Y$ ) e dell'attivazione di input ( $X$ ) direttamente nei registri FP32 della GPU.
Ottimizzazione In-Register: Una volta che il gradiente del tile è stato completamente accumulato nei registri, l'aggiornamento dell'ottimizzatore (ad es., AdamW) viene applicato immediatamente a quel tile, sempre all'interno dei registri.
Scarto Immediato: Il tile del gradiente viene consumato nell'istante stesso in cui viene prodotto. Non viene mai scritto in HBM come tensore.
Write-Back: Solo il tile del peso aggiornato e i tile dello stato dell'ottimizzatore (momenti) (opzionalmente quantizzati) vengono scritti nuovamente nella memoria globale.

Precisione e Fedeltà

Esattezza: Per qualsiasi ottimizzatore elemento per elemento separabile (dove l'aggiornamento di un peso dipende solo dalle coordinate di quel peso, ad es., AdamW, SGD, Lion), l'aggiornamento per tile è matematicamente identico a un passo standard del gradiente completo. L'unica deviazione è l'ordine della sommatoria in virgola mobile sugli token, che è intrinseco a qualsiasi implementazione GEMM e non costituisce un'approssimazione della logica dell'ottimizzatore.
Fedeltà a Bassa Precisione: Nei regimi di addestramento BF16 o INT8, FORGE preserva la piena fedeltà della precisione in due modi:
- Il gradiente non viene mai troncato a BF16 prima dell'ottimizzazione; viene calcolato e consumato in registri FP32.
- Per prevenire la "zona morta" del BF16 (dove piccoli aggiornamenti vengono arrotondati a zero), FORGE impiega il rounding stocastico per i write-back dei pesi, garantendo un'accumulazione non influenzata (unbiased) degli aggiornamenti sub-ULP.
- Per gli stati dell'ottimizzatore INT8, FORGE utilizza blocchi più piccoli (64 contro 2048 degli ottimizzatori standard a 8 bit) per minimizzare l'errore di quantizzazione.

Compatibilità con l'Addestramento Distribuito

FORGE è dimostrabilmente esatto sotto Parallelismo di Tensore (TP) e Parallelismo di Sequenza (SP) perché il gradiente locale per uno shard è completo al momento del backward. Tuttavia, è incompatibile con il Parallelismo di Dati (DP) o il Parallelismo di Contesto nella sua forma attuale, poiché questi richiedono un all-reduce dei gradienti parziali prima che l'aggiornamento possa essere applicato, il che richiederebbe la materializzazione del tensore del gradiente.

3. Contributi Chiave

Eliminazione Strutturale, non Compressione: FORGE identifica il gradiente materializzato come un artefatto dello scheduling dell'addestramento piuttosto che come una necessità matematica. Elimina il termine di memoria del gradiente $2P$ (numero di parametri), riducendo il pool del gradiente da $O(P)$ a $O(T_R T_C)$ (dimensione del tile), indipendentemente dalla dimensione del modello.
Esattezza Dimostrabile: Il documento prova che per gli ottimizzatori elemento per elemento, lo schema fuso per tile produce aggiornamenti identici a un AdamW standard su gradienti FP32. Questa esattezza si mantiene sotto parallelismo di tensore e di sequenza.
Fedeltà a Piena Precisione a Bassa Precisione: Mantenendo i gradienti nei registri FP32 fino all'ottimizzazione, FORGE evita il bias introdotto dal troncamento dei gradienti a BF16 prima dell'ottimizzazione. Combinato con il rounding stocastico, eguaglia la convergenza del master FP32 anche quando si utilizzano stati BF16 o INT8.
Generalità: Il metodo è agnostico rispetto all'architettura per i livelli lineari e supporta otto ottimizzatori completamente fusi elemento per elemento (AdamW, NAdam, RAdam, Lion, RMSprop, AdaGrad, SGD, SGD-momentum).

4. Risultati Sperimentali

Valutato su GPU NVIDIA H200 e B200 su modelli con range da 0.6B a 14B di parametri (Llama-3.1-8B, famiglia Qwen3):

Riduzione della Memoria: FORGE riduce il picco di memoria del 38–53% rispetto ad AdamW fuso. Per Llama-3.1-8B, il picco di memoria scende da ~75 GiB a 35.4 GiB (una riduzione del 53%).
Throughput: FORGE raggiunge un incremento di velocità di 1.52× nel tempo di step (99.1 ms vs 150.9 ms) eliminando il lancio del kernel separato dell'ottimizzatore e il round-trip in HBM per i gradienti.
Scalabilità:
- Su una singola H200, FORGE addestra Qwen3-14B (che causerebbe errori Out-Of-Memory con l'AdamW fuso standard) a 63.2 GiB.
- In una configurazione a 4 GPU A100-40GB, FORGE consente una dimensione del micro-batch di 4 (rientrando in 24.6 GiB), mentre l'AdamW fuso standard permette solo un micro-batch di 1. Ciò consente a FORGE di ottenere un throughput 2.42× superiore rispetto al baseline standard sullo stesso hardware.
Convergenza: FORGE segue la convergenza di AdamW fuso entro 0.001 nats (deviazione media) nei compiti di pre-addestramento continuo, confermando che le deviazioni a bassa precisione sono entro il rumore tipico delle esecuzioni.
Componibilità: FORGE si compone agevolmente con la quantizzazione FP8 e il checkpointing delle attivazioni, assorbendo i costi di dequantizzazione che solitamente penalizzano gli ottimizzatori standard.

5. Significato e Rivendicazioni

Il documento sostiene che FORGE rappresenti un cambio strutturale piuttosto che un'ottimizzazione incrementale. Eliminando il gradiente di base, trasforma l'overhead di memoria dell'ottimizzatore in spazio utile per batch più grandi, sequenze più lunghe o modelli più grandi su hardware esistente.

Confine di Capacità: Gli autori affermano che l'eliminazione del gradiente di base non è solo un miglioramento a fattore costante, ma un confine di capacità. Ad esempio, le configurazioni di addestramento standard per modelli da 14B su singole H200 sono impossibili a causa del gradiente di base, mentre FORGE le rende fattibili.
Ortogonalità: Il metodo è ortogonale ad altre tecniche di risparmio di memoria (come la quantizzazione dello stato o la ricomputazione delle attivazioni) e può essere combinato con esse per ridurre ulteriormente l'impronta di memoria.
Limitazioni: Il metodo è attualmente limitato agli ottimizzatori separabili elemento per elemento e non supporta il parallelismo di dati senza modifiche. Richiede inoltre che il gradiente di input venga letto prima dell'aggiornamento dei pesi, escludendo gli embedding condivisi (tied embeddings) o gli esperti MoE con molteplici consumatori di gradiente in un singolo step.

In sintesi, FORGE dimostra che il "gradiente di base" è un artefatto non necessario degli attuali cicli di addestramento. Fondendo l'ottimizzatore nel passaggio backward a livello di tile, ottiene significativi risparmi di memoria e accelerazioni mantenendo l'esattezza matematica per una vasta classe di ottimizzatori.