💬 NLP

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

Cet article présente Jet-RL, un cadre d'apprentissage par renforcement unifié en FP8 qui résout l'instabilité de l'entraînement et l'effondrement de la précision causés par les stratégies de précision mixte en appliquant une précision FP8 cohérente tant pour le déploiement (rollout) que pour l'entraînement, permettant ainsi d'obtenir des accélérations significatives tout en maintenant une convergence stable.

Auteurs originaux : Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

Publié 2026-01-27

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Haocheng Xi, Charlie Ruan, Peiyuan Liao, Yujun Lin, Han Cai, Yilong Zhao, Shuo Yang, Kurt Keutzer, Song Han, Ligeng Zhu

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Le gros problème : La « marche lente » de l'entraînement de l'IA

Imaginez que vous entraînez un robot très intelligent (un grand modèle de langage) pour résoudre des problèmes mathématiques complexes. Pour devenir vraiment bon, le robot doit apprendre à « penser à voix haute » (générer de longues chaînes de raisonnement) avant de donner une réponse.

Dans le monde de l'entraînement de l'IA, cette séance d'entraînement est appelée une Rollout (phase de déploiement).

Le goulot d'étranglement : Le papier souligne que cette phase de « pensée à voix haute » est incroyablement lente. Elle occupe 70 % du temps total passé à entraîner le robot. C'est comme si un étudiant passait 7 heures à étudier pour un examen, mais que 5 de ces heures consistaient simplement à fixer la page blanche, essayant de trouver la première phrase.
L'objectif : Nous voulons accélérer cette phase de « pensée » sans rendre le robot moins intelligent.

L'échec du raccourci : « L'erreur de la double comptabilité »

Pour accélérer les choses, les ingénieurs ont tenté une astuce simple : la Quantification.

L'analogie : Imaginez que le cerveau du robot travaille habituellement avec des nombres « HD » (haute définition) en 16 bits (BF16). Pour aller plus vite, ils ont essayé de passer le mode de « pensée » du robot en mode « SD » (basse définition) de 8 bits (FP8) juste pour la séance d'entraînement, tout en gardant la séance de « notation » en HD.
L'idée : « Utilisons le mode rapide et basse résolution pour les longues parties de réflexion, mais gardons le mode haute résolution pour les véritables mises à jour de l'apprentissage. »
Le désastre : Le papier a découvert que cette approche est un désastre pour les tâches longues ou difficiles.
- Le décalage : C'est comme demander à un étudiant d'écrire une dissertation de 10 pages au crayon (basse résolution), mais de la noter comme s'il l'avait écrite à l'encre (haute résolution). Les minuscules erreurs du croquis au crayon s'accumulent. Au moment où l'essai atteint 16 000 mots, la version « crayon » ne ressemble plus du tout à la version « encre ».
- Le résultat : Le robot s'embrouille. Il commence à halluciner, l'entraînement plante, et les performances du robot s'effondrent. Le papier appelle cela un « effondrement catastrophique de la précision ».

La solution : Jet-RL (L'approche de la « langue unifiée »)

Les auteurs proposent un nouveau système appelé Jet-RL. Au lieu de changer de langue entre la pratique et la notation, ils font en sorte que le robot parle la même langue (FP8) pour tout.

L'analogie : Imaginez que le robot écrit désormais ses essais de pratique au crayon, et que le professeur évalue aussi les essais en utilisant une grille de notation basée sur le crayon.
- Consistance : Parce que la « pensée » et l'« apprentissage » se déroulent exactement dans le même format de basse résolution, il n'y a pas de confusion. Le robot apprend exactement ce qu'il a pratiqué.
- Stabilité : Même pour des essais très longs (longues rollouts) ou des sujets très difficiles, le robot reste stable car les erreurs du « crayon » sont cohérentes tout au long du processus.

Comment ça marche (La magie technique)

Pour que cela fonctionne, l'équipe a dû être très habile dans la gestion des calculs :

Flux unifié : Ils ont veillé à ce que les données circulant à travers les puces informatiques soient traitées en FP8 (basse résolution) du tout premier pas de la pensée jusqu'à la toute dernière étape de la mise à jour du cerveau.
Groupement intelligent : Ils n'ont pas simplement réduit tout de manière aveugle. Ils ont regroupé les nombres ensemble (comme organiser des livres sur une étagère) afin que la version « basse résolution » reste suffisamment précise pour permettre l'apprentissage.
Pas de calibration : Habituellement, lorsque l'on change de format, il faut passer du temps à « calibrer » (tester et ajuster) pour s'assurer que cela fonctionne. Jet-RL saute cette étape car le système est conçu pour être cohérent dès le départ.

Les résultats : Rapide et précis

Le papier a testé cela sur plusieurs modèles et a constaté :

Vitesse : Cela a rendu la phase de « pensée » 33 % plus rapide et l'ensemble du processus d'entraînement 16 % plus rapide.
Précision : Contrairement au raccourci raté (qui provoquait l'échec du robot), Jet-RL a maintenu l'intelligence du robot presque identique à la version lente et haute définition. La baisse de performance était négligeable (moins de 1 %).
Fiabilité : Cela a fonctionné même lorsque le robot devait écrire des réponses très longues (16 000 mots) ou résoudre des problèmes mathématiques très difficiles, là où l'ancienne méthode aurait échoué.

Résumé

Jet-RL est une nouvelle façon d'entraîner l'IA qui empêche le robot de s'embrouiller en passant entre le « mode rapide » et le « mode intelligent ». En parlant la même « langue rapide » pour la pensée et l'apprentissage, cela rend l'entraînement de l'IA beaucoup plus rapide sans rendre l'IA moins intelligente. Cela transforme un raccourci instable et défectueux en une autoroute fiable et à grande vitesse.

Résumé Technique : Jet-RL

Énoncé du Problème

L'apprentissage par renforcement (RL) est devenu crucial pour améliorer les capacités de raisonnement complexe des grands modèles de langage (LLM), particulièrement à travers la génération de Chaîne de Pensée (Chain-of-Thought - CoT). Cependant, les pipelines d'entraînement par RL sont inefficaces sur le plan computationnel, la phase de rollout (génération de réponses) consommant plus de 70 % du temps total d'entraînement, surtout lorsque la longueur de génération dépasse 8 000 jetons.

Pour remédier à ce goulot d'étranglement, une stratégie courante consiste à utiliser la quantification FP8 pendant la phase de rollout tout en conservant une précision BF16 pour la phase d'entraînement (BF16-train + FP8-rollout). Bien que cette approche soit largement adoptée dans des frameworks comme VeRL et OpenRLHF, cet article identifie deux modes de défaillance critiques :

Instabilité des Longs-Rollouts : À mesure que la longueur de génération augmente (ex: >8k jetons), le décalage numérique entre la distribution d'entraînement BF16 et la distribution de rollout FP8 s'accumule. Cela conduit à une instabilité sévère de l'entraînement et à un effondrement catastrophique de la précision.
Échec sur les Tâches Difficiles : La stratégie échoue lorsqu'elle est appliquée à des tâches de raisonnement difficiles ou à des modèles de base plus faibles. Dans ces scénarios, la faible confiance du modèle amplifie les erreurs de quantification, provoquant une divergence rapide de la courbe d'apprentissage par rapport aux bases en pleine précision.

La cause profonde est identifiée comme étant une nature off-policy : l'acteur apprend à partir de données générées par une politique (rollout FP8) qui diffère numériquement de la politique mise à jour (entraînement BF16), violant ainsi l'hypothèse on-policy essentielle à un RL stable.

Méthodologie : Jet-RL

Les auteurs proposent Jet-RL, un framework conçu pour permettre un entraînement RL FP8 robuste et on-policy en imposant un flux de précision unifié entre l'entraînement et le rollout.

Architecture Centrale

Jet-RL élimine le décalage de précision en garantissant que le passage direct (forward pass) de l'entraînement et le rollout d'inférence utilisent l'identique flux de quantification FP8.

Graphe de Précision Unifié : Le framework modélise le modèle comme un graphe dirigé où les arêtes représentent la propagation des tenseurs. Jet-RL garantit que le graphe d'inférence ( $\mathcal{G}_{infer}$ ) est un sous-graphe du graphe direct de l'entraînement ( $\mathcal{G}^{fwd}_{train}$ ), partageant la même granularité de quantification et la même précision.
Poids Maîtres (Master Weights) : Tandis que le passage direct et le rollout opèrent en FP8, les poids maîtres sont conservés en BF16 pour stabiliser les mises à jour de l'entraînement, de manière similaire à l'entraînement en précision mixte standard.

Schéma de Quantification

Pour maintenir la stabilité tout en accélérant le calcul, Jet-RL emploie une stratégie de quantification mixte pour les couches linéaires (GEMM) :

Activations : Quantifiées via une quantification 1×128 par groupe.
Poids : Quantifiés via une quantification 128×128 par bloc.
Passage Arrière (Backward Pass) :
- Les activations sauvegardées pour le passage arrière sont stockées en FP8 (fusionnées avec l'opérateur précédent).
- Les gradients transportés entre les opérateurs sont conservés en BF16 pour prévenir le sous-flux de gradient (underflow) et la dégradation de la convergence.
- Les opérations GEMM de passage arrière (DGrad et WGrad) sont effectuées en FP16 pour l'accélération, utilisant des configurations de noyaux spécifiques (ex: 1×128 et 128×1 pour WGrad) afin de s'aligner sur les exigences matérielles.

Implémentation

Jet-RL est implémenté en utilisant vLLM pour le moteur d'inférence et VeRL pour le framework d'entraînement RL. Il exploite les noyaux GEMM FP8 de pointe de DeepGEMM et des noyaux Triton personnalisés pour la quantification, la transposition et les opérations fusionnées d'activation/RMSNorm.

Contributions Clés

Identification du Problème : L'article démontre que la stratégie standard BF16-train + FP8-rollout est fondamentalement défectueuse pour les longs-rollouts et les tâches difficiles en raison des décalages numériques off-policy, menant à l'effondrement de l'entraînement.
Framework Jet-RL : La proposition d'un flux de précision FP8 unifié qui impose une cohérence on-policy entre l'entraînement et le rollout, éliminant le besoin de calibration inter-étapes et résolvant le décalage de politique.
Performance et Stabilité : La démonstration que Jet-RL atteint des accélérations significatives tout en maintenant une convergence et une précision proches des bases BF16 à travers divers modèles et tâches.

Résultats Expérimentaux

Les auteurs ont évalué Jet-RL sur des modèles incluant Llama3.1-8B, Qwen2.5-7B et Qwen3-8B-Base sur des jeux de données tels que GSM8K, MATH et DeepMATH.

Précision et Convergence

Stabilité : Contrairement à la base BF16-train + FP8-rollout, qui n'a pas réussi à converger sur Qwen2.5-7B et a montré une dégradation significative sur les rollouts de 16k, Jet-RL a convergé avec succès dans tous les réglages.
Écart de Précision : Jet-RL a réduit l'écart de performance par rapport à l'entraînement BF16 à environ 1 % (ex: dégradation de 1,1 % sur Qwen3-8B-Base contre 2,9 % pour la base). Dans certains cas (Llama3.1-8B), Jet-RL a même surpassé la base BF16 de 2,0 %.
Robustesse : Jet-RL a maintenu ses performances sur des tâches difficiles (DeepMATH) et des contextes longs (16k) là où la stratégie de rollout FP8 naïve s'est effondrée.

Efficacité

Accélération du Rollout : La quantification FP8 a apporté une accélération de 1,07× à 1,33× dans la phase de rollout (jetons/s), les modèles plus grands (32B) affichant les gains les plus élevés.
Accélération de l'Entraînement : Le flux FP8 unifié a accéléré la phase de mise à jour de l'acteur de 1,54× et l'inférence du modèle de référence de 1,80×, résultant en une accélération de 1,41× pour la phase d'entraînement.
End-to-End : Le framework a atteint une accélération de bout en bout de 1,16× pour l'entraînement de modèles 8B, avec des gains attendus plus élevés pour les modèles plus grands.

Signification et Revendications

L'article affirme que Jet-RL offre une voie fiable et efficace pour appliquer le calcul FP8 afin d'accélérer l'entraînement RL à grande échelle. En traitant l'instabilité fondamentale causée par les décalages de précision, Jet-RL permet l'utilisation du FP8 non seulement pour l'inférence, mais pour l'ensemble de la boucle d'entraînement RL. Cela permet une accélération significative (jusqu'à 33 % pour le rollout et 41 % pour l'entraînement) sans sacrifier la stabilité de convergence ou la précision requise pour l'entraînement de modèles de raisonnement complexes. Ce travail établit qu'un flux de précision unifié est essentiel pour passer à l'échelle l'entraînement RL avec du matériel à basse précision.