""" ============================================================================= 第 8 讲 · 代码 1：玩具语言模型、监督微调（SFT）、把生成看成马尔可夫决策过程 ============================================================================= 对应讲义：lectures/lecture08_rlhf.md 第 1、2 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- RLHF 流程的第一步是 SFT：在人写的示范上做最大似然。这个文件先介绍 toy_lm.py 里的玩具世界， 再做 SFT，最后把"生成一个回答"写成强化学习的语言（状态、动作、奖励），并验证一个后面要反复用到的事实： **整个回答的 KL = 每一步 KL 的期望之和**。 演示 1：玩具世界：词表、三种"奖励"、真实质量最高 / 标注员最喜欢的回答。 演示 2：SFT：在 5000 个示范上训练，得到参考策略 π_ref；示范越少，学得越不像。 演示 3：生成 = 马尔可夫决策过程：一个回答一步一步的概率；全部回答的概率之和 = 1； KL 的链式法则（逐词 KL 之和 = 整个回答的 KL）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 01_toy_lm_sft.py # 约 10 秒（CPU） 需要 numpy 和 PyTorch。依赖本目录的 toy_lm.py。 """ import numpy as np import torch import torch.nn.functional as F from toy_lm import (ALL, ALL_DEMO_LP, ALL_ENC, ALL_GOLD, ALL_LEN, CONTENT, EOS, MAX_LEN, WORDS, all_logprobs, clone_policy, demo_sample, encode, gold_reward, kl, labeler_utility, pad, print_header, sample, show, state_features, train_sft) # ============================================================================= # 演示 1 # ============================================================================= def demo_world(): print_header("演示 1：玩具世界 —— 词表、真实质量、标注员的效用") print(f"词表：{'、'.join(WORDS)}。一个回答最多 {MAX_LEN} 个内容词，所有可能的回答共 {len(ALL)} 个。\n") print("真实质量（只用于评估）：要点前两个各 +1.0、之后每个 −0.6；例子第一个 +0.7、之后每个 −0.4；") print("客套 −0.15、华丽 −0.1、错误 −1.2；空回答 −1.0。") print("众包标注员的效用 = 真实质量 + 0.5 × 华丽的个数 + 0.1 × 长度（风格偏差和长度偏差）。\n") lab = np.array([labeler_utility(s) for s in ALL]) for title, score in [("真实质量最高的 5 个回答", ALL_GOLD), ("标注员效用最高的 5 个回答", lab)]: order = np.argsort(-score, kind="stable")[:5] print(f" {title}：") for k in order: print(f" {pad(show(ALL[k]), 34)} 真实质量 {ALL_GOLD[k]:+.2f} 标注员效用 {lab[k]:+.2f}") print(""" 观察： · 真实质量最高的是“要点 要点 例子”这类回答：两个要点、一个例子，简洁。 · 标注员最喜欢的回答把剩下的位置填满了“华丽” —— 真实质量反而低了 0.3。 标注员的偏差不大（每个华丽词 +0.5），但它**有方向**：只要有空间，偏差就会被放大（第 5 讲第 9 节）。 · 这两个“最高”之间的差别，就是后面 RLHF 会被拉向的方向。""") # ============================================================================= # 演示 2 # ============================================================================= def demo_sft(): print_header("演示 2：监督微调（SFT）—— 在人写的示范上做最大似然") rng = np.random.default_rng(0) demos = demo_sample(5000, rng) print("示范分布（“人怎样写回答”）：写了 k 个词后停笔的概率依次为 0.02、0.10、0.20、0.30、0.45、0.60；") print("不停笔时，下一个词是要点 / 例子 / 客套 / 华丽 / 错误的概率为 0.35 / 0.20 / 0.20 / 0.15 / 0.10。\n") print(f"在 {len(demos)} 个示范上训练（Adam，每步 256 个示范）：") ref = train_sft(demos, verbose=True) ref_lp = all_logprobs(ref) p_demo, p_ref = np.exp(ALL_DEMO_LP), np.exp(ref_lp) print(f"\n {pad('', 12)} {pad('平均长度', 9, 'right')} {pad('平均真实质量', 13, 'right')} {pad('空回答的概率', 13, 'right')}") print(f" {pad('示范分布', 12)} {p_demo @ ALL_LEN:>9.3f} {p_demo @ ALL_GOLD:>+13.3f} {p_demo[0]:>13.3f}") print(f" {pad('SFT 策略', 12)} {p_ref @ ALL_LEN:>9.3f} {p_ref @ ALL_GOLD:>+13.3f} {p_ref[0]:>13.3f}") print("\n 示范的数量对 SFT 的影响（训练步数相同）：") for n in [50, 500, 5000]: pol = train_sft(demo_sample(n, np.random.default_rng(n)), seed=1) print(f" {n:>5d} 个示范：KL(示范分布 ‖ SFT 策略) = {kl(ALL_DEMO_LP, all_logprobs(pol)):.4f}") gen = torch.Generator().manual_seed(0) print("\n SFT 策略生成的 5 个回答：") for s in sample(ref, 5, gen): print(f" {pad(show(s), 34)} 真实质量 {gold_reward(s):+.2f}") print(""" 观察： · SFT 学的是“人通常怎样写”，学得很像（KL 只有 0.01 左右）。可是人写的回答平均真实质量只有 +0.7 左右， 远低于最好的 +2.7：示范里有客套、有错误、有时太长或太短。SFT 不知道哪些好、哪些不好，它照单全收。 · 示范太少时，SFT 会“背下”这几十个示范（KL 大得多）——小数据的过拟合。 · 这就是为什么还需要偏好：示范告诉模型“怎样写得像人”，偏好告诉模型“哪种写法更好”。 InstructGPT 的 SFT 数据约 1.3 万个提示词；偏好（奖励模型）数据约 3.3 万个提示词，每个有 4~9 个回答的排序。""") return ref # ============================================================================= # 演示 3 # ============================================================================= @torch.no_grad() def demo_mdp(ref): print_header("演示 3：生成 = 马尔可夫决策过程；KL 的链式法则") seq = (1, 1, 2) print(f"回答“{show(seq)}”是怎样一步一步生成的（SFT 策略）：\n") print(f" {pad('步', 3)} {pad('状态（已写的词）', 22)} {pad('动作', 8)} {pad('π(动作 | 状态)', 15, 'right')} " f"{pad('累计 log π', 11, 'right')}") cnt = torch.zeros((1, CONTENT)) total = 0.0 steps = list(seq) + ([EOS] if len(seq) < MAX_LEN else []) for t, a in enumerate(steps): probs = F.softmax(ref(state_features(cnt, torch.tensor([t]))), dim=-1)[0] total += float(torch.log(probs[a])) prefix = show(seq[:t]) if t else "（还没写）" print(f" {t:>3d} {pad(prefix, 22)} {pad(WORDS[a], 8)} {float(probs[a]):>15.3f} {total:>11.3f}") if a != EOS: cnt[0, a - 1] += 1 print(f"\n 这个回答的概率 = exp({total:.3f}) = {np.exp(total):.4f}。奖励只在最后一步（写完之后）给出。") lp = all_logprobs(ref) print(f" 全部 {len(ALL)} 个回答的概率之和 = {np.exp(lp).sum():.6f}（应该是 1）") # KL 的链式法则：构造一个和 SFT 策略不一样的策略（把参数随机扰动一下） torch.manual_seed(5) other = clone_policy(ref) for p in other.parameters(): p.add_(0.3 * torch.randn_like(p)) lp_o = all_logprobs(other) seq_kl = kl(lp_o, lp) # 逐词：在每个回答经过的每个状态上，算两个策略"下一个词的分布"之间的 KL，再按 π_other(回答) 加权 lo = F.log_softmax(other(ALL_ENC["states"]), dim=-1) lr = F.log_softmax(ref(ALL_ENC["states"]), dim=-1) step_kl = (lo.exp() * (lo - lr)).sum(-1) # (回答数, 7)：每一步的 KL # 第 t 步"发生"的条件：这个回答在第 t 步还没结束（mask） per_seq = (step_kl * ALL_ENC["mask"]).sum(1).double().numpy() token_kl = float(np.exp(lp_o) @ per_seq) print(f""" 另一个策略 π（把 SFT 策略的参数随机扰动）相对 π_ref 的 KL： 按整个回答算：KL(π ‖ π_ref) = Σ_y π(y) [log π(y) − log π_ref(y)] = {seq_kl:.5f} 按每一步算：E_π[ Σ_t KL( π(·|前缀_t) ‖ π_ref(·|前缀_t) ) ]（每一步“下一个词的分布”之间的 KL） = {token_kl:.5f} 观察： · 生成一个回答就是在一个马尔可夫决策过程里走一遍：状态 = 已经写出的词，动作 = 下一个词， 写完（〈结束〉或写满）时得到奖励。回答的概率 = 每一步概率的乘积。 · 两种算法得到同一个 KL —— 这是 KL 的链式法则。所以 RLHF 实现里“每个词上扣 β·(log π − log π_ref)”， 和“对整个回答的 KL 加惩罚”是一回事；真实的实现（InstructGPT、TRL 库）正是在每个词上算 KL 惩罚的。 · 真实的语言模型不能穷举所有回答（词表几万、长度几千），只能用采样估计这些量； 玩具世界可以穷举，所以后面的每个数都是精确的。""") def main(): demo_world() ref = demo_sft() demo_mdp(ref) if __name__ == "__main__": main()