Nereus: parallelismo adattivo per il post-addestramento degli LLM

Abstract degli autori

Il post-addestramento dei modelli linguistici di grandi dimensioni (LLM) mediante apprendimento per rinforzo (RL) coordina più modelli nelle fasi di generazione, inferenza e addestramento su cluster di GPU. Durante un’esecuzione possono cambiare diversi fattori, tra cui la disponibilità delle risorse, la lunghezza delle sequenze, la pressione sulla memoria e i colli di bottiglia nelle diverse fasi. Di conseguenza, un piano di esecuzione inizialmente adatto può diventare lento o persino impossibile da attuare. Adattare un job i cui modelli condividono le GPU presenta però sfide notevoli: stabilire se un nuovo piano giustifichi il costo della transizione, riutilizzare lo stato distribuito del job e coordinare i trasferimenti tra GPU che coinvolgono modelli e fasi diversi. Nereus affronta queste sfide con un sistema runtime che tiene conto dei costi e adatta i job di post-addestramento con RL a piani di esecuzione efficienti. Il suo controller a basso overhead seleziona un piano globale compatibile con la memoria disponibile e autorizza la transizione in base a un modello di costo calibrato sul job in esecuzione. Per stimare ed eseguire una transizione, Nereus rappresenta lo stato distribuito di ciascuna replica di una coppia modello-fase (un modello in una fase) come una Elastic Model Unit. Usa quindi un grafo globale delle transizioni per stabilire l’ordine delle trasformazioni e dei trasferimenti tra GPU di queste unità. In una traccia costruita a partire da dati reali, l’adattamento online di TP/PP riduce la latenza media per step del 27,7% rispetto alla configurazione iniziale fissa di TP/PP con scalabilità tramite DP. In un’esecuzione di 1.000 step che arriva a 1.024 GPU, sei transizioni occupano lo 0,079% del tempo totale di esecuzione. Nereus aumenta il throughput end-to-end di PPO 8B di 2,14–7,27$\times$ rispetto a OpenRLHF e di 1,10–1,47$\times$ rispetto a Verl su cluster diversi.