A5（Alignment and Reasoning RL）

围绕 “让小模型学会做数学推理”，从 零样本提示 → SFT（用强模型推理轨迹） → Expert Iteration（自生成+验证筛选） → GRPO（带verified reward的策略梯度） 一步步搭起完整的 reasoning post-training 流水线。

下面按“它到底做哪些任务”拆解。

0) 总目标与学生要交付什么

总体目标：在 MATH 数据集上提升 Qwen2.5-Math-1.5B Base 的解题正确率（用“答案解析器/验证奖励”判对错），并理解不同对齐/推理训练方法的差别。

必须实现的四件事（A5 Overview）：

1. Zero-shot prompting baseline（在 MATH 上测零样本）
2. SFT：用更强 reasoning 模型（DeepSeek R1）的推理轨迹做监督微调
3. Expert Iteration（EI / STaR风格）：模型自己生成 CoT → 用 verified reward 筛出对的 → 再 SFT，迭代提升
4. GRPO：Group-Relative Policy Optimization（带 verified rewards 的 policy gradient）进一步提升

1) Baseline：零样本 MATH 测评流水线（math_baseline，4分）

这一块要求你搭建评测框架：prompt → 批量生成 → 答案解析 → 统计并落盘。

核心任务：

• 读取 MATH validation 集
• 用指定 prompt（r1_zero，要求输出 <think>...</think><answer>...</answer>）构造 prompts
• 用 vLLM 做批量推理，并在遇到 </answer> 时停止生成
• 用提供的 reward/answer parser 评判（format reward、answer reward、total reward）
• 把 prompts、生成、评分结果序列化保存，方便后续分析
• 统计三类样本数量：
  1. format=1 & answer=1
  2. format=1 & answer=0
  3. format=0 & answer=0，并做案例分析（是模型输出问题还是 parser 问题）

能力点：工程化评测（可复用评测函数 evaluate_vllm）、可审计日志、错误归因。

2) SFT：推理轨迹监督微调（SFT primitives + sft_experiment）

不让用 HF Trainer，要自己写训练循环；但允许用 Transformers 载入模型、做 forward。

2.1 SFT 训练所需的“基础组件实现”（多道小题）

这一段要求你写一系列训练与统计的基础函数（并配测试）：

• tokenize_prompt_and_output（2分）：prompt 和 response 分别 tokenize 再拼接，构造 response_mask，并返回 input_ids/labels/response_mask（shifted labels）
• compute_entropy（1分）：计算每个位置的 next-token 分布熵（数值稳定）
• get_response_log_probs（2分）：从模型 logits 得到每 token 的 log-prob（可选返回 entropy）
• masked_normalize（1分）：带 mask 的求和/归一化 primitive
• sft_microbatch_train_step（3分）：SFT 的 microbatch 训练步：用 mask 只在 response tokens 上算 NLL、处理 gradient accumulation、并 loss.backward()
• log_generations（1分）：训练中定期生成并记录：prompt、response、GT、reward、平均 entropy、长度统计（正确/错误长度对比）

2.2 sft_experiment（2分，算力昂贵）

• 在给定的 reasoning SFT 数据（R1 轨迹）上训练 Qwen2.5-Math-1.5B
• 做数据规模对比：{128,256,512,1024} + full，画验证准确率曲线；调 lr/batch 让 full 达到 ≥15% 验证准确率
• 过滤 SFT 数据：只保留“答案正确”的轨迹再训练，报告过滤后规模与曲线，并比较差异

能力点：mask 训练、梯度累积、训练中可视化与样例日志；理解“坏轨迹会害人”，过滤能显著改善。

3) Expert Iteration（2分）：自生成 → 验证筛选 → SFT 迭代

这部分实现并运行 Expert Iteration (EI)（类似 STaR）：

• 每轮：从 train 集采样问题 Db
• 用当前（旧）策略对每题采样 G 条输出
• 用 verified reward 判对错，只保留正确输出组成新的 SFT 数据
• 对这些“自举的正确轨迹”做 SFT，重复 \(n_{ei}=5\) 轮

运行要求（写在 experiment 题里）：

• 变动超参：rollout 数 G、SFT epochs、Db batch size ∈ {512,1024,2048}（不必全组合）
• 记录并绘制：验证准确率曲线；并且记录 response entropy 随训练变化
• 目标：最终模型验证准确率 ≥15%

能力点：把“验证奖励”当作自动标签器，实现自举；理解 rollout 多样性与筛选率的 trade-off。

4) GRPO：实现策略梯度训练 + 大量消融实验（核心）

A5 最重的部分是 GRPO：既要实现关键函数，又要跑一系列实验对比。

4.1 GRPO 关键实现题（从 advantage 到 loss 到 train loop）

• compute_group_normalized_rewards（2分）：对每题的 G 个 rollouts 计算 raw rewards，并做组内归一化：
  • 标准版：减均值再除以 std + eps（Eq.28）
  • 简化版：只减均值（Eq.31），作为可切换 ablation
    • compute_naive_policy_gradient_loss（1分）：朴素 PG loss：\(-A \cdot \log p\)（逐 token）
• compute_grpo_clip_loss（2分）：实现 GRPO-Clip 的 per-token loss（含概率比率、clip、统计 clip fraction 等）
• compute_policy_gradient_loss（1分）：wrapper，支持三种 loss：
  1. no_baseline（raw reward）
  2. reinforce_with_baseline（组归一化 advantage）
  3. grpo_clip（需要 old_log_probs 与 cliprange）
     • masked_mean（1分）：对 response_mask 做 masked mean，用于把 per-token loss 聚合到每条样本
• grpo_microbatch_train_step（3分）：GRPO microbatch 训练步：计算 loss → masked_mean → batch mean → 梯度累积缩放 → backward，并返回日志信息
• grpo_train_loop（5分）：把整套 GRPO 算法跑起来（按 Algorithm 3），要求观察到验证 reward 上升并展示例子 rollouts

4.2 GRPO 实验题（大量消融 + 目标更高）

• 学习率 sweep（2分）：对多种 lr 画曲线，目标：验证准确率 ≥25%
• baseline 消融（2分）：no_baseline vs reinforce_with_baseline 的对比
• 长度归一化讨论（1分）：masked_mean vs masked_normalize 的利弊分析（不跑实验先讲）
• 长度归一化实验（2分）：masked_mean vs masked_normalize 的端到端对比
• std 归一化消融（2分）：use_std_normalization True vs False（Eq.28 vs Eq.31）
• off-policy GRPO（实现）：多 epochs / 多 train steps per rollout batch + 使用 GRPO-Clip + old_log_probs 缓存
• off-policy sweep（4分）：对 epochs_per_rollout_batch 与 train_batch_size 做 sweep，并同时按 validation step 和 wall-clock time 作图，对比 on-policy；观察 entropy/长度等趋势
• no-clip ablation（2分）：在 off-policy 最佳设置下，比较 GRPO-Clip vs GRPO-No-Clip 稳定性/效果
• prompt ablation（2分）：r1_zero prompt vs question_only prompt，对 reward/entropy/长度/梯度等影响做解释

能力点：完整 RLHF/RLAIF（verified reward 版）训练闭环、稳定性机制（baseline、clip、长度归一化）、on-policy vs off-policy 的系统/样本效率 trade-off、prompt 与训练分布匹配。

5) Leaderboard（16分）：4小时内最大化验证准确率（强工程赛）

最后是一个“在固定时间预算内做最强”的挑战：

• 约束：只能用 MATH train/val；验证必须用 r1_zero prompt + 固定 reward_fn，temperature=1、max tokens=1024；报告 4 小时内的曲线截图
• 鼓励系统优化：让 GPU 不闲着、精度/compile/更好的并行等

能力点：LLMOps + RL 系统工程：吞吐、并行、缓存、日志、早停与调参策略。

一句话总结：A5 的任务谱系

• 评测框架：prompt → 生成 → 解析 → 分类统计 → 落盘（baseline）
• SFT 工程：tokenize+mask、logprob、熵、masked loss、梯度累积、训练中生成日志
• 自举对齐：Expert Iteration（采样 → 验证 → 筛选 → 再训）
• 策略梯度对齐：GRPO（优势估计、clip、baseline、长度归一化、off-policy、多组消融）
• 工程竞赛：固定时间预算下最大化准确率（leaderboard）