""" ============================================================================= CS486 · 第 5 讲模块：rl_lab.py —— 用强化学习训练推理的零件（GRPO、RFT、奖励、缓存） ============================================================================= 【这个模块要解决什么问题】 第 4 讲的所有方法都**不改变模型**：多采样几次、挑一挑、搜一搜。第 5 讲问：能不能把 "采样很多次、知道哪些对了"变成**训练信号**，让模型本身变强？这就是"用强化学习训练推理"： 策略（policy） = 语言模型本身（这里是第 1 讲的 TinyReasoner） 一个回合（episode）= 对一道题生成一整段回答（推理链 + "|" + 答案 + "$"） 奖励（reward） = 回答完之后打一个分：答案对了得 1，错了得 0（"可验证奖励"，RLVR） 训练目标 = 提高"会得高分的回答"的概率，降低"会得低分的回答"的概率 本模块提供实验需要的全部零件： 1. sample_fast：带 KV cache 的批量采样（比第 4 讲逐步重算整段序列的 sample 快约 5 倍，结果相同） 2. 回答的解析与分类：response_of、true_reward、classify（对 / 走过头 / 停早了 / 查找错 / 抄错 / 格式坏） 3. 策略梯度的零件：sequence_batch、token_logprobs、group_advantages（GRPO 的组内基线） 4. 奖励：真实奖励、第 4 讲的验证器（ORM / PRM-full）当奖励、随机奖励；格式检查（guard） 5. 训练：train_rl —— 一个循环同时支持 GRPO、RFT（只学自己做对的样本，STaR 式）、SFT（学标准答案）， 以及 KL 惩罚、长度惩罚、是否按组标准差归一化 6. "R1-Zero 式"实验用的基座模型：mixed_base（很少主动写推理链、但写了就会写对的模型） 7. 缓存与并行：load_runs（按配置指纹缓存在 checkpoints/，用第 2 讲的 run_parallel 并行训练） 8. 评测：eval_set、greedy_accuracy、pass_curve（pass@k 曲线）、think_profile（按跳数看"想不想、对不对"） 【复用】只读地导入第 1 讲的 hop_world、tiny_reasoner，第 2 讲的 exp_kit（run_parallel、config_id） 和 method_lab（fit_stream、encode_batch），第 4 讲的 ttc_lab（验证器、pass_at_k、step_labels）。 【运行】 python3 rl_lab.py # 自检：快速采样 = 第 4 讲的采样；回答分类；优势函数；冒烟训练 【依赖】PyTorch + 第 1、2、4 讲的模块。 """ import copy import math import os import random import re import sys import time import torch import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture04")) sys.path.insert(0, os.path.join(HERE, "..", "lecture02")) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import exp_kit as ek # noqa: E402 import hop_world as hw # noqa: E402 import method_lab as ml # noqa: E402 import tiny_reasoner as tr # noqa: E402 import ttc_lab as tl # noqa: E402 CKPT_DIR = os.path.join(HERE, "checkpoints") # --------------------------------------------------------------------------- # 1. 带 KV cache 的快速采样 # --------------------------------------------------------------------------- def _rope_at(rope, x, pos): """ 对"每行只有一个新 token"的 query / key 做旋转位置编码，每一行的位置可以不同。 【参数】rope：tiny_reasoner.RoPE；x：(B, H, 1, head_dim)；pos：(B,) 每行新 token 的位置 【返回】同形状、旋转过的张量 【思路】和 RoPE.forward 的公式完全相同，只是 cos / sin 按每行的位置单独取出来（rope.cos[pos]）。 """ cos = rope.cos[pos][:, None, None, :] sin = rope.sin[pos][:, None, None, :] x1, x2 = x[..., 0::2], x[..., 1::2] return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2) @torch.no_grad() def sample_fast(model, prompts, max_new, temperature, gen): """ 批量温度采样（带 KV cache）：对每个题面接着生成 max_new 个 token。 【思路】第 4 讲的 ttc_lab.sample 每生成一个 token 都把整段序列重新算一遍。强化学习每一步要采样 几百条回答、训练几百步，这样太慢。KV cache 的想法（推理系统里的标准做法）： 1. 先把题面整段算一遍（prefill），把每层的 key、value 存起来； 2. 之后每一步只算**新 token 这一个位置**：它的 query 和缓存里所有的 key 做注意力， 再把它自己的 key、value 追加到缓存里。 因为注意力是因果的，旧位置的 key / value 不会因为后面的 token 而改变，所以缓存是精确的， 结果和逐步重算相同（自检里验证：贪心和同一随机种子下的采样都逐字相同）。 难点是一批题面长短不一：第 i 行的第 t 个新 token 写在位置 len_i + t。我们给缓存留出 "最长题面 + max_new"个槽位，用掩码只让每一行看到"自己已经写过的位置"。 短题面在 prefill 时右边补的 PAD 也会算出 key / value，但它们所在的槽位在被看到之前就会被新 token 覆盖。 【参数】model：TinyReasoner；prompts：题面字符串列表；max_new：生成多少个 token； temperature：温度（0 = 贪心）；gen：torch.Generator（温度 > 0 时用来抽样） 【返回】list[str]：每个题面之后生成的 max_new 个字符（遇到 "$" 也不提前停，由 response_of 截断） 【边界】prompts 为空时返回空列表；题面长度 + max_new 不能超过 model.cfg.max_len。 【PyTorch 小知识】Kc[rows, :, cur] = k 是"高级索引"赋值：同时给第 i 行的第 cur[i] 个槽位写入。 masked_fill(~mask, -inf) 把看不到的位置的注意力分数设成负无穷，softmax 之后权重就是 0。 """ if not prompts: return [] B = len(prompts) lens = torch.tensor([len(s) for s in prompts]) L = int(lens.max()) S = L + max_new X = torch.full((B, S), hw.PAD_ID, dtype=torch.long) for i, s in enumerate(prompts): X[i, :len(s)] = torch.tensor(hw.encode(s)) cfg = model.cfg H, D = cfg.n_heads, cfg.d_model dh = D // H # ---- prefill：整段题面算一遍，存下每层的 key / value ---- x = model.embed(X[:, :L]) caches = [] for blk in model.blocks: q, k, v = blk.qkv(blk.ln1(x)).split(D, dim=-1) q, k, v = (t.view(B, L, H, dh).transpose(1, 2) for t in (q, k, v)) q, k = blk.rope(q), blk.rope(k) Kc, Vc = torch.zeros(B, H, S, dh), torch.zeros(B, H, S, dh) Kc[:, :, :L], Vc[:, :, :L] = k, v caches.append((Kc, Vc)) out = F.scaled_dot_product_attention(q, k, v, is_causal=True) x = x + blk.proj(out.transpose(1, 2).reshape(B, L, D)) x = x + blk.mlp(blk.ln2(x)) rows = torch.arange(B) last = model.head(model.ln_f(x))[rows, lens - 1] # 每行最后一个题面字符处的 logits cur = lens.clone() # 下一个 token 要写到的位置 slots = torch.arange(S) # ---- 逐个生成 ---- for t in range(max_new): if temperature == 0: nxt = last.argmax(-1) else: nxt = torch.multinomial((last / temperature).softmax(-1), 1, generator=gen).squeeze(1) X[rows, cur] = nxt if t == max_new - 1: break x = model.embed(nxt)[:, None, :] # (B, 1, D)：只算新 token 这一个位置 mask = (slots[None, :] <= cur[:, None])[:, None, None, :] # 每行只看到 ≤ 当前位置的槽位 for blk, (Kc, Vc) in zip(model.blocks, caches): q, k, v = blk.qkv(blk.ln1(x)).split(D, dim=-1) q, k, v = (u.view(B, 1, H, dh).transpose(1, 2) for u in (q, k, v)) q, k = _rope_at(blk.rope, q, cur), _rope_at(blk.rope, k, cur) Kc[rows, :, cur] = k[:, :, 0] Vc[rows, :, cur] = v[:, :, 0] scores = (q @ Kc.transpose(-2, -1)) / math.sqrt(dh) out = scores.masked_fill(~mask, float("-inf")).softmax(-1) @ Vc x = x + blk.proj(out.transpose(1, 2).reshape(B, 1, D)) x = x + blk.mlp(blk.ln2(x)) last = model.head(model.ln_f(x))[:, 0] cur = cur + 1 return [hw.decode(X[i, int(lens[i]):int(lens[i]) + max_new].tolist()) for i in range(B)] def sample_responses(model, probs, n_samples, temperature, gen, chunk=4096): """ 对每道题采样 n_samples 条回答（已经在第一个 "$" 处截断）。 【参数】probs：list[HopProblem]；n_samples：每道题几条；temperature、gen：同 sample_fast； chunk：一次最多解码多少条（控制内存） 【返回】list[list[str]]：第 i 道题的 n_samples 条回答 【边界】max_new = 这批题里最大的 k + 5（和第 4 讲一样：多给几个位置，让"走过头"能被看到）。 """ flat = [p.prompt for p in probs for _ in range(n_samples)] max_new = max(p.k for p in probs) + 5 outs = [] for i in range(0, len(flat), chunk): outs += sample_fast(model, flat[i:i + chunk], max_new, temperature, gen) resps = [response_of(o) for o in outs] return [resps[i * n_samples:(i + 1) * n_samples] for i in range(len(probs))] # --------------------------------------------------------------------------- # 2. 回答：截断、奖励、分类 # --------------------------------------------------------------------------- # 【Python 小知识】re.compile 预先编译一个正则表达式。^ 和 $ 分别匹配开头和结尾（这里的 \$ 是字面的 "$" 字符）； # [a-p]+ 表示"一个或多个 a–p 之间的字母"，\| 是字面的 "|"。 WELL_FORMED = re.compile(r"^[a-p]+\|[a-p]\$$") # 合格的思维链回答：至少一步推理 + "|" + 一个字母 + "$" THINKS = re.compile(r"^[a-p]+\|") # "写了推理链"：在 "|" 之前至少有一个字母 def response_of(output): """ 把生成的字符串截到第一个 "$"（含）为止——这就是"一个回合"的全部动作。 【思路】采样时我们总是生成固定个数的 token（方便批量处理），"$" 之后的字符模型并没有"打算"说， 不应该算进奖励，也不应该参与训练。没有 "$"（写到长度上限还没结束）时保留全部字符。 """ j = output.find(hw.END) return output if j < 0 else output[:j + 1] def true_reward(prob, resp): """真实奖励（可验证奖励）：答案字母对了得 1.0，否则 0.0。和第 1、4 讲评测用的标准完全相同。""" return float(hw.parse_answer(resp, "cot") == prob.answer) FAILURE_TYPES = ("correct", "overshoot", "early", "lookup", "miscopy", "malformed") def classify(prob, resp): """ 把一条回答归入一种结果类型。 【返回】FAILURE_TYPES 之一： correct ：答案对（不管推理链） overshoot ：每一步查找都对，但走了超过 k 步（第 4 讲发现的主要失败方式："走过头"） early ：每一步查找都对，但不到 k 步就停了（"停早了"） lookup ：推理链里至少有一步查找错了 miscopy ：推理链恰好 k 步且都对，"|" 后面的答案却和最后一步不同（"抄错了"） malformed ：没有 "|" 或 "|" 后面不是字母。贪心解码时几乎都是"一直没停下"：写到了长度上限（k + 5 个字符） 还没写 "|"——也就是走过头至少 5 步。表格里显示为"没写 |" 【边界】一条推理链可能同时"有查找错误"又"长度不对"——我们优先归为 lookup（错误发生得更早）。 """ if hw.parse_answer(resp, "cot") == prob.answer: return "correct" if hw.ANSWER_MARK not in resp or hw.parse_answer(resp, "cot") is None: return "malformed" trace = resp.split(hw.ANSWER_MARK, 1)[0] steps = tl.step_labels(prob, trace) if not all(steps): return "lookup" if len(trace) > prob.k: return "overshoot" if len(trace) < prob.k: return "early" return "miscopy" # --------------------------------------------------------------------------- # 3. 策略梯度的零件 # --------------------------------------------------------------------------- def sequence_batch(prompts, resps): """ 把 (题面, 回答) 拼成训练用的张量。 【返回】(X, M)：X 是 (B, T) 的 token 编号（右补 PAD）；M 是 (B, T) 的 0/1 掩码，回答部分为 1 【思路】和第 1 讲 make_batch 的损失掩码一样：只有回答里的 token 是"动作"，题面不是模型选的， 不参与策略梯度。 """ seqs = [hw.encode(p + r) for p, r in zip(prompts, resps)] T = max(len(s) for s in seqs) X = torch.full((len(seqs), T), hw.PAD_ID, dtype=torch.long) M = torch.zeros(len(seqs), T) for i, (p, s) in enumerate(zip(prompts, seqs)): X[i, :len(s)] = torch.tensor(s) M[i, len(p):len(s)] = 1.0 return X, M def token_logprobs(model, X): """ 每个位置上"实际出现的下一个 token"的对数概率，以及那个位置的预测分布的熵。 【返回】(logp, ent)：都是 (B, T-1)；logp[i, t] = log π(X[i, t+1] | X[i, ≤t]) 【思路】一次前向就得到整条回答里每个 token 的对数概率（"教师强制"：输入是已经采好的回答）。 策略梯度要的就是它：∇ log π(回答) = Σ_t ∇ log π(第 t 个 token)。 熵 H = −Σ p log p 衡量模型在这个位置"有多不确定"：RL 训练中熵往往快速下降（"熵坍缩"，讲义 §7）。 【PyTorch 小知识】gather(-1, idx) 按下标取出最后一维上的元素：这里取出"实际 token"那一列的对数概率。 """ lp = model(X[:, :-1]).log_softmax(-1) logp = lp.gather(-1, X[:, 1:, None]).squeeze(-1) ent = -(lp.exp() * lp).sum(-1) return logp, ent def group_advantages(r, std_norm=True, eps=1e-4): """ GRPO 的组内优势：同一道题的 G 个样本互相比较。 【思路】策略梯度需要一个"基线"（baseline）：奖励高于基线的回答被鼓励，低于基线的被抑制， 这样梯度的方差小得多（只用原始奖励时，"所有回答都得 0.9"也会把它们全部推高）。 PPO 用一个单独训练的价值网络估计基线；GRPO（Shao et al. 2024）的想法更简单： **同一道题采样 G 次，用这 G 个奖励的平均值当基线**，再除以它们的标准差（"标准化"）。 Dr. GRPO（Liu et al. 2025）指出除以标准差会带来偏差：一组奖励几乎相同（标准差很小）时， 微小的差别会被放大成很大的优势（讲义 §8.3，代码 03 演示 4）。std_norm=False 就是去掉这一步。 【参数】r：(B, G) 张量，每行是一道题的 G 个奖励；std_norm：是否除以组内标准差；eps：防止除以 0 【返回】(B, G) 的优势 【边界】一组奖励完全相同（全对或全错）时，优势全是 0——这一组对梯度**没有任何贡献** （讲义 §4.3 的"零方差组"；DAPO 的"动态采样"就是把这种组过滤掉、补采新题）。 """ a = r - r.mean(dim=1, keepdim=True) if std_norm: a = a / (r.std(dim=1, keepdim=True) + eps) return a # --------------------------------------------------------------------------- # 4. "R1-Zero 式"实验用的基座模型 # --------------------------------------------------------------------------- # 【思路】DeepSeek-R1-Zero 直接从预训练的基座模型开始做 RL，回答越来越长、自己"学会"了长思维链。 # 一个关键问题是：这种行为是 RL **创造**的，还是基座模型里本来就有、RL 只是把它**选出来**？ # 我们造几个基座模型来回答它。每个基座的"预训练语料"里，一部分题目带推理链（"abc|c$"）， # 其余的直接给答案（"|c$"，推理链长度为 0）。基座采样时于是会"有时想、有时不想"。 # # 直接用少量推理链数据从零训练（例如 2%）行不通：模型连查找本身都学不会（编写时的原型，见 README）。 # 所以用两阶段：先用 50% 带推理链的数据训练 2000 步（学会查找），再用很少推理链的数据继续训练 1000 步 # （把"主动写推理链"的频率压低，但写的时候仍然写得对）。这模仿了真实基座模型的状况： # **能力**来自大量数据，而某种**行为**（例如长思维链、自我检查）在语料里很少见。 MIXED_BASES = { # 名字: (第一阶段的推理链比例, 第二阶段的推理链比例；None 表示没有第二阶段) "noskill": (0.0, None), # 对照：语料里从来没有推理链——没有查找能力，也没有"想"的行为 "think0": (0.5, 0.0), # 学会了查找，但第二阶段的语料里完全没有推理链——"想"几乎绝迹 "think2": (0.5, 0.02), # 2% 的语料带推理链 "think10": (0.5, 0.10), # 10% 的语料带推理链 } def _mixed_pairs(rng, batch_size, p_think): """生成一批"预训练语料"：每道题以概率 p_think 带推理链，否则直接给答案（"|" + 答案 + "$"）。""" out = [] for _ in range(batch_size): p = hw.sample_train_problem(rng, (4, 8)) trace = p.trace if rng.random() < p_think else "" out.append((p.prompt, trace + hw.ANSWER_MARK + p.answer + hw.END)) return out def _train_mixed(name, quick=False): """ 训练一个混合语料的基座模型（两阶段），返回 (model, seconds)。 【思路】第一阶段用第 2 讲的 fit_stream（自定义数据流的训练循环，从零开始）； 第二阶段在第一阶段的模型上继续训练（学习率 3e-4，warmup 20 步，余弦退火）。 【边界】quick 模式下两阶段都只训练 30 步。 """ p1, p2 = MIXED_BASES[name] steps1, steps2 = (30, 30) if quick else (2000, 1000) tcfg1 = tr.TrainConfig(mode="cot", seed=0, steps=steps1) t0 = time.perf_counter() model, _ = ml.fit_stream(tr.ModelConfig(), tcfg1, lambda step, rng: _mixed_pairs(rng, 64, p1)) if p2 is not None: tcfg2 = tr.TrainConfig(mode="cot", seed=11, steps=steps2, lr=3e-4, warmup=20) torch.manual_seed(tcfg2.seed) rng = random.Random(tcfg2.seed) opt = torch.optim.AdamW(model.parameters(), lr=tcfg2.lr, weight_decay=tcfg2.weight_decay) for step in range(steps2): X, M = ml.encode_batch(_mixed_pairs(rng, 64, p2)) model.train() loss = tr.masked_loss(model, X, M) for g in opt.param_groups: g["lr"] = tr.lr_at(step, tcfg2) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() model.eval() return model, time.perf_counter() - t0 def _mixed_path(name): return os.path.join(CKPT_DIR, f"base_{name}_{ek.config_id({'name': name, 'spec': MIXED_BASES[name]})}.pt") def _save_mixed(name): """子进程入口：训练并保存一个混合基座。""" model, secs = _train_mixed(name) os.makedirs(CKPT_DIR, exist_ok=True) torch.save({"name": name, "spec": MIXED_BASES[name], "state_dict": model.state_dict(), "seconds": secs}, _mixed_path(name)) return secs def mixed_base(name, quick=False): """ 读取一个混合基座（没有缓存就训练）。 【参数】name：MIXED_BASES 里的名字 【返回】TinyReasoner（eval 模式） 【边界】quick 模式不读写缓存，几秒钟训练完，模型几乎没有学到东西。 """ if quick: return _train_mixed(name, quick=True)[0] path = _mixed_path(name) if not os.path.exists(path): _save_mixed(name) blob = torch.load(path, weights_only=False) model = tr.TinyReasoner(tr.ModelConfig()) model.load_state_dict(blob["state_dict"]) model.eval() return model def init_model(name, quick=False): """ RL 的起点：第 1 讲的标准模型（例如 "cot_L2"），或者 "mixed:<名字>" 形式的混合基座。 【返回】一个新的模型对象（训练会修改它，所以每次都从缓存重新读一份，互不影响） """ if name.startswith("mixed:"): return mixed_base(name.split(":", 1)[1], quick=quick) return tr.get_model(name, quick=quick, verbose=False)[0] # --------------------------------------------------------------------------- # 5. 训练 # --------------------------------------------------------------------------- DEFAULT_JOB = { "init": "cot_L2", # 起点模型 "method": "grpo", # "grpo"：组内优势的策略梯度；"rft"：只学自己做对的样本；"sft"：学标准答案 "reward": "true", # "true"：可验证奖励；"orm" / "prm_full"：第 4 讲的验证器；"random"：随机 0/1 "guard": False, # True：格式不合格（WELL_FORMED 不匹配）的回答奖励记为 0 "beta": 0.0, # KL 惩罚系数（相对于起点模型） "len_pen": 0.0, # 长度惩罚：奖励减去 len_pen × 回答长度（字符数，含 "|"、答案和 "$"） "std_norm": True, # 优势是否除以组内标准差（GRPO 原版 True；Dr. GRPO 去掉） "lr": 1e-3, # 学习率（AdamW，常数，不衰减） "steps": 200, # 更新步数 "B": 32, # 每步几道题 "G": 8, # 每道题采样几次（"组"的大小） "n_range": [4, 12], # 训练题的环大小范围 "k_max": 7, # 训练题的最大跳数（k 在 1 … min(n-1, k_max) 里均匀取） "temperature": 1.0, # 采样温度 "seed": 0, "eval": "none", # 训练中定期评测什么："none"、"n12"（n = 12 的测试集贪心准确率）、"think"（想不想、对不对） "eval_every": 25, } def make_job(**kw): """ 以 DEFAULT_JOB 为底，覆盖给定的字段，得到一个完整的任务配置。 【为什么要完整的配置】缓存文件名由配置的指纹（config_id）决定，配置里每个字段都写全， 才不会出现"默认值改了、却读到旧缓存"的问题。 【边界】给了 DEFAULT_JOB 里没有的字段时抛 KeyError（防止拼错字段名而悄悄被忽略）。 """ for k in kw: if k not in DEFAULT_JOB: raise KeyError(f"未知的字段 {k!r}；可用的有：{', '.join(DEFAULT_JOB)}") return {**DEFAULT_JOB, **kw} def _train_problems(rng, job): """按 job 的 n_range、k_max 抽一批训练题。""" lo, hi = job["n_range"] out = [] for _ in range(job["B"]): n = rng.randint(lo, hi) out.append(hw.make_problem(rng, n, rng.randint(1, min(n - 1, job["k_max"])))) return out def compute_rewards(job, probs, resps, verifier, gen): """ 给一批回答打分。 【参数】job；probs：每条回答对应的题；resps：回答；verifier：job["reward"] 是 orm / prm_full 时的验证器； gen：随机奖励用的随机数生成器 【返回】(proxy, true)：两个一维张量。proxy 是**训练用**的奖励（可能是验证器打的分），true 是真实奖励。 长度惩罚加在 proxy 上；格式检查（guard）把格式不合格的回答的 proxy 置 0（在长度惩罚之前）。 【研究联系】把"训练用的奖励"和"真实奖励"分开记录，是观察奖励黑客（reward hacking）的前提： 一旦两者分道扬镳——训练奖励在涨、真实奖励在跌——就说明策略在钻奖励的空子（讲义 §9，代码 04）。 """ true = torch.tensor([true_reward(p, r) for p, r in zip(probs, resps)]) kind = job["reward"] if kind == "true": proxy = true.clone() elif kind == "random": proxy = (torch.rand(len(resps), generator=gen) < 0.5).float() elif kind in ("orm", "prm_full"): exs = [tl.verifier_example(p, r) for p, r in zip(probs, resps)] if kind == "orm": proxy = torch.tensor(tl.score(verifier, exs, "orm")) else: proxy = torch.tensor([tl.prm_aggregate(s) for s in tl.score(verifier, exs, "prm_full")]) else: raise ValueError(f"未知的奖励 {kind!r}") if job["guard"]: proxy = proxy * torch.tensor([float(bool(WELL_FORMED.match(r))) for r in resps]) if job["len_pen"]: proxy = proxy - job["len_pen"] * torch.tensor([float(len(r)) for r in resps]) return proxy, true def train_rl(job, quick=False): """ 按 job 训练一个模型（可以在子进程里运行）。 【思路】每一步： 1. 抽 B 道训练题； 2. （grpo / rft）每道题用当前策略采样 G 条回答；（sft）直接用标准答案（推理链 + 答案）； 3. 打分：proxy（训练用）和 true（真实）； 4. 算损失并更新一次参数： grpo：优势 A = 组内标准化的 proxy，损失 = −Σ A_i · log π(回答_i 的每个 token) / 回答 token 总数 （每个回答的所有 token 共享同一个优势；按 token 总数平均，即 DAPO 的"token 级"平均）； 有 KL 惩罚时再加 β × KL(π ‖ π_起点) 的逐 token 估计（k3 估计量：r − log r − 1，r = π_起点/π）； rft ：只在 true = 1 的回答上做最大似然（"拒绝采样微调"，STaR / ReST-EM 的在线版本）； sft ：在标准答案上做最大似然（第 1 讲的训练方式，只是数据换成了 n ≤ 12 的题）。 5. 记录这一步的统计量；每隔 eval_every 步按 job["eval"] 评测一次。 为什么 GRPO 这里没有 PPO 的"裁剪"（clip）？PPO / GRPO 用同一批样本更新好几次， 需要裁剪"新旧策略的概率比"防止走得太远。我们每批样本只更新一次（完全在线，on-policy）， 概率比恒等于 1，裁剪不起作用，目标函数就退化成"带组内基线的 REINFORCE"（讲义 §3.4）。 【参数】job：make_job 生成的配置；quick：冒烟测试（只训练几步、评测集很小） 【返回】(model, history)。history 是字典： "steps"：每一步一个字典 {proxy, true, zero, ent, len, think} proxy / true：这一批回答的平均训练奖励 / 真实奖励；zero：零方差组的比例； ent：回答 token 的平均熵；len：回答的平均长度；think：写了推理链的回答的比例 "evals"：[(step, 指标字典), …] "examples"：[(step, [(标准答案, 回答, proxy), …]), …]——"训练奖励高、真实奖励为 0"的回答，用来看策略钻了什么空子 "seconds"：训练用时 【边界】rft 这一步如果没有任何正确的回答，就跳过更新（损失为 0）。 """ torch.set_num_threads(1) steps = 6 if quick else job["steps"] policy = init_model(job["init"], quick=quick) ref = copy.deepcopy(policy) if job["beta"] > 0 else None verifier = None if job["reward"] in ("orm", "prm_full"): verifier = tl.load_verifiers(quick=quick, verbose=False)[job["reward"]] torch.manual_seed(job["seed"]) opt = torch.optim.AdamW(policy.parameters(), lr=job["lr"], weight_decay=0.0) rng = random.Random(job["seed"] + 100) gen = torch.Generator().manual_seed(job["seed"] + 100) B, G = job["B"], job["G"] hist = {"steps": [], "evals": [], "examples": []} t0 = time.perf_counter() def do_eval(step): if job["eval"] == "n12": hist["evals"].append((step, {"acc": greedy_accuracy(policy, eval_set(12, per_k=5 if quick else 50))})) elif job["eval"] == "think": hist["evals"].append((step, think_profile(policy, think_set(per_k=4 if quick else 40), 4 if quick else 16))) do_eval(0) for step in range(steps): probs = _train_problems(rng, job) if job["method"] == "sft": prompts = [p.prompt for p in probs] resps = [p.trace + hw.ANSWER_MARK + p.answer + hw.END for p in probs] proxy = true = torch.ones(len(probs)) r = proxy.view(-1, 1) else: prompts = [p.prompt for p in probs for _ in range(G)] resps = [x for row in sample_responses(policy, probs, G, job["temperature"], gen) for x in row] flat_probs = [p for p in probs for _ in range(G)] proxy, true = compute_rewards(job, flat_probs, resps, verifier, gen) r = proxy.view(B, G) X, M = sequence_batch(prompts, resps) m = M[:, 1:] policy.train() logp, ent = token_logprobs(policy, X) if job["method"] == "grpo": adv = group_advantages(r, job["std_norm"]).view(-1) loss = -(logp * m * adv[:, None]).sum() / m.sum() elif job["method"] == "rft": keep = (true > 0.5).float() loss = -(logp * m * keep[:, None]).sum() / (m * keep[:, None]).sum().clamp(min=1.0) elif job["method"] == "sft": loss = -(logp * m).sum() / m.sum() else: raise ValueError(f"未知的方法 {job['method']!r}") if ref is not None: with torch.no_grad(): ref_logp, _ = token_logprobs(ref, X) d = ref_logp - logp # log(π_起点 / π) loss = loss + job["beta"] * ((d.exp() - d - 1.0) * m).sum() / m.sum() opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) opt.step() policy.eval() zero = float((r.std(dim=1) == 0).float().mean()) if r.shape[1] > 1 else 0.0 hist["steps"].append({ "proxy": proxy.mean().item(), "true": true.mean().item(), "zero": zero, "ent": ((ent.detach() * m).sum() / m.sum()).item(), "len": sum(len(x) for x in resps) / len(resps), "think": sum(bool(THINKS.match(x)) for x in resps) / len(resps), }) if job["method"] == "grpo" and job["reward"] != "true" and ((step + 1) % job["eval_every"] == 0 or step == 0): fooled = [(flat_probs[i].trace + hw.ANSWER_MARK + flat_probs[i].answer, resps[i], round(proxy[i].item(), 3)) for i in range(len(resps)) if true[i] == 0 and proxy[i] > 0.9] hist["examples"].append((step + 1, fooled[:3])) if (step + 1) % job["eval_every"] == 0 or step == steps - 1: do_eval(step + 1) hist["seconds"] = time.perf_counter() - t0 policy.eval() return policy, hist def run_path(job): """缓存文件名：方法 + 配置指纹。""" return os.path.join(CKPT_DIR, f"rl_{job['method']}_{ek.config_id(job)}.pt") def _train_and_save(job): """子进程入口：训练并保存（权重 + 历史）。""" model, hist = train_rl(job) os.makedirs(CKPT_DIR, exist_ok=True) torch.save({"job": job, "state_dict": model.state_dict(), "history": hist}, run_path(job)) return hist["seconds"] def load_runs(jobs, quick=False, verbose=True): """ 取得一组训练结果：缓存里有就读，没有就（并行地）训练。 【参数】jobs：make_job 生成的配置列表；quick：冒烟测试（每个只训练几步，不读写缓存，串行运行） 【返回】list[(model, history)]，顺序和 jobs 相同 【边界】 - 并行训练之前，先在主进程里把所有依赖（第 1 讲的 cot_L2、混合基座、第 4 讲的验证器）准备好并写进缓存： 否则几个子进程会同时训练同一个依赖、同时写同一个文件（第 4 讲 README 记录过这个问题）。 - 读取时核对缓存里存的配置和 job 完全一致。 """ if quick: return [train_rl(j, quick=True) for j in jobs] todo = [j for j in jobs if not os.path.exists(run_path(j))] if todo: inits = sorted({j["init"] for j in todo}) prepare_bases([n.split(":", 1)[1] for n in inits if n.startswith("mixed:")], verbose=verbose) for name in inits: if not name.startswith("mixed:"): cached = os.path.exists(os.path.join(tr.CKPT_DIR, f"{name}.pt")) tr.get_model(name, verbose=verbose and not cached) # 已有缓存时不重复打印"读取缓存" if any(j["reward"] in ("orm", "prm_full") for j in todo): tl.load_verifiers(verbose=verbose) if verbose: print(f" 需要训练 {len(todo)} 个模型（并行运行）…", flush=True) t0 = time.perf_counter() ek.run_parallel(_train_and_save, todo) if verbose: print(f" ⏱ 训练用时 {time.perf_counter() - t0:.0f} 秒（墙钟）", flush=True) out = [] for job in jobs: blob = torch.load(run_path(job), weights_only=False) assert blob["job"] == job, "缓存文件和配置不一致" model = tr.TinyReasoner(tr.ModelConfig()) model.load_state_dict(blob["state_dict"]) model.eval() out.append((model, blob["history"])) if verbose and len(todo) < len(jobs): print(f" 读取了 {len(jobs) - len(todo)} 个缓存的模型", flush=True) return out def prepare_bases(names, verbose=True): """ 并行训练还没有缓存的混合基座（代码 03 一次要用好几个）。 【返回】None；之后用 mixed_base(name) 读取 """ missing = [n for n in names if not os.path.exists(_mixed_path(n))] if missing: if verbose: print(f" 训练 {len(missing)} 个混合基座（两阶段，并行）…", flush=True) t0 = time.perf_counter() ek.run_parallel(_save_mixed, missing) if verbose: print(f" ⏱ 用时 {time.perf_counter() - t0:.0f} 秒（墙钟）", flush=True) # --------------------------------------------------------------------------- # 6. 评测 # --------------------------------------------------------------------------- def eval_set(n, ks=(3, 4, 5, 6), per_k=50): """ 固定的测试集：环大小 n，跳数 ks 各 per_k 道（种子 21，和第 4 讲 n = 12 的测试集相同）。 【返回】list[HopProblem] 【边界】n 最大 16（字母表只有 16 个字母）。 """ out = [] for k in ks: out += hw.make_eval_set(n, k, per_k, seed=21) return out def main_jobs(): """ 代码 01、02 共用的 9 个训练任务：三种方法（SFT、RFT、GRPO）× 三个种子，都从第 1 讲的 cot_L2 出发， 在 n = 4–12、k = 1–7 的题上训练 200 步，学习率都是 1e-3，每 25 步在 n = 12 的测试集上评测一次。 【为什么放在模块里】代码 02 要读代码 01 训练好的模型；两边用同一个函数生成配置， 配置指纹就一定相同，不会因为某一边改了一个字段而悄悄重新训练。 【返回】list[dict]，顺序：方法 sft、rft、grpo，每种方法内种子 0、1、2 """ return [make_job(method=m, seed=s, eval="n12") for m in ("sft", "rft", "grpo") for s in (0, 1, 2)] def pass_job(args): """ 子进程入口：在一个测试集上算一个模型的 pass@k 曲线（代码 02 用它并行地评测）。 【参数】args = (job, n, N, ks)：job 是 None（表示起点模型 cot_L2）或 main_jobs() 里的一个配置； n：测试集的环大小；N：每道题采样几次；ks：要算的 k 【返回】(curve, counts, distinct, ent)： curve / counts 同 pass_curve； distinct：每道题 N 个样本里不同答案的个数（平均）——"多样性"的直观度量； ent：模型在自己采样出的回答上，每个 token 的平均熵（nats） 【边界】子进程里只读缓存，不训练（调用前主进程要先用 load_runs 把模型训练好）。 """ torch.set_num_threads(1) job, n, N, ks = args if job is None: model = tr.get_model("cot_L2", verbose=False)[0] else: model = load_runs([job], verbose=False)[0][0] probs = eval_set(n) gen = torch.Generator().manual_seed(7) resps = sample_responses(model, probs, N, 1.0, gen) counts = [sum(true_reward(p, r) > 0.5 for r in rs) for p, rs in zip(probs, resps)] curve = {k: sum(tl.pass_at_k(N, c, k) for c in counts) / len(counts) for k in ks if k <= N} distinct = sum(len({hw.parse_answer(r, "cot") for r in rs}) for rs in resps) / len(resps) # 熵：每道题取前 16 个样本，算回答里每个 token 的预测分布的熵 sub_p = [p.prompt for p, rs in zip(probs, resps) for r in rs[:16]] sub_r = [r for rs in resps for r in rs[:16]] with torch.no_grad(): X, M = sequence_batch(sub_p, sub_r) _, ent = token_logprobs(model, X) m = M[:, 1:] ent = ((ent * m).sum() / m.sum()).item() return curve, counts, distinct, ent def think_set(per_k=40): """代码 03 的测试集：n = 8，k = 1 … 7 各 per_k 道（种子 31）。""" out = [] for k in range(1, 8): out += hw.make_eval_set(8, k, per_k, seed=31) return out def greedy_accuracy(model, probs): """贪心解码的准确率（第 1 讲的 evaluate；所有模型都用思维链格式作答）。""" return tr.accuracy(tr.evaluate(model, probs, "cot")) def pass_curve(model, probs, N, ks, temperature=1.0, seed=7): """ pass@k 曲线：每道题采样 N 次，用第 4 讲的无偏估计算每个 k 的 pass@k。 【返回】(curve, counts)：curve = {k: 平均 pass@k}；counts = 每道题 N 次里答对的次数 【边界】k > N 的项被跳过。 """ gen = torch.Generator().manual_seed(seed) resps = sample_responses(model, probs, N, temperature, gen) counts = [sum(true_reward(p, r) > 0.5 for r in rs) for p, rs in zip(probs, resps)] curve = {k: sum(tl.pass_at_k(N, c, k) for c in counts) / len(counts) for k in ks if k <= N} return curve, counts def think_profile(model, probs, N=16, seed=3): """ 按跳数统计：采样 N 次，有多少回答写了推理链（think）、多少答对（acc）、平均回答长度（len）。 【返回】dict：{"think": 总体比例, "acc": 总体准确率, "len": 平均长度, "acc_think": 写了推理链的回答的准确率, "acc_direct": 直接回答的准确率, "by_k": {k: {"think", "acc", "trace"}}}，trace 是写了推理链的回答里推理链的平均步数 【边界】没有任何"写了推理链"（或"直接回答"）的样本时，对应的条件准确率记为 0。 """ gen = torch.Generator().manual_seed(seed) resps = sample_responses(model, probs, N, 1.0, gen) tot = {"think": 0, "acc": 0, "acc_think": 0, "len": 0, "n": 0} by_k = {} for p, rs in zip(probs, resps): d = by_k.setdefault(p.k, {"think": 0, "acc": 0, "trace": 0, "n": 0}) for r in rs: t = bool(THINKS.match(r)) c = true_reward(p, r) > 0.5 tot["think"] += t tot["acc"] += c tot["acc_think"] += c and t tot["len"] += len(r) tot["n"] += 1 d["think"] += t d["acc"] += c d["trace"] += len(r.split(hw.ANSWER_MARK, 1)[0]) if t else 0 d["n"] += 1 out = {"think": tot["think"] / tot["n"], "acc": tot["acc"] / tot["n"], "len": tot["len"] / tot["n"], "acc_think": tot["acc_think"] / max(1, tot["think"]), "acc_direct": (tot["acc"] - tot["acc_think"]) / max(1, tot["n"] - tot["think"]), "by_k": {}} for k, d in sorted(by_k.items()): out["by_k"][k] = {"think": d["think"] / d["n"], "acc": d["acc"] / d["n"], "trace": d["trace"] / max(1, d["think"])} return out # --------------------------------------------------------------------------- # 自检 # --------------------------------------------------------------------------- def _self_test(): policy = tr.get_model("cot_L2", verbose=False)[0] rng = random.Random(0) probs = [hw.make_problem(rng, rng.randint(4, 12), rng.randint(1, 7)) for _ in range(64)] prompts = [p.prompt for p in probs] t0 = time.perf_counter() fast = sample_fast(policy, prompts, 12, 0, None) t1 = time.perf_counter() slow = tl.sample(policy, prompts, 12, 0, None) t2 = time.perf_counter() assert fast == slow a = sample_fast(policy, prompts, 12, 1.0, torch.Generator().manual_seed(5)) b = tl.sample(policy, prompts, 12, 1.0, torch.Generator().manual_seed(5)) assert a == b print(f"快速采样 = 第 4 讲的采样（64 道题，贪心和温度 1.0 都逐字相同） ⏱ {t1 - t0:.2f} 秒 vs {t2 - t1:.2f} 秒") p = hw.make_problem(random.Random(1), n=6, k=3) succ = {st[0]: st[2] for st in p.statements} over = p.trace + succ[p.answer] cases = { p.trace + "|" + p.answer + "$": "correct", over + "|" + over[-1] + "$": "overshoot", # 多走了一步，每一步查找都对 p.trace[:2] + "|" + p.trace[1] + "$": "early", # 只走了 2 步 p.trace + "|" + p.walk[1] + "$": "miscopy", # 推理链对，答案抄成了第 1 步的字母（n = 6 时它 ≠ 答案） p.trace[:-1]: "malformed", # 没有 "|" } for resp, want in cases.items(): assert classify(p, resp) == want, (resp, classify(p, resp), want) wrong_first = next(c for c in hw.LETTERS if c != p.walk[1]) assert classify(p, wrong_first + "|" + wrong_first + "$") in ("lookup", "correct") assert response_of("ab|b$xyz") == "ab|b$" and response_of("abc") == "abc" assert WELL_FORMED.match("ab|b$") and not WELL_FORMED.match("|b$") and not WELL_FORMED.match("ab|b|b") assert THINKS.match("ab|b$") and not THINKS.match("|b$") print("回答截断、分类（对 / 走过头 / 停早了 / 抄错 / 格式坏）、格式检查的自检通过") r = torch.tensor([[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.01, 0.0]]) a_std = group_advantages(r) a_raw = group_advantages(r, std_norm=False) assert a_std[1].abs().max() == 0 and a_raw[1].abs().max() == 0 # 全对的组：没有梯度 assert abs(a_raw[2, 2].item() - 0.0075) < 1e-6 # 几乎相同的奖励：原始优势很小 assert a_std[2, 2].item() > 1.4 # 除以标准差后被放大 print(f"组内优势：全对的组优势全为 0；奖励 [0, 0, 0.01, 0] 的组，原始优势 {a_raw[2, 2].item():.4f}，" f"除以标准差后 {a_std[2, 2].item():.2f}（Dr. GRPO 指出的放大效应）") job = make_job(steps=6, eval="none") t0 = time.perf_counter() _, hist = train_rl(job, quick=True) print(f"冒烟训练 GRPO 6 步：每步 {job['B']} 题 × {job['G']} 个样本，" f"第 1 步平均真实奖励 {hist['steps'][0]['true']:.3f} ⏱ {time.perf_counter() - t0:.1f} 秒") try: make_job(stpes=10) except KeyError: print("拼错字段名会报错（make_job(stpes=10) → KeyError）。自检通过。") if __name__ == "__main__": _self_test()