GRPO自定义奖励函数代码实现

最新推荐文章于 2025-08-12 15:44:22 发布

原创最新推荐文章于 2025-08-12 15:44:22 发布 · 1.5k 阅读

24 ·

CC 4.0 BY-SA版权

文章标签：

#llama #语言模型

Deepseek目前的方法仍然很有热度，比如说仅仅使用规则作为奖励指向，不关注中间的推理过程，又或者是使用GRPO代替PPO或者DPO，两者都是化繁为简的方法。

文章目录

前言
1.安装环境
2.设置Unsloth
3.数据准备
4.定义奖励函数
5.设置GRPOTrainer
6.测试模型

前言

比较一下当前在大模型中使用强化学习的常见算法和流程。

名称	类型	核心思想	数据来源	优化目标	是否需要奖励函数	是否需要人类偏好
PPO	算法	稳定高效地优化策略（RL）	环境或奖励模型	最大化期望奖励	✅	❌
DPO	算法	直接基于偏好对优化策略（无需 reward）	偏好对（A 比 B 好）	让好样本概率更高	❌	✅
GRPO	算法	使用排序/打分的奖励进行 RL（类似 DPO 的泛化）	相对排序数据	最大化奖励差异引导的策略提升	✅	❌
RLHF	流程	用人类偏好训练奖励模型 + 用 RL 优化策略	人类偏好 → 奖励模型	奖励最大化	✅ （奖励来自人类偏好）	✅
RLAIF	流程	用 AI 模型生成偏好或奖励，替代人工	AI 偏好/打分	奖励或偏好最大化	✅/❌ 可选	❌ （用 AI 代替）

GRPO：比如说同一个prompt，大模型给出了多个回答，给定了两个奖励函数，一个是生成的回答长度越长奖励越高，还有一个是回答正确了再给奖励，最终取所有回答的平均值，超过平均值的接受。基于规则选择，打分，超过平均的接受

目前流行的SFT+RL的框架有trl、OpenRLHF、LLama-Factory、Unsloth，看下来在Unsloth中使用trl使用GRPO有比较成熟的例子，并且可以使用Lora微调。

1.安装环境

使用以下方式安装环境和库

conda create --name unsloth_env \
    python=3.11 \
    pytorch-cuda=12.1 \
    pytorch cudatoolkit -c pytorch -c nvidia \
    -y
conda activate unsloth_env
pip install unsloth vllm

在实际使用的过程中，发现版本问题可能会造成报错。因此，推荐Unsloth 2025.3.8, Unsloth=2025.3.7，vLLM: 0.7.3，transformers：4.48.3

2.设置Unsloth

Unsloth提供了FastLanguageModel类，集成了转换器和Unsloth优化。

from unsloth import FastLanguageModel
import torch

max_seq_length = 1024  # 增加推理长度
lora_rank = 32  # 更大的秩 = 更聪明, 但更慢

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="model/Meta-Llama-3.1-8B-Instruct",
    max_seq_length=max_seq_length,
    load_in_4bit=True,  # 如果是LoRA 16bit就改成False
    fast_inference=True,  # 使得vLLM更快推理
    max_lora_rank=lora_rank,
    gpu_memory_utilization=0.6,  # 减少内存不足
)

model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,  # 选择的数 > 0 ! 建议 8, 16, 32, 64, 128
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],  # 如果内存不足移除QKVO
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",  # 启用上下文调优
    random_state=3407,
)

3.数据准备

格式化字符串

# 定义指定的prompt
SYSTEM_PROMPT = """
Respond in the following format:
<think>
...
</think>
<answer>
...
</answer>
"""

# 回答格式
XML_COT_FORMAT = """\
<think>
{think}
</think>
<answer>
{answer}
</answer>
"""

Dataset里一定要有prompt列，其余按照需求。

prompt	answer
[{“role”: “system”, “content”: SYSTEM_PROMPT}, {“role”: “user”, “content”: ”What is the result of (1 + 2) * 4?"}]	[{“role”: “assistant”, “content”: “[think]The sum of 1 and 2 is 3, which we multiply by 4 to get 12.[/think][answer](1 + 2) * 4 = 12[/answer]”}]

prompt	ground_truth
Solve the equation $2 x + 3 = 7$ .	2

4.定义奖励函数

huggingface中有各种详细的介绍案例。
completions是输出的回答，不用在数据集中给出。

奖励长回答

def reward_func(completions, **kwargs):
    """奖励功能：完成时间越长，得分越高。"""
    return [float(len(completion)) for completion in completions]

奖励格式回答

import re

def format_reward_func(completions, **kwargs):
    """奖励功能：检查完成是否具有特定格式。"""
    pattern = r"^<think>.*?</think><answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    return [1.0 if match else 0.0 for match in matches]

奖励正确答案

import re

def reward_func(completions, ground_truth, **kwargs):
    # 用于捕获\box{}内内容的正则表达式
    matches = [re.search(r"\\boxed\{(.*?)\}", completion) for completion in completions]
    contents = [match.group(1) if match else "" for match in matches]
    # 如果内容与基本事实相同，奖励1，否则奖励0
    return [1.0 if c == gt else 0.0 for c, gt in zip(contents, ground_truth)]

多任务奖励

例子：答案的正确性奖励+遵循格式奖励
最终用reward_funcs=[math_reward_func, coding_reward_func]结合

from datasets import Dataset
from trl import GRPOTrainer

# 定义一个同时包含数学和编码问题的数据集
dataset = Dataset.from_list(
    [
        {"prompt": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "system", "content": What is the result of (1 + 2) * 4?"}], "ground_truth": 12},
       {"prompt": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "system", "content": What is the result of (3 + 1) * 2?"}], "ground_truth": 12},
    ]
)

# 当模型的答案与正确答案匹配时奖励模型
def correctness_reward_func(prompts, completions, ground_truth, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    q = prompts[0][-1]["content"]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    print(
        "-" * 20,
        f"Question:\n{q}",
        f"\nAnswer:\n{answer[0]}",
        f"\nResponse:\n{responses[0]}",
        f"\nExtracted:\n{extracted_responses[0]}",
    )
    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, ground_truth)]

# 奖励模型提供数字答案
def int_reward_func(completions, **kwargs) -> list[float]:
    responses = [completion[0]["content"] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]

# 奖励严格遵循指定格式的模型
def strict_format_reward_func(completions, **kwargs) -> list[float]:
    pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]
    
# 同时使用特定任务的奖励功能
trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=[math_reward_func, coding_reward_func],
    train_dataset=dataset,
)

5.设置GRPOTrainer

from trl import GRPOConfig, GRPOTrainer

max_prompt_length = 256

training_args = GRPOConfig(
    learning_rate=5e-6,
    adam_beta1=0.9,
    adam_beta2=0.99,
    weight_decay=0.1,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="paged_adamw_8bit",
    logging_steps=1,
    per_device_train_batch_size=6,
    gradient_accumulation_steps=1,  # 增加到4，训练更顺畅
    num_generations=6,  # 内存不足时减少
    max_prompt_length=max_prompt_length,
    max_completion_length=max_seq_length - max_prompt_length,
    # num_train_epochs = 1, # 设置为1为完整的训练运行
    max_steps=250,
    save_steps=250,
    max_grad_norm=0.1,
    report_to="none",  # 可以使用 Weights & Biases
    output_dir="outputs",
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func,
    ],
    args=training_args,
    train_dataset=dataset,
)

6.测试模型

保存lora模型

model.save_lora("grpo_saved_lora")

用一个新问题来测试

from vllm import SamplingParams

text = tokenizer.apply_chat_template(
    [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": "Calculate pi."},
    ],
    tokenize=False,
    add_generation_prompt=True,
)

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1024,
)
output = (
    model.fast_generate(
        text,
        sampling_params=sampling_params,
        lora_request=model.load_lora("grpo_saved_lora"),
    )[0]
    .outputs[0]
    .text
)

print(output)