cs336作业五文档
这是一份名为 CS336 作业 5 (Alignment):对齐与推理强化学习 (RL) 的课程作业文档翻译。
以下是文档内容的详细中文翻译,保持了原文档的结构、代码块占位和专业术语。
[页 1]
CS336 作业 5 (Alignment):对齐与推理 RL
1. 作业概览
- 版本:1.0.0
- CS336 教研组
- 2025 年春季
在本作业中,通过训练语言模型解决数学问题,你将获得让模型进行推理 (Reasoning) 的实战经验。
你将实现的内容:
- Zero-shot prompting baseline(零样本提示基线) :针对 Hendrycks et al. [2021] 的 MATH 竞赛数学数据集。
- Supervised Finetuning (SFT,监督微调) :基于来自更强推理模型(DeepSeek R1, DeepSeek-AI et al. 2025)的推理轨迹数据。
- Expert Iteration (专家迭代) :利用已验证的奖励来提升推理性能。
- Group-Relative Policy Optimization (GRPO,组相对策略优化) :利用已验证的奖励来提升推理性能。
对于感兴趣的同学,我们在接下来的几天内还将发布一个完全可选的部分,内容是关于将语言模型与人类偏好进行对齐。
你将运行的内容:
- 测量 Qwen 2.5 Math 1.5B 的零样本提示性能(作为基线)。
- 使用 R1 的推理轨迹在 Qwen 2.5 Math 1.5B 上运行 SFT 。
- 在 Qwen 2.5 Math 1.5B 上运行带验证奖励的 Expert Iteration 。
- 在 Qwen 2.5 Math 1.5B 上运行带验证奖励的 GRPO 。
代码概览:
所有作业代码及本文档均可在 GitHub 上找到:
github.com/stanford-cs336/assignment5-alignment
请 git clone 该仓库。如果有更新,我们会通知你,你可以通过 git pull 获取最新版本。
cs336_alignment/*:这是你编写作业 5 代码的地方。注意这里除了少量起始代码外没有其他代码,你需要从头开始编写。
[页 2]
cs336_alignment/prompts/*:为了方便起见,我们提供了包含提示词(prompts)的文本文件,以最大限度减少从 PDF 复制粘贴提示词可能导致的错误。tests/*.py:包含你必须通过的所有测试。- 注意 :你只需要通过
tests/test_sft.py和tests/test_grpo.py中的测试——其余测试用于作业的非强制部分。 - 这些测试会调用在
tests/adapters.py中定义的钩子(hooks)。你需要实现这些适配器以将你的代码连接到测试。 - 编写更多测试或修改测试代码有助于调试,但你的实现必须通过原始提供的测试套件。
- 注意 :你只需要通过
README.md:包含设置环境的基本说明。
你可以使用的工具:
我们希望你从头构建大部分 RL 相关组件。
- 你可以使用像 vLLM 这样的工具从语言模型生成文本(见 §3.1)。
- 此外,你可以使用 HuggingFace Transformers 加载 Qwen 2.5 Math 1.5B 模型和分词器并运行前向传播(§4.1),但不能使用任何训练实用程序(例如
Trainer类)。
提交方式:
你需要向 Gradescope 提交以下文件:
writeup.pdf:回答所有书面问题。请对你的回复进行排版。code.zip:包含你编写的所有代码。
2. 语言模型推理 (Reasoning with Language Models)
2.1 动机
语言模型的一个显著用例是构建能够处理广泛自然语言处理任务的通用系统。在本作业中,我们将聚焦于语言模型的一个发展中的用例: 数学推理 (Mathematical Reasoning) 。这将作为我们的测试平台,用于建立评估、执行监督微调 (SFT),并尝试使用强化学习 (RL) 教导 LM 进行推理。
与之前的作业相比,这里有两点不同:
- 不使用之前的代码库和模型 :我们理想情况下想使用之前作业训练的基础模型,但微调那些模型不会得到满意的结果——它们太弱了,无法展示非平凡的数学推理能力。因此,我们将切换到一个我们可以访问的现代高性能语言模型 ( Qwen 2.5 Math 1.5B Base ),并在该模型之上完成大部分工作。
- 引入新的基准测试 :到目前为止,我们一直认为交叉熵 (Cross-entropy) 是许多下游任务的良好替代指标。然而,本作业的重点是弥合基础模型与下游任务之间的差距,因此我们必须使用独立于交叉熵的评估。我们将使用 Hendrycks et al. [2021] 的 MATH 12K 数据集,其中包含具有挑战性的高中竞赛数学问题。我们将通过将模型输出与参考答案进行比较来评估模型。
2.2 思维链推理与推理 RL (Chain-of-Thought Reasoning and Reasoning RL)
语言模型最近一个令人兴奋的趋势是使用思维链 (Chain-of-Thought, CoT) 推理来提高各种任务的性能。思维链是指逐步推理问题的过程,在得出最终答案之前生成中间推理步骤。
[页 3]
- LLM 的思维链推理 :早期的 CoT 方法微调语言模型以解决简单的数学任务(如算术),通过使用“草稿纸 (scratchpad)”将问题分解为中间步骤 [Nye et al., 2021]。其他工作提示强模型在回答前“一步一步思考 (think step by step)”,发现这显著提高了小学数学等任务的性能 [Wei et al., 2023]。
- 通过专家迭代 (Expert Iteration) 学习推理 :Self-Taught Reasoner (STaR) [Zelikman et al., 2022] 将推理框架化为一个自举循环:预训练模型首先采样多样化的思维链 (CoTs),只保留那些导致正确答案的链,然后在这些“专家”轨迹上进行微调。迭代此循环可以提高 LM 的推理能力和解答率。STaR 证明了这种使用基于字符串匹配的自动验证的专家迭代 [Anthony et al., 2017] 可以在没有人类编写的推理轨迹的情况下自举推理技能。
- 基于验证奖励的推理 RL、o1 和 R1 :最近的工作探索了使用更强大的强化学习算法配合验证奖励来提高推理性能。OpenAI 的 o1 (及后续的 o3/o4)、DeepSeek 的 R1 和 Moonshot 的 kimi k1.5 使用策略梯度方法在可以通过字符串匹配或单元测试验证正确性的数学和代码任务上进行训练,展示了竞赛数学和编码性能的显著提升。后续工作如 Open-R1 、SimpleRL-Zoo 和 TinyZero 证实,即使在 1.5B 参数这样小的模型上,纯粹的基于验证奖励的 RL 也能提高推理性能。
我们的设置:模型和数据集
在接下来的部分中,我们将考虑逐步复杂的方法来训练基础语言模型进行逐步推理以解决数学问题。
- 模型 :我们将使用 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 et al., 2025]:基于英国电视节目 Countdown 的简单合成任务,是小规模推理 RL 的流行测试平台。
- GSM8K [Cobbe et al., 2021a]:小学数学问题,比 MATH 简单,但适合调试正确性和熟悉推理 RL 流程。
- Tulu 3 SFT Math :使用 GPT-4o 和 Claude 3.5 Sonnet 生成的合成数学问题。由于是合成的,部分答案可能不完全正确。
- 其他链接的数学 SFT 数据集。
为了获得简短的真实标签(如 $1/2$),你可以使用如 Math-Verify 之类的数学答案解析器处理 ground-truth 列。
3. 测量 Zero-Shot MATH 性能
我们首先测量基础语言模型在 MATH 的 5K 示例测试集上的性能。建立这个基线对于理解后续每种方法如何影响模型行为非常有用。
[页 4]
除非另有说明,对于 MATH 上的实验,我们将使用来自 DeepSeek R1-Zero 模型 [DeepSeek-AI et al., 2025] 的以下提示词。我们将此称为 r1_zero 提示词:
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>.User:
Assistant:
<think>
该 r1_zero 提示词位于文本文件 cs336_alignment/prompts/r1_zero.prompt 中。
在提示词中,question 指的是我们插入的问题。预期模型扮演助手的角色,并开始生成思考过程(因为我们已经包含了左侧的 <think> 标签),用 </think> 结束思考过程,然后在 answer 标签内生成最终的符号答案,如 <answer> 4x + 10 </answer>。让模型生成 <answer> 标签的目的是便于我们解析模型的输出并将其与真实答案进行比较,并且我们可以在看到右侧答案标签 </answer> 时停止生成。
关于提示词选择的说明
事实证明,
r1_zero提示词并不是 RL 后最大化下游性能的最佳选择,因为提示词与 Qwen 2.5 Math 1.5B 模型的预训练方式不匹配。Liu et al. [2025] 发现,仅用问题提示模型(不做其他任何提示)一开始就有很高的准确率,例如在 100 步 RL 后就能匹配r1_zero提示词的效果。这表明该模型已经在类似的问题-答案对上进行了预训练。尽管如此,我们在本作业中选择r1_zero提示词,因为使用它进行 RL 可以在短时间内显示出明显的准确率提升,使我们能够快速完成 RL 机制的演练并进行正确性检查。作为现实检查,你稍后会在作业中直接与question_only提示词进行比较。
3.1 使用 vLLM 进行离线语言模型推理
为了评估语言模型,我们需要为各种提示生成续写(响应)。为了高效实现 RL,我们需要高性能的推理技术。因此,本作业推荐使用 vLLM 进行离线批量推理。vLLM 是一个高吞吐量且内存高效的推理引擎(支持 CUDA 内核优化、PagedAttention 等)。
使用 vLLM 生成的示例代码:
Python
1 | from vllm import LLM, SamplingParams |
[页 5]
在 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 - Llama 3.3 70B Instruct (可选):
/data/a5-alignment/models/Llama-3.3-70B-Instruct
3.2 Zero-shot MATH 基线
提示设置 :加载 MATH 验证集示例,使用 r1_zero 提示词提示模型回答。
评估指标 :数学问题的评估很棘手,因为模型可能输出 <answer> 1/2 </answer> 而真实答案是 0.5。我们需要一个答案解析函数。
我们将使用最近推理 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。指示 vLLM 在模型输出 </answer> 时停止:
Python
1 | sampling_params.stop = ["</answer>"] |
[页 6]
问题 (math_baseline): 4 分
(a) 编写一个脚本来评估 Qwen 2.5 Math 1.5B 在 MATH 上的 zero-shot 性能。
该脚本应:
- 从
/data/a5-alignment/MATH/validation.jsonl加载验证示例。 - 使用
r1_zero提示词将其格式化为字符串。 - 为每个示例生成输出。
- 计算评估指标。
- 将示例、模型生成和相应的评分序列化到磁盘以供后续分析。
建议实现一个 evaluate_vllm 方法(见文档中的函数签名)。
交付物 :评估基线 zero-shot MATH 性能的脚本。
(b) 在 Qwen 2.5 Math 1.5B 上运行评估脚本。
统计落入以下类别的模型生成数量:
- 正确(格式和答案奖励均为 1)。
- 格式奖励为 1 但答案奖励为 0。
- 格式奖励为 0 且答案奖励为 0。
观察至少 10 个格式奖励为 0 的案例,你认为是基础模型输出的问题还是解析器的问题?为什么?对于至少 10 个格式正确但答案错误的案例呢?
交付物 :关于模型和奖励函数性能的评论,包括每种类别的示例。
© Qwen 2.5 Math 1.5B zero-shot 基线在 MATH 上的表现如何?
交付物 :1-2 句话说明评估指标。
4. MATH 的监督微调 (Supervised Finetuning)
算法 1:监督微调 (SFT)
- 输入:初始策略模型 $\pi_{\theta_{init}}$;SFT 数据集 $\mathcal{D}$
- 循环步骤:采样批次 -> 计算交叉熵损失 -> 梯度更新
[页 7]
在本节中,我们将在 MATH 数据集上微调基础模型。由于目标是提高推理能力,我们将微调模型先生成思维链推理轨迹,然后再生成答案。我们提供了一个数据集,包含从 DeepSeek R1 [DeepSeek-AI et al., 2025] 获取的推理轨迹,位于 /data/a5-alignment/MATH/sft.jsonl。
注意:在实践中,SFT 通常作为 RL 微调的热启动。但由于我们的模型较小,叠加 SFT 和 RL 效果不明显,因此本作业将这两个阶段分开处理。
4.1 使用 Hugging Face 模型
加载模型和分词器 :
使用 AutoModelForCausalLM 和 AutoTokenizer。建议使用 bfloat16 和 flash_attention_2。代码示例见文档。
前向传播 :
运行前向传播获取 logits,然后计算 loss。
Python
1 | logits = model(input_ids).logits |
[页 8]
保存模型 :
使用 .save_pretrained()。请保存到 /data/yourusername 目录下。
梯度累积 (Gradient Accumulation) :
由于显存限制,使用梯度累积来支持更大的有效批次大小。
- 每 $k$ 步才调用
optimizer.step()和optimizer.zero_grad()。 - 在
loss.backward()之前将 loss 除以gradient_accumulation_steps。
4.2 SFT 辅助方法
分词提示词和输出
我们需要将问题 $q$ 和目标输出 $o$ 分别分词并拼接。还需要构建一个 response_mask(响应掩码),对于响应部分的 token 为 True,其他为 False。
[页 9]
问题 (tokenize_prompt_and_output): 提示词和输出分词 (2 分)
交付物 :实现 tokenize_prompt_and_output 方法。
- 输入:
prompt_strs,output_strs,tokenizer。 - 返回:
input_ids,labels(shifted input ids),response_mask。 - 测试 :实现
adapters.run_tokenize_prompt_and_output并运行pytest -k test_tokenize_prompt_and_output。
记录逐词熵 (Per-token entropies)
在 RL 中,跟踪熵以观察模型置信度变化很有用。熵定义为 $H(p) = -\sum p(x) \log p(x)$。
[页 10]
问题 (compute_entropy): 逐词熵 (1 分)
交付物 :实现 compute_entropy 方法。
- 输入:
logits(unnormalized)。 - 返回:每个 next-token 预测的熵。
- 提示:使用数值稳定的方法(如
logsumexp)。 - 测试 :
pytest -k test_compute_entropy。
从模型获取对数概率
对于 RL 和 SFT,我们需要获得 token 的对数概率:$\log p_\theta(y|x) = \log [\text{softmax}(f_\theta(x))]_y$。
问题 (get_response_log_probs): 响应对数概率 (2 分)
交付物 :实现 get_response_log_probs 方法。
- 输入:
model,input_ids,labels,return_token_entropy。 - 返回:条件对数概率
log_probs和可选的token_entropy。 - 测试 :
pytest -k test_get_response_log_probs。
[页 11]
SFT 微批次训练步
我们需要计算目标输出的负对数似然损失。首先实现一个带掩码的归一化函数。
问题 (masked_normalize): 带掩码归一化 (1 分)
交付物 :实现 masked_normalize 方法。
- 功能:在遵守掩码的情况下对 Tensor 元素求和并除以常数。
- 测试 :
pytest -k test_masked_normalize。
[页 12]
问题 (sft_microbatch_train_step): 微批次训练步 (3 分)
交付物 :实现 SFT 的单次微批次更新,包括交叉熵损失、带掩码求和及梯度缩放。
- 输入:
policy_log_probs,response_mask,gradient_accumulation_steps等。 - 返回:
loss(scalar),metadata(dict)。 - 提示:需在此函数中调用
loss.backward()。 - 测试 :
pytest -k test_sft_microbatch_train_step。
[页 13]
记录生成结果 (Logging generations)
编写 log_generations 函数,记录输入提示、模型响应、真实答案、奖励信息、平均熵、响应长度等。
问题 (log_generations): 记录生成 (1 分)
交付物 :实现 log_generations 函数。
4.3 SFT 实验
现在实现完整的 SFT 流程(算法 1)。
- 数据集:
/data/a5-alignment/MATH/sft.jsonl。 - 设置:使用 2 个 GPU。一个用于 Policy 模型训练,另一个运行 vLLM 进行评估。
- 代码提供了
init_vllm和load_policy_into_vllm_instance的辅助函数。
[页 14]
- 建议使用
wandb记录指标(区分train_step和eval_step)。 - 使用梯度裁剪(clip value 1.0)。
问题 (sft_experiment): 在 MATH 数据集上运行 SFT (2 分)
- 不同数据量 :在 reasoning SFT 示例上运行 SFT,变化唯一示例数量 {128, 256, 512, 1024} 及全量数据集。调整学习率和 Batch Size 以在使用全量数据时达到至少 15% 的验证准确率。
- 交付物 :不同数据集大小对应的验证准确率曲线。
- 数据过滤 :过滤 SFT 示例,仅包含产生正确答案的示例。在过滤后的数据集上运行 SFT。
- 交付物 :报告过滤后数据集的大小及验证准确率曲线。比较与之前实验的结果。
[页 15]
5. MATH 的专家迭代 (Expert Iteration)
上一节我们发现通过过滤 SFT 数据中的坏样本可以提升模型。本节我们将更进一步:对基础模型自身生成的推理轨迹应用此过滤过程。这就是专家迭代 (Expert Iteration, EI)。
算法 2:专家迭代 (EI)
- 采样一批问题。
- 对每个问题,从当前策略 $\pi_{\theta_{old}}$ 采样 $G$ 个输出。
- 计算奖励,过滤出奖励为 1 的正确输出,构成数据集 $\mathcal{D}_{sft}$。
- 在 $\mathcal{D}_{sft}$ 上运行 SFT 更新模型。
提示 :
- vLLM
SamplingParams设置min_tokens=4以防空字符串。 - 使用梯度裁剪。
问题 (expert_iteration_experiment): 运行专家迭代 (2 分)
在 MATH 数据集 (/data/a5-alignment/MATH/train.jsonl) 上运行 EI。
- 变化每个问题的 rollout 数量 $G$ 和 SFT 步骤的 epoch 数。
- 使用 $n_{ei_steps} = 5$。
- 交付物 :
- 不同配置的验证准确率曲线。
- 一个在 MATH 上达到至少 15% 验证准确率的模型。
- 简短讨论与 SFT 的比较及 EI 步骤间的性能变化。
- 模型响应熵随训练变化的图表。
[页 16]
6. 策略梯度 (Policy Gradients) 入门
本节简要介绍用于语言模型的策略梯度。DeepSeek R1 等强模型就是使用策略梯度训练的。
6.1 语言模型即策略
LM 定义了给定前缀 $s_t$ 下下一个 token $a_t$ 的概率分布。LM 是一个分类随机策略:$\pi_\theta(a_t|s_t) = [\text{softmax}(f_\theta(s_t))]_{a_t}$。
6.2 轨迹 (Trajectories)
轨迹 $\tau = (s_0, a_0, …, s_T, a_T)$。在 LLM 中,环境是确定性的:$s_{t+1} = s_t || a_t$。
[页 17]
6.3 奖励与回报 (Rewards and Return)
对于验证领域,如果轨迹匹配真实答案则奖励为 1,否则为 0。我们关注非折现回报 (undiscounted return)。
目标:最大化期望回报 $J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[R(\tau)]$。
6.4 原始策略梯度 (Vanilla Policy Gradient)
REINFORCE 策略梯度公式:
$$
\nabla_\theta J(\pi_\theta) = \mathbb{E}{\tau \sim \pi\theta} [\sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t|s_t) R(\tau)]
$$
(文档中提供了推导过程,包括对数导数技巧)。
[页 18]
梯度的样本估计 :
$$
\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)})
$$
6.5 策略梯度基线 (Baselines)
为了减少方差,我们从奖励中减去一个基线 $b(s_t)$。只要基线仅依赖于状态,梯度估计就是无偏的。
关于“策略梯度损失 (pg_loss)”的说明 :
在 PyTorch 中,我们定义 pg_loss 是为了通过反向传播获得梯度。它本身不是一个有意义的评估指标。在 RL 中,应始终记录和报告 训练和验证奖励 。
6.6 异策略 (Off-Policy) 策略梯度
REINFORCE 是同策略 (on-policy) 的,效率低。
异策略方法(如 PPO, GRPO)使用旧策略 $\pi_{\theta_{old}}$ 采样的轨迹来优化当前策略 $\pi_\theta$。这引入了重要性采样权重。
[页 20]
7. 组相对策略优化 (Group Relative Policy Optimization, GRPO)
7.1 GRPO 算法
优势估计 (Advantage Estimation)
GRPO 的核心思想是对每个问题采样多个输出,并利用它们计算基线。这避免了训练价值函数 $V_\phi(s)$ 的需要。
DeepSeekMath 和 R1 计算组归一化优势:
$$
A^{(i)} = \frac{r^{(i)} - \text{mean}({r}) }{\text{std}({r}) + \epsilon}
$$
Dr. GRPO 提出了一种简化版本,仅减去均值而不除以标准差,我们稍后会比较这两种方法。
GRPO 目标函数
结合了三个思想:
- 异策略策略梯度。
- 组归一化优势。
- 裁剪机制 (Clipping),类似于 PPO。
[页 21]
算法 3:GRPO
(表格形式展示了 GRPO 的训练循环:采样问题 -> 采样 G 个输出 -> 计算奖励和优势 -> 内部循环更新策略)。
GRPO-Clip 目标函数 :
$$
\text{per-token objective} = \min \left( \frac{\pi}{\pi_{old}} A, \text{clip}\left(\frac{\pi}{\pi_{old}}, 1-\epsilon, 1+\epsilon\right) A \right)
$$
(文档解释了裁剪如何防止策略变化过大)。
7.2 实现
计算优势(组归一化奖励)
[页 22]
问题 (compute_group_normalized_rewards): 组归一化 (2 分)
交付物 :实现 compute_group_normalized_rewards。
- 输入:
reward_fn,rollout_responses,group_size等。 - 功能:计算原始奖励,并在组内进行归一化(支持除以标准差或不除)。
- 测试 :
pytest -k test_compute_group_normalized_rewards。
朴素策略梯度损失
[页 23]
问题 (compute_naive_policy_gradient_loss): 朴素策略梯度 (1 分)
交付物 :实现 compute_naive_policy_gradient_loss。
- 公式:$-A_t \cdot \log p_\theta(o_t | q, o_{<t})$。
- 测试 :
pytest -k test_compute_naive_policy_gradient_loss。
GRPO-Clip 损失
问题 (compute_grpo_clip_loss): GRPO-Clip 损失 (2 分)
交付物 :实现 compute_grpo_clip_loss。
- 实现带裁剪的目标函数。
- 测试 :
pytest -k test_compute_grpo_clip_loss。
$$
- \min \left( \frac{\pi_\theta(o_t|q, o_{<t})}{\pi_{\theta_{old}}(o_t|q, o_{<t})} A_t, \text{clip} \left( \frac{\pi_\theta(o_t|q, o_{<t})}{\pi_{\theta_{old}}(o_t|q, o_{<t})}, 1-\epsilon, 1+\epsilon \right) A_t \right)
$$
[页 24]
策略梯度损失包装器 (Wrapper)
我们需要对比三种版本:
(a) no_baseline: 原始奖励作为优势。
(b) reinforce_with_baseline: 组归一化奖励作为优势。
© grpo_clip: GRPO-Clip 损失。
问题 (compute_policy_gradient_loss): 策略梯度包装器 (1 分)
交付物 :实现 compute_policy_gradient_loss,根据 loss_type 分发调用。
[页 25]
带掩码的均值 (Masked mean)
通常我们将 loss 在序列维度求均值。
问题 (masked_mean): 带掩码均值 (1 分)
交付物 :实现 masked_mean。
[页 26]
GRPO 微批次训练步
问题 (grpo_microbatch_train_step): 微批次训练步 (3 分)
交付物 :实现 GRPO 的单次微批次更新。
- 调用
compute_policy_gradient_loss。 - 使用
masked_mean聚合 loss。 - 处理梯度累积。
- 测试 :
pytest -k test_grpo_microbatch_train_step。
[页 27]
GRPO 训练循环
提供了一些起始超参数(LR=1e-5, group_size=8, etc.)。
注意:如果是 Off-policy 设置,需要计算一次 old_log_probs 并在内层循环复用(不要对 old log probs 求导)。
[页 28]
问题 (grpo_train_loop): GRPO 训练循环 (5 分)
交付物 :实现完整的 GRPO 训练循环。
- 开始训练并确认验证奖励提升。
- 提供验证奖励随步骤变化的图表和一些 rollout 示例。
[页 29]
8. GRPO 实验
注意 :如果发现超参数明显不好,可提前停止。
问题 (grpo_learning_rate): 调整学习率 (2 分)
- 扫描学习率,报告最终验证答案奖励。
- 交付物 :
- 验证奖励曲线。
- 一个达到至少 25% 验证准确率的模型。
- 简短讨论。
基线的影响
比较 no_baseline 和 reinforce_with_baseline(On-policy 设置)。
问题 (grpo_baselines): 基线的影响 (2 分)
- 交付物 :验证奖励曲线及简短讨论。
[页 30]
长度归一化 (Length normalization)
比较两种聚合 loss 的方式:
masked_mean:对序列中未掩码的 token 求平均。masked_normalize:求和后除以常数(如最大生成长度)。
(文档提供了一个代码示例展示梯度差异)。
问题 (think_about_length_normalization): 思考长度归一化 (1 分)
交付物 :比较两种方法的优缺点(无需运行实验)。
[页 31]
问题 (grpo_length_normalization): 长度归一化的影响 (2 分)
交付物 :通过端到端 GRPO 训练比较两种归一化方法。报告曲线并评论稳定性(如梯度范数)。
组标准差归一化
比较是否在优势计算中除以组标准差。
问题 (grpo_group_standard_deviation): 标准差归一化的影响 (2 分)
交付物 :比较 use_std_normalization 为 True 和 False 的情况。
[页 32]
Off-policy vs On-policy
On-policy 效率低。我们将实验 Off-policy(每个 rollout batch 进行多次梯度更新)。
问题 (grpo_off_policy): 实现 Off-policy GRPO
交付物 :修改代码以支持多 epoch 更新,并正确计算和使用 old_log_probs。使用 grpo_clip 损失。
问题 (grpo_off_policy_sweep): Off-policy GRPO 超参数扫描 (4 分)
交付物 :
- 固定
rollout_batch_size=256,扫描epochs_per_rollout_batch和train_batch_size。 - 与 On-policy 运行进行比较(验证步骤 vs 挂钟时间)。
- 比较熵的变化。
裁剪 (Clipping) 的消融实验
问题 (grpo_off_policy_clip_ablation): Off-policy GRPO-Clip 消融 (2 分)
交付物 :实现无裁剪版本 (GRPO-No-Clip) 并在最佳 Off-policy 设置下运行。比较结果。
[页 33]
提示词的影响 (Effect of prompt)
比较 r1_zero 提示词与简单的 question_only 提示词。
问题 (grpo_prompt_ablation): 提示词消融 (2 分)
交付物 :报告验证奖励曲线。解释为什么简单的提示词可能表现更好(或不同)。
Leaderboard: MATH 上的 GRPO
问题 (leaderboard): 排行榜 (16 分)
交付物 :在 2 张 H100 GPU 上训练 4 小时内 获得的最高验证准确率。
- 提供随时间变化的验证准确率截图。
- 约束 :
- 使用 Qwen 2.5 Math 1.5B Base。
- 仅使用 MATH 数据集。
- 验证时必须使用
r1_zero提示词和提供的r1_zero_reward_fn。 - 验证参数:Temp 1.0, Max tokens 1024。
系统优化提示 :你可以优化并行方式(vLLM 和 Policy 并行),降低精度,使用 torch.compile 等。
[页 34]
10. 结语
恭喜完成这门课的最后一次作业!
参考文献
[列出了相关的论文,包括 Hendrycks et al., DeepSeek R1, STaR, OpenAI o1, Kimi k1.5, PPO 等]
