AWS Machine Learning Blog

应对奖励信号挑战:使用GRPO在SageMaker AI上实现基于可验证奖励的强化学习

8.7内容质量
应对奖励信号挑战:使用GRPO在SageMaker AI上实现基于可验证奖励的强化学习

TL;DR · AI 摘要

AWS提出RLVR方法,结合GRPO提升强化学习中的奖励信号可靠性,适用于数学推理等可验证任务。

核心要点

  • RLVR方法通过可验证奖励提升训练透明度和稳定性
  • GRPO技术结合few-shot示例显著提高数学问题解决准确率
  • 适用于数学推理、代码生成等输出可验证的任务场景
#强化学习#AWS#GRPO#数学推理
打开原文

训练大型语言模型需要准确的反馈信号,但传统的强化学习(RL)在奖励信号的可靠性方面常常面临挑战。这些信号的质量直接影响模型的学习和决策过程。然而,构建稳健的反馈机制可能复杂且容易出错。现实世界的训练场景往往会引入隐藏的偏见、非预期的激励以及模糊的成功标准,从而干扰学习过程,导致模型行为不可预测或无法达到预期目标。

在本文中,你将学习如何实现具有可验证奖励的强化学习(Reinforcement Learning with Verifiable Rewards, RLVR),以在奖励信号中引入验证和透明性,从而提升训练性能。当输出可以被客观验证其正确性时,例如数学推理、代码生成或符号操作任务,这种方法效果最佳。你还将学习如何结合使用组相对策略优化(Group Relative Policy Optimization, GRPO)和少样本示例(few-shot examples)等技术来进一步提升效果。我们将使用 GSM8K 数据集(小学数学8K:包含小学数学问题的数据集)来提升数学问题求解的准确性,但此处使用的技术可灵活适配到多种其他应用场景。

技术概览

在深入实现细节之前,理解支撑该方法的RL概念是有帮助的。强化学习通过建立基于奖励信号的结构化反馈系统来解决模型训练中的挑战。这种范式使模型能够通过交互进行学习,接收引导其向最优行为发展的反馈。RL为模型提供了一个框架,使其能够根据输出质量的明确信号迭代改进响应,因此非常适合训练需要与用户交互并根据结果调整行为的模型。传统的RL凸显了一个重要考量:奖励信号的质量至关重要。当奖励函数不精确或不完整时,模型可能会进行“奖励欺骗”(reward hacking),即以非预期的方式最大化得分,而没有实现期望的行为。认识到这一局限性后,人们开发了更严格的方法,专注于创建可靠且定义明确的奖励函数。

RLVR通过模型调优者定义的基于规则的反馈来解决奖励欺骗问题。它使用程序化的奖励函数,自动根据特定标准对输出进行评分,从而无需依赖人工评分即可快速迭代。这些“可验证”的奖励来源于客观、可复现的规则,使RLVR非常适合需求不断演进的场景,因为它能学习通用的优化策略并快速适应新场景。GRPO是一种强化学习算法,它通过在组内而非全部数据中进行性能比较来提升AI模型的学习效果。它将训练数据组织成有意义的组,并相对于每组的基线进行性能优化,从而对每个类别给予适当的关注。这种组感知的优化减少了训练方差,加快了收敛速度,并能产生在各类别中表现一致的模型。将RLVR与GRPO结合,构建了一个由自动化奖励引导学习、组相对优化推动平衡性能的框架。

你为不同任务方面定义奖励函数,GRPO在训练中将这些函数视为不同的组,从而实现多个维度的同步改进。这种组合实现了快速适应和稳健性能,非常适合需要超越训练分布进行泛化的动态环境。加入少样本学习(few-shot learning)进一步从三个方面增强了这一框架。首先,少样本示例提供了模板,展示模型应如何输出高质量结果,从而缩小探索的搜索空间。其次,GRPO通过为每个提示生成多个候选响应,并基于每组内的相对表现进行学习。第三,可验证奖励立即确认哪些方法是成功的。这种组合加速了学习过程:模型从具体的示例开始,通过基于组的比较高效探索变体,并获得关于正确性的明确反馈。

解决方案概览

在本节中,你将学习如何在 SageMaker AI 上对 Qwen2.5-0.5B 模型进行微调,使用 Amazon SageMaker 训练任务。Amazon SageMaker 训练任务支持分布式多GPU和多节点配置,因此你可以按需启动高性能集群,更快地训练十亿参数级模型,并在任务完成后自动关闭资源。

注意: 虽然本例选择了 Qwen2.5-0.5B,但其他场景如代码生成可能需要更大的模型(例如 Qwen2.5-Coder-7B),并相应需要更大的训练实例。

Image 1
Image 1

前提条件

要在 Amazon SageMaker AI 上运行本文中的示例,你需要满足以下前提条件:

  • 一个包含 AWS 资源的 AWS 账户。
  • 一个 AWS Identity and Access Management (IAM) 角色,用于访问 SageMaker AI。要了解 IAM 如何与 SageMaker AI 配合使用,请参阅 AWS Identity and Access Management for Amazon SageMaker AI
  • 您可以从首选的开发环境(包括 PyCharm 或 Visual Studio Code 等交互式开发环境 (IDE))运行本文提供的笔记本,前提是您的 AWS 凭证已正确设置并配置为访问您的 AWS 账户。要设置本地环境,请参阅 Configuring settings for the AWS CLI。您也可以选择使用 Amazon SageMaker Studio 来简化在 SageMaker AI 上的开发流程。
  • 如果您要按照本文操作,您将需要一个 ml.p4d.24xlarge 实例进行训练。您需要访问这些 SageMaker 训练实例才能运行示例训练代码。如果您不确定,可以在 AWS 管理控制台 上查看 AWS 服务配额
  • 在“管理配额”下选择 Amazon SageMaker 作为 AWS 服务。
  • 选择 ml.p4d.24xlarge 用于训练作业使用情况,并请求提高账户级别的配额。

环境设置

您可以使用您偏好的 IDE(如 VS Code 或 PyCharm),但请确保如前提条件中所述,您的本地环境已配置为与 AWS 配合使用。

要使用 SageMaker Studio JupyterLab 空间,请完成以下步骤:

  1. 在 Amazon SageMaker AI 控制台中,选择导航窗格中的“Domains”,然后打开您的域。
  2. 在导航窗格的 Applications and IDEs 下,选择 Studio
  3. 在“User profiles”标签页中,找到您的用户配置文件,然后选择 LaunchStudio
  4. 在 Amazon SageMaker Studio 中,启动一个 ml.t3.medium JupyterLab 笔记本实例,存储空间至少为 50 GB。

不需要大型笔记本实例,因为微调作业将在一个单独的临时训练实例上运行,该实例具有 GPU 加速功能。

  1. 开始微调之前,请先克隆 GitHub 仓库,并导航至 3_distributed_training/reinforcement-learning/grpo-with-verifiable-reward 目录,然后启动 model-finetuning-grpo-rlvr.ipynb
  2. 使用 Python 3.12 或更高版本内核的笔记本

准备微调数据集

使用 RLVR 运行 GRPO 需要您拥有每个问题的最终答案以计算奖励。首先,通过提取每个问题的最终答案来准备数据。

code
dataset = GSM8K(split='train', include_answer=False, include_reasoning=True, few_shot=True, num_shots=8, seed=None, cot=True).dataset.shuffle(seed=42)
Dataset({
    features: ['question', 'answer', 'prompt', 'final_answer'],
    num_rows: 7473
})

此外,此示例使用了少量示例(8 个 shot)以提高模型训练性能。有关强化学习中少量示例的更多信息,请参阅论文 “Reinforcement Learning for Reasoning in Large Language Models with One Training Example”。尽管该研究论文主要关注单个示例,但本文将向您展示单个和多个示例的性能。

每个输入将包含 8 个示例,然后是需要解决的问题:

code
"Question: Mark has $50 and buys a toy that costs $35. How much money does he have left?
Solution: Let's think step by step. To find out how much money Mark has left, subtract the cost of the toy from the total amount of money Mark has. So, $50 - $35 = $15.
#### The final answer is 15

Question: Emily has 3 times as many pencils as Alice. If Alice has 15 pencils, how many pencils does Emily have?
Solution: Let's think step by step. To find out how many pencils Emily has, we multiply the number of pencils Alice has by 3. Alice has 15 pencils, so Emily has 15 * 3 = 45 pencils.
#### The final answer is 45

Question: Jack has collected 12 more marbles than Kevin. If Kevin has 27 marbles, how many marbles does Jack have?
Solution: Let's think step by step. To find how many marbles Jack has, we add 12 to the number of marbles Kevin has. So, Jack has 27 + 12 = 39 marbles.
#### The final answer is 39

Question: There are 24 students in a classroom. If each group must have 4 students, how many groups can be formed?
Solution: Let's think step by step. To find how many groups can be formed, we divide the number of students by the number of students per group. So, 24 / 4 = 6 groups can be formed.
#### The final answer is 6

Question: Samantha baked 40 cookies and wants to divide them equally into bags, with each bag containing 5 cookies. How many bags will Samantha need?
Solution: Let's think step by step. To find the number of bags needed, divide the total number of cookies by the number of cookies per bag. Thus, 40 divided by 5 equals 8.
#### The final answer is 8

问题:一包铅笔售价为 4 美元。如果你购买 7 包,总共要花多少钱? 解答:我们一步一步来思考。总费用是通过每包的价格乘以包数来计算的。因此,你需要花费 7 × 4 美元 = 28 美元。 #### 最终答案是 28

问题:一本书有 240 页,Sarah 每天读 20 页。她需要多少天才能读完这本书? 解答:我们一步一步来思考。Sarah 每天读 20 页,因此我们将总页数除以她每天读的页数。因此,她需要 240 ÷ 20 = 12 天才能读完这本书。 #### 最终答案是 12

问题:一个农民总共有 80 个苹果和橘子。如果他有 30 个苹果,那么他有多少个橘子? 解答:我们一步一步来思考。为了确定橘子的数量,我们需要从水果总数中减去苹果的数量。因此,橘子的数量是 80 - 30 = 50。 #### 最终答案是 50

问题:Mimi 在海滩上捡了 2 打贝壳。Kyle 找到的贝壳数量是 Mimi 的两倍,并把它们放进了口袋。Leigh 抓走了 Kyle 找到贝壳数量的三分之一。Leigh 有多少个贝壳? 解答:我们一步一步来思考。

code

Python

数据准备完成后,保留 10% 的数据作为验证集,并将训练集和验证集都推送到 S3。

## 可验证奖励函数

该 GRPO 数学推理实现采用双奖励系统,在训练期间提供客观、可验证的反馈。这种方法利用了数学问题的可验证性,创建了可靠的训练信号,而无需人工标注或主观评估。你将实现两个互补的奖励函数,它们共同引导模型朝向正确的回答格式和数学准确性:

#### 格式奖励函数

此函数通过以下方式帮助验证模型是否学会了正确构造响应:

*   **模式匹配**:搜索特定格式 `#### The final answer is [number]`
*   **一致性评分**:格式正确得 0.5 分,格式错误得 0.0 分
*   **训练信号**:鼓励模型遵循预期的答案结构

格式奖励函数

def format_reward_func_qa(completions, **kwargs): pattern = r"\n#### The final answer is \d+" completion_contents = [completion for completion in completions] matches = [re.search(pattern, content) for content in completion_contents] return [0.5 if match else 0.0 for match in matches]

code

Python

#### 正确性奖励函数

此函数通过以下方式提供核心数学验证:

*   **答案提取**:使用正则表达式从格式化响应中提取数值答案
*   **标准化**:删除常见格式字符(逗号、货币符号、单位)
*   **精度比较**:使用 1e-3 的容差来处理浮点精度
*   **二元评分**:正确答案得 1.0 分,错误答案得 0.0 分

正确性奖励函数

def correctness_reward_func_qa(completions, final_answer, **kwargs): rewards = []

for completion, ground_truth in zip(completions, final_answer): try: match = re.search(r'####.*?([\d,]+(?:\.\d+)?)', completion) if match: answer = match.group(1)

for remove_char in [',', '$', '%', 'g']: answer = answer.replace(remove_char, '')

if abs(float(answer)-float(ground_truth)) < 1e-3: rewards.append(1.0) else: rewards.append(0.0) else: rewards.append(0.0) except ValueError: rewards.append(0.0)

return rewards

code

Python

#### 将 RLVR 与 GRPO 集成

奖励函数通过 GRPOTrainer 集成到 GRPO 训练流程中:

rewards_funcs = [format_reward_func_qa, correctness_reward_func_qa]

trainer = GRPOTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=test_dataset, processing_class=tokenizer, peft_config=peft_config, reward_funcs=rewards_funcs, )

code

代码

在训练过程中,GRPO 使用这些奖励函数计算策略梯度。首先,模型会为每个数学问题生成多个回答。然后,为每个回答分别计算两个奖励函数的奖励值。格式奖励函数最多可为正确响应结构提供 0.5 分,正确性奖励函数最多可为数学准确性提供 1.0 分,每个回答的最高总奖励为 1.5 分。然后 GRPO 在组内比较回答,找出最佳响应。最后,在策略更新步骤中,损失函数使用奖励差异来更新模型参数。高奖励的回答会增加其出现的概率,而低奖励的回答会降低其概率。这种相对排名推动了优化过程。以下示例演示了如何对 Qwen2.5-0.5B 进行微调。脚本文件夹中提供了该方法的实现代码,你可以自定义它或更换基础模型。在这里,你将使用带有可验证奖励的 GRPO 和量化低秩适配(QLoRA)。QLoRA 是一种用于减少训练资源需求并加快训练过程的技术,它在准确性上有一些小的权衡。

模型参数

model_name_or_path: Qwen/Qwen2.5-0.5B tokenizer_name_or_path: Qwen/Qwen2.5-0.5B model_revision: main torch_dtype: bfloat16 attn_implementation: flash_attention_2 bf16: true tf32: true output_dir: /opt/ml/model/Qwen2.5-0.5B-RL-VR-GRPO

数据集参数

train_dataset_id_or_path: /opt/ml/input/data/train/dataset.json test_dataset_id_or_path: /opt/ml/input/data/val/dataset.json dataset_splits: 'train' max_seq_length: 2048 packing: true

code

# LoRA 参数
use_peft: true
load_in_4bit: true
lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
lora_modules_to_save: ["lm_head", "embed_tokens"] 
lora_r: 16
lora_alpha: 16

# 训练参数
num_train_epochs: 2
per_device_train_batch_size: 16
gradient_accumulation_steps: 2
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: True
learning_rate: 1.84e-4
lr_scheduler_type: cosine
warmup_ratio: 0.1

# 日志参数
logging_strategy: steps
logging_steps: 5
report_to:
- mlflow
save_strategy: "no"
seed: 42

代码

方案概述

本方案实现了用于数学推理任务的 Group Relative Policy Optimization (GRPO) 训练方法,并通过可验证奖励机制对 Qwen2.5-0.5B 模型进行微调。该方案采用双奖励系统,无需人工标注即可客观评估答案格式和数学正确性。

重要超参数:

  • learning_rate: 1.84e-4 – 为 GRPO 训练优化的学习率
  • num_train_epochs: 2 – 避免过拟合的训练轮数
  • per_device_train_batch_size: 16 且 gradient_accumulation_steps: 2 – 实际批量大小为 32
  • max_seq_length: 2048 – 支持 8-shot 提示的上下文窗口
  • lora_r: 16 和 lora_alpha: 16 – LoRA 秩和缩放参数
  • warmup_ratio: 0.1 配合余弦调度器 – 学习率调度方案
  • lora_target_modules – 针对注意力层和 MLP 层进行适配

接下来,您将使用 SageMaker AI 训练作业启动训练集群并运行模型微调。SageMaker AI Model Trainer ModelTrainer 在完全托管的基础设施上运行训练作业;处理环境设置、扩展和工件管理。它还允许您指定训练脚本、输入数据和计算资源,而无需手动配置服务器。库依赖项可以通过 scripts 文件夹中的 requirements.txt 文件管理。ModelTrainer 会自动检测此文件并在运行时安装列出的依赖项。

首先,设置您的环境。这里您需要指定训练使用的实例类型、实例数量以及训练容器的位置。

code
from sagemaker.core import image_uris
from sagemaker.core.helper.session_helper import Session
sagemaker_session = Session()

bucket_name = sagemaker_session.default_bucket()
default_prefix = sagemaker_session.default_bucket_prefix
configs = load_sagemaker_config()

instance_type = "ml.g6.48xlarge"
instance_count = 1
config_filename = "Qwen2.5-0.5B.yaml" 

image_uri = image_uris.retrieve(
    framework="pytorch",
    region=sagemaker_session.boto_session.region_name,
    version="2.7.1",
    instance_type=instance_type,
    image_scope="training"
)

Python

接下来,配置环境变量、代码位置和数据路径:

code
from sagemaker.train.configs import (
    CheckpointConfig,
    Compute,
    OutputDataConfig,
    SourceCode,
    StoppingCondition,
)
from sagemaker.train.distributed import Torchrun
from sagemaker.train.model_trainer import ModelTrainer

env = {}
env["FI_PROVIDER"] = "efa"
env["NCCL_PROTO"] = "simple"
env["NCCL_SOCKET_IFNAME"] = "eth0"
env["NCCL_IB_DISABLE"] = "1"
env["NCCL_DEBUG"] = "WARN"
env["HF_token"] = os.environ['hf_token']
env["CONFIG_PATH"] = f"recipes/{config_filename}"
env["MLFLOW_EXPERIMENT_NAME"]= "grpo-rlvr"
env["MLFLOW_TAGS"] =  '{"source.job": "sm-training-jobs", "source.type": "grpo-rlvr", "source.framework": "pytorch"}'
env["MLFLOW_TRACKING_URI"] =  MLFLOW_TRACKING_SERVER_ARN

# 定义要运行的脚本
source_code = SourceCode(
    source_dir="./scripts",
    requirements="requirements.txt",
    entry_script="run_finetuning.sh",
)

# 定义计算资源
compute_configs = Compute(
    instance_type=instance_type,
    instance_count=instance_count,
    keep_alive_period_in_seconds=3600,
)

# 定义训练作业名称
job_name = f"train-{config_filename.split('/')[-1].replace('.', '-').replace('yaml', 'rlvr')}"

# 定义输出数据配置路径
output_path = f"s3://{bucket_name}/{job_name}"
# 定义 ModelTrainer
model_trainer = ModelTrainer(
    training_image=image_uri,
    environment=env,
    source_code=source_code,
    base_job_name=job_name,
    compute=compute_configs,
	stopping_condition=StoppingCondition(max_runtime_in_seconds=18000),
    output_data_config=OutputDataConfig(s3_output_path=output_path),
    checkpoint_config=CheckpointConfig(
        s3_uri=output_path + "/checkpoint", local_path="/opt/ml/checkpoints"
    ),
)

Python

设置训练和验证数据的通道:

code
from sagemaker.train.configs import InputData

# 传递输入数据
train_input = InputData(
    channel_name="train",
    data_source=train_dataset_s3_path, # S3 路径存储训练数据
)

val_input = InputData(
    channel_name="val",
    data_source=val_dataset_s3_path, # S3 路径存储训练数据
)

# 检查输入通道配置
data = [train_input, val_input]

Python

然后开始训练:model_trainer.train(input_data_config=data) 以下是本示例的源代码目录结构:

code
scripts/
├── accelerate_configs/                       # Accelerate 配置文件
├── run_finetuning.sh      # 在 SageMaker 训练作业上使用 Accelerate 进行分布式训练的启动脚本
├── run_grpo.py               # GRPO 主训练脚本
├── utils/                   # 加载数据和创建提示的工具
├── recipes/                           # 预定义训练配置方案 (YAML)
└── requirements.txt                   # 运行时安装的 Python 依赖

Code

要跨多个 GPU 进行微调,示例训练脚本使用了 Huggingface Accelerate 和 DeepSpeed ZeRO-3,它们协同工作以更高效地训练大型模型。Huggingface Accelerate 通过自动处理设备分配、进程管理和混合精度设置来简化分布式训练的启动。DeepSpeed ZeRO-3 则通过在 GPU 之间划分优化器状态、梯度和参数来减少内存使用,从而使得十亿参数级别的模型能够被容纳并更快训练。你可以使用如下简单的命令,通过 Huggingface Accelerate 来运行你的 GRPO 训练器脚本:

code
NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
echo "Detected ${NUM_GPUS} GPUs on the machine"

# Launch fine-tuning with Accelerate + DeepSpeed (Zero3)
accelerate launch \
  --config_file accelerate_configs/deepspeed_zero3.yaml \
  --num_processes ${NUM_GPUS} \
  run_grpo.py \
  --config $CONFIG_PATH

结果

在 100 个测试样本上评估模型后,8-shot GRPO 训练的模型达到了 41% 的准确率,而基础模型仅为 11%,表明其在链式推理数学任务上实现了 3.7 倍的提升。

Image 2
Image 2

下图显示了一个与上下文长度相关的明显阈值,揭示了激活推理能力的最佳样本数量范围。尽管 0-shot(6%)和 2-shot(3%)配置表现不佳——甚至比基础模型还差——但在 4-shot 提示下性能显著提升至 33%,并在 8-shot 上下文时达到峰值 41%。这种非线性扩展模式表明,GRPO 训练创建的推理模式需要一定数量的样本来有效激活。模型似乎学会了从多个示例中利用组比较,这与 GRPO 的基于组的策略优化方法一致,即模型通过比较多个生成的解决方案来选择最佳推理路径。

Image 3
Image 3

将 RLVR 扩展到其他领域

尽管本文重点介绍了在 GSM8K 数据集上的数学推理,RLVR 方法可以推广到具有可客观验证输出的领域。以下两个方向展示了其多功能性:

基于执行奖励的代码生成

代码生成提供了通过执行进行自然验证的能力。当代码能够编译并运行无误时,可以给予部分奖励;而当输出通过全面的单元测试时则给予完整奖励。领域专家通过自然语言提示指定需求,而奖励模型则通过代码执行自动评估其正确性,从而避免了主观的人工评估。

基于语义验证的特定领域文本生成

对于医疗或技术写作等专业领域,关键词奖励可以引导模型使用适当的术语。部分奖励鼓励包含必要关键词,而完整奖励则要求在语义合适的上下文中包含完整的关键词集合。例如,在医疗文本生成中,可以通过将诊断关键词(“症状”、“诊断”)与治疗关键词(“疗法”、“药物”)以临床有效的模式组合来给予奖励,从而通过可衡量的目标教授领域词汇。这些例子说明,可验证奖励不仅可以用于数学推理,还可以扩展到可通过程序验证正确性的任务,为该训练方法的广泛应用奠定了基础。

清理资源

为了避免产生额外费用,请按照以下步骤清理资源:

  1. 删除任何未使用的 SageMaker Studio 资源
  2. 可选:删除 SageMaker Studio 域名
  3. 删除创建的 S3 存储桶
  4. 确认你的训练任务已停止运行!为此,请在 SageMaker 控制台中选择“训练”,然后检查“训练任务”。
Image 4
Image 4

要了解更多关于清理资源的信息,请参阅 Clean up

结论

在本示例中,你使用 GRPO(组相对策略优化)对 Qwen2.5-0.5B 模型进行了训练,训练数据集为 GSM8K:一个包含 8,500 道小学数学应用题的数据集,这些问题需要多步算术推理和自然语言理解。每个问题都像“_Janet 的鸭子每天产 16 个蛋…_”这样,并配有逐步解答,最终以数值答案结束,非常适合用于可验证奖励训练。

本实现展示了在数学推理任务中使用“可验证奖励强化学习”(RLVR)的有效性。经过 GRPO 训练的 Qwen2.5-0.5B 模型在 GSM8K 上的准确率达到 41%,相比基础模型的 11% 提升了 3.7 倍。评估结果验证了 RLVR 在具有客观可验证结果的领域中的潜力,为基于偏好的训练方法提供了替代方案。阈值行为表明 GRPO 能够从多个示例中学习组比较,这与其基于组的优化方法一致。这项工作为将可验证奖励系统应用于其他需要逻辑严谨性和数学准确性的领域奠定了基础。

有关 Amazon SageMaker AI 完全托管训练的更多信息,请参阅 SageMaker AI 文档的训练部分。本文相关支持代码可在 [GitHub](https://github.com/aws-samples/amazon-sagemaker-generativeai/tree/rl-vr/3_distributed_training/reinforcement-learning/grpo-with-verifiable-reward) 中找到。

  • * *

关于作者