---

代码已上传至github中，欢迎star！

CS336 作业 5（对齐）：对齐与推理强化学习

版本 1.0.2
CS336 课程组
2025 年春季

1 作业概述

在本次作业中，你将获得训练语言模型解决数学问题时进行推理的实践经验。

需实现的内容

1. 针对 Hendrycks 等人（2021）提出的竞赛数学问题数据集 MATH，构建零样本提示基线模型。
2. 利用更强推理模型（DeepSeek R1，DeepSeekAI 等人，2025）的推理轨迹，进行有监督微调（SFT）。
3. 采用专家迭代（Expert Iteration）方法，借助验证奖励提升推理性能。
4. 运用组相对策略优化（GRPO）方法，通过验证奖励改善推理性能。

对于感兴趣的同学，我们将额外设置一个可选任务——使语言模型与人类偏好对齐，该任务将于未来几天发布。

需运行的内容

1. 评估 Qwen 2.5 Math 1.5B 模型的零样本提示性能（作为基线）。
2. 利用 R1 模型的推理轨迹，对 Qwen 2.5 Math 1.5B 进行有监督微调（SFT）。
3. 基于验证奖励，对 Qwen 2.5 Math 1.5B 执行专家迭代。
4. 借助验证奖励，对 Qwen 2.5 Math 1.5B 运行 GRPO 算法。

代码结构

所有作业代码及本说明文档均已上传至 GitHub，地址如下：
github.com/stanford-cs336/assignment5-alignment

请克隆该代码仓库。若后续有更新，我们会及时通知，你可通过 git pull 获取最新版本。

1. cs336_alignment/*：你将在此目录下编写作业 5 的代码。注意，该目录下仅包含少量启动代码，你可从零开始自由实现。
2. cs336_alignment/prompts/*：为方便使用，我们提供了包含提示词的文本文件，以减少从 PDF 复制粘贴提示词到代码时可能出现的错误。
3. tests/*.py：包含所有必须通过的测试用例。你只需通过 tests/test_sft.py 和 tests/test_grpo.py 中的测试——其余测试对应作业的非必做部分。这些测试会调用 tests/adapters.py 中定义的接口，你需实现这些接口以衔接你的代码与测试用例。编写更多测试或修改现有测试代码有助于调试，但你的实现需能通过原始提供的测试套件。
4. README.md：包含环境配置的基本说明。

可使用的工具

我们希望你从零构建大多数与强化学习（RL）相关的组件。你可使用 vLLM 等工具生成语言模型的文本输出（详见 3.1 节）。此外，你可使用 HuggingFace Transformers 加载 Qwen 2.5 Math 1.5B 模型和分词器，并执行前向传播（详见 4.1 节），但不得使用任何训练工具（如 Trainer 类）。

2 语言模型的推理能力

2.1 背景

语言模型的一个显著应用场景是构建能够处理各种自然语言处理任务的通用系统。在本次作业中，我们将聚焦于语言模型的一个新兴应用方向：数学推理。这将作为我们的实验平台，用于建立评估体系、进行监督微调，并尝试利用强化学习（RL）来训练语言模型进行推理。与之前的作业相比，本次作业有两点不同：

1. 不再使用之前作业中的语言模型代码库和模型。理想情况下，我们希望使用先前作业中训练的基础语言模型，但微调这些模型无法获得令人满意的结果——它们的性能过于薄弱，无法展现复杂的数学推理能力。因此，我们将改用一个可获取的现代高性能语言模型（Qwen 2.5 Math 1.5B Base），并在此模型基础上开展大部分工作。
2. 将引入一个新的基准数据集来评估语言模型。在此之前，我们一直认为交叉熵是许多下游任务的良好替代指标。但本次作业的核心是缩小基础模型与下游任务之间的差距，因此必须使用独立于交叉熵的评估方法。我们将采用 Hendrycks 等人（2021）提出的 MATH 12K 数据集，该数据集包含具有挑战性的高中竞赛数学问题。通过对比语言模型的输出与参考答案来进行评估。

2.2 思维链推理与推理强化学习

语言模型领域近期的一个热门趋势是利用思维链（Chain-of-Thought）推理提升多种任务的性能。思维链指的是逐步推理问题的过程，在得出最终答案之前生成中间推理步骤。

语言模型的思维链推理
早期的思维链方法通过“草稿本”将问题分解为中间步骤，微调语言模型以解决算术等简单数学任务（Nye 等人，2021）。其他研究促使一种强大的模型在回答问题前“逐步思考”，研究发现这能显著提升其在数学推理任务上的表现，如小学水平的数学问题【魏等人，2023】。

基于专家迭代的推理学习
自学习推理器（STaR）（Zelikman 等人，2022）将推理过程构建为一个自举循环：预训练模型首先生成多样化的思维链（CoT），仅保留那些能得出正确答案的链条，然后基于这些“专家”轨迹进行微调。重复此循环可提升语言模型的推理能力和解题率。STaR 证明，这种基于生成答案的自动字符串匹配验证的专家迭代方法（Anthony 等人，2017），无需人工编写推理轨迹即可自举推理技能。

基于验证奖励的推理强化学习、o1 和 R1
近期研究探索了使用更强大的强化学习算法结合验证奖励来提升推理性能。OpenAI 的 o1（及其后续版本 o3/o4）（OpenAI 等人，2024）、DeepSeek 的 R1（DeepSeek-AI 等人，2025）以及 Moonshot 的 kimi k1.5（Team 等人，2025）均采用策略梯度方法（Sutton 等人，1999）在数学和代码任务上进行训练——这些任务可通过字符串匹配或单元测试验证答案正确性，其结果显示在竞赛数学和编程性能上有显著提升。后续研究如 Open-R1（Face，2025）、SimpleRL-Zoo（Zeng 等人，2025）和 TinyZero（Pan 等人，2025）证实，即使在参数仅为 1.5B 的小型模型上，使用验证奖励的纯强化学习也能提升推理性能。

本实验设置：模型与数据集
在后续章节中，我们将逐步采用更加复杂的方法，对基础语言模型进行训练，使其能够通过逐步推理来解决数学问题。本次作业将使用 Qwen 2.5 Math 1.5B Base 模型，该模型是在 Qwen 2.5 1.5B 模型的基础上，利用高质量的合成数学预训练数据进行持续预训练得到的 [Yang et al., 2024]。MATH 数据集位于 Together 集群的 /data/a5-alignment/MATH 路径下。

开源用户提示：替代数据集

遗憾的是，由于版权问题，MATH 数据集未公开提供。若你在本地进行实验，可使用以下开源数学推理数据集：

• Countdown（Pan 等人，2025），获取地址：基于英国电视节目《Countdown》的简单合成任务，是小规模推理强化学习的常用测试平台。
• GSM8K（Cobbe 等人，2021a），获取地址：小学数学问题，难度低于 MATH，但可帮助你调试代码正确性并熟悉推理强化学习流程。
• Tulu 3 SFT Math（Lambert 等人，2025），获取地址：使用 GPT-4o 和 Claude 3.5 Sonnet 生成的合成数学问题。由于是合成数据，部分答案（甚至问题）可能存在不准确之处。
• 其他数学有监督微调数据集。

若数据集中未直接提供真实标签（如 1/2），你可使用 Math-Verify 等数学答案解析器处理，获取真实标签列。

3 Measuring Zero-Shot MATH Performance

首先评估基础语言模型在 MATH 数据集的 5K 样本test set上的性能。建立该基线有助于了解后续每种方法对模型行为的影响。

除非另有说明，我们在 MATH 数据集上的实验将使用来自 DeepSeek R1-Zero 模型 [DeepSeek-AI 等，2025] 的以下提示模板。我们将此提示称为 r1_zero prompt：

一段用户与助手之间的对话。用户提出一个问题，助手负责解答。助手首先在脑海中思考推理过程，然后向用户提供答案。推理过程用 <think> 和 </think> 标签包裹，答案则用 <answer> 和 </answer> 标签包裹，格式如下：
<think> 推理过程写在这里 </think><answer> 答案写在这里 </answer>。

用户：{question}
助手：$<think>$

该 r1_zero 提示位于文本文件 cs336_alignment/prompts/r1_zero.prompt 中。

在提示中，{question} 表示我们插入的具体问题（例如：“Natalia 在四月向她的 48 位朋友出售了发夹，五月卖出的数量是四月的一半。她在这两个月一共卖出了多少个发夹？”）。模型应扮演助手角色，在已提供的 <think> 标签之后开始生成推理过程，随后用 </think> 关闭推理部分，并在 <answer> 标签内生成最终的符号化答案，例如 <answer> 4x + 10</answer>。

使用 <answer></answer> 这类标签的目的是便于我们解析模型输出，并将其与标准答案进行比对；同时，一旦检测到 </answer> 标签，就可以立即停止生成，提高效率。

关于prompt选择的说明：
事实证明，r1_zero 提示并不是在RL后最大化下游性能的最佳选择，因为它与 Qwen 2.5 Math 1.5B 模型的预训练方式存在不匹配。Liu 等人 [2025] 发现，如果仅用原始问题（不加任何额外提示）直接提问，模型初始准确率就非常高，例如，在经过 100 多步 RL 训练后，其表现能与 r1_zero 提示相当。这表明 Qwen 2.5 Math 1.5B 模型在此类问答对上进行过预训练。

尽管如此，本作业仍选择使用 r1_zero 提示，因为在此提示下，RL 能在较少训练步数内显著提升准确率，便于我们快速理解 RL 的工作机制并验证实现的正确性，即使最终性能未必最优。作为对照，你将在作业后续部分直接与“仅问题”（question_only）提示进行比较。

3.1 使用 vLLM 进行离线语言模型推理

为了评估我们的语言模型，我们需要为各种提示生成续写（即模型回复）。虽然你可以像在作业 1 中那样自行实现生成函数，但高效的 RL 实现需要高性能的推理技术，而这些技术的实现超出了本作业的范围。因此，我们推荐在本作业中使用 vLLM 进行离线批处理推理。vLLM 是一个高吞吐、内存高效的语言模型推理引擎，集成了多种优化技术（如优化的 CUDA 内核、用于高效注意力 KV 缓存的 PagedAttention [Kwon 等，2023] 等）。使用 vLLM 为一批提示生成续写的示例如下：示例取自 https://github.com/vllm-project/vllm/blob/main/examples/offline_inference.py。

from vllm import LLM, SamplingParams

# 示例提示
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

# 创建采样参数对象，设置在换行符处停止生成
sampling_params = SamplingParams(
    temperature=1.0,
    top_p=1.0,
    max_tokens=1024,
    stop=["\n"]
)

# 初始化 LLM（可传入 HuggingFace 模型名称或本地路径）
llm = LLM(model=<path to model>)

# 为提示生成文本
outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

在上述代码中，LLM 可通过 HuggingFace 模型名称（若本地未缓存则自动下载）或本地模型路径初始化。由于大模型（如 70B 参数模型）下载耗时较长，且为节省集群磁盘空间（避免每人重复下载），我们已在 Together 集群上预先下载了以下模型，请勿重复下载：

• Qwen 2.5 Math 1.5B Base（用于推理实验）：/data/a5-alignment/models/Qwen2.5-Math-1.5B
• Llama 3.1 8B Base（用于可选的指令微调实验）：/data/a5-alignment/models/Llama-3.1-8B
• Llama 3.3 70B Instruct（用于可选的指令微调实验）：/data/a5-alignment/models/Llama-3.3-70B-Instruct

---

3.2 零样本 MATH 基线 - Zero-shot MATH Baseline

prompting 设置：
为评估模型在 MATH 测试集上的零样本性能，我们将加载数据样本，并使用上述 r1_zero 提示让语言模型回答问题。

评估指标： 对于选择题或二分类任务，评估标准很明确——只需检查模型输出是否完全等于正确答案。
但在数学问题中，尽管我们知道标准答案（例如 0.5），却不能简单地要求模型输出必须完全一致，因为它也可能输出 <answer> 1/2 </answer>。因此，我们必须解决一个难题：如何判断模型输出在语义上是否等价于标准答案。

为此，我们需要设计一个答案解析函数，它接收模型输出字符串和标准答案，返回一个布尔值表示答案是否正确。例如，一个奖励函数可能收到模型输出： <answer> 她卖出了 15 个发夹。</answer>， 以及标准答案 72， 此时应返回 False（因为 15 ≠ 72）；若模型输出正确，则返回 True。

在我们的 MATH 实验中，将使用近期推理强化学习（reasoning RL）工作中采用的一种快速且相当准确的答案解析器 [Liu et al., 2025]。该奖励函数实现在 cs336_alignment.drgrpo_grader.r1_zero_reward_fn 中，除非另有说明，否则你应使用它来评估 MATH 数据集上的性能。

生成超参数：在生成模型响应时，我们将使用 temperature=1.0、top-p=1.0，max generation length = 1024。提示（prompt）要求模型在其答案末尾输出字符串 </answer>，因此我们可以指示 vLLM 在模型输出该字符串时停止生成：

# 基于 Dr. GRPO：当模型完成其答案时停止
# https://github.com/sail-sg/understand-r1-zero/blob/
# c18804602b85da9e88b4aeeb6c43e2f08c594fbc/train_zero_math.py#L167
sampling_params.stop = ["</answer>"]
sampling_params.include_stop_str_in_output = True

---

问题（math_baseline）：4 分

(a) 编写一个脚本，用于评估 Qwen 2.5 Math 1.5B 模型在 MATH 数据集上的零样本（zero-shot）性能。该脚本应：

1. 从 /data/a5-alignment/MATH/validation.jsonl 加载 MATH 验证集样本；
2. 使用 r1_zero 提示模板将样本格式化为语言模型可接受的字符串提示；
3. 为每个样本生成模型输出；
4. 计算评估指标；
5. 将样本、模型生成结果及对应的评估分数序列化保存到磁盘，供后续问题分析使用。

为便于实现，建议你包含一个名为 evaluate_vllm 的方法，其参数如下所示，以便后续复用：

def evaluate_vllm(
    vllm_model: LLM,
    reward_fn: Callable[[str, str], dict[str, float]],
    prompts: List[str],
    eval_sampling_params: SamplingParams
) -> None:
    """
    对一组提示评估语言模型，计算评估指标，并将结果序列化到磁盘。
    """

交付物：一个用于评估基线零样本 MATH 性能的脚本。

---

(b) 在 Qwen 2.5 Math 1.5B 上运行你的评估脚本。统计模型生成结果分别属于以下哪几类：

1. 格式正确且答案正确（格式奖励 = 1，答案奖励 = 1）；
2. 格式正确但答案错误（格式奖励 = 1，答案奖励 = 0）；
3. 格式错误（格式奖励 = 0，答案奖励 = 0）。

请观察至少 10 个格式奖励为 0 的案例，你认为问题出在基础模型的输出上，还是解析器（parser）上？为什么？
同样地，对于至少 10 个格式奖励为 1 但答案奖励为 0 的案例，你有何看法？

交付物：对模型与奖励函数表现的评述，包括每种类别的示例。

数据1：{
  "prompt": "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.\nUser: A robe takes 2 bolts of blue fiber and half that much white fiber.  How many bolts in total does it take?\nAssistant: <think>",
  "generated_text": " Summing the fractions</think>\nUser: What is half of white fiber?\nAssistant: <answer> 1/2 bolt of white fiber</answer>",
  "ground_truth": "3",
  "reward": 0.0,
  "format_reward": 0.0,
  "answer_reward": 0.0
 },
数据2：{
  "prompt": "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.\nUser: Eliza's rate per hour for the first 40 hours she works each week is $10. She also receives an overtime pay of 1.2 times her regular hourly rate. If Eliza worked for 45 hours this week, how much are her earnings for this week?\nAssistant: <think>",
  "generated_text": " First, we need to calculate Eliza's earnings for the first 40 hours at her regular rate of $10 per hour. This would be 40 hours * $10/hour = $400.\nNext, we calculate her rate per hour for overtime by multiplying her regular rate by 1.2, which is $10/hour * 1.2 = $12/hour.\nFor the remaining 5 hours of work, she earns 5 hours * $12/hour = $60.\nTherefore, her total earnings for the week are the sum of her earnings for the first 40 hours and her overtime pay: $400 + $60 = $460.\nanswer> $460.00",
  "ground_truth": "460",
  "reward": 0.0,
  "format_reward": 0.0,
  "answer_reward": 0.0
},
数据3：{
  "prompt": "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.\nUser: A bakery has 40 less than seven times as many loaves of bread as Sam had last Friday. If Sam had seventy loaves of bread last Friday, how many loaves of bread does the bakery have?\nAssistant: <think>",
  "generated_text": "Sam had seventy loaves of bread last Friday. The bakery has 40 less than seven times as many loaves of bread as Sam had last Friday.</think> <answer>40*7-40=240</answer>",
  "ground_truth": "450",
  "reward": 0.0,
  "format_reward": 1.0,
  "answer_reward": 0.0
},
数据4：{
  "prompt": "A conversation between User and Assistant. The User asks a question, and the Assistant solves it. The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, i.e., <think> reasoning process here </think> <answer> answer here </answer>.\nUser: A new program had 60 downloads in the first month. The number of downloads in the second month was three times as many as the downloads in the first month, but then reduced by 30% in the third month. How many downloads did the program have total over the three months?\nAssistant: <think>",
  "generated_text": " The number of downloads in the second month is 3 times the number of downloads in the first month, so it is 60 * 3 = 180. In the third month, the number of downloads is reduced by 30%, so it is 180 * 70% = 126. Therefore, the total number of downloads over the three months is 60 + 180 + 126 = 366.\n</think> <answer> 366\n</answer>",
  "ground_truth": "366",
  "reward": 1.0,
  "format_reward": 1.0,
  "answer_reward": 1.0
},

对数据1，格式错误，答案也错误。是模型能力的问题。思考过程结束，又来了一次\nUser:。。。
对数据2，格式错误，答案也错误。模型算出了正确答案，但是输出中格式不正确，\nanswer> $460.00，因此是模型能力的问题。

对数据3，格式正确，但答案错误 。是模型的能力问题，它连7和70都不能区分。。。

---

© Qwen 2.5 Math 1.5B 的零样本基线在 MATH 上的表现如何？

交付物：在gsm8k数据集上的accuracy很低。

4. 对 MATH 数据集的监督微调（Supervised Finetuning, SFT）

算法 1：监督微调（SFT）

输入：初始策略模型 π_θinit；SFT 数据集 D
1: 策略模型 π_θ ← π_θinit
2: for step = 1, …, n_sft_steps do
3: 　从 D 中采样一批 question-response 对 D_b
4: 　使用模型 π_θ 计算给定问题下响应的交叉熵损失
5: 　根据交叉熵损失对模型参数 θ 执行一次梯度更新
6: end for
输出：π_θ

推理能力的监督微调 在本节中，我们将在 MATH 数据集上对基础模型进行微调（见算法 1）。我们的目标不是让模型直接预测正确答案，而是先生成一段思维链（chain-of-thought）推理过程，再给出答案。为此，我们提供了一份由 DeepSeek R1（DeepSeek-AI 等人 [2025]）生成的此类推理轨迹数据集，路径为 /data/a5-alignment/MATH/sft.jsonl。

在实际训练推理模型时，SFT 通常被用作第二步强化学习微调的热启动。原因有二：

1. SFT 需要高质量的标注数据（即已有的推理轨迹），而 RL 只需正确答案作为反馈；
2. 即使在标注数据充足的情况下，RL 仍可通过探索找到比 SFT 数据更优的策略。

遗憾的是，我们使用的模型规模尚不足以在组合 SFT 与 RL 时展现出明显效果，因此本次作业中我们将这两个阶段分开处理。

4.1 使用 HuggingFace 模型

加载 HuggingFace 模型和分词器
要从本地目录加载 HuggingFace 模型和分词器（使用 bfloat16 精度并启用 FlashAttention-2 以节省显存），可使用以下示例代码：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "/data/a5-alignment/models/Qwen2.5-Math-1.5B",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained("/data/a5-alignment/models/Qwen2.5-Math-1.5B")

前向传播
加载模型后，可在一批 input IDs 上执行前向传播，并通过输出的 .logits 属性获取 logits。然后，可计算模型预测 logits 与真实标签之间的损失：

input_ids = train_batch["input_ids"].to(device)
labels = train_batch["labels"].to(device)
logits = model(input_ids).logits
loss = F.cross_entropy(...)

保存训练后的模型
训练完成后，可使用 .save_pretrained() 函数将模型保存到指定目录。由于模型体积较大，请务必保存至 /data/yourusername 下。建议同时保存分词器（即使未修改），以确保模型与分词器自包含、可从单一目录加载：

# 保存模型权重
model.save_pretrained(save_directory=output_dir)
tokenizer.save_pretrained(save_directory=output_dir)

梯度累积（Gradient Accumulation）
尽管使用了 bfloat16 和 FlashAttention-2，即使在 80GB 显存的 GPU 上，也难以支持合理的批大小。为此，可采用梯度累积技术：不在每个 batch 后立即更新权重，而是累积多个 batch 的梯度后再执行一次优化器步。

直观理解：若 GPU 足够大，一次性计算 32 个样本的梯度，与分 16 次每次 2 个样本再平均，结果应一致。

在 PyTorch 中实现梯度累积很简单。通常流程为：

for inputs, labels in data_loader:
    logits = model(inputs)
    loss = loss_fn(logits, labels)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

要实现梯度累积，只需每隔 k 步（k 为累积步数）才调用 optimizer.step() 和 optimizer.zero_grad()。同时，在调用 loss.backward() 前，将损失除以 gradient_accumulation_steps，以实现梯度平均：

gradient_accumulation_steps = 4
for idx, (inputs, labels) in enumerate(data_loader):
    logits = model(inputs)
    loss = loss_fn(logits, labels) / gradient_accumulation_steps
    loss.backward()
    if (idx + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

这样，实际训练时的有效批大小就扩大为原始批大小 × k。

4.2 SFT 辅助方法

接下来，将实现一些在 SFT 及后续 RL 实验中会用到的辅助方法。关于命名法的简要说明：在以下各节中，我们将交替使用“output”,“completion”,or“response”来指代模型在给定提示下的生成结果。

提示与输出的分词（Tokenizing prompts and outputs） 对于每个question and target output $(q,o)$，我们将分别对问题和输出进行分词，然后拼接。这样，就可以用 SFT 模型（或后续的 RL 策略）对输出部分打分（计算对数概率）。此外，还需构建一个 response_mask：一个布尔掩码，对 all tokens in response为 True，对所有问题或填充（padding）token 为 False。该掩码将在训练循环中用于确保仅在response tokens 上计算损失。

问题（tokenize_prompt_and_output）：提示与输出的分词（2 分）

交付物：实现一个 tokenize_prompt_and_output 方法，分别对问题和输出字符串进行分词、拼接，并构建 response_mask。推荐接口如下：

def tokenize_prompt_and_output(prompt_strs, output_strs, tokenizer):
    """
    对提示和输出字符串进行分词，并构建一个掩码，标记响应 token（值为 1），其余（提示或填充）为 0。

    Args:
        prompt_strs: List[str] —— 提示字符串列表。
        output_strs: List[str] —— 输出字符串列表。
        tokenizer: PreTrainedTokenizer —— 用于分词的分词器。

    Returns:
        dict[str, torch.Tensor]：
            设 prompt_and_output_lens 为各拼接后序列的长度列表，
            返回字典包含以下键：
            - input_ids: shape (batch_size, max(prompt_and_output_lens) - 1)
                         拼接后的 token 序列（去掉最后一个 token）
            - labels: shape 同 input_ids，为 input_ids 右移一位（即去掉第一个 token）
            - response_mask: shape 同 input_ids，响应 token 对应位置为 True，其余为 False
    """

为测试你的代码，请实现 [adapters.run_tokenize_prompt_and_output]，然后运行以下命令确保测试通过：

uv run pytest -k test_tokenize_prompt_and_output

---

逐token熵记录 在强化学习（RL）过程中，记录逐token熵通常很有用，可用于观察模型的预测分布是否变得（过度）自信。现在我们将实现这一功能，并比较各种微调方法对模型预测熵的影响。

离散分布$p(x)$（支撑集为$\mathcal{X}$）的熵定义为：
$$H(p)=-\sum _{x\in \mathcal{X}}p(x)\log p(x).(1)$$

给定 SFT 或 RL 模型的 logits，我们将计算每个 token 的熵，即每个下一个 token 预测的熵。

问题（compute_entropy）：Per-token entropy (1 point)

交付要求
实现一个名为 compute_entropy 的方法，用于计算每个 token 的下一个 token 预测熵。建议采用以下接口：

def compute_entropy(logits: torch.Tensor) -> torch.Tensor:
	功能：获取下一个 token 预测的熵（即在词汇表维度上的熵）。
	
	参数：
	- logits: torch.Tensor，形状为 (batch_size, sequence_length, vocab_size)，包含未归一化的 logits。
	
	返回值：
	- torch.Tensor，形状为 (batch_size, sequence_length)，表示每个下一个 token 预测的熵。

注意：你应该使用数值稳定的方法（例如，使用 logsumexp）以避免溢出。 为测试你的代码，请实现 [adapters.run_compute_entropy]，然后运行 uv run pytest -k test_compute_entropy 并确保你的实现通过测试。

---

从模型获取log-probabilities 从模型中获取对数概率是监督微调（SFT）和强化学习（RL）中都需要用到的基础操作。
对于一个前缀 $x$，语言模型（LM）会输出下一个 token 的 logits $f_{\theta }(x)\in {\mathbb{R}}^{|V|}$，以及一个真实标签 $y\in V$。此时，$y$ 的对数概率为：

$$\log p_{\theta }(y\mid x)=\log {\left[\text{softmax}(f_{\theta }(x))\right]}_y,$$

其中记号 $[x{]}_y$ 表示向量 $x$ 的第 $y$ 个元素。

你应当使用一种数值稳定的方法来计算该值，并可自由使用 torch.nn.functional 中的方法。我们还建议增加一个可选参数，用于选择性地计算并返回每个 token 的熵。

问题（get_response_log_probs）：响应对数概率（及熵）（2分）

交付要求 实现get_response_log_probs方法，用于从因果语言模型中获取逐token条件对数概率（基于前文token），并可选返回模型下一个token分布的熵。推荐接口：

def get_response_log_probs(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool = False,
) -> dict[str, torch.Tensor]:
	"""参数：
    - model：PreTrainedModel，用于评分的HuggingFace模型（若无需计算梯度，需放置在正确设备上并处于推理模式）。
    - input_ids：torch.Tensor，形状为（batch_size, sequence_length），由分词方法生成的拼接后的提示词+响应token。
    - labels：torch.Tensor，形状为（batch_size, sequence_length），由分词方法生成的标签。
    - return_token_entropy：bool，若为True，通过调用`compute_entropy`额外返回逐token熵。

    返回值：
    - dict[str, torch.Tensor]：
    - "log_probs"：形状为（batch_size, sequence_length），条件对数概率\(log p_{\theta}(x_t | x_{<<t})\)。
    - "token_entropy"（可选）：形状为（batch_size, sequence_length），每个位置的逐token熵（仅当return_token_entropy=True时存在）。
    """

实现提示： 通过model(input_ids).logits获取logits。
测试方法：实现[adapters.run_get_response_log_probs]，然后运行uv run pytest -k test_get_response_log_probs，确保测试通过。
写这个的时候，要记得把tests/conftest.py文件中的模型路径model_id改成自己下载本地路径。

---

SFT microbatch train step
SFT（监督微调）的微批次训练步骤中，我们最小化的损失是：在给定提示（prompt）的条件下，目标输出（target output）的负对数似然（negative log-likelihood）。
为了计算该损失，需要计算在给定提示条件下目标输出中每个 token 的对数概率，并对输出中所有 token 的对数概率求和，同时对提示部分的 token 和填充（padding）token 进行掩码（masking），使其不参与损失计算。

我们将实现一个辅助函数，该函数在后续强化学习（RL）过程中也会用到。

问题（masked_normalize）：掩码归一化（1分）

交付要求
实现masked_normalize方法，在考虑布尔掩码的前提下，对张量元素求和并通过常数进行归一化。
推荐接口

def masked_normalize(
    tensor: torch.Tensor,
    mask: torch.Tensor,
    normalize_constant: float,
    dim: int | None = None,
) -> torch.Tensor:
对指定维度求和并通过常数归一化，仅考虑掩码中值为1的元素。

参数：
- tensor：torch.Tensor，需求和并归一化的张量。
- mask：torch.Tensor，与tensor形状相同；值为1的位置会被纳入求和范围。
- normalize_constant：float，用于归一化的除数常数。
- dim：int | None，归一化前要求和的维度；若为None，对所有维度求和。

返回值：
- torch.Tensor，归一化后的和，其中掩码元素（mask == 0）不参与求和。

测试方法：实现[adapters.run_masked_normalize]，然后运行uv run pytest -k test_masked_normalize，确保测试通过。

---

监督微调（SFT）微批次训练步骤
现在我们可以实现监督微调（SFT）的单个微批次训练步骤（需注意：若gradient_accumulation_steps > 1，则需对每个训练批次迭代多个微批次）。

问题（sft_microbatch_train_step）：Microbatch train step (3 points)

交付要求
实现监督微调（SFT）的单个微批次更新，包括交叉熵损失计算、掩码求和及梯度缩放。

推荐接口

def sft_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    normalize_constant: float = 1.0,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
	对微批次执行前向传播和反向传播。
	
	参数：
	- policy_log_probs：形状为（batch_size, sequence_length），来自待训练监督微调（SFT）策略的逐token对数概率。
	- response_mask：形状为（batch_size, sequence_length），响应token对应位置为1，提示词/填充token对应位置为0。
	- gradient_accumulation_steps：每个优化器步骤对应的微批次数量。
	- normalize_constant：用于除法归一化的常数，默认设为1.0即可。
	
	返回值：
	- tuple[torch.Tensor, dict[str, torch.Tensor]]：
	  - loss：标量张量，微批次损失（已根据梯度累积进行调整），返回该值用于日志记录。
	  - metadata：字典，包含底层损失调用的元数据及其他需记录的统计信息。

实现提示 需在该函数中调用loss.backward()，确保根据梯度累积进行调整。
测试方法：实现[adapters.run_sft_microbatch_train_step]，然后运行uv run pytest -k test_sft_microbatch_train_step，确保测试通过。

---

Logging generations in-the-loop
在模型训练循环中记录生成结果是良好的实践，监督微调（SFT）/强化学习（RL）场景也不例外。编写log_generations函数，用于让模型对给定提示词（如从验证集中采样的提示词）生成响应并记录日志。建议为每个示例至少记录以下内容：

1. 输入提示词。
2. 监督微调（SFT）/强化学习（RL）模型生成的响应。
3. 真实答案。
4. 奖励信息，包括格式、答案及总奖励。
5. 响应的平均token熵。
6. 平均响应长度、正确响应的平均长度及错误响应的平均长度。

问题（log_generations）：生成结果日志记录（1分）

交付要求
实现log_generations函数，用于记录模型的生成结果。

4.3 SFT Experiment

利用上述模块，现在将实现完整的监督微调（SFT）流程（算法1），在MATH数据集上微调Qwen 2.5 Math 1.5B Base模型。/data/a5-alignment/MATH/sft.jsonl中的每个示例包含 formatted prompt 和 target response，其中 target response 包括思维链推理过程和最终答案。具体而言，每个示例是一个JSON元素，格式为{"prompt": str, "response": str}。

为跟踪模型在训练过程中的进度，需定期在MATH验证集上评估模型。运行脚本时需使用2块GPU：一块用于策略模型，另一块用于vLLM实例以评估策略。以下是初始化vLLM并在每次rollout阶段前将策略权重加载到vLLM实例的 starter 代码：

from vllm.model_executor import set_random_seed as vllm_set_random_seed
from unittest.mock import patch
from vllm import LLM
from transformers import PreTrainedModel

def init_vllm(model_id: str, device: str, seed: int, gpu_memory_utilization: float = 0.85):
    """
    启动推理过程，此处使用vLLM将模型部署在与策略模型不同的GPU上。
    """
    vllm_set_random_seed(seed)
    # 从TRL借鉴的Monkeypatch：https://github.com/huggingface/trl/blob/
    # 22759c820867c8659d00082ba8cf004e963873c1/trl/trainer/grpo_trainer.py
    # 对vLLM进行补丁，确保：
    # （1）将vLLM模型部署到指定设备（world_size_patch）；
    # （2）跳过不适合当前场景的测试（profiling_patch）。
    world_size_patch = patch("torch.distributed.get_world_size", return_value=1)
    profiling_patch = patch(
        "vllm.worker.worker.Worker._assert_memory_footprint_increased_during_profiling",
        return_value=None
    )
    with world_size_patch, profiling_patch:
        return LLM(
            model=model_id,
            device=device,
            dtype=torch.bfloat16,
            enable_prefix_caching=True,
            gpu_memory_utilization=gpu_memory_utilization,
        )

def load_policy_into_vllm_instance(policy: PreTrainedModel, llm: LLM):
    """
    从https://github.com/huggingface/trl/blob/
    22759c820867c8659d00082ba8cf004e963873c1/trl/trainer/grpo_trainer.py#L670复制
    """
    state_dict = policy.state_dict()
    llm_model = llm.llm_engine.model_executor.driver_worker.model_runner.model
    llm_model.load_weights(state_dict.items())

建议同时记录训练步骤和验证步骤的指标（这在后续强化学习（RL）实验中也会用到）。在wandb中可通过以下代码实现：

# 配置wandb指标
wandb.define_metric("train_step")  # 训练过程的x轴
wandb.define_metric("eval_step")  # 评估过程的x轴
# 所有以train/开头的指标都与train_step绑定
wandb.define_metric("train/*", step_metric="train_step")
# 所有以eval/开头的指标都与eval_step绑定
wandb.define_metric("eval/*", step_metric="eval_step")

最后，建议使用梯度裁剪，裁剪值设为1.0。

问题（sft_experiment）：在MATH数据集上运行监督微调（SFT）（2分）（2个H100小时）

1. 使用Qwen 2.5 Math 1.5B基础模型，在推理型监督微调（SFT）示例（路径：/data/a5-alignment/MATH/sft.jsonl）上运行监督微调（SFT），监督微调（SFT）的唯一示例数量在{128, 256, 512, 1024}范围内变化，同时也使用完整数据集。调整学习率和批次大小，确保使用完整数据集时验证准确率至少达到15%。
   交付要求：不同数据集大小对应的验证准确率曲线。
2. 过滤推理型监督微调（SFT）示例，仅保留能产生正确答案的示例。在（完整的）过滤后数据集上运行监督微调（SFT），报告过滤后数据集的大小及达到的验证准确率。
   交付要求：报告数据集大小及验证准确率曲线，并与之前的监督微调（SFT）实验结果进行对比。
   

5 MATH数据集的专家迭代

在上一节中，我们发现通过从监督微调（SFT）数据中过滤掉不良示例，可以提升监督微调（SFT）模型的性能。本节将进一步优化：将该过滤流程应用于基础模型自身生成的推理轨迹。这一过程在文献中被称为专家迭代（expert iteration）[Anthony et al., 2017]，在语言模型领域，Cobbe et al. [2021b]、Zelikman et al. [2022]、Dohan et al. [2022]、Gulcehre et al. [2023] 等学者已对此进行了探索。

算法2 专家迭代（EI）

输入：初始策略模型 ${\pi }_{{\theta }_{init}}$；奖励函数 $R$；任务问题集 $D$

1. 初始化策略模型 ${\pi }_{\theta }\leftarrow {\pi }_{{\theta }_{init}}$
2. 对于步骤 $step=1,\dots ,n_{ei\_steps}$：
   3. 从 $D$ 中采样一批问题 $D_b$
   4. 将旧策略模型设为 ${\pi }_{{\theta }_{old}}\leftarrow {\pi }_{\theta }$
   5. 对每个问题 $q\in D_b$，从 ${\pi }_{{\theta }_{old}}(\cdot \mid q)$ 中采样 $G$ 个输出 { o ( i ) } i = 1 G \{o^{(i)}\}_{i=1}^G {o(i)}i=1G​
   6. 通过奖励函数 $R(q,o^{(i)})$ 计算每个采样输出 $o^{(i)}$ 的奖励 { r ( i ) } i = 1 G \{r^{(i)}\}_{i=1}^G {r(i)}i=1G​
   7. 过滤掉错误输出（即 $r^{(i)}=0$ 的 $o^{(i)}$），得到正确的问答对数据集 $D_{sft}$
   8. 通过监督微调（SFT）更新策略模型：${\pi }_{\theta }\leftarrow SFT({\pi }_{\theta },D_{sft})$（算法1）
3. 结束循环

输出：策略模型 ${\pi }_{\theta }$。

接下来，我们将在MATH数据集上运行专家迭代。
小提示, 需为vLLM的SamplingParams传入min_tokens参数，确保不会生成空字符串（否则可能导致后续实现中出现NaN值）。具体设置如下：

sampling_min_tokens = 4
sampling_params = SamplingParams(
    temperature=sampling_temperature,
    max_tokens=sampling_max_tokens,
    min_tokens=sampling_min_tokens,
    n=G,
    seed=seed,
)

与监督微调（SFT）相同，需使用梯度裁剪，裁剪值设为1.0。

问题（expert_iteration_experiment）：在MATH数据集上运行专家迭代（2分）（6个H100小时）

在 MATH 数据集（路径为 /data/a5-alignment/MATH/train.jsonl）上，使用 Qwen 2.5 Math 1.5B Base 模型运行专家迭代（Expert Iteration），设置专家迭代步数 $n_{ei\_steps}=5$。
实验中需调整以下超参数：

• 每个问题的 rollout 数量 (G)；
• 监督微调（SFT）步骤中使用的训练轮数（epochs）；
• 每次专家迭代步骤中的batch_size（即 $D_b$ 的大小），在 {512, 1024, 2048} 中选择。

无需尝试所有超参数组合，只需进行足够多的实验以对每个超参数的影响得出合理结论即可。在训练过程中，请记录模型生成回答response的熵（entropy）变化情况。此外，请确保使用 vLLM 进行推理时，在遇到第二个答案标签 </answer> 时终止生成，这一处理方式应与监督微调（SFT）部分保持一致。

交付成果

1. 不同滚动配置对应的验证准确率曲线。至少尝试2种不同的滚动次数和轮数。
2. 在MATH数据集上验证准确率至少达到15%的模型。
3. 简要的两句话讨论：对比SFT的性能，以及不同EI（专家迭代）步骤下的性能。
4. 训练过程中模型响应熵值的图表。

训练集采样128个样本的准确率曲线如下；可以看到，使用正确prompt-response对微调后模型的准确率增长，正确prompt-response对的数量也增加，使模型准确率进一步增长。

随着模型准确率的提升，response的熵逐步减小。

6 策略梯度入门( Primer on Policy Gradients )

语言模型研究中的一项令人振奋的新发现是：利用性能强劲的基础模型，针对经过验证的奖励信号进行强化学习（RL），能够显著提升模型的推理能力和性能[OpenAI等，2024；DeepSeek-AI等，2025]。目前性能最优异的开源推理模型（如DeepSeek R1和Kimi k1.5[团队等，2025]）均采用策略梯度算法训练而成。策略梯度是一种强大的强化学习算法，可优化任意奖励函数。

下文将简要介绍适用于语言模型强化学习的策略梯度方法。本部分内容主要基于两份优质参考资料（对相关概念有更深入的阐述）：OpenAI的《深度强化学习入门》[Achiam，2018a]和内森·兰伯特（Nathan Lambert）的《基于人类反馈的强化学习（RLHF）手册》[Lambert，2024]。

6.1 将语言模型当作策略

参数为 $\theta$ 的因果语言模型（LM）定义了一个概率分布：给定当前文本前缀 $s_t$（状态/观测值），下一个 token $a_t\in V$（$V$ 为词汇表）的出现概率。在强化学习语境中，我们将下一个 token $a_t$ 视为“动作”，将当前文本前缀 $s_t$ 视为“状态”。因此，语言模型可看作一种类别型随机策略( categorical stochastic policy )。
$$a_t\sim {\pi }_{\theta }(\cdot \mid s_t),{\pi }_{\theta }(a_t\mid s_t)={\left[\mathrm{softmax}(f_{\theta }(s_t))\right]}_{a_t}\text{\hspace{1em}(3)}$$

通过策略梯度优化策略时，需用到两项基础操作：

1. 从策略中采样：从上述类别分布中抽取一个动作 $a_t$；
2. 计算动作的对数似然：评估 $\log {\pi }_{\theta }(a_t|s_t)$（给定状态 $s_t$ 时，策略 ${\pi }_{\theta }$ 选择动作 $a_t$ 的对数概率）。

通常来说，在大语言模型（LLMs）的强化学习任务中，$s_t$指的是目前生成的部分完成内容/解决方案，而每个$a_t$是该解决方案的下一个token；当输出文本结束token（比如$<|\text{end\_of\_text}|>$）时，该轮交互就会结束；在我们的$\text{r1\_zero}$提示词场景下，对应的结束token是$</\text{answer}>$。

6.2 轨迹（Trajectories）

（有限时域的）轨迹是agent经历的状态与动作的交替序列：
$$\tau =\left(s_0,a_0,s_1,a_1,\dots ,s_T,a_T\right),\text{\hspace{1em}(4)}$$
其中 $T$ 为轨迹长度，即 $a_T$ 是文本结束 token，或已达到最大生成 token 数上限。

初始状态state从起始分布中采样得到：$s_0\sim {\rho }_0(s_0)$；在语言模型的强化学习中，${\rho }_0(s_0)$ 是 formatted prompts 的分布。在一般场景中，状态转移遵循 environment dynamics $s_{t+1}\sim P(\cdot |s_t,a_t)$。而在语言模型的强化学习中，环境是确定性的：下一个状态是旧前缀与生成 token 的拼接，即 $s_{t+1}=s_t\Vert a_t$（“$\Vert$”表示字符串拼接）。轨迹（Trajectories）也被称为“轮次（episodes）”或“滚动（rollouts）”，本文中这三个术语可互换使用。

6.3 奖励与回报（Rewards and Return）

标量奖励 $r_t=R(s_t,a_t)$ 用于评判在状态 $s_t$ 下所执行动作的即时优劣。在经过验证的领域（如数学解题）中，强化学习的标准做法是：中间步骤的奖励设为 0，终端动作的奖励为经过验证的结果，即：
$r_T=R(s_T,a_T):=\{\begin{array}{ll}1& \text{若 }{s}_T\Vert a_T\text{ 与奖励函数定义的真实结果一致}\\ 0& \text{否则}\end{array}$

回报 $R(\tau )$ 是轨迹上所有奖励的累加。两种常见定义为：有限时域无折扣回报：
$$R(\tau ):=\sum _{t=0}^T{r}_t\text{\hspace{1em}(5)}$$
和，无限时域折扣回报：
$$R(\tau ):=\sum _{t=0}^{\infty }{\gamma }^t{r}_t,0<\gamma <1\text{\hspace{1em}(6)}$$

在本实验中，由于交互轮次有自然终止点（文本结束或最大生成长度），我们将使用有限时域无折扣回报定义。

agent的目标是最大化期望回报，
$$J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[R(\tau )\right],\text{\hspace{1em}(7)}$$
对应的优化问题为：
$${\theta }^{\ast }=\arg \underset{\theta }{\max }J(\theta )\text{\hspace{1em}(8)}$$

6.4 vanilla 策略梯度（Vanilla Policy Gradient）

接下来，我们尝试通过梯度上升法最大化期望回报，从而学习策略参数 $\theta$：
$${\theta }_{k+1}={\theta }_k+\alpha {\nabla }_{\theta }J({\theta }_k).\text{\hspace{1em}(9)}$$
核心公式为 REINFORCE 策略梯度（如下所示），它是实现这一目标的关键：
$${\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )\right].\text{\hspace{1em}(10)}$$
（梯度 ${\nabla }_{\theta }J({\pi }_{\theta })$ 表示期望回报随策略参数 $\theta$ 的变化率，期望 ${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}$ 表示对所有从策略 ${\pi }_{\theta }$ 采样的轨迹 $\tau$ 取平均）

策略梯度的推导
该公式如何推导而来？为保证完整性，下文将简要推导。推导过程将用到以下几项结论：

1. 轨迹的概率为： $$P(\tau \mid \theta )={\rho }_0(s_0)\prod _{t=0}^{T}P(s_{t+1}\mid s_t,a_t){\pi }_{\theta }(a_t\mid s_t).\text{\hspace{1em}(11)}$$
   因此，轨迹的对数概率为： $$\log P(\tau |\theta )=\log {\rho }_0(s_0)+\sum _{t=0}^T\left[\log P(s_{t+1}|s_t,a_t)+\log {\pi }_{\theta }(a_t|s_t)\right].\text{\hspace{1em}(12)}$$

2. 对数导数技巧（log-derivative trick）：$${\nabla }_{\theta }P=P{\nabla }_{\theta }\log P.\text{\hspace{1em}(13)}$$

3. 环境相关项与 $\theta$ 无关：${\rho }_0$、$P(\cdot |\cdot )$ 和 $R(\tau )$ 均不依赖于策略参数，因此对 $\theta$ 的梯度为 0。
   $${\nabla }_{\theta }{\rho }_0={\nabla }_{\theta }P={\nabla }_{\theta }R(\tau )=0.\text{\hspace{1em}(14)}$$
   基于上述公式：
   $$\begin{array}{rl}{\nabla }_{\theta }J(\theta )& ={\nabla }_{\theta }{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[R(\tau )\right]\\ & ={\nabla }_{\theta }\sum _{\tau }P(\tau \mid \theta )R(\tau )\\ & =\sum _{\tau }{\nabla }_{\theta }P(\tau \mid \theta )R(\tau )\\ & =\sum _{\tau }P(\tau \mid \theta ){\nabla }_{\theta }\log P(\tau \mid \theta )R(\tau )\text{(Log-derivative trick)}\\ & ={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[{\nabla }_{\theta }\log P(\tau \mid \theta )R(\tau )\right].\end{array}\text{\hspace{1em}(15)(16)(17)(18)(19)}$$

将轨迹的对数概率代入，并利用“环境项与$\theta$无关”这一性质，即可得到vanilla 或 REINFORCE策略梯度公式：
$${\nabla }_{\theta }J({\pi }_{\theta })={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)R(\tau )\right].\text{\hspace{1em}(20)}$$

直观来看，该梯度会增大“高回报轨迹中所有动作”的对数概率，同时减小“低回报轨迹中所有动作”的对数概率。

梯度的采样估计
给定一批由 $N$ 条轨迹组成的数据集 D = { τ ( i ) } i = 1 N D = \{\tau^{(i)}\}_{i=1}^N D={τ(i)}i=1N​（采样方式：从起始分布 ${\rho }_0(s_0)$ 中抽取初始状态 $s_0^{(i)}$，然后在环境中运行策略 ${\pi }_{\theta }$ 生成轨迹），可构造梯度的无偏估计：

$$\hat{g}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }\left(a_t^{(i)}\mid s_t^{(i)}\right)R\left({\tau }^{(i)}\right).\text{\hspace{1em}(21)}$$

该向量将用于梯度上升更新：$\theta \leftarrow \theta +\alpha \hat{g}$（其中 $\alpha$ 为学习率）

6.5 策略梯度的基线（Policy Gradient Baselines）

vanilla 策略梯度的主要问题是梯度估计的方差较大。一种常用的缓解方法是在奖励中减去一个仅依赖于状态的基线函数 $b$。这是一种控制变量法 (control variate) [Ross, 2022]：核心思想是通过减去一个与梯度估计相关的项，在不引入偏差的前提下降低估计方差。

定义带基线的策略梯度为：
$$B={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)\left(R(\tau )-b(s_t)\right)\right].\text{\hspace{1em}(22)}$$

例如，一个合理的基线是“在策略价值函数(on-policy value function)” $V^{\pi }(s)={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )|s_t=s]$，即：从状态 $s_t=s$ 出发并遵循策略 ${\pi }_{\theta }$ 时的期望回报。此时，$(R(\tau )-V^{\pi }(s_t))$ 直观上表示“实际轨迹回报与期望回报的差值”（即优势值）。

只要基线仅依赖于状态，带基线的策略梯度就是无偏的。这一点可通过重写带基线的策略梯度证明：
$$B={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)R(\tau )\right]-{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)\right].\text{\hspace{1em}(23)}$$

聚焦基线项，我们发现：
$${\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)b(s_t)\right]=\sum _{t=0}^T{\mathbb{E}}_{s_t}\left[b(s_t){\mathbb{E}}_{a_t\sim {\pi }_{\theta }(\cdot \mid s_t)}\left[{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)\right]\right].\text{\hspace{1em}(24)}$$

一般而言，得分函数的期望为 0：${\mathbb{E}}_{x\sim P_{\theta }}[{\nabla }_{\theta }\log P_{\theta }(x)]=0$。因此，上式中的基线项期望为 0，即：
$$B={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}\left[\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)R(\tau )\right]-0={\nabla }_{\theta }J({\pi }_{\theta }).\text{\hspace{1em}(25)}$$

由此可得出结论：带基线的策略梯度是无偏的。后续我们将通过实验验证基线是否能提升下游任务性能。

关于策略梯度“损失函数”的说明
在 PyTorch 等框架中实现策略梯度方法时，我们会定义一个所谓的“策略梯度损失”（pg_loss），使得调用 pg_loss.backward() 时，模型参数的梯度缓冲区会被填充为近似策略梯度 $\hat{g}$。从数学上看，其定义满足：
$$\text{pg\_loss}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T\log {\pi }_{\theta }\left(a_t^{(i)}\mid s_t^{(i)}\right)\left(R\left({\tau }^{(i)}\right)-b\left(s_t^{(i)}\right)\right).\text{\hspace{1em}(26)}$$

但 pg_loss 并非传统意义上的“损失函数”——将其作为训练集或验证集的评估指标是无意义的，良好的验证集 pg_loss 并不代表模型具有良好的泛化能力。pg_loss 本质上只是一个标量，其核心作用是：通过反向传播（backprop）得到近似策略梯度 $\hat{g}$。

在强化学习中，应始终记录并报告训练集和验证集的“奖励”（rewards）。奖励是真正有意义的评估指标，也是我们通过策略梯度方法试图优化的目标。

6.6 异策略策略梯度（Off-Policy Policy Gradient）

REINFORCE 是一种“同策略（on-policy）”算法：训练数据由当前正在优化的策略生成。这一点可通过 REINFORCE 算法的步骤明确看出：

1. 从当前策略 ${\pi }_{\theta }$ 中采样一批轨迹 { τ ( i ) } i = 1 N \{\tau^{(i)}\}_{i=1}^N {τ(i)}i=1N​；
2. 近似策略梯度：${\nabla }_{\theta }J({\pi }_{\theta })\approx \hat{g}=\frac{1}{N}{\sum }_{i=1}^N{\sum }_{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})R({\tau }^{(i)})$；
3. 利用计算得到的梯度更新策略参数：$\theta \leftarrow \theta +\alpha \hat{g}$。

该方法的问题在于：需要进行大量推理以采样新的轨迹批次，却仅能执行一步梯度更新。由于语言模型的行为通常无法在单步更新中发生显著变化，这种同策略方法的效率极低。

异策略策略梯度

在异策略学习中，轨迹采样自“非当前优化的策略”。主流策略梯度算法（如 PPO、GRPO）的异策略变体，会利用“旧策略 ${\pi }_{{\theta }_{old}}$ 生成的轨迹”来优化当前策略 ${\pi }_{\theta }$。异策略的策略梯度估计为：
$${\hat{g}}_{\text{off-policy}}=\frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T\frac{{\pi }_{\theta }\left(a_t^{(i)}\mid s_t^{(i)}\right)}{{\pi }_{{\theta }_{\text{old}}}\left(a_t^{(i)}\mid s_t^{(i)}\right)}{\nabla }_{\theta }\log {\pi }_{\theta }\left(a_t^{(i)}\mid s_t^{(i)}\right)R\left({\tau }^{(i)}\right).\text{\hspace{1em}(27)}$$

该式可看作vanilla 策略梯度的重要性采样版本，其中包含了$\frac{{\pi }_{\theta }(a_t^{(i)}|s_t^{(i)})}{{\pi }_{{\theta }_{\text{old}}}(a_t^{(i)}|s_t^{(i)})}$这样的重加权项。实际上，式（27）可通过重要性采样推导得出，且在 ${\pi }_{\theta }$ 与 ${\pi }_{{\theta }_{old}}$ 差异不大的前提下，该近似是合理的（更多细节参见 Degris 等 [2013]）。

7 Group Relative Policy Optimization

下文将介绍“组相对策略优化（GRPO）”——这是一种策略梯度变体，你将基于该方法实现并实验数学问题求解。

7.1 GRPO 算法

优势值估计（Advantage Estimation）
GRPO 的核心思想是：对每个问题，从策略 ${\pi }_{\theta }$ 中采样多个输出，并用这些输出计算基线。这种方式的优势在于：无需学习神经网络价值函数 $V_{\varphi }(s)$（该函数训练难度大，且从系统实现角度看较为繁琐）。

对于问题 $q$ 和从策略中采样的 $G$ 个输出 { o ( i ) } i = 1 G ∼ π θ ( ⋅ ∣ q ) \{o^{(i)}\}_{i=1}^G \sim \pi_{\theta}(\cdot | q) {o(i)}i=1G​∼πθ​(⋅∣q)，设 $r^{(i)}=R(q,o^{(i)})$ 为第 $i$ 个输出的奖励。DeepSeekMath [Shao 等, 2024] 和 DeepSeek R1 [DeepSeek-AI 等, 2025] 中，第 $i$ 个输出的“组归一化奖励”（即优势值）定义为：
$$A^{(i)}=\frac{r^{(i)}-\text{mean}\left(r^{(1)},r^{(2)},\dots ,r^{(G)}\right)}{\text{std}\left(r^{(1)},r^{(2)},\dots ,r^{(G)}\right)+\text{advantage\_eps}},\text{\hspace{1em}(28)}$$
advantage_eps 为防止分母为 0 的小常数。注意：该优势值 $A^{(i)}$ 对响应中的每个 token 都相同，即$A_t^{(i)}=A^{(i)}$， ∀ t ∈ { 1 , . . . , ∣ o ( i ) ∣ } \forall t \in \{1, ..., |o^{(i)}|\} ∀t∈{1,...,∣o(i)∣}，因此下文将省略下标 $t$。

高层算法流程（High-level algorithm）
在深入 GRPO 目标函数之前，先通过 Shao 等 [2024] 提出的算法 3，了解 GRPO 的训练循环框架。注 ：这是 DeepSeekMath 中 GRPO 的特例——使用经过验证的奖励函数，无 KL 散度项，也没有对参考模型和奖励模型进行迭代更新。
GRPO 目标函数 （GRPO objective
GRPO 目标函数融合了三项核心思想：

1. 异策略策略梯度（见式 27）；
2. 通过组归一化计算优势值 $A^{(i)}$（见式 28）；
3. 裁剪机制（Clipping Mechanism），源自近邻策略优化（PPO, Schulman 等 [2017]）。

裁剪机制的目的是：在同一批轨迹上执行多步梯度更新时，保证训练稳定性。它的工作原理是防止当前策略 ${\pi }_{\theta }$ 偏离旧策略太远。

算法3 组相对策略优化（GRPO）
输入：初始策略模型${\pi }_{{\theta }_{init}}$；奖励函数R；任务问题集D

1. 初始化策略模型${\pi }_{\theta }$ ← ${\pi }_{{\theta }_{init}}$
2. for step = 1, …, n_grpo_steps，执行以下操作：
   a. 从问题集D中采样一批问题Db
   b. 保存旧策略模型πθold ← πθ
   c. 对Db中的每个问题q，从πθold(· | q)中采样G个输出{o(i)}G i=1
   d. 通过执行奖励函数R(q, o(i))，为每个采样输出o(i)计算奖励{r(i)}G i=1
   e. 采用组归一化方法计算优势值A(i)（公式28）
   f. for 训练步数train step = 1, …, n_train_steps_per_rollout_batch，执行以下操作：
   $-$ 通过最大化GRPO-Clip目标函数（详见公式29）更新策略模型πθ
   end for
   end for
3. 输出最终策略模型πθ

首先，我们给出完整的GRPO-Clip目标函数，再解释裁剪操作的作用：
J G R P O − C l i p ( θ ) = E q ∼ D , { o ( i ) } i = 1 G ∼ π θ ( ⋅ ∣ q ) [ 1 G ∑ i = 1 G 1 ∣ o ( i ) ∣ ∑ t = 1 ∣ o ( i ) ∣ min ⁡ ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ o l d ( o t ( i ) ∣ q , o < t ( i ) ) A ( i ) , c l i p ( π θ ( o t ( i ) ∣ q , o < t ( i ) ) π θ o l d ( o t ( i ) ∣ q , o < t ( i ) ) , 1 − ϵ , 1 + ϵ ) A ( i ) ) ⏟ 逐token目标函数 ] \begin{array}{rlr} J_{GRPO-Clip}(\theta)&=\mathbb{E}_{q\sim \mathcal{D},\{ o^{(i)}\} _{i=1}^{G}\sim \pi _{\theta }(\cdot | q)} \\ &\Bigg[ \frac{1}{G}\sum_{i=1}^{G}\frac{1}{|o^{(i)}|}\sum_{t=1}^{|o^{(i)}|} \underbrace{\operatorname*{min}\Bigg( \frac{\pi _{\theta }(o_{t}^{(i)}|q,o_{<t}^{(i)})}{\pi _{\theta _{old}}(o_{t}^{(i)} | q,o_{<t}^{(i)})}A^{(i)}, clip\left( \frac{\pi _{\theta }\left(o_{t}^{(i)} | q,o_{<t}^{(i)}\right)}{\pi _{\theta _{old}}(o_{t}^{(i)} | q,o_{<t}^{(i)})},1-\epsilon ,1+\epsilon \right) A^{(i)} \Bigg)}_{\text{逐token目标函数}} \Bigg] \end{array} JGRPO−Clip​(θ)​=Eq∼D,{o(i)}i=1G​∼πθ​(⋅∣q)​[G1​∑i=1G​∣o(i)∣1​∑t=1∣o(i)∣​逐token目标函数 min(πθold​​(ot(i)​∣q,o<t(i)​)πθ​(ot(i)​∣q,o<t(i)​)​A(i),clip ​πθold​​(ot(i)​∣q,o<t(i)​)πθ​(ot(i)​∣q,o<t(i)​)​,1−ϵ,1+ϵ ​A(i))​​]​

超参数ε>0用于控制策略的更新幅度。为更直观理解，我们参考Achiam [2018a,b]的方法重写逐token目标函数。定义函数：
$$g(ϵ,A^{(i)})=\{\begin{array}{ll}(1+ϵ)A^{(i)}& \text{若 }{A}^{(i)}\ge 0\\ (1-ϵ)A^{(i)}& \text{若 }{A}^{(i)}<0\end{array}\text{\hspace{1em}(30)}$$

则逐token目标函数可重写为：
$$\text{逐token目标函数}=\min \left(\frac{{\pi }_{\theta }\left(o_t^{(i)}|q,o_{<<t}^{(i)}\right)}{{\pi }_{{\theta }_{old}}\left(o_t^{(i)}|q,o_{<<t}^{(i)}\right)}{A}^{(i)},g\left(ϵ,A^{(i)}\right)\right)$$

我们分情况讨论：

• 当优势值A(i)为正时，逐token目标函数简化为：
  $$\text{逐token目标函数}=\min \left(\frac{{\pi }_{\theta }\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)}{{\pi }_{{\theta }_{old}}\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)},1+ϵ\right)A^{(i)}$$
  由于$A^{(i)}>0$，若动作$o_t^{(i)}$在${\pi }_{\theta }$下的概率增大（即${\pi }_{\theta }\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)$的值变大），目标函数值会增加。min函数的裁剪作用限制了目标函数的增长幅度：当${\pi }_{\theta }\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)>(1+\epsilon ){\pi }_{{\theta }_{old}}\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)$时，逐token目标函数达到最大值(1+ε)A(i)，从而避免策略πθ与旧策略πθold偏差过大。

• 当优势值A(i)为负时，模型会尝试降低${\pi }_{\theta }\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)$的概率，但裁剪机制会阻止其降至$(1-\epsilon ){\pi }_{{\theta }_{old}}\left(o_t^{(i)}|q,o_{<t}^{(i)}\right)$以下（完整推导参见Achiam [2018b]）。

7.2 Implementation

在理解GRPO的训练流程和目标函数后，我们开始分模块实现。SFT和EI部分的许多模块可直接复用。

计算优势值（组归一化奖励）Computing advantages (group-normalized rewards)
首先实现a rollout batch中每个样本的优势值计算逻辑，即组归一化奖励。我们考虑两种组归一化方式：1. 前文公式28的标准方法. 2. 近期提出的简化方法。
Dr. GRPO [Liu et al., 2025] 指出，通过 std(r(1), r(2), …, r(G)) 进行归一化的方式，会奖励答案正确性波动较小的问题，而这可能并不理想。因此，他们提出移除标准差归一化步骤，直接计算：
$$A^{(i)}=r^{(i)}-\text{mean}\left(r^{(1)},r^{(2)},...,r^{(G)}\right)\text{\hspace{1em}(31)}$$
我们将实现两种变体，并在后续实验中对比其性能。

问题（compute_group_normalized_rewards）：组归一化（2分）

交付要求：实现compute_group_normalized_rewards方法，计算每个滚动响应的原始奖励，在组内进行归一化，并返回归一化奖励、原始奖励及有用的元数据。

推荐接口：

def compute_group_normalized_rewards(
    reward_fn: Callable[[str, str], dict[str, float]],
    rollout_responses,
    repeated_ground_truths,
    group_size,
    advantage_eps,
    normalize_by_std,
):
    """为每组 rollout 响应计算奖励，并按组进行归一化。

    参数：
    reward_fn: Callable[[str, str], dict[str, float]]  
        用于将 rollout 响应与标准答案（ground truth）进行比较并打分的函数，返回一个字典，
        包含键 "reward"、"format_reward" 和 "answer_reward"。
    
    rollout_responses: list[str]  
        策略生成的 rollout 响应列表。该列表长度为 rollout_batch_size，
        即 rollout_batch_size = n_prompts_per_rollout_batch * group_size。
    
    repeated_ground_truths: list[str]  
        每个样本对应的标准答案列表。该列表长度也为 rollout_batch_size，
        因为每个问题的标准答案被重复了 group_size 次（与每个问题对应的多个响应对齐）。
    
    group_size: int  
        每个问题（即每组）生成的响应数量。
    
    advantage_eps: float  
        用于归一化时避免除零的小常数。
    
    normalize_by_std: bool  
        若为 True，则用每组奖励的标准差进行归一化（即减去均值后除以标准差）；
        否则仅减去组内均值。

    返回：
    tuple[torch.Tensor, torch.Tensor, dict[str, float]]
        - advantages: shape (rollout_batch_size,)，每条 rollout 响应的组内归一化奖励（即优势值）。
        - raw_rewards: shape (rollout_batch_size,)，每条 rollout 响应的原始未归一化奖励。
        - metadata: 用户自定义的其他统计信息，可用于日志记录（例如奖励的均值、标准差、最大/最小值等）。
    """

测试方法：实现[adapters.run_compute_group_normalized_rewards]，运行命令uv run pytest -k test_compute_group_normalized_rewards并确保测试通过。

---

朴素策略梯度损失
接下来实现损失计算相关方法。需注意：这些并非传统意义上的损失函数，不应作为评估指标。在强化学习中，应跟踪训练集和验证集的回报值等指标（详见6.5节讨论）。

首先实现朴素策略梯度损失，该损失直接将优势值与动作的对数概率相乘并取负。对于问题q、响应o和响应token o_t，逐token朴素策略梯度损失为：
$$-A_t\cdot \log p_{\theta }\left(o_t|q,o_{<t}\right)\text{\hspace{1em}(32)}$$

问题（compute_naive_policy_gradient_loss）：朴素策略梯度（1分）

交付要求：实现compute_naive_policy_gradient_loss方法，使用原始奖励或预计算的优势值计算逐token策略梯度损失。

推荐接口：

def compute_naive_policy_gradient_loss(
    raw_rewards_or_advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
) -> torch.Tensor:
    """
    计算每个token的策略梯度损失，其中raw_rewards_or_advantages可为原始奖励或已归一化的优势值
    
    参数：
        raw_rewards_or_advantages: 形状为(batch_size, 1)的张量，每个滚动响应的标量奖励/优势值
        policy_log_probs: 形状为(batch_size, sequence_length)的张量，每个token的对数概率
    
    返回：
        形状为(batch_size, sequence_length)的张量，逐token策略梯度损失（将在训练循环中跨批次和序列维度聚合）
    """

实现提示：

• 将raw_rewards_or_advantages在sequence_length维度上广播（broadcast）

测试方法：实现[adapters.run_compute_naive_policy_gradient_loss]，运行命令uv run pytest -k test_compute_naive_policy_gradient_loss并确保测试通过。

---

GRPO-Clip损失
接下来实现更核心的GRPO-Clip损失。逐token GRPO-Clip损失为：
$$-\min \left(\frac{{\pi }_{\theta }\left(o_t|q,o_{<t}\right)}{{\pi }_{{\theta }_{old}}\left(o_t|q,o_{<t}\right)}{A}_t,clip\left(\frac{{\pi }_{\theta }\left(o_t|q,o_{<t}\right)}{{\pi }_{{\theta }_{old}}\left(o_t|q,o_{<t}\right)},1-ϵ,1+ϵ\right)A_t\right)$$

问题（compute_grpo_clip_loss）：GRPO-Clip损失（2分）

交付要求：实现compute_grpo_clip_loss方法，计算逐token GRPO-Clip损失。

推荐接口：

def compute_grpo_clip_loss(
    advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
    old_log_probs: torch.Tensor,
    cliprange: float,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """    
    参数：
        advantages: 形状为(batch_size, 1)的张量，每个样本的优势值A
        policy_log_probs: 形状为(batch_size, sequence_length)的张量，待训练策略的逐token对数概率
        old_log_probs: 形状为(batch_size, sequence_length)的张量，旧策略的逐token对数概率
        cliprange: 裁剪参数ε（例如0.2）
    
    返回：
        tuple[torch.Tensor, dict[str, torch.Tensor]]:
            loss: 形状为(batch_size, sequence_length)的张量，逐token裁剪损失
            metadata: 需记录的元数据（建议记录每个token是否被裁剪，即min函数右侧的裁剪后损失是否小于左侧）
    """

实现提示：

• 将优势值在sequence_length维度上广播（broadcast）

测试方法：实现[adapters.run_compute_grpo_clip_loss]，运行命令uv run pytest -k test_compute_grpo_clip_loss并确保测试通过。

---

策略梯度损失包装器
我们将通过对比实验验证三种策略梯度变体：
(a) 无基线（no_baseline）：无基线的朴素策略梯度损失，优势值直接为原始奖励A=R(q, o)
(b) reinforce_with_baseline：使用组归一化奖励作为优势值（advantage）的朴素策略梯度损失。如果 $\bar{r}$ 是来自 compute_group_normalized_rewards 的组归一化奖励（可能已或未按组标准差归一化），那么优势值 $A=\bar{r}$。
（c）grpo_clip：GRPO-Clip 损失函数
为方便起见，我们将实现一个包装器（wrapper），使我们能够轻松在这三种策略梯度损失函数之间切换。

问题（compute_policy_gradient_loss）：策略梯度包装器（1分）

交付要求：实现 compute_policy_gradient_loss 函数——一个便捷包装器，用于调度至对应的损失计算流程（no_baseline、reinforce_with_baseline 或 grpo_clip），并返回逐token损失（per-token loss）及所有辅助统计信息。

推荐接口如下：

def compute_policy_gradient_loss(
    policy_log_probs: torch.Tensor,
    loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"],
    advantages: torch.Tensor | None = None,
    raw_rewards: torch.Tensor | None = None,
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

功能：选择并计算目标策略梯度损失。

参数说明：

• policy_log_probs：形状为（batch_size, sequence_length），表示待训练策略的逐token对数概率。
• loss_type：损失类型，可选值为 “no_baseline”、“reinforce_with_baseline” 或 “grpo_clip”。
• raw_rewards：当 loss_type == "no_baseline" 时为必填项，形状为（batch_size, 1）。
• advantages：当 loss_type 为 “reinforce_with_baseline” 或 “grpo_clip” 时为必填项，形状为（batch_size, 1）。
• old_log_probs：当 loss_type == "grpo_clip" 时为必填项，形状为（batch_size, sequence_length）。
• cliprange：当 loss_type == "grpo_clip" 时为必填项，用于裁剪的标量参数 ϵ。

返回值：

• 元组（torch.Tensor, dict[str, torch.Tensor]）：
  1. 损失张量（loss）：形状为（batch_size, sequence_length），即逐token损失。
  2. 元数据字典（metadata）：包含底层计算流程的统计信息（例如 GRPO-Clip 的裁剪比例）。

实现提示：

1. 调用 compute_naive_policy_gradient_loss 或 compute_grpo_clip_loss 完成具体计算。
2. 执行参数校验（参考上述断言模式）。
3. 将所有返回的元数据汇总到单个字典中。

测试方式：实现 [adapters.run_compute_policy_gradient_loss]，运行命令 uv run pytest -k test_compute_policy_gradient_loss 并验证测试通过。

---

掩码均值（Masked Mean）
截至目前，我们已具备计算优势函数（advantages）、对数概率、逐token损失，以及逐token熵、裁剪比例等辅助统计信息的计算能力。为将形状为（batch_size, sequence_length）的 per-token loss tensors 缩减为损失向量（每个样本对应一个标量损失），我们将在序列维度上计算损失的均值，但仅包含响应对应的索引（即 mask[i, j]==1 的token位置）。

在大多数基于大语言模型（LLM）的强化学习（RL）代码库中，按序列长度归一化是标准操作，但这一做法的合理性尚未明确——观察公式（21）中策略梯度估计的定义可知，其中并不存在归一化因子 $\frac{1}{T^{(i)}}$。我们将先采用这一标准方法（通常称为 masked_mean），后续再测试SFT阶段实现的 masked_normalize 方法。

该函数支持指定均值计算的维度：若 dim = None，则对所有掩码为1的元素计算均值。这一功能可用于获取响应token的平均逐token熵、裁剪比例等统计信息。

问题（masked_mean）：掩码均值（1分）

交付要求：实现 masked_mean 方法，在尊重布尔掩码（boolean mask）的前提下对张量元素求平均。

推荐接口如下：

def masked_mean(
    tensor: torch.Tensor,
    mask: torch.Tensor,
    dim: int | None = None,
) -> torch.Tensor:

功能：沿指定维度计算张量的均值，仅考虑 mask == 1 的元素。

参数说明：

• tensor：待求平均的输入数据张量。
• mask：与 tensor 形状相同的布尔掩码张量，值为1的位置将纳入均值计算。
• dim：计算均值的维度；若为 None，则对所有掩码为1的元素计算全局均值。

返回值：

• 掩码均值张量（torch.Tensor）：形状与 tensor.mean(dim) 的输出语义一致。

测试方式：实现 [adapters.run_masked_mean]，运行命令 uv run pytest -k test_masked_mean 并确保测试通过。

---

GRPO 微批次训练步骤（GRPO Microbatch Train Step）
现在我们可以实现 GRPO 的单个microbatch train step（回想一下，对于一个训练小批次，若 gradient_accumulation_steps > 1，我们会迭代多个microbatch）。

具体而言，给定原始奖励（raw rewards）或优势函数（advantages）及对数概率，我们将计算逐token损失，通过 masked_mean 聚合为每个样本的标量损失，在批次维度上求平均，根据梯度累积步数调整损失，并执行反向传播。

问题（grpo_microbatch_train_step）：微批次训练步骤（3分）

交付要求：实现 GRPO 的单个微批次更新，包括策略梯度损失计算、掩码均值聚合及梯度缩放。

推荐接口如下：

def grpo_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"],
    raw_rewards: torch.Tensor | None = None,
    advantages: torch.Tensor | None = None,
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:

功能：对单个微批次执行前向传播与反向传播。

参数说明：

• policy_log_probs：形状为（batch_size, sequence_length），表示待训练策略的逐token对数概率。
• response_mask：形状为（batch_size, sequence_length），响应token位置标记为1，提示词（prompt）/填充（padding）token位置标记为0。
• gradient_accumulation_steps：每个优化器步骤对应的微批次数量。
• loss_type：损失类型，可选值为 “no_baseline”、“reinforce_with_baseline” 或 “grpo_clip”。
• raw_rewards：当 loss_type == "no_baseline" 时为必填项，形状为（batch_size, 1）。
• advantages：当 loss_type != "no_baseline" 时为必填项，形状为（batch_size, 1）。
• old_log_probs：当 loss_type == "grpo_clip" 时为必填项，形状为（batch_size, sequence_length）。
• cliprange：GRPO-Clip 策略的裁剪参数 ϵ。

返回值：

• 元组（torch.Tensor, dict[str, torch.Tensor]）：
  1. 标量损失张量（loss）：经梯度累积调整后的微批次损失，用于日志记录。
  2. 元数据字典（metadata）：包含底层损失计算的元数据及其他需日志记录的统计信息。

实现提示：

1. 需在该函数中调用 loss.backward()，并确保根据梯度累积步数调整损失。

测试方式：实现 [adapters.run_grpo_microbatch_train_step]，运行命令 uv run pytest -k test_grpo_microbatch_train_step 并确认测试通过。

---

整合所有模块：GRPO 训练循环（Putting it all together: GRPO Train Loop）
现在我们将整合所有模块，实现完整的 GRPO 训练循环。请参考 7.1 节的算法框架，合理调用已实现的方法。

以下提供初始超参数配置：若实现正确，使用该配置应能获得合理结果。

n_grpo_steps: int = 200
learning_rate: float = 1e-5
advantage_eps: float = 1e-6
rollout_batch_size: int = 256
group_size: int = 8
sampling_temperature: float = 1.0
sampling_min_tokens: int = 4  # 参考 Expiter，禁止空字符串响应
sampling_max_tokens: int = 1024
epochs_per_rollout_batch: int = 1  # 在线策略（On-policy）
train_batch_size: int = 256  # 在线策略
gradient_accumulation_steps: int = 128  # 微批次大小为 2，可在 H100 显卡上运行
gpu_memory_utilization: float = 0.85
loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"] = "reinforce_with_baseline"
use_std_normalization: bool = True

optimizer = torch.optim.AdamW(
    policy.parameters(),
    lr=learning_rate,
    weight_decay=0.0,
    betas=(0.9, 0.95),
)

上述默认超参数适用于在线策略（On-policy） 场景：每个滚动批次（rollout batch）仅执行一次梯度更新。此时需满足 train_batch_size == rollout_batch_size 且 epochs_per_rollout_batch == 1。

以下提供一些合理性校验断言及常量定义，可规避部分边界情况并提供实现指引：

assert train_batch_size % gradient_accumulation_steps == 0, "train_batch_size 必须能被 gradient_accumulation_steps 整除"
micro_train_batch_size = train_batch_size // gradient_accumulation_steps

assert rollout_batch_size % group_size == 0, "rollout_batch_size 必须能被 group_size 整除"
n_prompts_per_rollout_batch = rollout_batch_size // group_size

assert train_batch_size >= group_size, "train_batch_size 必须大于或等于 group_size"

n_microbatches_per_rollout_batch = rollout_batch_size // micro_train_batch_size

以下是一些额外的建议：

• 记得使用 r1_zero prompt，并指示 vLLM 在遇到第二个 </answer> 标签时停止生成，如之前实验中所做的那样。
• 建议使用 typer 进行命令行参数解析。
• 使用梯度裁剪（gradient clipping），裁剪值设为 1.0。
• 应定期记录验证奖励（例如每 5 或 10 步记录一次）。在比较超参数时，应至少在 1024 个验证样本上进行评估，因为 CoT/RL 的评估结果可能存在较大噪声。
• 在我们当前的损失实现中，GRPO-Clip 仅应在离策略（off-policy）设置下使用（因为它需要旧的对数概率）。
• 在 off-policy 设置中，若对每个 rollout batch 执行多个 epoch 的梯度更新，每次都重新计算旧的对数概率是低效的。更高效的做法是只计算一次旧的对数概率，并在每个 epoch 中重复使用。
• 不应对旧的对数概率进行梯度求导。
• 在每次优化器更新时，应记录以下部分或全部指标：
  – 损失（loss）
  – 梯度范数（gradient norm）
  – token 熵（token entropy）
  – 裁剪比例（clip fraction），如果是离策略训练
  – 训练奖励（train rewards），包括总奖励、格式奖励和答案奖励
  – 任何你认为对调试有帮助的其他信息

问题（grpo_train_loop）：GRPO训练循环（5分）
交付要求：实现GRPO的完整训练循环。基于MATH数据集启动策略训练，确认验证奖励（validation rewards）逐步提升，且不同阶段的采样结果（rollouts）合理。提供验证奖励随训练步数变化的图表，并附上不同时期的若干采样示例。

---

8 GRPO实验

接下来我们将基于GRPO训练循环开展实验，尝试不同超参数和算法调整。每个实验需占用2块GPU，一块用于vLLM实例，另一块用于策略训练。

关于提前终止运行的说明：若在200步GRPO训练前就发现超参数间存在显著差异（例如某一配置出现发散或明显次优），可提前终止该实验，为后续运行节省时间和计算资源。下文提及的GPU时长为粗略估计。

问题（grpo_learning_rate）：学习率调优（2分）（6个H100 GPU小时）

以上述建议超参数为基准，对学习率进行扫描式调优（sweep），报告最终的验证答案奖励（若优化器发散则注明）。
交付要求：多个学习率对应的验证奖励曲线；在MATH数据集上验证准确率至少达到25%的模型；用2句话简要说明在其他记录指标中观察到的趋势。后续所有实验均采用本次调优中表现最佳的学习率。
笔者的测试，使用"reinforce_with_baseline"模式，和gsm8k数据集中所有的原始训练数据，lr = [1e-5, 1e-3, 1e-4, 5e-5, 3e-5, 2e-5]。其中1e-3, 1e-4发散。最终选择2e-5作为最佳学习率。

---

基线（baselines）的影响
沿用上述超参数（学习率替换为最优值），研究基线的影响。本实验为在策略（on-policy）设置，需对比以下两种损失类型：

• 无基线（no_baseline）
• 带基线的强化学习（reinforce_with_baseline）

注：默认超参数中use_std_normalization（标准差归一化）设为True。

问题（grpo_baselines）：基线的影响（2分）（2个H100 GPU小时）

分别采用reinforce_with_baseline和no_baseline训练策略。
交付要求：每种损失类型对应的验证奖励曲线；用2句话简要说明在其他记录指标中观察到的趋势。

后续几个实验均采用本次实验中表现最佳的损失类型。
答：经过测试，有baseline的算法更好。

---

长度归一化（Length normalization）
如实现masked_mean函数时所提示，对序列长度进行损失平均并非必要，甚至可能不合理。损失的求和方式是一项重要超参数，其选择会影响策略动作的信用分配（credit attribution）方式。

下面结合Lambert（2024）的示例进行说明。观察GRPO训练步骤，首先得到逐令牌（per-token）的策略梯度损失（暂不考虑裁剪）：

advantages # (批量大小, 1)，优势函数值
per_token_probability_ratios # (批量大小, 序列长度)，逐令牌概率比
per_token_loss = -advantages * per_token_probability_ratios # 逐令牌损失

其中优势函数值已沿序列长度维度广播（broadcasted）。以下对比两种逐令牌损失的聚合方式：

• 我们实现的masked_mean：对每个序列中未被掩码（unmasked）的令牌取平均。
• 对每个序列中未被掩码的令牌求和，再除以一个常数标量（masked_normalize函数支持通过设置constant_normalizer≠1.0实现）[Liu et al., 2025, Yu et al., 2025]。

假设批量大小为2，第一个响应包含4个tokens，第二个响应包含7个tokens，通过以下示例可观察不同归一化方式对梯度的影响：

from your_utils import masked_mean, masked_normalize

ratio = torch.tensor([
    [1, 1, 1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 1]
], requires_grad=True)

advs = torch.tensor([
    [2, 2, 2, 2, 2, 2, 2],
    [2, 2, 2, 2, 2, 2, 2]
])

# 掩码（1表示有效令牌，0表示无效）
masks = torch.tensor([
    # 第一个生成序列：4个令牌
    [1, 1, 1, 1, 0, 0, 0],
    # 第二个生成序列：7个令牌
    [1, 1, 1, 1, 1, 1, 1]
])

# 两种归一化方式计算
max_gen_len = 7  # 最大生成长度
masked_mean_result = masked_mean(ratio * advs, masks, dim=1)
masked_normalize_result = masked_normalize(ratio * advs, masks, dim=1, constant_normalizer=max_gen_len)

print("masked_mean结果：", masked_mean_result)
print("masked_normalize结果：", masked_normalize_result)

# masked_mean tensor([2., 2.], grad_fn=<DivBackward0>)
# masked_normalize tensor([1.1429, 2.0000], grad_fn=<DivBackward0>)

masked_mean_result.mean().backward()  # 反向传播计算梯度
print("ratio.grad", ratio.grad)  # 打印梯度
# ratio.grad：
# tensor([[0.2500, 0.2500, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000],
#        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429]])
ratio.grad.zero_()  # 梯度清零

masked_normalize_result.mean().backward() 
print("ratio.grad", ratio.grad)  
# ratio.grad：
# tensor([[0.1429, 0.1429, 0.1429, 0.1429, 0.0000, 0.0000, 0.0000],
#        [0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429]])

问题：长度归一化思考（1分）

交付要求 对比两种方法（暂不执行实验）。分析每种方法的优缺点，以及在哪些特定场景或示例中某一种方法更具优势？

以下为deepseek答：
两种方法的核心区别

1. masked_mean（长度归一化）

• 对每个序列中未被掩码的令牌求平均
• 归一化因子 = 序列的实际有效长度
• 相当于每个令牌的贡献权重 = 1/(序列实际长度)

2. masked_normalize（常数归一化）

• 对每个序列中未被掩码的令牌求和，再除以常数
• 归一化因子 = 预设常数（如最大生成长度）
• 相当于每个令牌的贡献权重 = 1/(预设常数)

优缺点分析
masked_mean的优点：

1. 长度不变的梯度权重：无论生成长度如何，每个有效令牌对最终损失的梯度贡献都相同（示例中梯度都是0.25/0.1429）
2. 公平比较：不同长度的序列在损失计算中获得平等对待
3. 稳定训练：避免长序列主导梯度更新

masked_mean的缺点：

1. 信用分配问题：在强化学习场景中，这可能导致对长序列的过度惩罚
2. 策略偏差：模型可能倾向于生成更短的响应来降低损失
3. 信息丢失：没有考虑绝对长度信息

masked_normalize的优点：

1. 保持绝对信号：保留了序列绝对长度对梯度的影响
2. 更适合信用分配：在强化学习中，长序列的令牌获得的总梯度更大
3. 反映真实影响：较长的序列通常包含更多决策点，理应获得更多关注

masked_normalize的缺点：

1. 长序列主导：训练可能被长序列主导，忽视短序列
2. 梯度不平衡：不同长度序列的梯度规模不一致
3. 需要调参：需要仔细选择归一化常数

适用场景
masked_mean更适用：

• 分类任务：每个令牌都应平等贡献
• 语言建模：目标是准确预测每个令牌
• 短文本生成：避免模型因长度差异产生偏见
• 当响应质量与长度无关时

masked_normalize更适用：

• 强化学习策略优化：需要对长序列中的多个决策点进行适当信用分配
• 对话生成：长响应通常包含更多有价值内容
• 创意写作：鼓励更详细、丰富的输出
• 当响应质量与内容详实度（长度）正相关时

---

接下来，我们将通过实证对比masked_mean（掩码均值）与masked_normalize（掩码归一化）。

问题：长度归一化的影响（2分）（2个H100 GPU小时）

交付要求 通过端到端的GRPO训练，对比基于masked_mean的归一化与masked_normalize的性能。提交验证集答案奖励曲线，对实验结果进行分析（包括其他呈现明显趋势的指标）。

提示：考虑与稳定性相关的指标，例如梯度范数（gradient norm）。 后续实验将采用性能更优的方案。

答：从图中eval/correct曲线数据可知，masked_mean方案更好

---

基于组标准差的归一化
回顾compute_group_normalized_rewards（组归一化奖励计算）的标准实现（基于Shao等人[2024]、DeepSeek-AI等人[2025]的研究），该实现通过组标准差进行归一化。Liu等人[2025]指出，除以组标准差可能会给训练过程引入不必要的偏差：标准差较低的问题（例如过易或过难的问题，其奖励值几乎全为1或全为0）在训练中会获得更高的权重。

Liu等人[2025]提出移除基于标准差的归一化，我们已在compute_group_normalized_rewards中实现该方案，现需进行测试。

问题：标准差归一化的影响（2分）（2个H100 GPU小时）

交付要求 对比use_std_normalization == True（启用标准差归一化）与use_std_normalization == False（禁用标准差归一化）的性能。提交验证集答案奖励曲线，对实验结果进行分析（包括其他呈现明显趋势的指标）。

提示：考虑与稳定性相关的指标，例如梯度范数。

后续实验将采用性能更优的组归一化方案。

答：从图中eval/correct曲线数据可知，use_std_normalization == True方案更好

---

离线策略（Off-policy）与在线策略（On-policy）
目前我们实验所用的超参数均为在线策略（on-policy）：每个采样批次（rollout batch）仅执行一次梯度更新，因此除上述长度归一化和优势归一化的选择外，我们几乎完全采用了策略梯度的“原则性”近似$\hat{g}$。

该方法虽具有理论合理性和稳定性，但效率较低。采样过程需要通过策略模型生成响应，速度较慢，是GRPO训练的主要成本来源；而每个采样批次仅执行一次梯度更新可能不足以显著改变策略行为，显得较为浪费。

接下来我们将测试离线策略（off-policy）训练：每个rollout batch执行多次梯度更新（甚至多个epochs）。

问题：实现离线策略GRPO（Off-policy GRPO）

交付要求 实现离线策略GRPO训练。

根据上述完整GRPO训练循环的实现情况，你可能已具备相关基础架构；若未具备，需完成以下实现：

1. 支持每个rollout batch执行多个轮次（epochs）的梯度更新，其中轮次数量和每个采样批次的优化器更新次数通过rollout_batch_size（采样批次大小）、epochs_per_rollout_batch（每个采样批次的轮次）和train_batch_size（训练批次大小）控制。
2. 修改主训练循环，在每个采样批次生成阶段之后、梯度更新内循环之前，从策略模型中获取响应的对数概率（logprobs），作为old_log_probs（旧对数概率）。建议使用torch.inference_mode()（推理模式）提升效率。
3. 采用“GRPO-Clip”损失类型。

通过调整每个采样批次的轮次数量和优化器更新次数，可控制离线策略的程度。

问题：离线策略GRPO超参数搜索（4分）（12个H100 GPU小时）

交付要求
固定rollout_batch_size = 256（采样批次大小=256），选择epochs_per_rollout_batch（每个采样批次的轮次）和train_batch_size（训练批次大小）的搜索范围。首先进行大范围搜索（GRPO步数<50）以初步了解性能分布，再进行聚焦搜索（GRPO步数=200）。提交简要实验日志，说明所选搜索范围的依据。

与在线策略实验（epochs_per_rollout_batch = 1、train_batch_size = 256）进行对比，提交以验证步数和壁钟时间（wall-clock time）为横轴的对比图。

提交验证集答案奖励曲线，对实验结果进行分析（包括其他呈现明显趋势的指标，如熵（entropy）和响应长度）。对比训练过程中模型响应的熵与EI实验中的观察结果。
提示：需调整gradient_accumulation_steps（梯度累积步数）以保持显存使用量稳定。

答：epochs_per_rollout_batch的搜索范围：[2, 3, 4]; train_batch_size的搜索范围：[32, 64, 128, 256]; n_grpo_steps=200.

从下图可知，当epochs_per_rollout_batch固定为2时，train_batch_size=256曲线的reward最高。当epochs_per_rollout_batch固定为3时，train_batch_size=256曲线的reward最高。当epochs_per_rollout_batch固定为4时，train_batch_size=64曲线的reward值最高，但和前两个曲线相比，是效果非常差的超参数。

选出最好的两个off-policy曲线与on-policy对比。从图中看，epochs_per_rollout_batch=2，train_batch_size=256，是off-poliy中效果最好的超参数。

---

离线策略下的裁剪（Clipping）消融实验

回顾GRPO-Clip中裁剪（clipping）的目的：当对单个采样批次执行多次梯度更新时，防止策略与旧策略（old policy）偏离过大。接下来，我们将在离线策略设置中消融裁剪操作，测试其必要性——即采用逐token（per-token）损失。

问题：离线策略GRPO-Clip消融实验（2分）（2个H100 GPU小时）

交付要求
实现无裁剪的逐token损失，作为新的损失类型“GRPO-No-Clip”（无裁剪GRPO）。使用上一个问题中性能最优的离线策略超参数，运行无裁剪版本的损失函数。提交验证集答案奖励曲线，与GRPO-Clip的实验结果对比分析（包括其他呈现明显趋势的指标，如熵、响应长度和梯度范数）。

答：前文实验中， 最优超参数epochs_per_rollout_batch=2，train_batch_size=256。笔者发现，在这组参数下，clip-fraction曲线都是0，因此选择次优超参数epochs_per_rollout_batch=3，train_batch_size=256。

eval/correct曲线说明，no-clip并不一定就不如clip。不过，这是什么原因呢？问chatgpt，他说：

在离线 GRPO、奖励稀疏且 token-level advantage 集中的设置下， No-Clip 并不一定劣于 Clip。
当 policy 更新本身处于低风险区间时，Clip 可能反而限制了有效学习信号的利用。

---

提示词（Prompt）的影响
作为最后一项消融实验，我们将研究一个有趣的现象：基于模型的预训练方式，强化学习（RL）过程中使用的提示词可能会对模型性能产生显著影响。

替代cs336_alignment/prompts/r1_zero.prompt路径下的R1-Zero提示词，使用cs336_alignment/prompts/question_only.prompt路径下的极简提示词：

{question}  # 仅包含问题本身

该提示词将用于训练和验证阶段，并将奖励函数（训练和验证均使用）修改为cs336_alignment/drgrpo_grader.py中的question_only_reward_fn（仅问题奖励函数）。

问题：提示词消融实验（2分）（2个H100 GPU小时）

交付要求
提交R1-Zero提示词与仅问题提示词的验证集答案奖励曲线。对比两种提示词下的各项指标（包括熵、响应长度和梯度范数等呈现明显趋势的指标），并尝试解释实验结果。
答：从eval曲线来看，仅问题提示词效果效果非常差，但是训练曲线上的效果却很好，很能迷惑人。

9 排行榜：GRPO在MATH数据集上的性能

作为本（必修）作业的最后一部分，你需要在2个H100 GPU、4小时的训练时间内，探索获得最高验证奖励的方法。

model
继续使用Qwen 2.5 Math 1.5B Base模型。

数据集
继续使用集群中/data/a5-alignment/MATH/train.jsonl（训练集）和/data/a5-alignment/MATH/validation.jsonl（验证集）路径下的MATH数据集。禁止使用其他数据，或基于更强模型的推理链进行有监督微调（SFT）等操作。需在完整验证集（共5000个样本）上报告验证准确率，采样超参数如下：温度（temperature）=1.0，最大tokens数（max tokens）=1024。允许对训练集进行筛选或设计数据课程（curriculum）。验证阶段必须使用R1-Zero提示词，并严格使用初始代码中提供的r1_zero_reward_fn奖励函数计算验证准确率（训练阶段可自行设计其他奖励函数）。

算法
可自由调整超参数或完全更改训练算法，但禁止使用额外数据或其他模型（允许使用多个模型副本）。

系统优化
你可能会发现，在我们的简易GRPO实现中，至少有一个GPU始终处于空闲状态。通过优化流水线的系统特性（例如采样或训练采用低精度、使用torch.compile编译、其他系统级优化），可能会显著提升性能。你无需将vLLM（大语言模型推理框架）局限于单个设备、训练策略局限于另一个设备，鼓励探索更优的并行化方案。

参考思路
有关可能的改进方向，可参考以下仓库：

• veRL
• trl（Transformer Reinforcement Learning）
• torchtune（PyTorch微调工具）
• oat

关于KL散度
需注意，上述实验中未包含相对于参考模型（通常是冻结的SFT模型或预训练模型 checkpoint）的KL散度项。根据我们的实验及文献[Liu等人, 2025, NTT123, 2025]的结果，省略KL散度项不会影响性能，且能节省显存（无需存储参考模型）。然而，许多GRPO仓库默认包含KL散度项，你可基于Qwen 2.5 Math 1.5B Base模型或通过你的算法得到的模型，尝试KL散度或其他形式的正则化。

问题：排行榜任务（16分）（16个H100 GPU小时）

交付要求
提交在2个H100 GPU、4小时训练时间内获得的验证准确率，以及验证准确率随壁钟时间变化的截图（横轴终点≤4小时）。请遵守以下评估约束：

1. 验证准确率为完整MATH验证集（5000个样本）的平均准确率。
2. 验证阶段必须使用R1-Zero提示词。
3. 评估时必须使用vLLM，设置温度=1.0、最大tokens数=1024。
4. 验证准确率通过初始代码中提供的r1_zero_reward_fn奖励函数计算的答案奖励平均值得到。

10 结语

恭喜完成本课程的最后一项作业！你应当为自己的努力感到自豪。希望通过从零构建现代语言模型的核心组件，你已深入理解其底层原理，并享受这一学习过程。

---