在 Amazon SageMaker HyperPod 上微调 GPT-OSS-20B 实现多语言推理
Fine-tuning GPT-OSS-20B for multilingual reasoning on Amazon SageMaker HyperPod
在 Amazon SageMaker HyperPod 上对 GPT-OSS-20B 做微调,可修复其"用英文思考"的问题:基座模型在 50 例多语言推理测试中仅 22% 用目标语言推理(英语 100%、法语 0%、德语 10%、西班牙语 0%、意大利语 0%)。
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 个验证示例。模型接收一个 reasoning language: 系统提示指令,必须用该目标语言进行推理和回答,覆盖其强烈的英语先验。评估使用 50 个测试用例:5 种语言(EN/FR/DE/ES/IT)上的 10 个数学/逻辑问题。
两个指标很重要:
- 推理语言正确——思维链是否使用了请求的语言?
- 答案语言正确——最终答案是否使用了请求的语言?
SFT 与 GRPO 如何配合
思考这两种技术最清晰的方式是:SFT 教模型一种它不具备的行为。GRPO 优化模型已经有时能够产生的行为。GRPO 的工作原理是为每个提示采样一组响应,并将策略推向高于平均水平的响应。如果期望的行为在基础分布中的概率为零,那么 GRPO 几乎没有可强化的内容。这一个想法解释了本文中几乎所有的结果。
训练栈
我们在由 Amazon Elastic Kubernetes Service(Amazon EKS)编排的 Amazon SageMaker HyperPod 上运行。对于 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 head pod 以及若干 worker pod,分别运行 veRL 训练器和 vLLM rollout worker。这两组 pod 都挂载 Amazon FSx for Lustre,用于存放基础模型、训练数据和检查点。节点间集合通信使用 EFA。Rollout 从 vLLM 流式返回给 veRL 训练器,训练器再将检查点写回 Amazon FSx for Lustre。
图 1:在 Amazon SageMaker HyperPod 上使用 veRL、Ray、KubeRay、vLLM 和 Amazon FSx for Lustre 进行 SFT 与 GRPO 微调的端到端训练数据路径(Qwen 对照实验)
下面更详细地描述每一跳。提交者从 Amazon SageMaker Studio 或 AWS Command Line Interface(AWS CLI)启动运行。Amazon SageMaker HyperPod 上的 Amazon EKS 控制平面将 pod 放置到 GPU 节点上,并启动一个 Ray head 以及用于 veRL 训练器和 vLLM rollout worker 的 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处有一个名为fsx-claim的PersistentVolumeClaim。 - 支持 EFA 的实例,以及可正常工作的 NVIDIA Collective Communications Library(NCCL)和
aws-ofi-nccl栈,用于多节点 SFT。 - 一个容器镜像仓库(例如 Amazon Elastic Container Registry(Amazon ECR)),用于存放从参考仓库构建的训练镜像。
- 一个 HuggingFace 访问令牌,用于
openai/gpt-oss-20b和HuggingFaceH4/Multilingual-Thinking数据集。 - 你的工作站上可用的 AWS Command Line Interface(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,用于检查点转换和评估日志检查:
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 运行并检查结果。
基础模型的差距
基础 GPT-OSS-20B 在 50 个用例的评估中得分为 22% 推理语言正确(11/50)。按语言划分:英语 100%,法语 0%,德语 10%,西班牙语 0%,意大利语 0%。无论指令如何,该模型都默认使用英语推理。
SFT 配置
SFT 设置使用 LoRA r=8、α=16,在 4× ml.g6e.12xlarge(16× L40S 48GB)上训练 1,000 步,采用 Fully Sharded Data Parallel(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 checkpoint 转换为 TRL 的 GRPO trainer 和推理脚本所期望的 Parameter-Efficient Fine-Tuning (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)。
- 各语言答案准确率:FR 60%,DE 60%,ES 30%,IT 60%,EN 100%。
要点:这里 SFT 不是可选项。基础模型完全无法遵循推理语言指令,而 SFT 将推理语言从 22% 提升到 100%。剩余的差距是答案语言,为 62%。这是更难的一半,因为最终答案很短,英语先验在那里最强。
用 GRPO 改进答案语言
答案语言是 GRPO 的明显目标。SFT 已经 100% 产生正确的推理,而模型有时会用正确的语言给出答案,因此有信号可以强化。
在 SFT checkpoint 之上运行 GRPO
我们使用 TRL 的 GRPOTrainer 运行 GRPO,从 SFT checkpoint 开始。设置:加载 SFT LoRA,调用 merge_and_unload,并添加一个新的 LoRA(r=8,α=16,约 0.068% 可训练参数,20B 中的 13.6M)。训练使用每个 prompt K=8 次生成,device_map=auto,学习率 1e-6,10 个 epoch。奖励函数对答案语言(±5.0)、推理语言(±1.5)和简洁性(+0.5 或 −1.0)评分,合计范围为 −7.5 到 +7.0。
从转换后的 SFT checkpoint 启动 GRPO:
envsubst '$REGISTRY $IMAGE $TAG $HF_TOKEN' \
< kubernetes/train-grpo-singlenode.yaml | kubectl apply -f -
kubectl logs -f grpo-singleGRPO 将 checkpoint 写入 /fsx/checkpoints/grpo-singlenode/。根据你的留出评估选择最佳 checkpoint(本次运行中答案语言指标在早期达到峰值),并将其转换为 PEFT 格式用于推理和评估。
结果(来自 SFT checkpoint):
- 推理语言正确:100%(保持:GRPO 没有破坏 SFT 所学的内容)。
- 答案语言正确:62% 到 70%(+8 个百分点,35/50)。
- 两者均正确:70%。
- 各语言答案:FR 70%,DE 80%,ES 60%,IT 40%,EN 100%。
GRPO 在答案语言指标上带来了实实在在的 +8 个百分点提升,同时将推理保持在 100%。提升出现在有提升空间的地方。答案语言是 SFT 尚未完全收敛的有时正确的行为,GRPO 对其进行了强化。
并排比较 base、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 -对每个变体运行相同的法语推理 prompt:
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 Frenchbase 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 对照组测试该模式
为了在另一个模型和不同任务上检验“让技术与能力差距相匹配”这一论点,我们在 Qwen2.5-7B-Instruct 上进行了对照实验。第一个任务与之前相同,是多语言推理任务(该任务存在提升空间)。第二个任务是 GSM8K,模型在该任务上已达到峰值准确率。这是一项内部实验,公开仓库尚未发布。
基础设施:2× g6e.48xlarge(16× L40S),EFA 100Gbps,FSx Lustre,Ray 2.44 + KubeRay 1.4.2,veRL v0.6.1,vLLM 0.8.4。
实验组:A = 基线,B = 仅 SFT,C = SFT+GRPO。
多语言推理(50 个案例)
| 指标 | A 基线 | B SFT | C SFT+GRPO | Δ GRPO |
| 结构化格式 | 0% | 92% | 98% | +6 |
| 推理语言 | 0% | 96% | 98% | +2 |
| 答案语言 | 0% | 38% | 44% | +6 |
| 两者均正确 | 0% | 38% | 44% | +6 |
GSM8K:一个没有提升空间的任务
在 GSM8K 上,基线 Qwen2.5-7B 的得分已达 90.5%(181/200)。第 30 步的 GRPO 得分同样为 90.5%。当大多数采样组全部正确时,组相对优势会坍缩为零,因此没有学习信号。这正是该论点的另一半:GRPO 需要提升空间(pass@k >> pass@1)才有可优化的内容。没有提升空间就没有收益。
最佳实践
- 从 SFT 开始。GRPO 无法教会模型一种基础概率为零的行为。它只能强化模型已经偶尔会产生的回答。SFT 将推理语言从 22% 提升到 100%,本工作中 GRPO 的每一项收益都建立在该 SFT 检查点之上。
- SFT 和 GRPO 都使用 LoRA。LoRA r=8/α=16 具有参数效率(可训练参数占 0.068%),并且在 Qwen 对照组上,它在 GRPO 过程中保持数值稳定(梯度范数 0.23,每 GPU 约 11.7 GB)。当你希望 SFT 和 GRPO 共享相同的适配器形状时,这是一个强有力的默认选择。
- 让技术与差距相匹配。用 SFT 补齐缺失的行为。用 GRPO 优化已有的、偶尔正确的行为。GRPO 在此处的现实上限是 +6 到 +8 分。
- 在花费 GPU 小时之前检查奖励提升空间。如果大多数组全部正确(GSM8K 基线 90.5%),组相对优势会坍缩为零,GRPO 得不到梯度。在启动 GRPO 运行之前,请确认存在提升空间(pass@k >> pass@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 - 缩减或删除 Amazon SageMaker HyperPod 集群或其节点组。
- 导出需要保留的检查点后,删除 Amazon FSx for Lustre 文件系统。
- 删除用于训练数据或检查点归档的任何 Amazon Simple Storage Service(Amazon S3)存储桶或对象。
结论
我们展示了如何在 Amazon SageMaker HyperPod 上为多语言推理微调 GPT-OSS-20B,以及如何在 SFT 和 GRPO 之间做出选择。SFT 完成了大部分工作:推理语言从 22% 提升到 100%。在存在提升空间的地方,GRPO 在其之上带来了真实而适度的优化(答案从 62% 提升到 70%)。并行的 Qwen2.5-7B 对照组证实了这一模式。在多语言推理上,GRPO 带来了 +6% 的提升,因为那里有改进空间。在 GSM8K 上,模型已达到峰值准确率(90.5%),GRPO 没有带来任何收益。
决策框架:
- 模型完全无法完成的行为:使用 SFT。
- 模型偶尔能完成、且有提升空间的行为:使用 先 SFT 再 GRPO,可再提升几分。
- 模型已达到峰值准确率,或目标的基础概率为零:GRPO 不会有帮助。
要自行尝试,请逐步运行 awsome-distributed-ai 仓库中的测试用例。有关 Amazon SageMaker HyperPod 的更多信息,请参阅服务详情页面和 Amazon SageMaker HyperPod 文档。有关在 Amazon SageMaker HyperPod 上使用 Ray 的演练,请参阅相关文章使用 Ray 和 Amazon SageMaker Studio 在 Amazon SageMaker HyperPod 上运行分布式训练。该链接将在文章发布后添加。欢迎在评论中分享你的结果。
关于作者
Mijanur Palash
Mijanur 是 AWS AI/ML 基础设施团队的一名 Specialist Solutions Architect。他专注于为汽车、制药和生成式 AI 客户在 Amazon SageMaker HyperPod 上运行分布式训练和强化学习工作负载。工作之余,Palash 喜欢徒步和摄影。
Mark Vinciguerra
Mark 曾是 AWS 生成式 AI 的 Worldwide Specialist Solutions Architect,专注于 Amazon SageMaker HyperPod 上的分布式训练和推理。他与战略客户合作处理跨 EKS 和 Slurm 编排的大规模基础模型工作负载,并为 HyperPod、Ray 和 SageMaker Unified Studio 贡献参考架构和研讨会。
来源:AWS Machine Learning Blog · aws.amazon.com
