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

TL;DR · AI 摘要
Fine-tuning GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod Artificial Intelligence Fine-tuning GPT-...
核心要点
- 主题聚焦:Fine-tuning GPT-OSS-20B for multilingual reasoni
- 来源:AWS Machine Learning Blog,建议结合原文判断细节。
- AI 分析暂不可用,本条为保底评分与摘要。
在Amazon SageMaker HyperPod上对GPT-OSS-20B进行多语言推理的微调 | 人工智能
在Amazon SageMaker HyperPod上对GPT-OSS-20B进行多语言推理的微调
GPT-OSS-20B是一个开源权重推理模型,其默认使用英语进行思考。当要求它用法语、德语、西班牙语或意大利语进行推理时,它虽然会接受指令,但最终仍会用英语生成推理过程。在包含50个示例的多语言推理测试中,基础模型仅有22%的推理结果符合目标语言要求。这22%几乎全部来自英语示例(英语100%、法语0%、德语10%、西班牙语0%、意大利语0%)。
对于需要为非英语用户部署推理模型的团队而言,这种差距具有实际影响。当用户用法语提问时,他们期望模型用法语进行推理并给出答案,而不是在推理过程中切换到英语。弥合这一差距属于微调问题。这引发了一个所有团队在开始定制开源权重模型时都会面临的问题:是选择监督微调(SFT),还是基于群体相对策略优化的强化学习(GRPO),或者两者结合?
在本文中,我们将详细介绍如何在Amazon SageMaker HyperPod上对GPT-OSS-20B进行多语言推理的微调。我们将展示单独使用SFT已足够的情形,以及GRPO在哪些情况下能带来可衡量的价值提升。为了更广泛地理解GRPO的适用场景,我们还针对两个具有不同能力差距的任务进行了Qwen2.5-7B的对照实验。简要结论是:SFT完成主要工作,GRPO在存在优化空间时进行精细化调整,而将技术手段与能力差距匹配是决定最终效果的关键。
解决方案概览
本节将涵盖多语言推理任务与数据集、SFT与GRPO的互补关系,以及我们在Amazon SageMaker HyperPod上运行的训练架构。
任务与数据集
我们使用的是HuggingFaceH4/Multilingual-Thinking数据集——一个包含英语、法语、西班牙语、德语和意大利语的思维链推理数据集,包含约949个训练示例和50个验证示例。模型会接收一个推理语言指令(system-prompt),需要在该目标语言中完成推理和回答,从而克服其固有的英语偏好。评估包含50个测试用例:涵盖5种语言(EN/FR/DE/ES/IT)的10个数学/逻辑问题。
两个关键指标:
- 推理语言正确率 – 推理过程是否使用请求的语言?
- 答案语言正确率 – 最终答案是否使用请求的语言?
SFT与GRPO的协同关系
理解这两种技术的最佳方式:SFT教会模型原本不具备的行为,GRPO则优化模型已能部分实现的行为。GRPO通过为每个提示采样一组响应,并将策略向表现优于平均水平的响应倾斜来实现优化。如果目标行为在基础分布中的概率为零,GRPO几乎无法发挥作用。这一核心观点几乎解释了本文所有实验结果。
训练架构
我们使用 Amazon SageMaker HyperPod 和 Amazon Elastic Kubernetes Service (Amazon EKS) 进行编排。在 GPT-OSS-20B 上进行初始的 SFT 和 GRPO 训练时,我们使用了 PyTorch TRL 库。在 Qwen 控制实验中,我们切换到 veRL 作为训练器,使用 Ray 和 KubeRay 进行分布式编排,使用 vLLM 进行快速 rollout 生成。检查点和数据集存储在 Amazon FSx for Lustre 上,节点间通信使用 Elastic Fabric Adapter (EFA)。所有微调均采用低秩适配 (LoRA),这种方法参数高效,且适用于 SFT 和 GRPO。
下图展示了训练数据的端到端传输路径。用户的提交请求会进入 Amazon SageMaker HyperPod 集群的 Amazon EKS 控制平面。Amazon EKS 会调度 Ray 头节点 Pod 和运行 veRL 训练器及 vLLM rollout 工作节点的 worker Pods。这两组 Pod 均挂载 Amazon FSx for Lustre,用于存储基础模型、训练数据和检查点。节点间通信使用 EFA。vLLM 生成的 rollout 会回传至 veRL 训练器,训练器将检查点写回 Amazon FSx for Lustre。
图 1:使用 veRL、Ray、KubeRay、vLLM 和 Amazon FSx for Lustre 在 Amazon SageMaker HyperPod 上进行 SFT 和 GRPO 微调的端到端训练数据路径(Qwen 控制实验)
以下对每个环节进行更详细的描述。提交者从 Amazon SageMaker Studio 或 AWS 命令行界面 (AWS CLI) 启动训练。Amazon SageMaker HyperPod 上的 Amazon EKS 控制平面将 Pod 部署到 GPU 节点,并启动 Ray 头节点及运行 veRL 训练器和 vLLM rollout 工作节点的 worker Pod。Amazon FSx for Lustre 为模型基础、训练数据和检查点提供共享的高吞吐存储。训练完成后,训练器将 LoRA 适配器写回 Amazon FSx for Lustre 用于下游评估或服务。
先决条件
要跟随本教程,您需要:
- 带有 Amazon EKS 编排的 Amazon SageMaker HyperPod 集群。GPT-OSS SFT 使用了 4× ml.g6e.12xlarge(16× L40S 48GB)。Qwen 控制实验使用了 2× g6e.48xlarge(16× L40S)。
- 挂载在集群上的 Amazon FSx for Lustre 文件系统,通过名为 fsx-claim 的 PersistentVolumeClaim 挂载在 /fsx 路径。
- 启用 EFA 的实例,以及用于多节点 SFT 的可用 NVIDIA 集群通信库 (NCCL) 和 aws-ofi-nccl 堆栈。
- 用于存储从参考仓库构建的训练镜像的容器注册表,例如 Amazon 弹性容器注册表 (Amazon ECR)。
- 用于 openai/gpt-oss-20b 和 HuggingFaceH4/Multilingual-Thinking 数据集的 HuggingFace 访问令牌。
- 在工作站上安装 AWS 命令行界面 (AWS CLI)、kubectl、docker 和 envsubst。
配置训练环境
以下所有命令、清单和脚本均位于 awsome-distributed-ai GitHub 仓库的 3.test_cases/pytorch/trl/gpt-oss-lora-grpo/ 路径下。克隆仓库并进入测试用例目录:
git clone https://github.com/awslabs/awsome-distributed-ai
cd awsome-distributed-ai/3.test_cases/pytorch/trl/gpt-oss-lora-grpo复制环境模板,然后编辑 env_vars 以设置您的 ECR 注册表、镜像名称、标签和 HuggingFace 令牌:
cp env_vars.example env_vars
source env_vars构建训练镜像并推送到您的注册表。该镜像包含 PyTorch、TRL、HuggingFace Transformers、PEFT 以及训练和推理脚本:
cd artifacts && ./build_push.sh && cd ..将 kubectl 指向你的集群并应用 FSx 存储管理器 Pod。存储管理器为你提供共享文件系统的 shell 环境,用于检查点转换和 eval-log 检查:
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 -使用 SFT 和 LoRA 对 GPT-OSS-20B 进行微调
我们首先测量基线模型,然后配置 SFT 训练并分析结果。
基线模型差距
基于 50 个案例的评估显示,基线 GPT-OSS-20B 的推理语言正确率为 22%(11/50)。按语言划分:英语 100%,法语 0%,德语 10%,西班牙语 0%,意大利语 0%。无论指令语言如何,模型默认使用英语进行推理。
SFT 配置
SFT 配置使用 LoRA r=8, α=16,在 4× ml.g6e.12xlarge(16× L40S 48GB)上训练 1,000 步,采用全分片数据并行(FSDP)全分片、bf16 精度和梯度检查点。任务作为由 Kueue 管理的弹性 HyperPodPyTorchJob 运行,通过 hyperpodrun 启动。
应用 SFT 配置文件:
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
< hyperpod-eks/train-lora-hyperpod-elastic-g6e.yaml | kubectl apply -f -查看训练进度:
kubectl get pods -l app=gpt-oss-lora-elastic
kubectl logs -f -l app=gpt-oss-lora-elastic每一步都会打印损失值和梯度范数。正常运行示例如下:
{'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}训练完成后,将 FSDP 检查点转换为 TRL 的 GRPO 训练器和推理脚本所需的参数高效微调(PEFT)格式:
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-peftSFT 结果与分阶段学习曲线
SFT 后,模型在所有五种语言中 100%(50/50)的推理语言正确。SFT 通过两个阶段实现这一目标:
- 在约第 200 步,模型学会了在目标语言进行推理(推理语言指标达到 100%)。
- 在约第 1000 步,模型进一步学会了在目标语言给出答案。
最终 SFT 评估(第 1000 步):
- 推理语言正确率:100%(50/50)
- 答案语言正确率:62%(31/50)
- 两者都正确:62%(31/50)
- 按语言答案准确率:法语 60%,德语 60%,西班牙语 30%,意大利语 60%,英语 100%
结论:此处 SFT 是必需的。基线模型完全无法遵循推理语言指令,而 SFT 将推理语言正确率从 22% 提升至 100%。剩余差距在于答案语言 62% 的正确率。这是更具挑战性的部分,因为最终答案较短且英语先验最强。
使用 GRPO 优化答案语言
答案语言是 GRPO 的明显目标。SFT 已经实现了 100% 的推理正确率,模型有时也能用目标语言给出答案,因此存在可强化的信号。
在 SFT 检查点上运行 GRPO
我们使用TRL的GRPOTrainer从SFT检查点启动GRPO训练。配置方案:加载SFT的LoRA权重,调用merge_and_unload方法,新增一个LoRA模块(r=8,α=16,约0.068%可训练参数,20B模型中13.6M参数)。训练参数设置为每提示生成K=8个样本,device_map=auto,学习率1e-6,训练10个周期。奖励函数对答案语言(±5.0)、推理语言(±1.5)和简洁性(+0.5或-1.0)进行评分,综合得分范围为-7.5至+7.0。
从转换后的SFT检查点启动GRPO:
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
< kubernetes/train-grpo-singlenode.yaml | kubectl apply -f -
kubectl logs -f grpo-singleGRPO将检查点写入/fsx/checkpoints/grpo-singlenode/目录。根据预留评估集选择最佳检查点(本次运行中答案语言指标早期达到峰值),将其转换为PEFT格式用于推理和评估。
结果(基于SFT检查点):
- 推理语言正确率:100%(保留了SFT训练的所有成果)
- 答案语言正确率:62%至70%(提升8分,35/50)
- 双语言正确率:70%
- 按语言统计:法语70%,德语80%,西班牙语60%,意大利语40%,英语100%
GRPO在保持推理能力100%的前提下,使答案语言指标提升了8分。这种提升发生在SFT尚未完全收敛的领域——答案语言是SFT有时能正确处理但未完全掌握的指标,GRPO通过强化训练进一步优化了这一能力。
对比基础模型、SFT模型和GRPO模型的推理效果
要观察定性差异,部署三个共享相同基础模型但加载不同适配器的推理Pod:
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 -对每个变体运行相同的法语推理提示:
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基础模型Pod尽管收到指令仍用英语推理。SFT模型Pod用法语推理但常以英语作答。GRPO模型Pod实现了法语推理和法语作答。要复现之前生成数值的完整50案例评估,运行评估清单并从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为更全面理解GRPO的价值场景,我们在Qwen2.5-7B模型上针对两个具有不同能力差距的任务进行了受控实验。
用Qwen2.5-7B控制实验验证模式
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
96%
+2
Answer language
38%
44%
Both correct
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
- 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.
- 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.
- 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.
- 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:
- 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
- Scale down or delete the Amazon SageMaker HyperPod cluster or its node groups.
- Delete the Amazon FSx for Lustre file system once you have exported any checkpoints you want to keep.
- 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.
- 在某些情况下,当存在余量时:先使用 SFT,然后使用 GRPO 以获得更高的分数。
- 模型已达到峰值精度,或目标的初始概率为零:GRPO 无法提供帮助。
要亲自尝试,可以浏览 awsome-distributed-ai 仓库中的测试用例。了解更多关于 Amazon SageMaker HyperPod 的信息,请查看服务详情页面和 Amazon SageMaker HyperPod 文档。要了解如何在 Amazon SageMaker HyperPod 上使用 Ray,请参阅相关文章《使用 Ray 和 Amazon SageMaker Studio 在 Amazon SageMaker HyperPod 上进行分布式训练》。该链接将在该文章发布后添加。请在评论区分享你的结果。
作者简介
'"`