Alibaba Qwen dévoile TRACE : le Rollout en FP4 jusqu'à 5,4 fois plus rapide

L'équipe Alibaba Qwen a présenté TRACE, une méthode d'entraînement de grands modèles par renforcement qui fait passer la phase de Rollout en FP4 : les poids, les activations (W/A) et le KV-cache. Selon les auteurs, le Rollout est jusqu'à 5,4 fois plus rapide qu'en BF16, et l'entraînement reste stable. Ce dernier point compte ici plus que la vitesse : beaucoup savent faire tourner la génération en quatre bits, mais peu parviennent à empêcher le RL de s'effondrer dans ces conditions.

Pourquoi le

Rollout est devenu la partie la plus gourmande du RL

Pour les modèles de langage, la boucle de RL se déroule ainsi : le modèle génère de nombreuses réponses, souvent de longues chaînes de raisonnement, celles-ci sont évaluées, puis les poids sont mis à jour à partir de ces trajectoires. La génération se fait token par token, elle est limitée par la bande passante mémoire, et le KV-cache grossit avec la longueur de la réponse. Plus le raisonnement est long et plus il y a d'échantillons par prompt, plus c'est le Rollout, et non l'étape de descente de gradient, qui consomme du temps et de la mémoire vidéo.

D'où l'envie logique de comprimer cette étape. Le passage du BF16 au FP4 réduit, en théorie, le volume de calcul comme la consommation de mémoire. Moins d'octets par poids et par élément du cache, cela signifie davantage de requêtes parallèles par carte et moins d'attente des données.

Où le FP4 met l'entraînement en difficulté

Cette compression a un prix : une grille numérique grossière. Quatre bits ne permettent de coder que 16 valeurs, si bien que deux grandeurs très proches se retrouvent facilement de part et d'autre d'une frontière de quantification et tombent sur des points de code différents. Si l'entraîneur et le moteur de Rollout calculent avec des précisions différentes, ils obtiennent des résultats nettement différents à partir d'entrées quasi identiques. Un écart minuscule au niveau du nombre se transforme en divergence entre ce que le modèle a « généré » et ce qu'il « pense » au moment de la mise à jour des poids.

C'est dangereux pour le RL. L'algorithme suppose que les trajectoires sont produites par une politique proche de celle qui est entraînée. Quand le générateur et l'entraîneur divergent, le gradient s'appuie sur des données que le modèle actuel aurait produites autrement. Cela se manifeste en général par du bruit, de la dérive et une instabilité de l'entraînement. Avec le FP4, le problème est plus aigu qu'avec le FP8, car les effets de bord sont plus nombreux et la marge de précision plus faible. C'est précisément ce problème que TRACE cherche à résoudre.

Ce que l'on sait de TRACE

Paramètres clés du travail :

• Développeur : l'équipe Qwen (Alibaba).
• Domaine d'application : l'étape de Rollout dans l'entraînement de grands modèles par RL.
• Format : FP4 pour les poids et les activations (W/A), plus FP4 pour le KV-cache.
• Référence : BF16.
• Effet annoncé : accélération du Rollout jusqu'à 5,4 fois.
• Objectif : préserver la stabilité de l'entraînement avec une génération en basse précision.

La formule « jusqu'à 5,4 fois » désigne un maximum, pas un résultat moyen. Le gain réel dépend de la longueur de génération, de la taille du batch, de l'architecture du modèle et du matériel. Plus le contexte est long et le KV-cache lourd, plus l'effet de la quantification du cache est marqué, car c'est souvent lui qui devient le goulot d'étranglement sur les longs raisonnements.

Ce que doivent vérifier les équipes qui font tourner du RL sur leurs

propres clusters

Pour ceux qui entraînent eux-mêmes des modèles de raisonnement, la principale conclusion est qu'il faut traiter le Rollout comme un problème d'optimisation à part entière, et non comme un processus annexe. Si la génération occupe l'essentiel du temps d'un pas, comprimer la précision rapporte plus que régler l'optimiseur.

Avant de passer au FP4, il faut mesurer l'écart entre les probabilités de tokens de l'entraîneur et celles du moteur de Rollout. S'il est déjà sensible en BF16 ou en FP8, il augmentera à quatre bits. Il faut aussi vérifier séparément la qualité sur les longues trajectoires : l'erreur du KV-cache s'accumule avec la longueur de la réponse.

Enfin, le gain dépend du support matériel. Le FP4 natif existe dans les nouvelles générations d'accélérateurs, par exemple l'architecture Blackwell, tandis que sur des cartes plus anciennes, la voie à 4 bits n'apportera pas une telle accélération. Pour les laboratoires disposant d'un parc de GPU limité, des travaux comme TRACE présentent un intérêt : ils permettent d'obtenir davantage d'itérations de RL avec le même matériel. La question suivante est de savoir dans quelle mesure le résultat se reproduit sur des modèles d'autres tailles et sur d'autres tâches, des mathématiques au code.