Jet-RL: Enabling On-Policy FP8 Reinforcement Learning with Unified Training and Rollout Precision Flow

Riepilogo Tecnico: Jet-RL

Definizione del Problema

L'Apprendimento per Rinforzo (RL) è diventato fondamentale per potenziare le capacità di ragionamento complesso dei Large Language Models (LLM), in particolare attraverso la generazione di Chain-of-Thought (CoT). Tuttavia, le pipeline di addestramento RL sono computazionalmente inefficienti, con la fase di rollout (generazione delle risposte) che consuma oltre il 70% del tempo totale di addestramento, specialmente quando la lunghezza della generazione supera gli 8.000 token.

Per affrontare questo collo di bottiglia, una strategia comune prevede l'uso della quantizzazione FP8 durante la fase di rollout, mantenendo la precisione BF16 per la fase di addestramento (BF16-train + FP8-rollout). Sebbene questo approccio sia ampiamente adottato in framework come VeRL e OpenRLHF, questo articolo identifica due modalità di fallimento critiche:

1. Instabilità nei Long-Rollout: Man mano che la lunghezza della generazione aumenta (ad es., >8k token), il disallineamento numerico tra la distribuzione di addestramento BF16 e la distribuzione di rollout FP8 si accumula. Ciò porta a una grave instabilità dell'addestramento e a un collasso catastrofico dell'accuratezza.
2. Fallimento in Task Complessi: La strategia fallisce quando applicata a task di ragionamento difficili o a modelli base più deboli. In questi scenari, la minore confidenza del modello amplifica gli errori di quantizzazione, causando una divergenza rapida della curva di addestramento rispetto ai baseline a precisione piena.

La causa principale è identificata nella sua natura off-policy: l'attore apprende da dati generati da una policy (rollout FP8) che differisce numericamente dalla policy che viene aggiornata (addestramento BF16), violando l'assunto on-policy essenziale per un RL stabile.

Metodologia: Jet-RL

Gli autori propongono Jet-RL, un framework progettato per consentire un addestramento RL FP8 robusto e on-policy, imponendo un flusso di precisione unificato tra addestramento e rollout.

Architettura Core

Jet-RL elimina il disallineamento di precisione garantendo che sia il forward pass di addestramento che il rollout di inferenza utilizzino lo stesso identico flusso di quantizzazione FP8.

• Grafo di Precisione Unificato: Il framework modella il modello come un grafo diretto dove gli archi rappresentano la propagazione dei tensori. Jet-RL assicura che il grafo di inferenza ($\mathcal{G}_{infer}$) sia un sottografo del grafo forward di addestramento ($\mathcal{G}^{fwd}_{train}$), condividendo la stessa granularità di quantizzazione e precisione.
• Master Weights: Mentre il forward pass e il rollout operano in FP8, i master weights sono mantenuti in BF16 per stabilizzare gli aggiornamenti dell'addestramento, similmente al tipico addestramento a precisione mista.

Schema di Quantizzazione

Per mantenere la stabilità pur accelerando il calcolo, Jet-RL impiega una strategia di quantizzazione mista per i Linear Layers (GEMM):

• Attivazioni: Quantizzate utilizzando una quantizzazione 1×128 per-group.
• Pesi: Quantizzati utilizzando una quantizzazione 128×128 per-block.
• Backward Pass:
  • Le attivazioni salvate per il backward pass sono memorizzate in FP8 (fuse con l'operatore precedente).
  • I gradienti trasportati tra gli operatori sono mantenuti in BF16 per prevenire l'underflow del gradiente e la degradazione della convergenza.
  • Le operazioni di backward GEMM (DGrad e WGrad) sono eseguite in FP8 per l'accelerazione, utilizzando configurazioni di kernel specifiche (ad es., 1×128 e 128×1 per WGrad) per allinearsi ai requisiti hardware.

Implementazione

Jet-RL è implementato utilizzando vLLM per il motore di inferenza e VeRL per il framework di addestramento RL. Sfrutta i kernel GEMM FP8 all'avanguardia di DeepGEMM e kernel Triton personalizzati per quantizzazione, trasposizione e operazioni fuse di attivazione/RMSNorm.

Contributi Chiave

1. Identificazione del Problema: L'articolo dimostra che la strategia standard BF16-train + FP8-rollout è fondamentalmente difettosa per i long-rollout e i task complessi a causa dei disallineamenti numerici off-policy, che portano al collasso dell'addestramento.
2. Framework Jet-RL: La proposta di un flusso di precisione FP8 unificato che impone la coerenza on-policy tra addestramento e rollout, eliminando la necessità di calibrazione tra i passaggi e risolvendo il mismatch della policy.
3. Performance e Stabilità: La dimostrazione che Jet-RL ottiene incrementi significativi di velocità mantenendo convergenza e accuratezza vicine ai baseline BF16 attraverso diversi modelli e task.

Risultati Sperimentali

Gli autori hanno valutato Jet-RL su modelli tra cui Llama3.1-8B, Qwen2.5-7B e Qwen3-8B-Base su dataset come GSM8K, MATH e DeepMATH.

Accuratezza e Convergenza

• Stabilità: A differenza del baseline BF16-train + FP8-rollout, che non è riuscito a convergere su Qwen2.5-7B e ha mostrato una significativa degradazione su rollout da 16k, Jet-RL è confluito con successo in tutte le configurazioni.
• Gap di Accuratezza: Jet-RL ha ridotto il gap di performance rispetto all'addestramento BF16 a circa l'1% (ad es., degradazione dell'1,1% su Qwen3-8B-Base contro il 2,9% del baseline). In alcuni casi (Llama3.1-8B), Jet-RL ha persino superato il baseline BF16 del 2,0%.
• Robustezza: Jet-RL ha mantenuto le performance su task difficili (DeepMATH) e contesti lunghi (16k) dove la strategia di rollout FP8 naive è crollata.

Efficienza

• Accelerazione del Rollout: La quantizzazione FP8 ha fornito un'accelerazione da 1,07× a 1,33× nella fase di rollout (token/s), con i modelli più grandi (32B) che hanno registrato i guadagni maggiori.
• Accelerazione dell'Addestramento: Il flusso FP8 unificato ha accelerato la fase di aggiornamento dell'attore di 1,54× e l'inferenza del modello di riferimento di 1,80×, risultando in un'accelerazione di 1,41× per la fase di addestramento.
• End-to-End: Il framework ha ottenuto un'accelerazione end-to-end di 1,16× per l'addestramento di modelli da 8B, con aspettative di guadagni superiori per modelli più grandi.

Significato e Rivendicazioni

L'articolo sostiene che Jet-RL fornisce una via affidabile ed efficiente per applicare il calcolo FP8 per accelerare l'addestramento RL su larga scala. Affrontando l'instabilità fondamentale causata dai disallineamenti di precisione, Jet-RL consente l'uso di FP8 non solo per l'inferenza, ma per l'intero ciclo di addestramento RL. Ciò consente una significativa accelerazione (fino al 33% di velocità nel rollout e al 41% nell'addestramento) senza sacrificare la stabilità della convergenza o l'accuratezza necessaria per addestrare modelli di ragionamento complessi. Il lavoro stabilisce che un flusso di precisione unificato è essenziale per scalare l'addestramento RL con hardware a bassa precisione.