AWS Machine Learning来源发布时间

在 Amazon SageMaker HyperPod 上针对多语言推理微调 GPT-OSS-20B

GPT-OSS-20B 是一个以英语进行思考的开放权重推理模型。当你要求它用法语、德语、西班牙语或意大利语进行推理时,它会确认该指令,但随后仍然用英语生成其思维链。

GPT-OSS-20B 是一个以英语进行思考的开放权重推理模型。当你要求它用法语、德语、西班牙语或意大利语进行推理时,它会确认该指令,但随后仍然用英语生成其思维链。在一个包含 50 个示例的多语言推理测试中,基础模型只有 22% 的时间使用正确的目标语言进行推理。这 22% 几乎全部是英语示例(英语 100%,法语 0%,德语 10%,西班牙语 0%,意大利语 0%)。

对于为非英语用户部署推理模型的团队来说,这一差距非常重要。用法语提问的用户期望模型用法语进行推理和回答,而不是在思考过程中切换到英语。弥合这一差距是一个微调问题。这就引出了每个团队在开始定制开放权重模型时都会面临的问题:监督微调(SFT)、使用组相对策略优化(GRPO)的强化学习,还是两者兼用?

在这篇文章中,我们将逐步介绍我们如何对 GPT-OSS-20B 进行微调,以实现多语言推理 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 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 head pod 以及运行 veRL 训练器和 vLLM rollout 工作进程的工作 pod。两组 pod 都挂载了 Amazon FSx for Lustre,用于存放基础模型、训练数据和检查点。节点间集合通信使用 EFA。Rollout 从 vLLM 流式传回 veRL 训练器,训练器将检查点写回 Amazon FSx for Lustre。

图 1:在 Amazon SageMaker HyperPod 上结合 veRL、Ray、KubeRay、vLLM 以及用于 Qwen 对照实验的 Amazon FSx for Lustre 进行 SFT 和 GRPO 微调的端到端训练数据路径

下面更详细地描述每个环节。提交者从 Amazon SageMaker Studio 或 AWS Command Line Interface(AWS CLI)启动运行。Amazon SageMaker HyperPod 上的 Amazon EKS 控制平面将 pod 放置到 GPU 节点上,并启动一个带有工作 pod 的 Ray head,用于 veRL 训练器和 vLLM rollout 工作进程。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)。
  • 一个挂载到集群上并带有 PersistentVolumeClaim 名称为 fsx-claim at 的 Amazon FSx for Lustre 文件系统 /fsx.
  • 启用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 命令行界面(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 registry、镜像名称、标签和 HuggingFace token:

cp env_vars.example env_vars
source env_vars

构建训练镜像并将其推送到你的 registry。该镜像打包了 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 检查点转换为 TRL 的 GRPO 训练器和推理脚本所期望的 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-peft

SFT 结果与分阶段学习曲线

SFT 之后,模型 100% 的情况下能以正确的语言进行推理 (50/50),在全部五种语言中都是如此。SFT 分两个阶段学会了这一点:

  • 大约在第 200 步,模型学会了用目标语言进行 推理 (推理语言指标达到 100 percent)。
  • 大约在第 1000 步,模型又学会了用目标语言给出 答案 。

SFT 最终评估(第 1000 步):

  • 推理语言正确: 100 percent (50/50).
  • 答案语言正确: 62 percent (31/50).
  • 两者都正确: 62 percent (31/50).
  • 各语言的答案准确率:FR 60 percent,DE 60 percent,ES 30 percent,IT 60 percent,EN 100 percent。

要点: 在这里 SFT 不是可选的。基础模型完全无法遵循推理语言的指令,而 SFT 将推理语言从 22 percent 提升到了 100 percent。剩下的差距是答案语言的 62 percent。这是更难的那一半,因为最终答案很短,而英语先验在那里最强。

用 GRPO 改进答案语言

答案语言是 GRPO 显而易见的目标。SFT 已经 100 percent 的时间里产生正确推理,而且模型 有时 会用正确的语言给出答案,因此存在可以强化的信号。

在 SFT 检查点之上运行 GRPO

我们使用 TRL 的 GRPOTrainer从 SFT 检查点开始运行了 GRPO。设置如下:加载 SFT LoRA,调用 merge_and_unload,并添加一个新的 LoRA(r=8,α=16,约 0.068 percent 的可训练参数,20B 中的 13.6M)。训练每个提示使用 K=8 个生成结果、 device_map=auto、学习率 1e-6 以及 10 个 epoch。奖励函数对答案语言(±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-single

GRPO将检查点写入到 /fsx/checkpoints/grpo-singlenode/。针对您的留出评估集挑选最佳检查点(本次运行中回答语言指标较早达到峰值),并将其转换为PEFT格式以进行推理和评估。

结果(来自SFT检查点):

  • 推理语言正确: 100 percent (留出集:GRPO没有破坏SFT所教的内容)。
  • 回答语言正确: 62 to 70 percent (+8 points, 35/50)。
  • 两者均正确: 70 percent.
  • 按语言回答结果:FR 70 percent, DE 80 percent, ES 60 percent, IT 40 percent, EN 100 percent。

GRPO在回答语言指标上实现了实际的+8分提升,同时将推理保持在100 percent。提升来自有提升空间的地方。回答语言是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 -

将相同的法语推理提示运行到每个变体:

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

base 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 = base,B = 仅SFT,C = SFT+GRPO。

多语言推理(50 cases)

Metric A base B SFT C SFT+GRPO Δ GRPO
Structured format 0% 92% 98% +6
Reasoning language 0% 96% 98% +2
回答语言 0% 38% 44% +6
两者均正确 0% 38% 44% +6

GSM8K,一个没有提升空间的任务

在GSM8K上,基础Qwen2.5-7B已经得分 90.5 percent (181/200)。第30步的GRPO同样得分90.5 percent。当大多数采样组全部正确时,组相对优势坍缩为零,因此没有学习信号。这是论文主张的另一半: GRPO 需要有提升空间(pass@k >> pass@1)才有可优化的东西。 没有提升空间就意味着没有收益。

最佳实践

  1. 从SFT开始。 GRPO 无法教授基础概率为零的行为。它只能强化模型已经偶尔会生成的回答。SFT 将推理语言从 22 percent 提升到 100 percent,而本工作中每一次 GRPO 的收益都是建立在该 SFT 检查点之上的。
  2. 在SFT和GRPO中都使用LoRA。 LoRA r=8/α=16 是参数高效的(0.068 percent 可训练),并且在 Qwen 对照实验中它在整个 GRPO 过程中保持了数值稳定(梯度范数 0.23,每块 GPU 约 11.7 GB)。当你希望 SFT 和 GRPO 共享相同的适配器形状时,这是一个很好的默认选择。
  3. 使技术与差距相匹配。 对缺失的行为使用SFT。对改进现有的、偶尔正确的行为使用GRPO。GRPO在本文中的现实上限是+6到+8分。
  4. 在花费GPU时长之前,先检查奖励的提升空间(headroom)。 如果大多数组全部正确(GSM8K 基线为 90.5 percent),组相对优势会塌缩为零,GRPO 将得不到梯度。在启动 GRPO 运行之前,请确认存在提升空间(pass@k >> pass@1)。

清理资源

为避免持续产生费用,请先删除工作负载资源,然后删除底层基础设施:

  1. 删除训练和推理的 pod 及作业:
    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. 缩减或删除 Amazon SageMaker HyperPod 集群或其节点组。
  3. 在导出所有想要保留的检查点之后,删除 Amazon FSx for Lustre 文件系统。
  4. 删除用于训练数据或检查点归档的所有 Amazon Simple Storage Service (Amazon S3) 存储桶或对象。

结论

我们展示了如何在 Amazon SageMaker HyperPod 上对 GPT-OSS-20B 进行多语言推理微调,以及如何在 SFT 和 GRPO 之间做出选择。SFT 完成了大部分工作:推理语言比例从 22 percent 提升到 100 percent。在存在提升空间的情况下,GRPO 在其上增加了真实但适度的改进(答案准确率从 62 percent 提升到 70 percent)。一个并行的 Qwen2.5-7B 对照实验证实了这一模式。GRPO 在有改进空间的多语言推理上增加了 +6 percent。而在模型已达到峰值准确率(90.5 percent)的 GSM8K 上,它没有带来任何提升。

决策框架:

  • 模型完全无法做到的行为:使用 SFT.
  • 模型有时能做、且存在提升空间的行为:使用 SFT 然后进行 GRPO ,可再多得几个百分点。
  • 模型已达到峰值准确率,或目标的基础概率为零: GRPO 不会有帮助.

要亲自尝试,请逐步完成 awsome-distributed-ai 代码库中的测试用例。有关 Amazon SageMaker HyperPod 的更多信息,请参阅 服务详情页面 以及 Amazon SageMaker HyperPod 文档。有关在 Amazon SageMaker HyperPod 上使用 Ray 的演练,请参阅相关文章 在 Amazon SageMaker HyperPod 上结合 Ray 与 Amazon SageMaker Studio 运行分布式训练。该链接将在那篇文章发布后添加。欢迎在评论中分享您的结果。

 


关于作者

原始出处

AWS Machine Learning

内容说明

原始发布及相关权利归来源方。

机器翻译 · 请以原文为准