GRPO 与 GSM8K 训练：代码解析与社区讨论综述

近年来，针对大语言模型如何通过强化学习（RL）进行微调，从而提高回答准确率和格式一致性，成为了社区关注的热点。本文将结合 GRPO Llama-1B 的示例代码 以及讨论社区中的交流，对代码整体结构、奖励函数设计、训练流程以及实际应用中的问题与优化策略进行详细梳理和解析。

1. 背景介绍

在该示例代码中，作者使用了 GRPO（可理解为一种基于奖励优化（Reward Policy Optimization）的算法）对 GSM8K 数据集进行强化学习调优。目标是让模型输出符合预定格式的回答，同时确保答案的正确性。本次调优主要依赖于 Hugging Face 的 trl 库，并使用了 PEFT（参数高效微调）和 Flash-Attention 等优化组件来加速模型训练。

2. 数据预处理与提示模板

2.1 系统提示与答案格式

代码中定义了一个固定的系统提示（SYSTEM_PROMPT）：

“Respond in the following format:
<reasoning> ... </reasoning>
<answer> ... </answer>”

这一提示起到了规范输出格式的作用，要求模型在回答时输出推理过程和最后答案，并分别用 XML 风格的标签 <reasoning> 与 <answer> 包裹。除了系统提示外，还定义了一个模板 XML_COT_FORMAT，用于格式化链式思考（Chain-of-Thought，CoT），该模板将推理过程与答案嵌入到预定格式中。

2.2 数据集构造

数据集使用的是 OpenAI 的 GSM8K——一个数学问题数据集。函数 get_gsm8k_questions 加载了数据集的指定 split（如“train”或“test”），并将每个数据样本处理成一个字典格式，其中包括：

• prompt：一个对话列表，最开始是系统角色提供的格式提示，然后由用户角色给出具体问题。
• answer：通过辅助函数 extract_hash_answer 将原答案文本（多以 “####” 标记）提取出来。

这种数据预处理方式不仅保证了数据的一致性，同时也为后续强化学习过程中使用奖励函数评估输出提供了标准化输入。

3. 奖励函数设计

在 RL 调优中，奖励函数（Reward Function）起到了至关重要的作用。该代码中设计了多种奖励函数来衡量模型输出与预期结果之间的差距，具体包括：

3.1 正确性奖励

• correctness_reward_func
  该函数首先提取用户问题最后的 prompt，然后解析模型输出中的 <answer> 部分，并与正确答案（经过预处理后的结果）进行严格对比。如果答案完全一致，则给予较高奖励（例如 2.0 分），否则返回 0 分。
  此外，函数会打印调试信息，展示问题、目标答案、模型完整输出和从中提取的答案，便于检查可能产生的问题。

3.2 格式奖励

为了鼓励模型输出符合预期的 XML 格式（即同时包含 <reasoning> 与 <answer> 标签），设计了两种奖励函数：

• strict_format_reward_func
  使用严格的正则表达式（regex）检测输出内容是否精准匹配预定格式（包括换行符等细节），匹配成功给予固定奖励（例如 0.5 分）。

• soft_format_reward_func
  使用相对宽松的正则表达式，允许一定的变动性，但依然要求输出包含 <reasoning> 和 <answer> 标签。匹配成功同样给予奖励。

3.3 数字格式奖励

• int_reward_func
  针对 GSM8K 数据集中答案本质上为正整数，该函数检查提取出来的答案是否为数字，若是则给予一定奖励（例如 0.5 分），否则为 0。

3.4 XML 内部结构奖励

• xmlcount_reward_func
  函数内部通过 count_xml 对输出文本进行简单计数：
  • 检查是否恰好出现 <reasoning>\n、\n</reasoning>\n、\n<answer>\n 和 \n</answer>\n 标记，并相应地加分；
  • 对于额外的或者不合理的分段，则进行轻微扣分。

这种设计鼓励模型在回答时严格遵守预定格式，但同时也允许一定的灵活性。

4. 模型加载与训练配置

4.1 模型加载

代码中根据模型名称（例如是否包含 “Llama” 来区别使用模型文件夹）确定输出路径，并通过 Hugging Face 的 AutoModelForCausalLM 与相关配置加载模型。本示例中默认使用的是 Qwen 系列模型，如 "Qwen/Qwen2.5-1.5B-Instruct"，同时设置模型的量化类型（torch.bfloat16）并采用 Flash-Attention 来加速计算。

4.2 Tokenizer 配置

加载 Tokenizer 后，将其 pad_token 设为模型的 eos_token，保证输入长度对齐，适应生成任务要求。

4.3 训练器与 GRPO 配置

利用 GRPOConfig 对训练超参数进行设定，例如：

• 学习率（例如 5e-6）
• 权重衰减、warmup 比例、梯度积累步数
• 每个 batch 的样本数、每次生成的样本数（num_generations）等
  此外，还设置了多种奖励函数，将它们传递给 GRPOTrainer。这样在训练过程中，模型不仅通过语言建模损失进行更新，还会依据奖励函数反馈进一步优化输出，使其同时满足正确性与格式要求。

4.4 可选的 PEFT 配置

代码中使用的 LoRA（低秩适配）配置（当前被注释掉）则为减少微调过程中模型参数更新量提供了可能。这部分配置在大规模多 GPU 训练场景下曾经引起一些兼容性问题，相关讨论在社区中也不断被提及。

5. 社区讨论与优化探索

在代码的 GitHub 讨论区，社区成员就以下几个问题展开了热烈讨论：

5.1 同一代码在不同硬件平台上的差异

• H100 vs. 7900XTX
  部分评论提到，同一份代码在 H100 与 7900XTX 上运行时，生成结果存在差异。从进一步对比 pip freeze 中的 TRL 代码版本（不同 commit，如 1c35… 与 2ce36…），可以看出库内更新可能引入了细微差异。社区成员也讨论了如何锁定特定版本并调试潜在的兼容性问题。

5.2 奖励函数参数与超参数调优

• 提到类似 “beta” 的参数（实际上指的是 KL 系数）对训练效果的影响。通过调整诸如学习率、warmup_ratio 以及生成样本数等超参数，部分团队在 GSM8K 上的准确率得到了显著提升（例如从 41.6% 到 51% 左右）。
• 有成员尝试通过 Optuna 自动化调整整个奖励函数的权重（例如修改 correct_reward、strict_match_reward 等），使得训练过程免受手动调节困扰，从而进一步挖掘出最佳组合。这种方法为奖励设计提供了灵活性，并帮助找到更优的超参数组合。

5.3 格式化输出问题

• 社区中有讨论为什么模型有时会输出诸如 “(\boxed{ANSWE})” 的字符，这种内容并非代码中设计奖励时关注的格式部分。这提示开发者在奖励设计时，需要额外考虑如何屏蔽与奖励无关的符号，确保奖励反馈更准确。

5.4 vLLM 的应用场景讨论

• 部分用户分享了如何在单 GPU 或多 GPU 环境下使用 vLLM 加速生成，有助于快速评估模型在推理端的输出质量。利用 vLLM 可以节省时间，使得周期性评估（例如使用 TrainerCallback）成为可能。

6. 总结与展望

本文通过对 GRPO 调优代码的详细解析，我们可以看出：

• 使用预定义格式化提示和严格的奖励函数设计，有助于引导模型在回答问题时做到既准确又格式统一。
• 多种奖励函数的组合（正确性、格式、数字检查以及 XML 内部结构计数）可以针对具体任务（如 GSM8K 数学问答）提供多维度优化。
• 社区在相同代码基础上，通过不同硬件环境的测试、微调超参数以及引入 Optuna 自动调参，均在探索如何进一步提高模型性能。这些讨论为同类项目提供了宝贵的实践经验和优化思路。

未来，随着社区在 RL 调优和奖励设计上的不断探索，或许能发现更自动化、鲁棒性更强的方法来统一和提升大语言模型的输出质量。开源社区的力量也将在这样的试验和探索中发挥越来越重要的作用。

下面给出对grpo_demo.py文件的详细解读，帮助你理解每个部分的作用。

1. 文件总体作用

该文件主要用于使用 GRPO（Graded Reward Policy Optimization）算法对因果语言模型（Causal LM）进行强化学习式微调。目标是让模型在回答问题时满足指定格式（例如，通过 XML 标签将推理过程和答案分隔开来），同时通过一系列奖励函数对生成结果进行打分，从而促使模型输出既符合格式要求又能正确回答问题的内容。

2. 模块导入与常量定义

首先导入了所需的库和模块：

• 正则表达式模块 (re)：用于验证输出格式。
• PyTorch (torch)：支持 GPU 加速以及张量计算。
• Datasets 库 (load_dataset, Dataset)：用于加载 GSM8k 数据集。
• Transformers 模块 (AutoTokenizer, AutoModelForCausalLM)：用于加载预训练模型和分词器。
• PEFT 模块 (LoraConfig)：支持 LoRA（低秩适应）微调方法。
• TRL（Transformer Reinforcement Learning）模块 (GRPOConfig, GRPOTrainer)：用于强化学习策略微调。

此外，还定义了两个常量：

• SYSTEM_PROMPT
  一个系统提示字符串，明确要求模型按照特定 XML 格式（包含 <reasoning> 和 <answer> 标签）来回复答案。

• XML_COT_FORMAT
  一个格式化模板，用于将推理过程和答案包装成 XML 样式输出。

3. 数据处理函数

3.1. 提取 XML 格式答案

def extract_xml_answer(text: str) -> str:
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()

• 作用：从模型的输出文本中提取 <answer> 标签内部的答案部分，该函数假定输出文本严格包含 <answer> 和 </answer> 标记。

3.2. 提取 Hash 标记的答案

def extract_hash_answer(text: str) -> str | None:
    if "####" not in text:
        return None
    return text.split("####")[1].strip()

• 作用：用于从数据集中提取答案。部分数据可能使用 “####” 将答案包裹起来，函数提取 “####” 后面的内容，如果不存在 “####” 则返回 None。

3.3. GSM8k 数据集加载与预处理

def get_gsm8k_questions(split = "train") -> Dataset:
    data = load_dataset('openai/gsm8k', 'main')[split] # type: ignore
    data = data.map(lambda x: {  # type: ignore
        'prompt': [
            {'role': 'system', 'content': SYSTEM_PROMPT},
            #{'role': 'user', 'content': 'What is the largest single-digit prime number?'},
            #{'role': 'assistant', 'content': XML_COT_FORMAT.format(
            #    reasoning="9 is divisble by 3 and 8 is divisible by 2, but 7 is prime.",
            #    answer="7"
            #)},
            {'role': 'user', 'content': x['question']}
        ],
        'answer': extract_hash_answer(x['answer'])
    }) # type: ignore
    return data # type: ignore

• 作用：

  1. 加载 gsm8k 数据集（主要为数学推理题目）。
  2. 使用 map 方法遍历数据，将每个样本预处理成一个包含 prompt 和 answer 的字典：
     • prompt 为一个消息列表，其中包含系统角色（给出回答格式要求）和用户角色（问题内容）。
     • answer 则利用 extract_hash_answer 将原始答案进行抽取。

• 备注：部分代码被注释掉，注释内展示了如何使用 1-shot 提示（通过提供一个示范问答）来引导模型，但目前默认仅使用单一用户问题。

4. 奖励函数设计

为了引导模型产生符合要求的回答，文件中定义了多个奖励函数，每个函数根据输出的不同特性给予奖励分数。

4.1. 正确性奖励函数

def correctness_reward_func(prompts, completions, answer, **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, answer)]

• 作用：
  对生成的答案进行正确性检验：
  • 首先从生成的输出中提取 XML 格式内的答案。
  • 将提取的答案与标准答案进行比较，相同则奖励 2.0 分，不同则奖励 0 分。
  • 同时打印调试信息，便于查看问题、标准答案和模型生成的完整输出细节。

4.2. 数字格式奖励函数

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]

• 作用：
  检查模型输出的答案是否完全由数字组成。如果是，则奖励 0.5 分；否则不奖励。

4.3. 严格格式奖励函数

def strict_format_reward_func(completions, **kwargs) -> list[float]:
    """Reward function that checks if the completion has a specific format."""
    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]

• 作用：
  利用正则表达式严格校验生成文本的格式，要求：

  • 文本以 <reasoning> 开始，紧跟换行和非贪婪模式文本，再到 </reasoning>，

  • 接着 <answer> 标签及相应内容，最后以 </answer> 结束。

    
    

    若完全匹配指定格式则奖励 0.5 分，否则奖励 0 分。

4.4. 宽松格式奖励函数

def soft_format_reward_func(completions, **kwargs) -> list[float]:
    """Reward function that checks if the completion has a specific format."""
    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
    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]

• 作用：
  与严格格式不同，这里使用一个宽松一点的正则表达式，只要包含 <reasoning> 和 <answer> 标签及内容即可，奖励同样为 0.5 分。

4.5. 基于 XML 标签计数的奖励函数

def count_xml(text) -> float:
    count = 0.0
    if text.count("<reasoning>\n") == 1:
        count += 0.125
    if text.count("\n</reasoning>\n") == 1:
        count += 0.125
    if text.count("\n<answer>\n") == 1:
        count += 0.125
        count -= len(text.split("\n</answer>\n")[-1])*0.001
    if text.count("\n</answer>") == 1:
        count += 0.125
        count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001
    return count

• 作用：
  以更细粒度的方式检查模型输出中的 XML 标签是否正确：
  • 对每个关键标签（<reasoning>\n、\n</reasoning>\n、\n<answer>\n、\n</answer>）各自计分，保证恰好出现一次。
  • 同时如果在答案部分存在额外字符，则按字符长度扣分，用以惩罚格式不严谨的回答。

随后函数 xmlcount_reward_func 会对完整生成结果应用 count_xml 得分：

def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

5. 模型、训练配置及训练过程

5.1. 模型名称和输出目录配置

#model_name = "meta-llama/Llama-3.2-1B-Instruct"
model_name = "Qwen/Qwen2.5-1.5B-Instruct"

if "Llama" in model_name:
    output_dir = "outputs/Llama-1B-GRPO"
    run_name = "Llama-1B-GRPO-gsm8k"
else:
    output_dir="outputs/Qwen-1.5B-GRPO"
    run_name="Qwen-1.5B-GRPO-gsm8k"

• 说明：
  文件中可以选择不同的预训练模型（例如 Llama 或 Qwen）。当前选用的是 Qwen 模型，因此训练输出目录和运行名称均以 “Qwen” 命名。

5.2. 训练参数配置

training_args = GRPOConfig(
    output_dir=output_dir,
    run_name=run_name,
    learning_rate=5e-6,
    adam_beta1 = 0.9,
    adam_beta2 = 0.99,
    weight_decay = 0.1,
    warmup_ratio = 0.1,
    lr_scheduler_type='cosine',
    logging_steps=1,
    bf16=True,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_generations=16,
    max_prompt_length=256,
    max_completion_length=786,
    num_train_epochs=1,
    save_steps=100,
    max_grad_norm=0.1,
    report_to="wandb",
    log_on_each_node=False,
)

• 作用：
  利用 GRPOConfig 定义训练过程中的超参数，包括：
  • 学习率、Adam 优化器的 beta 参数、权重衰减；
  • warmup比例和 cosine 学习率调度器；
  • 混合精度（bf16）；
  • 批量大小、梯度累积步数；
  • 每个 prompt 生成回答时的生成样本数（num_generations：16）；
  • 最大 prompt 与回答长度；
  • 训练轮数、保存步数以及梯度裁剪等配置；
  • 日志输出（例如报告到 wandb）。

5.3. LoRA 参数配置（目前未启用）

peft_config = LoraConfig(
    r=16,
    lora_alpha=64,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"],
    task_type="CAUSAL_LM",
    lora_dropout=0.05,
)

• 说明：
  如果使用 PEFT 进行低秩适应微调，可以通过此配置定义：
  • 低秩维度 r、比例系数 lora_alpha；
  • 针对模型中的哪些模块（如查询、键、值投影等）进行 LoRA 应用；
  • 指定任务类型和 dropout 概率。
    需要注意的是，在后续初始化 GRPOTrainer 时，此配置被注释掉，说明当前没有启用 PEFT。

5.4. 模型和 Tokenizer 加载

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    device_map=None
).to("cuda")
        
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

• 作用：
  • 使用 AutoModelForCausalLM 从 Hugging Face Hub 加载预训练模型，并指定使用 bfloat16 数据类型以及 flash attention 优化（若支持）。
  • 加载对应的分词器并将 pad_token 设置为 eos_token，防止模型输入时因缺少 pad_token 而报错。

5.5. GRPOTrainer 实例化与训练

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func],
    args=training_args,
    train_dataset=dataset,
    #peft_config=peft_config
)
trainer.train()

• 作用：
  • 实例化 GRPOTrainer 时传入：
  • 模型 与 分词器（此处分词器作为 processing_class，负责数据预处理）。
  • 奖励函数列表：包括对格式（XML 标签）和答案正确性等多维度的奖励函数，组成混合奖励信号。
  • 训练参数（training_args）与预处理后的 训练数据集（dataset）。
  • 可选的 peft_config（当前被注释掉）。
  • 通过调用 trainer.train() 启动训练过程。

总结

该文件通过以下几个步骤实现了强化学习式的微调流程：

1. 数据准备：
   从 GSM8k 数据集中提取问题及答案，同时添加系统提示以强制模型输出带有 <reasoning> 与 <answer> 标签的格式。

2. 奖励函数设计：
   定义多个奖励函数，从格式正确性、答案是否为整数字、XML 标签的完整性及答案正确性等多个角度评分，帮助模型逐步调整输出格式与内容。

3. 训练配置：
   基于 GRPOConfig 指定训练中的所有超参数，并加载预训练模型与分词器，保证训练过程在 GPU 上高效运行。

4. 训练过程：
   使用 GRPOTrainer 将模型、数据、奖励函数、训练参数等组合在一起，开始强化学习式微调，从而使模型在回答问题时既满足特定格式，又尽可能给出正确答案。

通过这些步骤，文件展示了一种基于奖励信号引导预训练语言模型生成结构化、规范化回答的方案，并且可以通过调整奖励函数和超参数来个性化训练效果。