LoGRA: addestramento per rinforzo dei LLM su larga scala con sketch dei gradienti a basso rango

Abstract degli autori

L’apprendimento per rinforzo (RL) ha migliorato notevolmente le capacità dei modelli linguistici di grandi dimensioni (LLM), ma la memoria necessaria resta un ostacolo a un’adozione più ampia. Presentiamo LoGRA, un approccio al post-addestramento mediante RL che riduce il consumo di memoria conservando i segnali utili all’apprendimento in sketch dei gradienti a basso rango. Queste rappresentazioni compatte consentono sia di aggiornare il modello sia di sincronizzare le policy in modo efficiente. Per evitare che aggiornamenti troppo ampi compromettano l’apprendimento, affianchiamo alla compressione dei gradienti un controllo dell’ampiezza del passo basato sulla KL prevista: il metodo stima i cambiamenti della policy prima di applicare ciascun aggiornamento e ne regola di conseguenza l’ampiezza. Nei compiti di ragionamento, LoGRA riduce il consumo medio di memoria durante l’addestramento fino al 45,7\% senza compromettere le prestazioni. Consente inoltre di addestrare in modo stabile un modello da 27 miliardi di parametri per oltre 1.100 passi su un singolo nodo con otto GPU, dove l’addestramento con Adam standard esaurisce la memoria. Rende così praticabile l’addestramento mediante RL che prima non era fattibile per limiti di memoria. Il codice è disponibile nella \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt library}.