Fine-tuning de GPT-OSS-20B pour le raisonnement multilingue sur Amazon SageMaker HyperPod

GPT-OSS-20B est un modèle de raisonnement à poids ouverts qui raisonne en anglais. Demandez-lui de raisonner en français, en allemand, en espagnol ou en italien, et il reconnaît l'instruction, puis produit néanmoins sa chaîne de raisonnement en anglais. Sur un test de raisonnement multilingue de 50 exemples, le modèle de base n'a raisonné dans la langue cible correcte que 22 pour cent du temps. Ces 22 pour cent étaient presque entièrement constitués des exemples en anglais (anglais 100 pour cent, français 0 pour cent, allemand 10 pour cent, espagnol 0 pour cent, italien 0 pour cent).

Pour les équipes qui déploient des modèles de raisonnement pour des utilisateurs non anglophones, cet écart compte. Un utilisateur qui pose une question en français s'attend à ce que le modèle raisonne et réponde en français, et non à ce qu'il passe à l'anglais en pleine réflexion. Combler cet écart est un problème de fine-tuning. Cela soulève une question à laquelle toute équipe est confrontée dès qu'elle commence à personnaliser des modèles à poids ouverts : fine-tuning supervisé (SFT), apprentissage par renforcement avec Group Relative Policy Optimization (GRPO), ou les deux ?

Dans cet article, nous présentons comment nous avons affiné GPT-OSS-20B pour le raisonnement multilingue sur Amazon SageMaker HyperPod. Nous montrons quand le SFT seul suffisait et quand le GRPO apportait une valeur mesurable en complément. Pour mieux comprendre quand le GRPO est utile de manière générale, nous avons également mené une expérience contrôlée avec Qwen2.5-7B sur deux tâches présentant des écarts de capacité différents. La réponse courte : le SFT fait l'essentiel du travail, le GRPO affine quand il y a une marge de progression, et c'est l'adéquation entre la technique et l'écart de capacité qui détermine le résultat.

Vue d'ensemble de la solution

Cette section couvre la tâche de raisonnement multilingue et le jeu de données, la manière dont le SFT et le GRPO se complètent, et la pile d'entraînement que nous avons exécutée sur Amazon SageMaker HyperPod.

La tâche et le jeu de données

Le jeu de données est HuggingFaceH4/Multilingual-Thinking — un jeu de données de raisonnement en chaîne traduit en français, en espagnol, en allemand et en italien en plus de l'anglais, avec environ 949 exemples d'entraînement et 50 exemples de validation. Le modèle reçoit une instruction de reasoning language: prompt système et doit à la fois raisonner et répondre dans cette langue cible, passant outre son fort a priori en faveur de l'anglais. L'évaluation porte sur 50 cas de test : 10 problèmes de mathématiques/logique dans 5 langues (EN/FR/DE/ES/IT).

Deux métriques comptent :

• Raisonnement dans la bonne langue – la chaîne de raisonnement est-elle dans la langue demandée ?
• Réponse dans la bonne langue – la réponse finale est-elle dans la langue demandée ?

Comment le SFT et le GRPO s'articulent

La manière la plus claire de penser ces deux techniques : le SFT enseigne un comportement que le modèle n'a pas. Le GRPO optimise un comportement que le modèle peut déjà parfois produire. Le GRPO fonctionne en échantillonnant un groupe de réponses par prompt et en orientant la politique vers celles qui sont au-dessus de la moyenne. Si le comportement souhaité a une probabilité nulle dans la distribution de base, il y a très peu de choses à renforcer pour le GRPO. Cette idée unique explique presque tous les résultats de cet article.

La pile d'entraînement

Nous avons travaillé sur Amazon SageMaker HyperPod avec l'orchestration Amazon Elastic Kubernetes Service (Amazon EKS). Pour les premières exécutions de SFT et de GRPO sur GPT-OSS-20B, nous avons utilisé la bibliothèque PyTorch TRL. Pour l'expérience de contrôle avec Qwen, nous sommes passés à veRL comme entraîneur, avec Ray et KubeRay pour l'orchestration distribuée et vLLM pour une génération rapide des rollouts. Les checkpoints et les jeux de données étaient stockés sur Amazon FSx for Lustre, et la fabric inter-nœuds utilisait Elastic Fabric Adapter (EFA). Tout le fine-tuning a utilisé Low-Rank Adaptation (LoRA), qui est économe en paramètres et bien adapté à la fois au SFT et au GRPO.

Le diagramme suivant montre le chemin des données d'entraînement de bout en bout. Une soumission de l'utilisateur arrive sur le plan de contrôle Amazon EKS du cluster Amazon SageMaker HyperPod. Amazon EKS planifie un pod head Ray ainsi que des pods workers qui exécutent l'entraîneur veRL et les workers de rollout vLLM. Les deux ensembles de pods montent Amazon FSx for Lustre pour le modèle de base, les données d'entraînement et les checkpoints. Les collectifs inter-nœuds utilisent EFA. Les rollouts sont diffusés depuis vLLM vers l'entraîneur veRL, et l'entraîneur réécrit les checkpoints sur Amazon FSx for Lustre.

Figure 1 : Chemin des données d'entraînement de bout en bout pour le fine-tuning SFT et GRPO sur Amazon SageMaker HyperPod avec veRL, Ray, KubeRay, vLLM et Amazon FSx for Lustre pour les expériences de contrôle Qwen

Ce qui suit décrit chaque saut plus en détail. Le soumetteur lance l'exécution depuis Amazon SageMaker Studio ou l'AWS Command Line Interface (AWS CLI). Le plan de contrôle Amazon EKS sur Amazon SageMaker HyperPod place les pods sur des nœuds GPU et démarre un head Ray avec des pods workers pour l'entraîneur veRL et les workers de rollout vLLM. Amazon FSx for Lustre fournit un stockage partagé à haut débit pour le modèle de base, les données d'entraînement et les checkpoints. Lorsque l'entraînement se termine, l'entraîneur écrit l'adaptateur LoRA sur Amazon FSx for Lustre pour une évaluation ou un déploiement en aval.

Prérequis

Pour suivre ce tutoriel, vous avez besoin de :

• Un cluster Amazon SageMaker HyperPod avec orchestration Amazon EKS. Le SFT de GPT-OSS a utilisé 4× ml.g6e.12xlarge (16× L40S 48GB). L'expérience de contrôle Qwen a utilisé 2× g6e.48xlarge (16× L40S).
• Un système de fichiers Amazon FSx for Lustre monté sur le cluster avec un PersistentVolumeClaim nommé fsx-claim à /fsx.
• des instances compatibles EFA et une bibliothèque NVIDIA Collective Communications Library (NCCL) fonctionnelle, ainsi qu'un aws-ofi-nccl stack pour le SFT multi-nœuds.
• Un registre de conteneurs tel qu'Amazon Elastic Container Registry (Amazon ECR) pour stocker l'image d'entraînement construite à partir du dépôt de référence.
• Un jeton d'accès HuggingFace pour openai/gpt-oss-20b et le HuggingFaceH4/Multilingual-Thinking jeu de données.
• L'AWS Command Line Interface (AWS CLI), kubectl, docker, et envsubst disponibles sur votre poste de travail.

Configurer l'environnement d'entraînement

Toutes les commandes, manifestes et scripts référencés ci-dessous se trouvent dans le awsome-distributed-ai dépôt GitHub sous 3.test_cases/pytorch/trl/gpt-oss-lora-grpo/. Clonez le dépôt et placez-vous dans le répertoire du cas de test :

git clone https://github.com/awslabs/awsome-distributed-ai
cd awsome-distributed-ai/3.test_cases/pytorch/trl/gpt-oss-lora-grpo

Copiez le modèle d'environnement, puis modifiez env_vars pour définir votre registre ECR, le nom de l'image, le tag et votre jeton HuggingFace :

cp env_vars.example env_vars
source env_vars

Construisez l'image d'entraînement et poussez-la vers votre registre. L'image regroupe PyTorch, TRL, HuggingFace Transformers, PEFT ainsi que les scripts d'entraînement et d'inférence :

cd artifacts && ./build_push.sh && cd ..

Pointez kubectl vers votre cluster et appliquez le pod du gestionnaire de stockage FSx. Le gestionnaire de stockage vous donne un shell dans le système de fichiers partagé pour la conversion des checkpoints et l'inspection des journaux d'évaluation :

aws eks update-kubeconfig --name <YOUR_CLUSTER_NAME> --region us-east-2
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/fsx-storage-manager.yaml | kubectl apply -f -

Fine-tuner GPT-OSS-20B avec SFT et LoRA

Nous commençons par mesurer le modèle de base, puis nous configurons l'exécution SFT et examinons les résultats.

L'écart du modèle de base

Le GPT-OSS-20B de base sur l'évaluation de 50 cas a obtenu 22 pour cent de raisonnement-langue correct (11/50). Par langue : anglais 100 pour cent, français 0 pour cent, allemand 10 pour cent, espagnol 0 pour cent, italien 0 pour cent. Le modèle raisonne par défaut en anglais quelle que soit l'instruction.

Configuration SFT

La configuration SFT utilisait LoRA r=8, α=16, entraînée pendant 1,000 étapes sur 4× ml.g6e.12xlarge (16× L40S 48GB) avec Fully Sharded Data Parallel (FSDP) full-shard, la précision bf16 et le gradient checkpointing. Le travail s'exécute comme une tâche élastique HyperPodPyTorchJob gérée par Kueue, lancée avec hyperpodrun.

Appliquez le manifeste SFT :

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < hyperpod-eks/train-lora-hyperpod-elastic-g6e.yaml | kubectl apply -f -

Suivez la progression de l'entraînement :

kubectl get pods -l app=gpt-oss-lora-elastic
kubectl logs -f -l app=gpt-oss-lora-elastic

La loss et la norme du gradient sont affichées à chaque étape. Une exécution saine ressemble à ceci :

{'loss': 1.035, 'grad_norm': 0.163, 'learning_rate': 0.0002, 'mean_token_accuracy': 0.716, 'epoch': 3.25}
{'loss': 1.003, 'grad_norm': 0.171, 'lr': 0.0002, 'acc': 0.722, 'epoch': 3.32}
{'loss': 0.976, 'grad_norm': 0.159, 'lr': 0.0002, 'acc': 0.726, 'epoch': 3.57}

Une fois l'entraînement terminé, convertissez le checkpoint FSDP au format Parameter-Efficient Fine-Tuning (PEFT) attendu par le trainer GRPO de TRL et par les scripts d'inférence :

kubectl exec -it fsx-storage-manager -- python /app/convert_fsdp_checkpoint.py \
  --checkpoint /fsx/checkpoints/checkpoint-1000 \
  --output    /fsx/checkpoints/converted-peft/lora-checkpoint-1000-peft

Résultats SFT et courbe d'apprentissage par étapes

Après SFT, le modèle raisonne dans la langue correcte 100 pour cent du temps (50/50) dans les cinq langues. Le SFT l'a appris en deux étapes :

• Autour de l'étape 200, le modèle a appris à raisonner dans la langue cible (la métrique langue-de-raisonnement atteint 100 pour cent).
• Autour de l'étape 1000, le modèle a en outre appris à donner la réponse dans la langue cible.

Évaluation SFT finale (étape 1000) :

• Langue-de-raisonnement correcte : 100 pour cent (50/50).
• Langue-de-réponse correcte : 62 pour cent (31/50).
• Les deux correctes : 62 pour cent (31/50).
• Précision des réponses par langue : FR 60 pour cent, DE 60 pour cent, ES 30 pour cent, IT 60 pour cent, EN 100 pour cent.

À retenir : Le SFT n'est pas optionnel ici. Le modèle de base ne peut absolument pas suivre l'instruction de langue de raisonnement, et le SFT fait passer la langue de raisonnement de 22 pour cent à 100 pour cent. L'écart restant concerne la langue de réponse à 62 pour cent. C'est la moitié la plus difficile, car la réponse finale est courte et le prior anglais y est le plus fort.

Affiner la langue de réponse avec GRPO

La langue de réponse était la cible évidente pour GRPO. Le SFT produisait déjà un raisonnement correct 100 pour cent du temps, et le modèle parfois produisait la réponse dans la bonne langue, il y avait donc un signal à renforcer.

Exécuter GRPO sur la base du checkpoint SFT

Nous avons exécuté GRPO avec le GRPOTrainerde TRL, en partant du checkpoint SFT. La configuration : charger le LoRA du SFT, appeler merge_and_unload, et ajouter un nouveau LoRA (r=8, α=16, environ 0.068 pour cent de paramètres entraînables, 13.6M sur 20B). L'entraînement a utilisé K=8 générations par prompt, device_map=auto, un taux d'apprentissage de 1e-6, et 10 époques. La fonction de récompense notait la langue de réponse (±5.0), la langue de raisonnement (±1.5) et la brièveté (+0.5 ou −1.0), pour une plage combinée de −7.5 à +7.0.

Lancer GRPO à partir du checkpoint SFT converti :

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/train-grpo-singlenode.yaml | kubectl apply -f -
kubectl logs -f grpo-single

GRPO écrit les checkpoints dans /fsx/checkpoints/grpo-singlenode/. Choisissez le meilleur checkpoint par rapport à votre eval de validation (la métrique de langue de réponse a culminé tôt dans cette exécution) et convertissez-le au format PEFT pour l'inférence et l'évaluation.

Résultat (à partir du checkpoint SFT) :

• Réponses correctes en langue de raisonnement : 100 pour cent (sur les données held-out : GRPO n'a pas cassé ce que SFT a appris).
• Réponses correctes en langue de réponse : 62 à 70 pour cent (+8 points, 35/50).
• Les deux corrects : 70 pour cent.
• Réponses par langue : FR 70 pour cent, DE 80 pour cent, ES 60 pour cent, IT 40 pour cent, EN 100 pour cent.

GRPO a apporté un véritable gain de +8 points sur la métrique de langue de réponse tout en maintenant le raisonnement à 100 pour cent. Le gain est venu là où il y avait de la marge de progression. La langue de réponse était le comportement parfois correct sur lequel SFT n'avait pas complètement convergé, et GRPO l'a renforcé.

Comparer l'inférence de base, SFT et GRPO côte à côte

Pour voir la différence qualitative, déployez trois pods d'inférence qui partagent le même modèle de base mais chargent des adaptateurs différents :

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-base.yaml    | kubectl apply -f -
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-trained.yaml | kubectl apply -f -
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/inference-g6e-grpo.yaml    | kubectl apply -f -

Exécutez la même invite de raisonnement en français à travers chaque variante :

kubectl exec -it inference-base    -- python /app/inference_g6e.py \
  --reasoning_language French

kubectl exec -it inference-trained -- python /app/inference_g6e.py \
  --checkpoint /fsx/checkpoints/converted-peft/lora-checkpoint-1000-peft \
  --reasoning_language French

kubectl exec -it inference-grpo    -- python /app/inference_grpo_new.py \
  --use_grpo \
  --grpo_checkpoint /fsx/checkpoints/converted-peft/checkpoint-1470-peft \
  --reasoning_language French

Le pod de base raisonne en anglais malgré l'instruction. Le pod SFT raisonne en français mais répond souvent en anglais. Le pod GRPO raisonne et répond en français. Pour reproduire l'eval complet des 50 cas qui a généré les chiffres précédents, exécutez le manifeste d'évaluation et lisez le fichier de résultats depuis le gestionnaire de stockage FSx :

envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
  < kubernetes/eval-grpo.yaml | kubectl apply -f -
kubectl logs -f eval-grpo
kubectl exec fsx-storage-manager -- \
  cat /fsx/experiments/grpo_eval_checkpoint1470_*.txt

Pour comprendre plus largement quand GRPO apporte de la valeur, nous avons mené une expérience contrôlée sur Qwen2.5-7B à travers deux tâches présentant des écarts de capacité différents.

Tester le schéma avec un contrôle Qwen2.5-7B

Pour tester la thèse « adapter la technique à l'écart de capacité » sur un modèle différent et des tâches différentes, nous avons mené des expériences contrôlées sur Qwen2.5-7B-Instruct. La première tâche est la même tâche de raisonnement multilingue (où il existe une marge de progression). La deuxième tâche est GSM8K, où le modèle est déjà à sa précision maximale. Il s'agit d'une expérience interne, et le dépôt public n'a pas encore été publié.

Infrastructure : 2× g6e.48xlarge (16× L40S), EFA 100Gbps, FSx Lustre, Ray 2.44 + KubeRay 1.4.2, veRL v0.6.1, vLLM 0.8.4.

Bras : A = base, B = SFT seul, C = SFT+GRPO.

Raisonnement multilingue (50 cas)

GSM8K, une tâche sans marge de progression

Sur GSM8K, le modèle de base Qwen2.5-7B obtient déjà 90,5 pour cent (181/200). GRPO à l'étape 30 obtient également 90,5 pour cent. Lorsque la plupart des groupes échantillonnés sont entièrement corrects, l'avantage relatif au groupe s'effondre à zéro, donc il n'y a aucun signal d'apprentissage. C'est l'autre moitié de la thèse : GRPO a besoin de marge de progression (pass@k >> pass@1) pour avoir quelque chose à optimiser. Pas de marge de progression signifie pas de gain.

Bonnes pratiques

1. Commencer par le SFT. GRPO ne peut pas enseigner un comportement de probabilité de base nulle. Il ne peut que renforcer les réponses que le modèle produit déjà parfois. Le SFT a fait passer la langue du raisonnement de 22 pour cent à 100 pour cent, et chaque victoire de GRPO dans ce travail s'est appuyée sur ce point de contrôle SFT.
2. Utiliser LoRA pour le SFT et le GRPO. LoRA r=8/α=16 était économe en paramètres (0,068 pour cent entraînables), et sur le contrôle Qwen il est resté numériquement stable tout au long du GRPO (norme de gradient 0,23, environ 11,7 Go par GPU). C'est un choix par défaut solide lorsque vous voulez que le SFT et le GRPO partagent la même forme d'adaptateur.
3. Adapter la technique à l'écart. Utilisez le SFT pour un comportement manquant. Utilisez le GRPO pour affiner un comportement existant, parfois correct. Le plafond réaliste de GRPO ici était de +6 à +8 points.
4. Vérifiez la marge de progression des récompenses avant de dépenser des GPU-heures. Si la plupart des groupes sont entièrement corrects (90,5 pour cent de base sur GSM8K), l'avantage relatif au groupe s'effondre à zéro et GRPO n'obtient aucun gradient. Vérifiez qu'il existe une marge de progression (pass@8 >> pass@1) avant de lancer une exécution GRPO.

Nettoyer les ressources

Pour éviter des frais continus, supprimez les ressources de la charge de travail, puis l'infrastructure sous-jacente :

1. Supprimez les pods et les jobs d'entraînement et d'inférence :

   kubectl delete hyperpodpytorchjob lora-hyperpod-elastic --ignore-not-found
   kubectl delete pod grpo-single inference-base inference-trained inference-grpo \
     eval-grpo fsx-storage-manager --ignore-not-found

2. Réduisez la taille ou supprimez le cluster Amazon SageMaker HyperPod ou ses groupes de nœuds.
3. Supprimez le système de fichiers Amazon FSx for Lustre une fois que vous avez exporté les checkpoints que vous souhaitez conserver.
4. Supprimez tous les compartiments ou objets Amazon Simple Storage Service (Amazon S3) utilisés pour les données d'entraînement ou les archives de checkpoints.

Conclusion

Nous avons montré comment affiner GPT-OSS-20B pour le raisonnement multilingue sur Amazon SageMaker HyperPod, et comment choisir entre SFT et GRPO. SFT a fait l'essentiel du travail : la langue de raisonnement est passée de 22 pour cent à 100 pour cent. GRPO a apporté un véritable raffinement, mais modeste, par-dessus (réponse de 62 pour cent à 70 pour cent) là où une marge de progression existait. Un contrôle parallèle avec Qwen2.5-7B a confirmé ce schéma. GRPO a ajouté +6 pour cent sur le raisonnement multilingue, où il y avait de la place pour s'améliorer. Il n'a apporté aucun gain sur GSM8K, où le modèle était déjà à sa précision maximale (90,5 pour cent).

Cadre de décision :

• Comportement que le modèle ne peut pas du tout réaliser : utilisez SFT.
• Comportement qu'il réalise parfois, avec une marge de progression : utilisez SFT puis GRPO pour quelques points de plus.
• Modèle déjà à sa précision maximale, ou cible ayant une probabilité de base nulle : GRPO n'aidera pas.

Pour essayer vous-même, parcourez les cas de test dans le dépôt awsome-distributed-ai . Pour en savoir plus sur Amazon SageMaker HyperPod, consultez la page de détails du service et la documentation Amazon SageMaker HyperPod. Pour une présentation de Ray sur Amazon SageMaker HyperPod, consultez l'article connexe Exécution de l'entraînement distribué sur Amazon SageMaker HyperPod avec Ray et Amazon SageMaker Studio. Le lien sera ajouté après la publication de cet article. Partagez vos résultats dans les commentaires.

À propos des auteurs