gpt-oss Apprentissage par renforcement
Vous pouvez maintenant entraîner OpenAI gpt-oss avec RL et GRPO via Unsloth. Unsloth propose désormais le inférence la plus rapide (3x plus rapide), utilisation VRAM la plus faible (50% de moins) et contexte le plus long (8x plus long) pour gpt-oss RL par rapport à toute implémentation - sans dégradation de précision. Étant donné que l'apprentissage par renforcement (RL) sur gpt-oss n'est pas encore compatible vLLM, nous avons dû réécrire le code d'inférence depuis Transformers pour offrir 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), notamment en termes d'utilisation de la VRAM, en utilisant 50% de VRAM en moins par rapport à toute autre implémentation RL. Nous prévoyons de supporter notre fonctionnalité de partage de poids à 50% une fois que vLLM deviendra compatible avec le RL.
Notebook gratuit : Notebook Colab gpt-oss-20b GRPO Ce notebook crée automatiquement des noyaux de multiplication matricielle plus rapides et utilise 4 nouvelles fonctions de récompense Unsloth. Nous montrons aussi comment contrer le piratage de la 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 pour gratuit sur Colab. Nous avons introduit le déchargement des embeddings qui réduit l'utilisation d'1 Go également via offload_embeddings. La nouvelle inférence d'Unsloth s'exécute plus rapidement sur n'importe quel GPU incluant A100, H100 et les anciennes T4. gpt-oss-120b tient parfaitement sur un GPU de 120 Go de VRAM.
Unsloth est le seul framework à supporter le RL en 4 bits pour gpt-oss. Tous les gains de performance sont dus à l'unique partage de poids, et notre fonctionnalité Standby dans, Standby et aux noyaux personnalisés d'Unsloth.
Rappel : Flash Attention 3 (FA3) est inapproprié pour gpt-oss l'entraînement puisqu'il ne prend actuellement pas en charge la passe arrière pour les attention sinks, ce qui provoque des pertes d'entraînement incorrectes. Si vous n' utilisez Unsloth, FA3 peut être activé par défaut, donc veuillez vérifier qu'il n'est pas utilisé ! Désactiver FA3 entraînera également une utilisation mémoire O(N^2) , donc Unsloth est le seul framework RL à offrir utilisation mémoire 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 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 incluant des algorithmes personnalisés comme Unsloth et notre fonctionnalité Standby dans, en utilisant des drapeaux spéciaux au sein de torch.compile (comme des noyaux combinés). Notre nouveau code d'inférence pour gpt-oss a été évalué par rapport à une base 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 en pleine précision BF16 fonctionne, rendant l'utilisation de la mémoire plus élevée de 800%+. La plupart des frameworks activent FA3 (Flash Attention 3) par défaut (ce qui réduit l'utilisation de la VRAM et augmente la vitesse) mais cela provoque une perte d'entraînement incorrecte. Voir Problème 1797 dans le dépôt FA3. Vous devez désactiver FA3 cependant, car il empêchera l'entraînement sur de longs contextes puisque FA3 utilise une utilisation mémoire O(N), tandis que l'attention naïve gonflera avec une utilisation O(N^2). Donc, pour permettre que les attention sinks soient différentiables, nous avons implémenté Unsloth Flex Attention.
Nous avons évalué l'inférence RL de gpt-oss en comparant BitsandBytes 4-bit 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, surtout en utilisation de VRAM.
Le meilleur avec le 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'anciennes T4 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 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 :
Le masquage avec padding à gauche pendant l'inférence a également été un problème délicat à traiter dans gpt-oss. Nous avons constaté que nous devions non seulement tenir compte du préremplissage du cache KV lors de la génération de tokens, mais aussi tenir compte d'une quantité unique de tokens de padding dans chaque prompt pour les générations par lot, ce qui changeait la manière dont nous devions stocker le masque de bloc. Un exemple peut être vu ci-dessous :
Masque causal normal :
Pour l'inférence dans le cas général (décodage)
Si nous utilisons naïvement la même stratégie de masquage, cela échouera :
Pour la génération (phase de décodage), nous ne nous intéressons généralement qu'à la dernière ligne de la matrice d'attention, car il n'y a qu'un seul token de requête qui regarde tous les tokens clés précédents. Si nous appliquons naïvement le masque causal (q_idx ≥ k_idx), cela échoue car notre unique requête a l'indice 0, tandis qu'il y a n_k tokens clés. Pour corriger cela, nous avons besoin d'un décalage dans la création du masque pour décider quels tokens prendre en compte. Mais une approche naïve est lente, puisque les décalages changent à chaque étape, obligeant à régénérer le masque et le 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 lot. Les séquences diffèrent en longueur, donc le padding complique la création du masque. Flex Attention a rencontré beaucoup de défis et les masques dynamiques sont délicats. Pire encore, si ce n'est pas compilé, il revient à une attention eager non fusionnée qui est lente et gourmande en mémoire (quadratique vs linéaire en longueur de séquence).
Citation de https://github.com/meta-pytorch/attention-gym/issues/15#issuecomment-2284148665
Vous devez appeler cela avec _compile=True. Nous cartographions essentiellement votre masque de bloc sur une matrice complète Q_LEN x KV_LEN afin de produire le masque de bloc. Sans compilation, nous devons matérialiser cette chose complète, et cela peut causer 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 vs le décodage avec le KV Cache, le lot et les tokens de padding par séquence, rester torch.compile convivial, 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 opère uniquement sur la sortie d'attention et les valeurs logsumexp que FlashAttention fournit facilement. Compte tenu de 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 fortement de l'implémentation en mode eager dans transformers. Il est important de noter que cette divergence ne peut pas être attribuée à une accumulation d'erreurs, puisque les entrées pour chaque méthode sont identiques à chaque couche. Pour une validation supplémentaire, nous avons également comparé les résultats avec Unsloth FlexAttention.
Cela nécessite une enquête plus approfondie pour comprendre pourquoi seules les dernières couches montrent une différence aussi drastique 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 c'est incorrect pour gpt-oss. L'utilisation de FA3 rendra la perte d'entraînement complètement fausse car FA3 ne prend pas en charge les passes arrière gpt-oss pour les attention sinks. Beaucoup de personnes l'ignorent encore, donc soyez prudents !
⚠️ Pouvons-nous contrer le piratage de la récompense ?
L'objectif ultime de l'apprentissage par renforcement est de maximiser une récompense (par exemple la vitesse, le revenu, une métrique). Mais l'apprentissage par renforcement peut tricher. Lorsque l'algorithme RL apprend une astuce ou exploite quelque chose pour augmenter la récompense, sans réellement accomplir la tâche finale, cela s'appelle "Manipulation de la 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 pour le déploiement dans le monde réel. D'autres bons exemples proviennent de Wikipédia.
Dans notre cahier RL gpt-oss gratuit nous explorons comment contrer la manipulation de la récompense dans un contexte de génération de code et présentons des solutions tangibles aux modes d'erreur courants. Nous avons vu le modèle modifier la fonction de temporisation, externaliser vers d'autres bibliothèques, mettre en cache les résultats et carrément tricher. Après contre-mesures, le résultat est que notre modèle génère de véritables noyaux de multiplication de matrices optimisés, pas des astuces ingénieuses.
🏆Manipulation de la récompense
Quelques exemples courants de manipulation de la récompense lors de l'apprentissage par renforcement incluent :
Paresse
Le RL apprend à utiliser Numpy, Torch, d'autres bibliothèques, qui appellent des noyaux CUDA optimisés. Nous pouvons empêcher l'algorithme RL d'appeler du code optimisé en vérifiant si le code généré importe d'autres bibliothèques Python non standard.
Mise en cache et tricherie
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 mises en cache en effaçant le cache avec une grande matrice factice. Nous devons également mesurer soigneusement les performances avec plusieurs boucles et itérations.
Tricherie
Le RL apprend à modifier la fonction de temporisation pour qu'elle affiche 0 comme temps écoulé. Nous pouvons empêcher l'algorithme RL d'utiliser des variables globales ou mises en cache en restreignant ses locaux et globaux. Nous allons également utiliser exec pour créer la fonction, donc nous devons sauvegarder la sortie dans un dict vide. Nous interdisons également 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écompense, ces défis peuvent être surmontés.
Le RL peut être adapté à des tâches telles que la création automatique de noyaux ou de stratégies. Ce tutoriel montre comment entraîner gpt-oss avec GRPO et Unsloth pour battre automatiquement 2048.
Nos notebooks incluent des guides étape par étape sur la façon de naviguer dans l'ensemble du processus déjà.
notebook 2048 (Exemple officiel d'OpenAI)
Ce que vous allez construire :
Entraîner gpt-oss-20b afin que le modèle puisse gagner automatiquement à 2048
Créer un environnement 2048 minimal avec lequel le modèle peut interagir
Définir fonctions de récompense qui :
Vérifier que la stratégie générée se compile et s'exécute,
Prévenir la manipulation des récompenses (interdire les imports externes), et
Récompenser le succès réel du jeu
Exécuter l'inférence et exporter le modèle (MXFP4 4‑bit ou FP16 fusionné)
Matériel : L'exemple 2048 fonctionne 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 d'adapter un modèle 20B dans une VRAM modeste
Mis à jour
Ce contenu vous a-t-il été utile ?