CS336：从零手搓LLM-后训练部分

本文总结了斯坦福 CS336 课程第五次作业的核心内容——从零实现 SFT、Expert Iteration 和 GRPO，训练一个 1.5B 参数的语言模型在 MATH 数据集上学会数学推理。

一、实验动机

语言模型有一个令人兴奋的能力——思维链推理（Chain-of-Thought Reasoning）：在给出最终答案前，先生成一步步的推理过程，显著提升复杂问题的解决准确率。

本实验的目标是：使用 Qwen 2.5 Math 1.5B 这个预训练的数学基础模型，通过三种递进的训练方法，让它学会在 MATH（高中竞赛数学题）上做推理：

二、实验环境

本次实验在 双卡 NVIDIA H20（各 96GB） 上完成，环境为 Python 3.12 + torch 2.5.1+cu124 + transformers 4.51.3 + vLLM 0.7.2。

由于 1.5B 模型的 rollout（采样）与训练都需要 GPU，采用了双卡分工架构：

每步训练流程：vLLM 采样 → 算 reward → policy 在 cuda:1 上算 loss 并反向传播 → 更新后的权重同步回 vLLM（load_policy_into_vllm_instance）→ 下一轮采样。这一"采样-训练分离"的架构是 RL 训练 LLM 的标准做法。

三、基础组件

3.1 提示词格式

所有实验统一使用 R1-Zero 提示词，模型被训练输出特定格式：

<think>
一步步的推理过程...
</think>
<answer>
\boxed{最终答案}
</answer>

这种格式便于后续自动化解析答案和判分。

3.2 答案判分

判分由 r1_zero_reward_fn 完成，返回三个分数：

{
    "format_reward": 0.0 | 1.0,    # 格式是否正确（含 <think>/<answer> 标签）
    "answer_reward": 0.0 | 1.0,    # 答案是否数学等价于标准答案
    "reward":        0.0 | 1.0     # 总奖励 = 格式分 × 答案分
}

答案等价判断经过三层 fallback：

1. 字符串规范化匹配：去掉空格、统一 \frac 写法、补齐花括号等

2. SymPy 表达式等价：将 \frac{1}{2} 和 0.5 转为数学表达式比较

3. LaTeX 解析：用 math_verify 做更深层的语义等价

3.3 梯度累积

1.5B 模型单卡显存有限，采用 梯度累积（Gradient Accumulation）技术：将 batch 拆成多个 microbatch，每次只做 loss.backward()，等累积够 gradient_accumulation_steps 步后再统一更新参数。GRPO 中实测 gradient_accumulation_steps=64（microbatch=4）是最优平衡——128 太慢、16 会 OOM。

四、Zero-shot 基线（§3）

在开始任何训练之前，先评测未微调的 Qwen 2.5 Math 1.5B Base 在 MATH 验证集（5000 条）上的表现，作为整份作业的对照基线：

两个关键观察：

1. 准确率仅 2.70%：base 模型是纯数学语料上持续预训练得到的，从未做过指令微调，完全不理解 User: ... Assistant: <think> 这种对话模板该输出什么。

2. 格式正确率只有 16.82%：超过 83% 的输出连 </think> <answer> 格式都不符合。从样本看，base 模型常输出错误的标签（如 <end think>、<answer>...<answer> 配对错误等）。这是模型输出本身的问题——r1_zero_reward_fn 严格匹配 </think> <answer>，而 base 模型根本没被训练过按这种结构作答。

五、SFT（监督微调）

原理

SFT 是最简单直接的方法——把 DeepSeek R1 生成的"问题 → 推理链 → 答案"数据喂给 Qwen 模型，用交叉熵损失训练：

只计算输出部分的 loss（用 response_mask 屏蔽 prompt 和 padding）。

关键实现

在 sft_utils.py 中实现了以下核心函数：

实验配置

• 基础模型：Qwen 2.5 Math 1.5B Base

• SFT 数据：DeepSeek-V4-Flash API 生成的 R1 风格推理轨迹（见下）

• 优化器：AdamW（betas=(0.9, 0.95)，weight_decay=0.0），lr=1e-4

SFT 数据：用 DeepSeek 生成真实推理轨迹

官方提供的是 DeepSeek R1 轨迹，但环境里缺失。我们改用 DeepSeek-V4-Flash API 现场生成推理轨迹：

关键点：真实推理轨迹（带多步思考、验证）远优于 MATH 参考 solution。后者是简洁教科书式解答，缺乏 R1 式深度思考，是 SFT 效果好的核心原因之一。

训练曲线（lr=1e-4，600 步，每 50 步 eval，抽样 256 条）

最终全量评测结果（5000 条验证集）

最终验证准确率 55.82%，远超作业要求的 ≥15% 目标（约 3.7 倍）。

两个关键 bug 的发现与修复（重要工程经验）

调试过程中定位到两个导致"SFT 无效"的根本问题：

1. Bug A：eval 不同步 policy 权重到 vLLM。 train_sft.py 训练的是内存中的 policy，但 eval 通过 llm（vLLM 实例）生成，而 vLLM 里的模型是初始 base 模型、从未同步更新，导致 eval 测的永远是 base 模型（~2.7%），与训练完全脱节。修复：eval 前调用 load_policy_into_vllm_instance(policy, llm)。

2. Bug B：lr=1e-5 过低。 修复 Bug A 前，lr=1e-5 下模型权重变化仅 0.07%（几乎没训练）。修复：lr 提高到 1e-4，权重真正更新，准确率才释放。

实验结论

• 用全量数据 SFT 可达到 55.82% 验证准确率（远超 15% 目标）

• 真实推理轨迹远优于教科书式解答，是 SFT 效果的关键

• lr 的敏感性极高：1e-5 几乎不动，1e-4 才真正起效

六、Expert Iteration（专家迭代）

原理

Expert Iteration 是一个自举循环：

flowchart TD
    A["当前模型 π_θ"] --> B["对每道题采样 G 条回答"]
    B --> C["用 reward_fn 判对错"]
    C --> D["筛掉错误回答\n只保留 reward=1 的 (q, o) 对"]
    D --> E["用保留的数据做 SFT\n更新 π_θ"]
    E --> A

核心洞察：不需要人类标注推理链，模型自己生成、自己筛选、自己学习。

实验配置

从 Qwen 2.5 Math 1.5B Base 完全自举（bootstrapping）。固定超参：n_ei_steps=5、rollout_batch_size Db=1024、lr=1e-4、batch=32、grad_accum=16。按 PDF §5 要求（至少 2 种 rollout 数 + 2 种 epoch 数），共跑了 3 组配置：

三组配置的迭代结果

配置 B（G=16, epoch=1）—— 最优：

配置 A（G=8, epoch=1）：峰值 33.20%（step 4），第 5 步回落至 23.83%
配置 C（G=8, epoch=2）：峰值 34.77%（step 5）

消融结论

三个关键发现：

1. rollout 数 G 是关键因子：G 从 8 增到 16，峰值从 33.20% 升到 42.58%（+9.4pp）。每题采样更多 rollout，能覆盖更多正确答案，过滤后保留的 SFT 数据质量更高、覆盖更广。

2. epoch 数影响较小：sft_epochs 从 1 增到 2，仅 +1.6pp（受限于过滤数据量）。

3. 熵持续下降（三组一致）：avg entropy 从 ~0.45-0.58 单调降到 ~0.20，说明模型从"乱猜"收敛到"稳定推理"，逐步学会 <think>/<answer> 格式和推理路径——这是 EI 的典型特征。

与 SFT 对比

EI 虽然不如 SFT（SFT 用了 DeepSeek 高质量轨迹），但从 base 完全自举到 42.58%，证明了专家迭代机制的有效性。

七、GRPO（组相对策略优化）

这是本作业最核心的部分——用强化学习的方法训练语言模型。

7.1 语言模型 = 策略（Policy）

在 RL 视角下：

7.2 策略梯度（Policy Gradient）

朴素策略梯度的核心公式：

直观理解：好回答（reward=1）的所有 token 被"鼓励"（概率增大），坏回答的 token 被"压制"（概率减小）。

代码实现仅一行：

def compute_naive_policy_gradient_loss(raw_rewards, policy_log_probs):
    return -raw_rewards * policy_log_probs   # 负号 = 最小化 → 等价于梯度上升

7.3 Baseline 降低方差

直接用原始 reward（只有 0/1）作为梯度权重，方差很大。引入 baseline 后：

GRPO 的巧思：不用额外训练 Critic 网络，而是用同一道题的 G 条回答的组内均值当作 baseline。

# 组归一化（Group Normalization）
avg = mean(rewards_in_group)                    # baseline = 组内均值
advantages = [(r - avg) for r in rewards]       # 相对优势
# 可选：再除以组内标准差
advantages = [a / (std + eps) for a in advantages]

7.4 GRPO-Clip 防止策略突变

在 off-policy 训练中（同一批采样数据上做多次梯度更新），新旧策略差异过大可能导致训练崩溃。Clip 机制限制策略更新幅度：

def compute_grpo_clip_loss(advantages, policy_log_probs, old_log_probs, cliprange):
    ratio = torch.exp(policy_log_probs - old_log_probs)   # π_new / π_old
    loss  = -torch.min(
        ratio * advantages,                                # 无限制项
        torch.clamp(ratio, 1-cliprange, 1+cliprange) * advantages  # Clip 项
    )
    return loss

7.5 完整 GRPO 训练流程

flowchart TD
    A["1. vLLM 采样：对每道题生成 G 条回答"] --> B["2. 保存 old_log_probs（旧策略的 log-prob）"]
    B --> C["3. reward_fn 打分 → run_compute_group_normalized_rewards()"]
    C --> D["4. 组归一化：advantages = (r - mean) / std"]
    D --> E["5. tokenize + get_response_log_probs 获取 policy_log_probs"]
    E --> F["6. compute_grpo_clip_loss() 计算 per-token loss"]
    F --> G["7. masked_mean() 只在 response 位置求均值"]
    G --> H["8. loss /= gradient_accumulation_steps → loss.backward()"]
    H --> I["9. optimizer.step()"]
    I -->|下一轮| A

7.6 实现的 6 个 GRPO 核心函数

全部位于 grpo_utils.py，通过 14 个测试用例验证：

八、GRPO 实验：完整消融结果

GRPO 统一超参：rollout_batch_size=256、group_size=8、train_batch_size=256、gradient_accumulation_steps=64、eval_every=5、eval_examples=1024。从 base 冷启动（不接 SFT/EI warm-start）。

8.1 学习率 sweep（§8 grpo_learning_rate）

结论：lr=2e-5 为最优（最终 answer reward ~60.7%，远超 25% 目标），后续所有消融统一用它。lr=2e-5 满程曲线稳定攀升：step50 47.9% → step100 52.9% → step150 58.2% → step200 60.7%。

8.2 Baseline 消融（§8 grpo_baselines）

固定 lr=1e-5，各 80 步：

结论：有 baseline 明显更优（26.4% vs 17.1%），符合预期——baseline 降低方差、收敛更稳。

8.3 长度归一化消融（§8 grpo_length_normalization）

固定 lr=2e-5，各 200 步，对比 masked_mean（默认）与 masked_normalize（÷1024）：

结论：masked_mean 显著优于 masked_normalize。masked_mean 按 response 均权、更稳定、不奖励冗长；masked_normalize 按 token 均权，对长度更敏感且梯度 norm 随长度放大，在数学推理上反而不利于收敛。

8.4 标准差归一化消融（§8 grpo_group_standard_deviation）

固定 lr=2e-5，各 200 步，对比 use_std_normalization=True/False：

结论：开启 std 归一化（除以组内 std）明显更优。std 归一化将组内优势缩放到统一尺度，消除不同问题难度带来的 reward 方差差异，使梯度更新更稳定。

8.5 Off-policy 扫描（§8 grpo_off_policy_sweep）

off-policy 用 grpo_clip loss。先 4 组粗扫（各 50 步），再对最优配置精扫 200 步：

粗扫（epochs × train_batch_size）：

200 步精扫（最优配置 epochs=2, batch=128）：answer 峰值 53.4% @105，末值 47.5%。

结论：off-policy 精扫 53.4% 仍低于 on-policy 基线 61.2%，印证 off-policy 上限低于 on-policy。epochs=2、batch=128 是最稳配置；epochs=4 或 batch=64 后期更易发散。

8.6 Clip 消融（§8 grpo_off_policy_clip_ablation）

off-policy 最优配置下对比 grpo_clip 与 grpo_no_clip：

结论：有 clip 略优（峰值 53.4% vs 51.7%），clip 在训练后期起到防发散作用。

8.7 Prompt 消融（§8 grpo_prompt_ablation）

off-policy 最优配置下对比 r1_zero 与 question_only prompt：

结论（重要发现）：question_only prompt 显著优于 r1_zero（68.65% vs 53.4%），甚至超过 on-policy 冷启动 GRPO 的 61.2%。说明 r1_zero 强制"step by step"指令在 off-policy 配置下反而拖累性能，去掉强制推理模板后模型表现更好。

8.8 SFT warm-start 串联实验

验证"SFT 作为 warm-start 串联 RL/EI"是否有增益，链路 SFT ckpt → EI → GRPO：

结论：在本实验的 1.5B 小模型 + 特定超参（G=16、lr=2e-5、5 步 EI、200 步 GRPO） 下，SFT warm-start 串联路线没有带来增益——SFT→EI 掉到 52.73%（低于 SFT 的 55.82%），EI→GRPO 57.13%（低于冷启动 60.7%）。

⚠️ 这个结论不能推广为"串联本身是错的"。它只在当前小模型场景下成立，原因有三：

1. 起点分布已偏：SFT 用外部 DeepSeek 轨迹拉到 55.82%，但 EI 自举采样出的数据分布与之有偏差，从 SFT 起点继续自举反而把模型往偏的方向拉。

2. 探索空间被压缩：SFT 后模型熵已经很低（~0.20），采样多样性不足，筛选出的增量数据有限；接近收敛的模型上 GRPO 组归一化的相对优势信号增益也很小。

3. 算力预算太短：5 步 EI + 200 步 GRPO 对 1.5B 模型，串联带来的"格式冷启动"收益本就不大，而代价（分布偏差、探索受限）却更大。

这印证的是 R1-Zero 论文"SFT 不是必需的"（冷启动 RL 可行）这一点，而不是"串联有害"。在大模型实践中（如 DeepSeek-R1 正式版、OpenAI o1），多阶段串联（冷启动 SFT → RL → 再 SFT → 再 RL）恰恰是标准做法——因为大模型冷启动 RL 的前期探索成本极高，先用少量 SFT 教会格式与初始 CoT 结构收益巨大。

九、最终成绩汇总

完整的训练路径：Zero-shot (2.70%) → SFT (55.82%) → EI (42.58%) → GRPO (60.7%)，GRPO 从 base 冷启动达到了最优成绩 60.7%。

十、核心代码片段

Advantage 的计算

def run_compute_group_normalized_rewards(reward_fn, rollout_responses, 
                                          repeated_ground_truths, group_size,
                                          advantage_eps, normalize_by_std):
    raw_rewards, group_normalized_rewards, curr_group = [], [], []
    for i, response in enumerate(rollout_responses):
        raw_reward = reward_fn(response, repeated_ground_truths[i])["reward"]
        raw_rewards.append(raw_reward)
        curr_group.append(raw_reward)
        if (i + 1) % group_size == 0:
            avg = sum(curr_group) / group_size
            curr_norm = [n - avg for n in curr_group]
            if normalize_by_std:
                curr_norm = [n / (advantage_eps + statistics.stdev(curr_group)) 
                            for n in curr_norm]
            group_normalized_rewards.extend(curr_norm)
            curr_group = []
    return (torch.tensor(group_normalized_rewards), 
            torch.tensor(raw_rewards), {})

完整的微批次训练步

def grpo_microbatch_train_step(policy_log_probs, response_mask,
                                gradient_accumulation_steps, loss_type,
                                raw_rewards, advantages, old_log_probs, cliprange):
    loss, meta = compute_policy_gradient_loss(
        policy_log_probs, loss_type, raw_rewards, advantages, old_log_probs, cliprange
    )
    loss = masked_mean(loss, response_mask)      # 只在 response token 求均值
    loss /= gradient_accumulation_steps           # 梯度累积调整
    loss.backward()                               # 反向传播
    return loss, {}

十一、关键 Insight

1. GRPO 冷启动就能超越 SFT：GRPO 从 base 冷启动达到 60.7%，超过用了 DeepSeek 高质量轨迹的 SFT（55.82%）。这与直觉"先 SFT 再 RL"相反——实验证明在 1.5B 小模型 + 本次超参下，SFT 与 RL 可以视为两个独立阶段：SFT warm-start 串联路线（SFT→EI→GRPO）反而掉点。但这印证的是 R1-Zero"冷启动 RL 可行"的可行性结论，不等于"串联永远有害"——大模型实践中（如 DeepSeek-R1 正式版）用少量 SFT 冷启动来教格式、降低 RL 前期探索成本，多阶段串联仍是标准做法（详见 §8.8）。

2. 组归一化（Group Normalization）是 GRPO 的核心创新：同一道题的 G 条回答互相对比，自动生成 baseline，避免了传统 Actor-Critic 方法需要额外训练 Critic 网络的负担。消融实验证实：有 baseline 26.4% vs 无 baseline 17.1%，开启 std 归一化 61.2% vs 关闭 53.5%。

3. 学习率是 GRPO 的第一敏感超参：lr=5e-6 收敛过慢（9.9%）、lr=1e-5 达标（32%）、lr=2e-5 最优（61.2%）、lr=5e-5 先冲高后崩溃（47.9%→0.5%）。窗口非常窄，SFT/EI 用 1e-4、GRPO 用 2e-5 的差异体现了不同训练范式的梯度尺度差异。

4. Prompt 工程出乎意料地关键：去掉 r1_zero 强制"step by step"模板的 question_only prompt 反而拿到 68.65%（全实验最高），说明过度指定的推理模板在 off-policy 配置下会拖累性能。

5. 答案解析器（Reward Function）决定了训练质量：如果 reward 函数不能准确判断答案对错，整个 RL 训练的方向就会偏。三层 fallback（字符串 → SymPy → LaTeX）的设计体现了工程中对鲁棒性的追求。

6. 系统工程与 bug 排查占了大头：梯度累积、vLLM 高效推理、双 GPU 调度（一个 GPU 跑采样、一个 GPU 跑训练）都是让 1.5B 模型训练可行的关键。SFT 阶段的两个 bug（eval 不同步 vLLM 权重、lr 过低）一度让训练"看起来完全无效"，这类"训练脱节"问题是 LLM 训练中的高频陷阱。

十二、总结

这次作业完整覆盖了从监督微调到强化学习的整个对齐训练管线，四阶段全链路实测：

通过亲手实现 6 个 GRPO 核心函数和完整的 SFT 工具链，并完成 §8 的全部消融实验（学习率、baseline、长度归一化、std 归一化、off-policy、clip、prompt），深入理解了策略梯度、Advantage 估计、Clip 机制、组归一化等 RL 核心概念——这些正是 DeepSeek R1、OpenAI o1 等前沿推理模型背后的关键技术。

最终成绩：GRPO 冷启动 60.7% 为最优，显著超越作业要求的 25% 目标（约 2.4 倍）。