Apprentissage par renforcement gpt-oss

Vous pouvez désormais entraîner OpenAI gpt-oss avec RL et GRPO via Unsloth. Unsloth offre désormais la inférence la plus rapide (3x plus rapide), la plus faible utilisation de VRAM (50 % de moins) et le contexte le plus long (8x plus long) pour le RL gpt-oss par rapport à toute implémentation — sans dégradation de la précision. Comme l’apprentissage par renforcement (RL) sur gpt-oss n’est pas encore compatible avec vLLM, nous avons dû réécrire le code d’inférence à partir du code Transformers pour fournir une inférence 3x plus rapide pour gpt-oss à ~21 tokens/s. Pour BF16, Unsloth atteint également l’inférence la plus rapide (~30 tokens/s), surtout par rapport à l’utilisation de VRAM, en utilisant 50 % moins de VRAM que toute autre implémentation RL. Nous prévoyons de prendre en charge notre fonctionnalité de partage de poids à 50 % une fois que vLLM deviendra compatible avec le RL.

• Notebook gratuit : notebook Colab GRPO gpt-oss-20b Ce notebook crée automatiquement des noyaux de multiplication matricielle plus rapides et utilise 4 nouvelles fonctions de récompense Unsloth. Nous montrons également comment contrer le piratage de récompense qui est l’un des plus grands défis du RL.\

  

Avec Unsloth, vous pouvez entraîner gpt-oss-20b avec GRPO sur 15 Go de VRAM et gratuitement sur Colab. Nous avons introduit le déchargement des embeddings qui réduit également l’utilisation de 1 Go via offload_embeddings. La nouvelle inférence d’Unloth s’exécute plus rapidement sur n’importe quel GPU, y compris les A100, H100 et les anciens T4. gpt-oss-120b tient bien sur un GPU avec 120 Go de VRAM.

Unsloth est le seul framework à prendre en charge le RL en 4 bits pour gpt-oss. Tous les gains de performance sont dus au partage de poids, Flex Attention, Standby et aux noyaux personnalisés.

Rappel : Flash Attention 3 (FA3) est inadapté à gpt-oss l’entraînement car il ne prend actuellement pas en charge la passe arrière pour les attention sinks, provoquant des pertes d’entraînement incorrectes. Si vous n’êtes pas en train d’utiliser Unsloth, FA3 peut être activé par défaut, alors veuillez vérifier qu’il n’est pas utilisé ! La désactivation de FA3 entraînera également une utilisation mémoire en O(N^2) , donc Unsloth est le seul framework RL à offrir une utilisation mémoire en O(N) pour gpt-oss via notre implémentation Flex attention.

⚡Rendre l’inférence beaucoup plus rapide

L’inférence est cruciale dans l’entraînement RL, car nous en avons besoin pour générer des solutions candidates avant de maximiser une certaine fonction de récompense (voir ici pour une explication plus détaillée). Pour atteindre la vitesse d’inférence la plus rapide pour gpt-oss sans vLLM, nous avons réécrit le code d’inférence de Transformers et intégré de nombreuses innovations, y compris des algorithmes personnalisés comme Unsloth Flex Attention, en utilisant des indicateurs spéciaux dans torch.compile (comme les noyaux combinés). Notre nouveau code d’inférence pour gpt-oss a été évalué par rapport à une base de référence déjà optimisée (2x plus rapide que Transformers natif).

vLLM ne prend pas en charge le RL pour gpt-oss car il manque l’entraînement BF16 et le support LoRA pour gpt-oss. Sans Unsloth, seul l’entraînement via BF16 en pleine précision fonctionne, rendant l’utilisation mémoire 800 % et plus supérieure. La plupart des frameworks activent FA3 (Flash Attention 3) par défaut (ce qui réduit l’utilisation de VRAM et augmente la vitesse) mais cela provoque une perte d’entraînement incorrecte. Voir Issue 1797 dans le dépôt FA3. Vous devez toutefois désactiver FA3, car cela empêchera l’entraînement à long contexte puisque FA3 utilise une mémoire en O(N), tandis que l’attention naïve gonflera avec une utilisation en O(N^2). Donc, pour rendre les attention sinks différentiables, nous avons implémenté Unsloth Flex Attention.

Nous avons évalué l’inférence RL de gpt-oss en mesurant BitsandBytes 4 bits et avons également effectué des tests séparés pour BF16. L’inférence 4 bits d’Unsloth est ~4x plus rapide, et BF16 est aussi plus efficace, en particulier pour l’utilisation de la VRAM.

La meilleure partie du RL gpt-oss d’Unsloth est qu’il peut fonctionner sur n’importe quel GPU, même ceux qui ne prennent pas en charge BF16. Nos notebooks Colab gratuits gpt-oss-20b utilisent d’anciens GPU T4 de 15 Go, donc les exemples d’inférence fonctionnent bien !

🛠️ Problèmes et particularités de gpt-oss Flex Attention

Nous avons dû modifier notre implémentation pour les attention sinks comme décrit ici pour permettre à la génération de fonctionner avec un padding à gauche. Nous avons dû obtenir le logsumexp et appliquer l’activation sigmoïde pour modifier les poids d’attention comme ci-dessous :

$$A(X) = \sigma \bigg( \frac{1}{\sqrt{d}}QK^T \bigg)V \\ A(X) = \frac{\exp{\frac{1}{\sqrt{d}}QK^T}}{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}}V \\ \text{LSE} = \log{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}} \\ A_{sinks}(X) = A(X) \odot \sigma (\text{LSE} - \text{sinks})$$

Le masquage avec padding à gauche pendant l’inférence était aussi un problème délicat à gérer dans gpt-oss. Nous avons constaté que nous devions non seulement prendre en compte le préremplissage du cache KV lors des générations de tokens, mais aussi tenir compte d’une quantité unique de tokens de padding dans chaque prompt pour les générations par lots, ce qui modifiait la manière dont nous devions stocker le block mask. Un exemple de cela peut être vu ci-dessous :

Masque causal normal :

   k0 k1 k2 k3 k4   <-- clés
q0  X
q1  X  X
q2  X  X  X
q3  X  X  X  X
q4  X  X  X  X  X   <-- dernière ligne de requête (la plus importante pour le décodage)

Pour l’inférence dans le cas général (décodage)

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X  X  X  X  X

Si nous utilisons naïvement la même stratégie de masquage, cela échouera :

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X   (notez que q4 a q_idx=0 car c’est la première requête dans la configuration actuelle)

Pour la génération (phase de décodage), nous ne nous soucions généralement que de la dernière ligne de la matrice d’attention, puisqu’il n’y a qu’un seul token de requête assistant à tous les tokens de clé précédents. Si nous appliquons naïvement le masque causal (q_idx ≥ k_idx), cela échoue car notre requête unique a l’indice 0, alors qu’il y a n_k tokens de clé. Pour corriger cela, nous avons besoin d’un décalage dans la création du masque pour décider quels tokens assister. Mais une approche naïve est lente, puisque les décalages changent à chaque étape, forçant la régénération du masque et du noyau. Nous avons résolu cela avec des optimisations de cache et de compilation.

La partie la plus difficile est la génération par lots. Les séquences diffèrent en longueur, donc le padding complique la création du masque. Flex Attention avait beaucoup de défis et les masques dynamiques sont délicats. Pire encore, s’il n’est pas compilé, il revient à l’attention eager qui est lente et gourmande en mémoire (quadratique contre linéaire selon la longueur de séquence).

Citation de https://github.com/meta-pytorch/attention-gym/issues/15#issuecomment-2284148665

Vous devez appeler ceci avec _compile=True. Nous projetons essentiellement votre block mask sur une matrice complète Q_LEN x KV_LEN afin de produire le block mask. Sans compilation, nous devons matérialiser tout cela, et cela peut provoquer des OOM sur de longues séquences.

De plus, vous devez exécuter flex_attention = torch.compile(flex_attention). Sans compilation, flex revient à une implémentation eager non fusionnée qui est excellente pour le débogage, mais elle est beaucoup plus lente et matérialise la matrice complète des scores.

En fin de compte, le masque doit gérer dynamiquement le préremplissage par rapport au décodage avec le cache KV, les tokens de lot et de padding par séquence, rester torch.compile compatible, et prendre en charge les fenêtres glissantes.

🔍 Enquête sur Flash Attention

Une autre direction intéressante que nous avons explorée était d’essayer d’intégrer Flash Attention. Ses avantages sont largement reconnus, mais une limitation est qu’il ne prend pas en charge les attention sinks pendant la passe arrière pour gpt-oss. Pour contourner cela, nous avons restructuré le mécanisme d’attention afin qu’il fonctionne uniquement sur la sortie d’attention et les valeurs logsumexp que FlashAttention fournit facilement. Étant donné ces avantages, cela semblait être un choix évident à essayer.

Cependant, nous avons rapidement commencé à remarquer des problèmes. Alors que les premières couches se comportaient comme prévu, les couches ultérieures, en particulier les couches 18 à 24, produisaient des sorties qui divergeaient significativement de l’implémentation en mode eager dans transformers. Il est important de noter que cet écart ne peut pas être attribué à une accumulation d’erreurs, puisque les entrées de chaque méthode sont identiques à chaque couche. Pour une validation supplémentaire, nous avons également comparé les résultats à Unsloth FlexAttention.

Cela nécessite une enquête plus approfondie pour comprendre pourquoi seules les dernières couches montrent une différence aussi importante entre l’implémentation flash attention et les autres.

Flash Attention 3 ne prend pas en charge la passe arrière pour les attention sinks

FA3 est souvent activé par défaut pour la plupart des packages d’entraînement (pas Unsloth), mais cela est incorrect pour gpt-oss. Utiliser FA3 rendra la perte d’entraînement complètement erronée car FA3 ne prend pas en charge les passes arrière gpt-oss pour les attention sinks. Beaucoup de gens ne sont toujours pas au courant de cela, alors soyez prudent !

⚠️ Pouvons-nous contrer le piratage de récompense ?

L’objectif ultime du RL est de maximiser une certaine récompense (par exemple la vitesse, le revenu, une métrique). Mais le RL peut tricher. Lorsque l’algorithme RL apprend une astuce ou exploite quelque chose pour augmenter la récompense, sans réellement effectuer la tâche à la fin, cela s’appelle « Piratage de récompense".

C’est la raison pour laquelle les modèles apprennent à modifier les tests unitaires pour réussir des défis de codage, et ce sont des obstacles critiques au déploiement dans le monde réel. D’autres bons exemples viennent de Wikipédia.

Dans notre notebook RL gpt-oss gratuit nous explorons comment contrer le piratage de récompense dans un cadre de génération de code et présentons des solutions concrètes aux modes d’erreur courants. Nous avons vu le modèle modifier la fonction de chronométrage, externaliser vers d’autres bibliothèques, mettre les résultats en cache, et tricher ouvertement. Après correction, le résultat est que notre modèle génère de véritables noyaux de multiplication matricielle optimisés, et non des tricheries astucieuses.

🏆Piratage de récompense

Quelques exemples courants de piratage de récompense pendant le RL incluent :

Paresse

Le RL apprend à utiliser Numpy, Torch, d’autres bibliothèques, ce qui appelle des noyaux CUDA optimisés. Nous pouvons empêcher l’algorithme RL d’appeler du code optimisé en inspectant si le code généré importe d’autres bibliothèques Python non standard.

Mise en cache et triche

Le RL apprend à mettre en cache le résultat de la sortie et le RL apprend à trouver la sortie réelle en inspectant les variables globales Python.

Nous pouvons empêcher l’algorithme RL d’utiliser des données en cache en effaçant le cache avec une grande fausse matrice. Nous devons également faire des benchmarks soigneusement avec plusieurs boucles et tours.

Triche

Le RL apprend à modifier la fonction de chronométrage pour lui faire renvoyer 0 comme temps écoulé. Nous pouvons empêcher l’algorithme RL d’utiliser des variables globales ou mises en cache en restreignant ses variables locales et variables globales. Nous allons également utiliser exec pour créer la fonction, donc nous devons enregistrer la sortie dans un dictionnaire vide. Nous interdisons aussi l’accès aux variables globales via types.FunctionType(f.__code__, {})\

Tutoriel : Comment entraîner gpt-oss avec RL

Les LLM ont souvent du mal avec des tâches impliquant des environnements complexes. Cependant, en appliquant l’apprentissage par renforcement (RL) et en concevant une fonction de récompensepersonnalisée, ces défis peuvent être surmontés.

Le RL peut être adapté à des tâches telles que la génération automatique de noyaux ou la création de stratégies. Ce tutoriel montre comment entraîner gpt-oss avec GRPO et Unsloth pour battre 2048 de manière autonome.

Nos notebooks incluent déjà des guides étape par étape sur la manière de parcourir tout le processus.

Notebook 2048 (Exemple officiel d’OpenAI)	Notebook de génération de noyaux

Ce que vous allez construire :

• Entraîner gpt-oss-20b pour que le modèle puisse gagner automatiquement à 2048

• Créer un environnement 2048 minimal avec lequel le modèle peut interagir

• Définir des fonctions de récompense qui :

  1. Vérifient que la stratégie générée compile et s’exécute,

  2. Empêchent le reward hacking (interdisent les imports externes), et

  3. Récompensent la réussite effective du jeu

• Exécuter l’inférence et exporter le modèle (MXFP4 4 bits ou FP16 fusionné)

Matériel : L’exemple 2048 s’exécute sur un Colab T4 gratuit, mais l’entraînement sera lent. A100/H100 est beaucoup plus rapide. Le chargement en 4 bits + LoRA vous permet de faire tenir un modèle 20B dans une VRAM modeste