Ch3: RL Scaling — GRPO 自我进化 🔥

目标: 用 GRPO 训练模型从 2 位数自我进化到 3-5 位数乘法 复刻 DeepSeek-R1 的 "aha moment": 观察模型自发延长推理链

核心思想

Ch2 的 RoPE 模型已经 100% 掌握了 2 位数乘法 (9801 个组合)。但它只在 max_len=64 上训练过 — 从未见过 3 位数以上的题目。

GRPO (Group Relative Policy Optimization) 让模型通过强化学习自我进化: 1. 给模型一道题 (比如 123*456=) 2. 模型尝试生成 8 个候选答案 (temperature sampling) 3. 用 reward function 评分: 答案对不对? 推理过程对不对? 4. 做得好的候选获得正向优势, 做得差的获得负向优势 5. 用 PPO-clip 更新策略, 让模型倾向于生成更好的回答

为什么不需要 Critic? 传统 PPO 需要一个 Value Network 估计状态价值。GRPO 用组内统计替代:

advantage = (reward - group_mean) / group_std

数学上更简洁, 内存省一半, DeepSeek 验证对数学推理效果极好。

SFT vs RL (GRPO) 底层原理对比

深入理解为什么要在昂贵的 SFT 之后加上更昂贵的 RL 阶段：

维度	SFT (监督微调)	RL (GRPO 强化学习)
通俗比喻	“照猫画虎” 老师给标准答案，学生照着抄，写错一个字就被扣分。	“摸着石头过河” 没有标准步骤，学生试出8种解法，看哪种能拿到最终奖励。
输入数据	必须是完美的人工标注。例如：`34=S1:34=12;Z12<EOS>`	只需要问题 (Prompt)。例如：`3*4=`。模型自己生成后面的所有过程。
前向传播	Teacher Forcing (强制纠偏)：强行喂给模型绝对正确的上下文。	Autoregressive (自回归生成)：硬着头皮基于自己生成的 Token 预测下一步。
目标函数	交叉熵 (Cross Entropy)： $Loss = -\log P(正确Token)$	策略梯度 (Policy Gradient)： $Loss = -\log P(预测Token) \times Advantage$
误差来源	来自于“每一步的标准答案”。 “老师说这里该填 7，你填了 6，产生误差”。	来自于“最终的成败与相对优势”。 “我知道这条命算错了，所以我要削弱整条路径的概率”。
参数更新 ($W_Q, W_K$)	“向标准答案对齐”：强行修改参数让概率向预定目标靠拢。	“强化高奖励的注意力模式”：根据 Advantage 分数，动态决定是正向强化还是反向惩罚刚才算出的注意力分布。
优化上限	被标注数据锁死。最高达到标注数据的拟合极限，没见过的题容易崩。	超越人类标注（上限极高）。模型通过试错，自己探索出极其有效（甚至反常识）的注意力路径。
计算成本	低。每条数据只过一次前向和反向。	极高 (SFT的 8-10倍)。生成 8 条长回答，算 Reward 和 KL 散度后再反向传播。

错位交叉训练的异同 (核心解惑)

在底层 Pytorch 代码中，SFT 和 RL 都使用了“错位交叉熵 (Shifted Cross-Entropy)” 来计算 Loss，但目的截然不同：

SFT 阶段：数据源是人类写的绝对正确答案。错位交叉强迫模型把预测概率拉向这些固定 Token，这本质上是死记硬背。
RL 阶段：
上半场 (出题与做题)：只给模型 246*722=，模型靠随机性自己蹦出 8 条完整的轨迹（有些对，有些错）。系统自动批改，算出每条的相对优势 $Advantage$。
下半场 (发工资更新)：把模型自己刚刚生成的 8 条轨迹原封不动塞回模型，计算错位交叉 Loss。此时不再是强制背诵，而是看 $Advantage$ 脸色：对于 $A > 0$ 的轨迹，利用 Loss 正向放大这批 Token 的生成概率；对于 $A < 0$ 的轨迹，则反向惩罚这批 Token 的概率。

一句话总结: RL 出题只给算式，走出来的路不管是好是坏，都拿来做错位交叉训练，只是在计算梯度前，乘上了一个由“最终考试成绩”决定的权重 ($Advantage$)。这就是 RL 的威力之源。

训练流程

Phase 1: SFT 预热 (~20 epochs)
  ├── 加载 Ch2 best.pt (2位数 100%)
  ├── 扩展 max_len 64→1500 (RoPE 天然支持)
  └── 用 2位数+3位数混合 CoT 做 SFT (教长格式)

Phase 2: RL-2位数 (巩固, 建立 reward 基线)
  └── 达到 >98% → 进阶

Phase 3: RL-3位数混合 (70% 2位 + 30% 3位)
  └── 3位数 >90% → 进阶

Phase 4: RL-4位数混合 (40/40/20%)
  └── 4位数 >80% → 进阶

Phase 5: RL-5位数混合 (20/30/30/20%)
  └── 终极挑战

快速开始

# 1. 完整训练 (SFT预热 + RL)
python train.py --ch2_ckpt ../ch2_rope/checkpoints_v3/best.pt

# 2. 跳过 SFT, 直接 RL (如果已有 SFT 检查点)
python train.py --skip_sft --sft_ckpt checkpoints/sft_warmup.pt

# 3. 评估
python eval.py --checkpoint checkpoints/best_rl.pt --samples 500

# 4. 生成数据 (独立运行)
python generate_data.py --mode both

文件清单

文件	职责
`model.py`	扩展 Ch2 模型: max_len=1500, generate_with_logprobs
`generate_data.py`	多位数 CoT 数据生成 (SFT + RL 题库)
`reward.py`	多维度奖励: 答案正确性 + 格式 + 过程奖励
`grpo_trainer.py`	GRPO 核心训练器 (纯 PyTorch)
`curriculum.py`	课程学习调度器 (自动升级难度)
`monitor.py`	Aha Moment 监控 + 训练曲线
`train.py`	主训练入口
`eval.py`	多位数评估脚本

奖励函数设计

总奖励 ∈ [-1.0, +1.5]

1. 最终答案 Z{answer}
   ✅ 正确 → +1.0
   ❌ 错误 → -0.5

2. 格式合规 (S/A/Z 结构)
   ✅ 完整 → +0.2
   ❌ 缺失 → -0.3

3. 过程奖励 (每步 ±0.1)
   S步骤: a*val=result 对不对?
   A步骤: x+y=z (反转) 对不对?

"Aha Moment" 指标

指标	含义	期望变化
`avg_response_length`	平均回复长度	📈 模型生成更长推理链
`accuracy_by_digits`	按位数准确率	📈 逐渐攻克更多位数
`avg_reward`	平均奖励	📈 持续上升
`kl_divergence`	与参考模型距离	📊 适度增长

Token 长度

位数	示例	Token 长度
2×2	99×99	~50
3×3	999×999	~109
4×4	9999×9999	~166
5×5	99999×99999	~268

模型参数

与 Ch2 相同, 仅扩展 max_len: - d_model=256, n_heads=8, n_layers=6, d_ff=1024 - 4.74M 参数 (RoPE 零参数) - max_len: 64 → 1500

阶段性总结 (Ch3 末期状态)

在长达数百步的 GRPO 强化学习中，我们观察到了激动人心的智能涌现与工程瓶颈：

1. 算法层面：精度打磨与“数零综合症”

压舱石稳固：1d、2d、3d 乘法在 RL 中稳定在 95%~100% 正确率，毫无灾难性遗忘。
4d 的痛苦挣扎：4d 乘法突破了 SFT 40% 的天花板，在 30%~60% 之间剧烈震荡。KL 散度的持续飙升证明模型正在经历深度的参数重组。
典型死法：逻辑拆解已完全正确（不再乱编乘法规则），所有的错误都收敛到了 Copy Error（抄错前面的中间结果） 和 Zero-Counting Collapse（眼花数错 0 的个数）。这证明 RL 正在死磕模型最微观的 Attention 锐度。

2. 工程层面：算力黑洞与 $O(N^3)$ 噩梦

为了跑通长达 300 Token 的 5d 乘法，我们将 max_len 强行拉到了 1500。
尽管我们在 5090 上开启了 BF16 自动混合精度 救场，但朴素的自回归生成使得每一次吐出新词，都要重新计算前面所有词的 Q、K、V 并重新点积。这导致生成阶段的计算量呈立方级 $O(N^3)$ 爆炸。
纯 Python 的 for 循环生成已经让顶级显卡不堪重负。算法的 Scaling 撞上了工程架构的墙。

3. Ch3 最终停机快照 (Step 150)

在转移至 Ch4 前，我们截取了最后的训练状态：

  📊 Step   150 | Phase 4 | reward +1.394 | acc 63.3% | loss 0.0002
  📏 avg_len 213 | KL 0.0164 | clip 0.00
  🎲 训练: 2d:30✓/2✗ | 3d:30✓/2✗ | 4d:21✓/43✗
  🎯 按位数: 2d:100% | 3d:95% | 4d:40% | 5d:15%

此时模型正处于最痛苦的蜕变期，虽然 4d 胜率仍在 40% 挣扎，但低级错误已消失，全部转为注意力精准度问题。这就是我们继续向上 Scaling 的绝佳起跑线！

下一步：开启 Ch4 (基础设施升级) 在下一章，我们将不动模型参数，纯粹从底层架构动刀，引入大模型推理的灵魂技术——KV Cache（键值缓存）。我们将亲手把 $O(N^3)$ 的无底洞降维回 $O(N^2)$，让生成速度飙升数十倍，为 5d 甚至更高位数的 RL 铺平算力大道！

NanoTransformer Courses