Mantenere i token in flusso: lezioni da 16 librerie RL open-source

Per coloro che non hanno tempo di leggere 5.000 parole sulle complessità dell'RL asincrono (capiamo, avete modelli da addestrare): se preferite andare dritti al punto, ecco la tabella di confronto completa (non è richiesta alcuna lettura, non vi giudicheremo). Ma seriamente, se restate con noi, potreste imparare una o due cose sul perché le vostre GPU restano inattive il 60% del tempo.

L'addestramento asincrono per l'apprendimento per rinforzo (RL) è emerso come il paradigma dominante per il post-addestramento su larga scala. Diverse tendenze nel moderno post-addestramento hanno reso i cicli di addestramento sincroni quasi impossibili da scalare. L'ecosistema open-source è converguto su una risposta architettonica comune: disaggregare l'inferenza dall'addestramento su pool di GPU separati, connetterli con un buffer di rollout e permettere a entrambi i lati di funzionare in modo concorrente.

Stiamo sviluppando un nuovo trainer asincrono per TRL, una delle librerie più utilizzate per il post-addestramento dei modelli. Per guidare la nostra progettazione, abbiamo esaminato sedici librerie open-source costruite da zero attorno all'addestramento asincrono e le abbiamo confrontate lungo sette assi: primitivi di orchestrazione, progettazione del buffer, protocolli di sincronizzazione dei pesi, gestione della "vecchiaia" (staleness), gestione dei rollout parziali, supporto LoRA e backend di addestramento distribuito. Questo articolo distilla i principi di progettazione che abbiamo estratto da tale indagine.

Oltre all'RL, la necessità di un'infrastruttura asincrona è sempre più evidente. Ad esempio, la distillazione on-policy, in cui uno studente genera sequenze e un insegnante le valuta, rispecchia GRPO ma scambia la funzione di ricompensa con un passaggio forward dell'insegnante. Riconoscendo questa somiglianza strutturale, tutto ciò che è presente in questa indagine si applica ugualmente alla distillazione asincrona. Torneremo su questo punto più ampio nella Sezione 5.

Le limitazioni dell'attuale GRPOTrainer di TRL

L'attuale GRPOTrainer di TRL implementa il ciclo completo di GRPO (campionamento dei prompt, generazione, assegnazione del punteggio di ricompensa, calcolo del vantaggio, aggiornamento del gradiente e sincronizzazione dei pesi) in una singola chiamata sincrona training_step(). Questo design è semplice e corretto, ma non può sovrapporre la generazione all'addestramento, lasciando un significativo utilizzo della GPU inutilizzato.

Osservando il GRPOTrainer, abbiamo le seguenti fasi sequenzialmente all'interno di ogni passo di addestramento:

• Campionamento dei prompt
• Generazione
• Assegnazione del punteggio di ricompensa
• Calcolo del vantaggio
• Aggiornamento del gradiente
• Sincronizzazione dei pesi

Ogni fase si blocca fino al completamento prima che inizi la successiva. TRL offre l'opzione di configurazione steps_per_generation per riutilizzare un singolo set di rollout su più passi di gradiente (riutilizzo temporale), ammortizzando il costo di generazione. Ma la chiamata di generazione stessa rimane completamente sincrona e bloccante; il trainer non può iniziare il calcolo del gradiente finché ogni completamento nel batch non è terminato. La libreria supporta anche l'esecuzione di vLLM in modalità server come processo separato. Ciò libera la GPU di addestramento durante la generazione, ma rimangono due dure barriere di sincronizzazione: le chiamate HTTP finché tutti i completamenti non sono restituiti, e la sincronizzazione dei pesi blocca sia il trainer che vLLM durante il trasferimento.

Topologie di deployment: co-localizzata vs. disaggregata

Prima di discutere l'addestramento asincrono, è essenziale comprendere le due topologie di deployment per l'addestramento RL con un motore di inferenza separato:

• Modalità co-localizzata: inferenza e addestramento condividono le stesse GPU.
• Modalità disaggregata: inferenza e addestramento vengono eseguiti su pool di GPU separati.

Il vantaggio maggiore della modalità disaggregata è che inferenza e addestramento possono essere eseguiti in modo concorrente. Mentre il trainer calcola i gradienti sul batch N, il pool di inferenza sta già generando rollout per il batch N+K, abilitando l'addestramento asincrono. Tuttavia, questo beneficio ha un costo: sono necessarie GPU aggiuntive.

Concorrenza, asincronicità e parallelismo sono concetti distinti che spesso vengono confusi. In questo articolo, quando parliamo di "addestramento asincrono", intendiamo qualcosa di specifico: generazione e addestramento che operano in parallelo, con un'efficace sovrapposizione; il pool di inferenza produce il batch successivo di rollout mentre il pool di addestramento calcola i gradienti sul batch corrente. Questa è fondamentalmente una capacità della modalità disaggregata. La modalità co-localizzata può beneficiare di ottimizzazioni come la gestione della memoria sleep/wake o un rapido resharding in-place per accelerare l'inferenza, ma non può raggiungere una vera sovrapposizione simultanea; inferenza e addestramento si alternano ancora sulle stesse GPU. Ogni libreria in questa indagine che implementa una significativa sovrapposizione asincrona utilizza la modalità disaggregata come fondamento.

Il costo della generazione nei modelli di ragionamento

Nell'addestramento RL per i modelli di ragionamento, la generazione autoregressiva domina il tempo di esecuzione effettivo. Un singolo rollout per un compito matematico o di codifica può produrre 8K–64K token di ragionamento "chain-of-thought" (vedi le lunghezze dei rollout di QED-Nano).

Per contestualizzare concretamente, consideriamo i benchmark di vLLM su una singola GPU H100 80GB (bf16, senza quantizzazione, modalità throughput offline). Un modello da 7B (DeepSeek-R1-Distill-Qwen-7B) raggiunge circa 6.300 token di output/s di throughput aggregato; un modello da 32B (DeepSeek-R1-Distill-Qwen-32B) scende a circa 1.200 token di output/s. Questi sono throughput totali su tutte le richieste concorrenti, il numero che il motore di inferenza può elaborare al secondo, indipendentemente da quante sequenze condividono la GPU.

Ora consideriamo un tipico passo di addestramento GRPO: G=8 completamenti per prompt × 64 prompt/batch = 512 rollout. Quanto tempo richiede la generazione?

• Con un modello da 7B (2K token/rollout): 512 rollout × 2K token/rollout / (6300 token/s) = ~162 secondi = ~2,7 minuti.
• Con un modello da 32B (2K token/rollout): 512 rollout × 2K token/rollout / (1200 token/s) = ~853 secondi = ~14 minuti.
• Con un modello da 32B (32K token/rollout): 512 rollout × 32K token/rollout / (1200 token/s) = ~13653 secondi = ~227 minuti = ~3,8 ore.

Anche all'estremità inferiore (2K token generati con un modello da 7B), la sola generazione consuma diversi minuti per passo di addestramento. All'estremità superiore, dove i modelli di ragionamento all'avanguardia operano sempre più spesso, una singola fase di generazione può richiedere ore su una GPU. Scalare a 8 GPU di inferenza divide questi tempi per circa 8× (assumendo una scalabilità lineare del throughput), ma anche così, i rollout da 32K token su un modello da 32B richiedono ancora circa 28 minuti per passo.

Il problema dello "straggler"

Il problema dello "straggler" (elemento più lento) aggrava ulteriormente questo aspetto. Negli algoritmi basati su gruppi come GRPO, si campionano G completamenti per prompt. Il batch non può procedere finché il completamento più lento non è terminato. Le lunghezze degli output "chain-of-thought" sono altamente variabili; un singolo prompt potrebbe produrre completamenti che vanno da 1K a 32K token. Il batch è vincolato dal completamento più lungo, e il batching continuo lo mitiga solo parzialmente: le sequenze più corte liberano slot per nuovo lavoro, ma l'ultima sequenza in un gruppo GRPO blocca comunque il calcolo della ricompensa del gruppo e il passo di addestramento.

La soluzione architettonica: cicli disaccoppiati

Ogni libreria in questa indagine è converguta indipendentemente sullo stesso principio architettonico: separare fisicamente le GPU di inferenza dalle GPU di addestramento, e spingere i pesi in modo asincrono, in modo che la generazione non si fermi mai e l'addestramento non aspetti mai. Il pool di inferenza funziona continuamente, alimentando i rollout completati in un buffer. Il pool di addestramento estrae dal buffer, calcola gli aggiornamenti del gradiente e periodicamente spinge nuovi pesi al pool di inferenza per mantenerlo sincronizzato. I due cicli funzionano al proprio ritmo, disaccoppiati dal buffer.

Questa configurazione è altamente scalabile, ma introduce una nuova classe di problemi: la "vecchiaia" (rollout generati sotto una vecchia policy), l'overhead di sincronizzazione dei pesi, la gestione dei rollout parziali, ecc. Il resto di questo articolo analizza in dettaglio come le attuali librerie open-source affrontano questi problemi.

Sette assi di confronto

Per dare un senso all'ecosistema in rapida espansione delle librerie RL asincrone, proponiamo sette assi ortogonali di confronto. Ogni asse cattura una decisione di progettazione fondamentale che modella le prestazioni, la complessità e i compromessi della libreria. Questi assi forniscono una lente attraverso la quale analizzare e comprendere le diverse approcci implementativi e le loro implicazioni pratiche.

• Primitivi di orchestrazione
• Progettazione del buffer
• Protocolli di sincronizzazione dei pesi
• Gestione della "vecchiaia" (staleness)
• Gestione dei rollout parziali
• Supporto LoRA
• Backend di addestramento distribuito

Primitivi di orchestrazione

Come il sistema coordina i suoi componenti distribuiti? La scelta del framework di orchestrazione determina il modello di programmazione, la semantica dei guasti e il limite di scalabilità. Piuttosto che elencare i dettagli di implementazione per singola libreria, il panorama si scompone nettamente in quattro tipi di orchestrazione, paradigmi di coordinamento fondamentali che differiscono per livello di astrazione, modello di guasto e requisiti di deployment:

• Centralizzata: un singolo controller gestisce più worker, che sono essenzialmente server RPC stateless. Semplice da implementare, ma il controller può diventare un collo di bottiglia e un singolo punto di guasto. Esempi includono Apex-DQN e Reverb.
• Decentralizzata: i worker si coordinano direttamente o tramite memoria condivisa. Più robusta ai guasti del controller ma complessa da implementare. Un esempio è Dopamine.
• Mesh-Native: nativa di JAX/TPU, con partizionamento statico di calcolo e dati su una mesh di dispositivi. Altamente performante ma meno flessibile. V-PPO ne è un esempio.
• Basata su attori: un esempio è Ray. Utilizza il passaggio di messaggi tra attori indipendenti. Flessibile, scalabile e robusta ai guasti. Esempi sono Rllib e Acme.

Nota su Tunix: Tunix (Google) utilizza un modello mesh nativo di JAX con ThreadPoolExecutor per la sovrapposizione asincrona e jax.device_put per il trasferimento dei pesi tra mesh. È architettonicamente abbastanza distinto dall'ecosistema PyTorch da rendere il confronto diretto sull'orchestra zione non significativo; vive nel mondo XLA/TPU con i propri primitivi di coordinamento.

La tabella sopra rivela un modello sorprendente: otto delle sedici librerie esaminate utilizzano Ray come loro spina dorsale di orchestrazione. Questa non è una coincidenza; riflette una profonda aderenza architettonica tra il modello ad attori e la struttura dell'addestramento RL. Un sondaggio di Anyscale (l'azienda dietro Ray) sulle librerie RL open-source per LLM conferma questa convergenza. L'addestramento RL su larga scala coinvolge fondamentali capacità di gestione distribuita.