跳到正文
AWS Machine Learning Blog· Mijanur Palash·· 4 小时前AI 评分56

在 Amazon SageMaker HyperPod 上微调 GPT-OSS-20B 的多语言推理能力

Fine-tuning GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod

AI 导读

AWS 团队在 Amazon SageMaker HyperPod 上对 openai/gpt-oss-20b 做多语言推理微调,用 SFT 将 50 例评测中的目标语言链式思考正确率从 22% 提升到 100%,再叠加 GRPO 把答案语言正确率从 62% 提到 70%。

正文 · 原文

GPT-OSS-20B is an open-weight reasoning model that thinks in English. Ask it to reason in French, German, Spanish, or Italian and it acknowledges the instruction, then produces its chain-of-thought in English anyway. On a 50-example multilingual reasoning test, the base model reasoned in the correct target language only 22 percent of the time. That 22 percent was almost entirely the English examples (English 100 percent, French 0 percent, German 10 percent, Spanish 0 percent, Italian 0 percent).

For teams deploying reasoning models for non-English users, that gap matters. A user who asks a question in French expects the model to reason and answer in French, not to switch to English mid-thought. Closing that gap is a fine-tuning problem. It raises a question every team faces once they start customizing open-weight models: supervised fine-tuning (SFT), reinforcement learning with Group Relative Policy Optimization (GRPO), or both?

In this post, we walk through how we fine-tuned GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod. We show when SFT alone was sufficient and when GRPO added measurable value on top. To understand more broadly when GRPO helps, we also ran a controlled Qwen2.5-7B experiment across two tasks with different capability gaps. The short answer: SFT does most of the work, GRPO refines when there is headroom to refine, and matching the technique to the capability gap is what decides the outcome.

Solution overview

This section covers the multilingual reasoning task and dataset, how SFT and GRPO complement each other, and the training stack we ran on Amazon SageMaker HyperPod.

The task and dataset

The dataset is HuggingFaceH4/Multilingual-Thinking — a chain-of-thought reasoning dataset translated into French, Spanish, German, and Italian alongside English, roughly 949 training and 50 validation examples. The model receives a reasoning language: system-prompt instruction and must both reason and answer in that target language, overriding its strong English prior. The evaluation is 50 test cases: 10 math/logic problems across 5 languages (EN/FR/DE/ES/IT).

Two metrics matter:

  • Reasoning-language correct – is the chain-of-thought in the requested language?
  • Answer-language correct – is the final answer in the requested language?

How SFT and GRPO fit together

The clearest way to think about the two techniques: SFT teaches a behavior the model doesn’t have. GRPO optimizes a behavior the model can already sometimes produce. GRPO works by sampling a group of responses per prompt and pushing the policy toward the above-average ones. If the desired behavior has zero probability in the base distribution, there is very little for GRPO to reinforce. That single idea explains almost every result in this post.

The training stack

We ran on Amazon SageMaker HyperPod with Amazon Elastic Kubernetes Service (Amazon EKS) orchestration. For the initial SFT and GRPO runs on GPT-OSS-20B, we used the PyTorch TRL library. For the Qwen control experiment, we switched to veRL as the trainer, with Ray and KubeRay for distributed orchestration and vLLM for fast rollout generation. Checkpoints and datasets lived on Amazon FSx for Lustre, and the inter-node fabric used Elastic Fabric Adapter (EFA). All fine-tuning used Low-Rank Adaptation (LoRA), which is parameter-efficient and well-suited to both SFT and GRPO.

The following diagram shows the training data path end-to-end. A submission from the user lands on the Amazon EKS control plane of the Amazon SageMaker HyperPod cluster. Amazon EKS schedules a Ray head pod plus worker pods that run the veRL trainer and vLLM rollout workers. Both sets of pods mount Amazon FSx for Lustre for the base model, training data, and checkpoints. Inter-node collectives use EFA. Rollouts stream from vLLM back to the veRL trainer, and the trainer writes checkpoints back to Amazon FSx for Lustre.

End-to-end fine-tuning architecture on Amazon SageMaker HyperPod

Figure 1: End-to-end training data path for SFT and GRPO fine-tuning on Amazon SageMaker HyperPod with veRL, Ray, KubeRay, vLLM, and Amazon FSx for Lustre for Qwen control experiments

The following describes each hop in more detail. The submitter starts the run from Amazon SageMaker Studio or the AWS Command Line Interface (AWS CLI). The Amazon EKS control plane on Amazon SageMaker HyperPod places pods on GPU nodes and starts a Ray head with worker pods for the veRL trainer and the vLLM rollout workers. Amazon FSx for Lustre provides shared, high-throughput storage for the base model, the training data, and the checkpoints. When training completes, the trainer writes the LoRA adapter back to Amazon FSx for Lustre for downstream evaluation or serving.

Prerequisites

To follow along, you need:

  • An Amazon SageMaker HyperPod cluster with Amazon EKS orchestration. GPT-OSS SFT used 4× ml.g6e.12xlarge (16× L40S 48GB). The Qwen control used 2× g6e.48xlarge (16× L40S).
  • An Amazon FSx for Lustre file system mounted on the cluster with a PersistentVolumeClaim named fsx-claim at /fsx.
  • EFA-enabled instances and a working NVIDIA Collective Communications Library (NCCL) and aws-ofi-nccl stack for multi-node SFT.
  • A container registry such as Amazon Elastic Container Registry (Amazon ECR) to hold the training image built from the reference repository.
  • A HuggingFace access token for openai/gpt-oss-20b and the HuggingFaceH4/Multilingual-Thinking dataset.
  • The AWS Command Line Interface (AWS CLI), kubectl, docker, and envsubst available on your workstation.

Set up the training environment

All of the commands, manifests, and scripts referenced below live in the awsome-distributed-ai GitHub repository under 3.test_cases/pytorch/trl/gpt-oss-lora-grpo/. Clone the repository and change into the test case directory:

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

Copy the environment template, then edit env_vars to set your ECR registry, image name, tag, and HuggingFace token:

cp env_vars.example env_vars
source env_vars

Build the training image and push it to your registry. The image bundles PyTorch, TRL, HuggingFace Transformers, PEFT, and the training and inference scripts:

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

Point kubectl at your cluster and apply the FSx storage manager pod. The storage manager gives you a shell into the shared file system for checkpoint conversion and eval-log inspection:

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-tune GPT-OSS-20B with SFT and LoRA

We begin by measuring the base model, then configure the SFT run and examine the results.

The base-model gap

Base GPT-OSS-20B on the 50-case eval scored 22 percent reasoning-language correct (11/50). Per-language: English 100 percent, French 0 percent, German 10 percent, Spanish 0 percent, Italian 0 percent. The model defaults to English reasoning regardless of the instruction.

SFT configuration

The SFT setup used LoRA r=8, α=16, trained for 1,000 steps on 4× ml.g6e.12xlarge (16× L40S 48GB) with Fully Sharded Data Parallel (FSDP) full-shard, bf16 precision, and gradient checkpointing. The job runs as an elastic HyperPodPyTorchJob managed by Kueue, launched with hyperpodrun.

Apply the SFT manifest:

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

Watch training progress:

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

Loss and gradient norm print each step. A healthy run looks like this:

{'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}

When training completes, convert the FSDP checkpoint to the Parameter-Efficient Fine-Tuning (PEFT) format that TRL’s GRPO trainer and the inference scripts expect:

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

SFT results and the staged learning curve

After SFT, the model reasons in the correct language 100 percent of the time (50/50) across all five languages. SFT learned this in two stages:

  • Around step 200, the model learned to reason in the target language (reasoning-language metric hits 100 percent).
  • Around step 1000, the model additionally learned to give the answer in the target language.

Final SFT eval (step 1000):

  • Reasoning-language correct: 100 percent (50/50).
  • Answer-language correct: 62 percent (31/50).
  • Both correct: 62 percent (31/50).
  • Per-language answer accuracy: FR 60 percent, DE 60 percent, ES 30 percent, IT 60 percent, EN 100 percent.

Takeaway: SFT is not optional here. The base model can’t follow the reasoning-language instruction at all, and SFT takes reasoning-language from 22 percent to 100 percent. The remaining gap is answer-language at 62 percent. That’s the harder half, because the final answer is short and the English prior is strongest there.

Refine answer-language with GRPO

Answer-language was the obvious target for GRPO. SFT already produced correct reasoning 100 percent of the time, and the model sometimes produced the answer in the right language, so there was a signal to reinforce.

Run GRPO on top of the SFT checkpoint

We ran GRPO with TRL’s GRPOTrainer, starting from the SFT checkpoint. The setup: load the SFT LoRA, call merge_and_unload, and add a new LoRA (r=8, α=16, about 0.068 percent trainable parameters, 13.6M of 20B). Training used K=8 generations per prompt, device_map=auto, learning rate 1e-6, and 10 epochs. The reward function scored answer-language (±5.0), reasoning-language (±1.5), and brevity (+0.5 or −1.0), for a combined range of −7.5 to +7.0.

Launch GRPO from the converted SFT checkpoint:

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

GRPO writes checkpoints to /fsx/checkpoints/grpo-singlenode/. Pick the best checkpoint against your held-out eval (the answer-language metric peaked early in this run) and convert it to PEFT format for inference and evaluation.

Result (from the SFT checkpoint):

  • Reasoning-language correct: 100 percent (held: GRPO didn’t break what SFT taught).
  • Answer-language correct: 62 to 70 percent (+8 points, 35/50).
  • Both correct: 70 percent.
  • Per-language answer: FR 70 percent, DE 80 percent, ES 60 percent, IT 40 percent, EN 100 percent.

GRPO delivered a real +8 point gain on the answer-language metric while holding reasoning at 100 percent. The gain came where there was headroom to gain. Answer-language was the sometimes-correct behavior that SFT had not fully converged on, and GRPO reinforced it.

Compare base, SFT, and GRPO inference side by side

To see the qualitative difference, deploy three inference pods that share the same base model but load different adapters:

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 -

Run the same French-reasoning prompt through each variant:

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

The base pod reasons in English despite the instruction. The SFT pod reasons in French but often answers in English. The GRPO pod reasons and answers in French. To reproduce the full 50-case eval that generated the numbers earlier, run the evaluation manifest and read the result file from the FSx storage manager:

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

To understand more broadly when GRPO adds value, we ran a controlled experiment on Qwen2.5-7B across two tasks with different capability gaps.

Test the pattern with a Qwen2.5-7B control

To test the “match technique to capability gap” thesis on a different model and different tasks, we ran controlled experiments on Qwen2.5-7B-Instruct. The first task is the same multilingual reasoning task (where headroom exists). The second task is GSM8K, where the model is already at peak accuracy. This is an internal experiment, and the public repo has not yet been released.

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.

Arms: A = base, B = SFT-only, C = SFT+GRPO.

Multilingual reasoning (50 cases)

Metric A base B SFT C SFT+GRPO Δ GRPO
Structured format 0% 92% 98% +6
Reasoning language 0% 96% 98% +2
Answer language 0% 38% 44% +6
Both correct 0% 38% 44% +6

GSM8K, a task with no headroom

On GSM8K, base Qwen2.5-7B already scores 90.5 percent (181/200). GRPO at step 30 also scores 90.5 percent. When most sampled groups are all-correct, the group-relative advantage collapses to zero, so there is no learning signal. This is the other half of the thesis: GRPO needs headroom (pass@k >> pass@1) to have anything to optimize. No headroom means no gain.

Best practices

  1. Start with SFT. GRPO cannot teach a behavior with zero base probability. It can only reinforce responses the model already sometimes produces. SFT took reasoning-language from 22 percent to 100 percent, and every GRPO win in this work was built on that SFT checkpoint.
  2. Use LoRA for both SFT and GRPO. LoRA r=8/α=16 was parameter-efficient (0.068 percent trainable), and on the Qwen control it stayed numerically stable through GRPO (grad norm 0.23, about 11.7 GB per GPU). It’s a strong default when you want SFT and GRPO to share the same adapter shape.
  3. Match the technique to the gap. Use SFT for missing behavior. Use GRPO for refining an existing, sometimes-correct behavior. GRPO’s realistic ceiling here was +6 to +8 points.
  4. Check for reward headroom before spending GPU-hours. If most groups are all-correct (GSM8K 90.5 percent base), the group-relative advantage collapses to zero and GRPO gets no gradient. Verify there is headroom (pass@k >> pass@1) before starting a GRPO run.

Clean up resources

To avoid ongoing charges, delete the workload resources and then the underlying infrastructure:

  1. Delete the training and inference pods and jobs:
    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. Scale down or delete the Amazon SageMaker HyperPod cluster or its node groups.
  3. Delete the Amazon FSx for Lustre file system once you have exported any checkpoints you want to keep.
  4. Delete any Amazon Simple Storage Service (Amazon S3) buckets or objects used for training data or checkpoint archives.

Conclusion

We showed how to fine-tune GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod, and how to decide between SFT and GRPO. SFT did most of the work: reasoning-language went from 22 percent to 100 percent. GRPO added a real, modest refinement on top (answer 62 percent to 70 percent) where headroom existed. A parallel Qwen2.5-7B control confirmed the pattern. GRPO added +6 percent on multilingual reasoning, where there was room to improve. It added no gain on GSM8K, where the model was already at peak accuracy (90.5 percent).

Decision framework:

  • Behavior the model can’t do at all: use SFT.
  • Behavior it does sometimes, with headroom: use SFT then GRPO for a few points more.
  • Model already at peak accuracy, or target has zero base probability: GRPO won’t help.

To try this yourself, walk through the test cases in the awsome-distributed-ai repository. For more on Amazon SageMaker HyperPod, see the service detail page and the Amazon SageMaker HyperPod documentation. For a walkthrough of Ray on Amazon SageMaker HyperPod, see the related post Running distributed training on Amazon SageMaker HyperPod with Ray and Amazon SageMaker Studio. The link will be added after that post is published. Share your results in the comments.


About the authors

来源: AWS Machine Learning Blog · aws.amazon.com