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

Technische Samenvatting: Jet-RL

Probleemstelling

Reinforcement Learning (RL) is cruciaal geworden voor het verbeteren van de complexe redeneercapaciteiten van Large Language Models (LLM's), met name door middel van Chain-of-Thought (CoT) generatie. De RL-trainingspipelines zijn echter computationeel inefficiënt, waarbij de rollout-fase (het genereren van responsen) meer dan 70% van de totale trainingstijd in beslag neemt, vooral wanneer de generatielengte de 8.000 tokens overschrijdt.

Om deze bottleneck aan te pakken, is een veelgebruikte strategie het gebruik van FP8-kwantisatie tijdens de rollout-fase, terwijl de BF16-precisie behouden blijft tijdens de trainingsfase (BF16-train + FP8-rollout). Hoewel deze aanpak breed wordt toegepast in frameworks zoals VeRL en OpenRLHF, identificeert dit artikel twee kritieke faalmodi:

1. Long-Rollout Instabiliteit: Naarman de generatielengte toeneemt (bijv. >8k tokens), cumuleert de numerieke mismatch tussen de BF16-trainingsdistributie en de FP8-rolloutdistributie. Dit leidt tot ernstige trainingsinstabiliteit en een catastrofale nauwkeurigheidsinstorting.
2. Falen bij Uitdagende Taken: De strategie faalt bij het toepassen op moeilijke redeneertaken of zwakkere basismodellen. In deze scenario's versterkt de lagere betrouwbaarheid van het model de kwantisatiefouten, waardoor de trainingscurve sneller divergeert vergeleken met full-precision baselines.

De grondoorzaak wordt geïdentificeerd als een off-policy natuur: de actor leert van data gegenereerd door een beleid (FP8 rollout) dat numeriek verschilt van het beleid dat wordt bijgewerkt (BF16 training), wat de on-policy assumptie schendt die essentieel is voor stabiele RL.

Methodologie: Jet-RL

De auteurs stellen Jet-RL voor, een framework ontworpen om robuuste, on-policy FP8 RL-training mogelijk te maken door een geëenheidigde precisieflow af te dwingen tussen training en rollout.

Kernarchitectuur

Jet-RL elimineert de precisie-mismatch door ervoor te zorgen dat zowel de training forward pass als de inference rollout gebruikmaken van de identieke FP8-kwantisatieflow.

• Unified Precision Graph: Het framework modelleert het model als een gerichte graaf waarbij randen tensorpropagatie vertegenwoordigen. Jet-RL zorgt ervoor dat de inference graaf ($\mathcal{G}_{infer}$) een subgraaf is van de training forward graaf ($\mathcal{G}^{fwd}_{train}$), waarbij dezelfde kwantisatiegranulariteit en precisie worden gedeeld.
• Master Weights: Hoewel de forward pass en de rollout opereren in FP8, worden de master weights in BF16 behouden om trainingsupdates te stabiliseren, vergelijkbaar met standaard mixed-precision training.

Kwantisatieschema

Om stabiliteit te behouden terwijl de berekening wordt versneld, gebruikt Jet-RL een gemengde kwantisatiestrategie voor Lineaire Lagen (GEMMs):

• Activations: Gekwantiseerd met 1×128 per-group kwantisatie.
• Weights: Gekwantiseerd met 128×128 per-block kwantisatie.
• Backward Pass:
  • Activations die worden bewaard voor de backward pass worden opgeslagen in FP8 (gefuseerd met de vorige operator).
  • Gradiënten die tussen operators worden getransporteerd, worden in BF16 gehouden om gradient underflow en convergentie-degradatie te voorkomen.
  • De backward GEMM-operaties (DGrad en WGrad) worden uitgevoerd in FP8 voor versnelling, waarbij specifieke kernelconfiguraties worden gebruikt (bijv. 1×128 en 128×1 voor WGrad) om aan hardwarevereisten te voldoen.

Implementatie

Jet-RL is geïmplementeerd met behulp van vLLM voor de inference engine en VeRL voor het RL-trainingsframework. Het maakt gebruik van state-of-the-art FP8 GEMM-kernels van DeepGEMM en aangepaste Triton-kernels voor kwantisatie, transpositie en gefuseerde activation/RMSNorm operaties.

Belangrijkste Bijdragen

1. Probleemidentificatie: Het artikel toont aan dat de standaard BF16-train + FP8-rollout strategie fundamenteel gebrekkig is voor long-rollout en uitdagende taken vanwege off-policy numerieke mismatches, wat leidt tot trainingsinstorting.
2. Jet-RL Framework: De voorstel van een uniforme FP8 precisieflow die on-policy consistentie afdwingt tussen training en rollout, waardoor de noodzaak voor inter-step kalibratie wordt geëlimineerd en de policy mismatch wordt opgelost.
3. Prestaties en Stabiliteit: De demonstratie dat Jet-RL significante versnellingen bereikt terwijl het de convergentie en nauwkeurigheid dicht bij BF16-baselines behoudt over diverse modellen en taken.

Experimentele Resultaten

De auteurs evalueerden Jet-RL op modellen waaronder Llama3.1-8B, Qwen2.5-7B en Qwen3-8B-Base over datasets zoals GSM8K, MATH en DeepMATH.

Nauwkeurigheid en Convergentie

• Stabiliteit: In tegen tegenstelling tot de BF16-train + FP8-rollout baseline, die niet kon convergeren op Qwen2.5-7B en significante degradatie vertoonde op 16k rollouts, convergeerde Jet-RL succesvol in alle settings.
• Nauwkeurigheidsgat: Jet-RL verkleinde het prestatieverschil ten opzien van BF16-training tot ongeveer 1% (bijv. 1,1% degradatie op Qwen3-8B-Base versus 2,9% voor de baseline). In sommige gevallen (Llama3.1-8B) presteerde Jet-RL zelfs 2,0% beter dan de BF16 baseline.
• Robuustheid: Jet-RL behield de prestaties op moeilijke taken (DeepMATH) en lange contexten (16k) waar de naïeve FP8-rolloutstrategie instortte.

Efficiëntie

• Rollout Versnelling: FP8-kwantisatie zorgde voor een versnelling van 1,07× tot 1,33× in de rollout-fase (tokens/s), waarbij grotere modellen (32B) de hoogste winsten lieten zien.
• Trainingsversnelling: De uniforme FP8-flow versnelde de actor update fase met 1,54× en de reference model inference met 1,80×, wat resulteerde in een 1,41× versnelling voor de trainingsfase.
• End-to-End: Het framework bereikte een 1,16× end-to-end versnelling voor de training van 8B modellen, met de verwachting van hogere winsten voor grotere modellen.

Betekenis en Claims

Het artikel claimt dat Jet-RL een betrouwbaar en efficiënt pad biedt voor het toepassen van FP8-berekeningen om grootschalige RL-training te versnellen. Door de fundamentele instabiliteit veroorzaakt door precisie-mismatches aan te pakken, maakt Jet-RL het mogelijk om FP8 niet alleen voor inference te gebruiken, maar voor de gehele RL-trainingslus. Dit staat een significante versnelling toe (tot 33% rollout versnelling en 41% trainingsversnelling) zonder de convergentiestabiliteit of nauwkeurigheid die nodig is voor het trainen van complexe redeneermodellen op te offeren. Het werk stelt vast dat een uniforme precisieflow essentieel is voor het schalen van RL-training met low-precision hardware.