""" ============================================================================= 第 10 讲 · 代码 4：在迷你 GPT 上做 DPO 和 best-of-n ============================================================================= 对应讲义：lectures/lecture10_post_training.md 第 5、6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 03 号文件在只有 6 个回答的玩具策略上看清了 DPO 和过度优化。这里在真正的语言模型上做同样的事： 参考模型 π_ref：01 号文件的 SFT 模型（会听 "write a positive review"） 偏好数据 ：从 π_ref 采样两条回答，模拟的标注者按金标准奖励选一条（post_utils） 目标 ：让模型写出金标准奖励更高的评论（更正面，但不啰嗦） —— 这就是 DPO 论文第一个实验（IMDb 受控情感生成）的玩具版 演示 1：DPO 训练（β = 0.1 和 β = 0.5），比较金标准奖励、奖励模型分数、长度、多样性、和参考模型的 KL 演示 2：best-of-n：用奖励模型（02 号文件）挑 n 条里最好的 —— n 越大，真实奖励一定越高吗？ 需要 checkpoints/sft.pt（01 号文件）和 checkpoints/rm.pt（02 号文件）；找不到时会自动生成（比较慢）。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_dpo_lm.py # 完整运行（约 5~10 分钟，不含 SFT 和奖励模型） python 04_dpo_lm.py --quick # 冒烟测试 【关于结果】编写讲义时这个文件只运行了 --quick 冒烟测试，没有完整运行。 "预期会看到"描述的是应该出现的趋势，具体数字以你运行的结果为准。 """ import os os.environ.setdefault("OMP_NUM_THREADS", "1") import copy import math import sys import time import numpy as np import torch import torch.nn.functional as F import post_utils as P # 先导入它：它会把第 7、8、9 讲的文件夹加进搜索路径 import lm_utils as U # noqa: E402 from toy_reviews import Vocab, sentiment_score # noqa: E402 torch.set_num_threads(1) def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def response_logprobs(model, vocab, responses, batch_size=64): """ log π(y | x)：对指令 "write a positive review"，每条回复（加上结尾的 ）的对数概率。 【返回】torch.Tensor (N,)，保留梯度（训练 DPO 时需要） """ prompt = P.format_prompt(vocab, P.INSTRUCTIONS["pos"]) eos = vocab.stoi[Vocab.EOS] out = [] for start in range(0, len(responses), batch_size): chunk = responses[start:start + batch_size] conts = [vocab.encode(r, add_bos=False, add_eos=False) for r in chunk] out.append(U.sequence_logprob(model, [prompt] * len(chunk), conts, eos=eos)) return torch.cat(out) def train_dpo(policy, ref, pairs, vocab, beta, epochs=2, lr=1e-4, batch_size=32, seed=0): """ DPO 训练。 loss = −mean log σ( β·[ (log π_θ(y_w) − log π_ref(y_w)) − (log π_θ(y_l) − log π_ref(y_l)) ] ) 【参数】 policy : 要训练的模型（从 ref 复制而来） ref : 参考模型（不训练） pairs : post_utils.make_preference_pairs 的输出（只用 chosen / rejected，不看金标准奖励） beta : β，越大越"保守" epochs, lr, batch_size, seed : 训练超参数（DPO 通常用很小的学习率） 【返回】policy 【思路】（讲义第 5 节）参考模型的对数概率在训练中不变，先算好存起来； 每一步只需要对 policy 算一次前向和反向传播。不需要采样、不需要奖励模型、不需要价值函数。 """ with torch.no_grad(): ref_w = response_logprobs(ref, vocab, [p["chosen"] for p in pairs]) ref_l = response_logprobs(ref, vocab, [p["rejected"] for p in pairs]) opt = torch.optim.AdamW(policy.parameters(), lr=lr, weight_decay=0.0) rng = np.random.default_rng(seed) policy.train() for epoch in range(epochs): t0, losses, accs = time.perf_counter(), [], [] order = rng.permutation(len(pairs)) for start in range(0, len(order), batch_size): idx = order[start:start + batch_size] lw = response_logprobs(policy, vocab, [pairs[i]["chosen"] for i in idx]) ll = response_logprobs(policy, vocab, [pairs[i]["rejected"] for i in idx]) margin = beta * ((lw - ref_w[idx]) - (ll - ref_l[idx])) loss = -F.logsigmoid(margin).mean() opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) opt.step() losses.append(loss.item()) accs.append((margin > 0).float().mean().item()) print(f" 第 {epoch + 1} 轮：DPO 损失 {np.mean(losses):.3f}，隐式奖励把\"选中\"排在前面的比例 {np.mean(accs):.1%}，" f"用时 {time.perf_counter() - t0:.0f} 秒", flush=True) policy.eval() return policy def distinct2(responses): """多样性：所有回复里不同的相邻两词（2-gram）占全部 2-gram 的比例（第 5 讲）。""" grams = [tuple(r[i:i + 2]) for r in responses for i in range(len(r) - 1)] return len(set(grams)) / max(1, len(grams)) @torch.no_grad() def evaluate(model, ref, vocab, w_rm, n=300, seed=123): """ 从模型采样 n 条回复，统计各项指标。 【返回】dict：gold（金标准奖励的平均）、rm（奖励模型分数的平均）、pos（正面的比例）、 len（平均长度）、over30（超过 30 个词的比例）、div（distinct-2）、kl（KL(π‖π_ref) 的估计）、example 【思路】KL(π ‖ π_ref) = E_{y~π}[ log π(y) − log π_ref(y) ]，用采样到的回复做蒙特卡洛估计。 【边界】KL 本身不会是负数，但样本少、两个模型又很接近时，蒙特卡洛估计可能略小于 0（采样误差）。 """ ys = P.generate(model, vocab, P.INSTRUCTIONS["pos"], n, seed=seed) kl = (response_logprobs(model, vocab, ys) - response_logprobs(ref, vocab, ys)).mean().item() return {"gold": np.mean([P.gold_reward(y) for y in ys]), "rm": P.rm_score(w_rm, ys, vocab).mean(), "pos": np.mean([sentiment_score(y) > 0 for y in ys]), "len": np.mean([len(y) for y in ys]), "over30": np.mean([len(y) > P.LENGTH_LIMIT for y in ys]), "div": distinct2(ys), "kl": kl, "example": ys[0]} def print_row(name, r): print(f" {name:<16s}{r['gold']:>8.2f}{r['rm']:>8.2f}{r['pos']:>8.1%}{r['len']:>7.1f}{r['over30']:>8.1%}" f"{r['div']:>8.2f}{r['kl']:>8.2f}") def main(quick=False): vocab = Vocab() ref = P.load_or_sft(vocab, quick=quick) w_rm = P.load_or_train_rm(vocab, sft_model=ref, quick=quick) pairs = P.make_preference_pairs(ref, vocab, 100 if quick else 3000, seed=0) n_eval = 50 if quick else 300 # ---------------- 演示 1 ---------------- print_header(f"演示 1：DPO（{len(pairs)} 对偏好数据）") rows = {"SFT（参考模型）": evaluate(ref, ref, vocab, w_rm, n=n_eval)} for beta in [0.1, 0.5]: print(f"\n β = {beta}：") policy = train_dpo(copy.deepcopy(ref), ref, pairs, vocab, beta, epochs=1 if quick else 2) rows[f"DPO，β = {beta}"] = evaluate(policy, ref, vocab, w_rm, n=n_eval) print(f"\n {'模型':<16s}{'金标准':>8s}{'奖励模型':>8s}{'正面':>8s}{'长度':>7s}{'>30 词':>8s}{'多样性':>8s}{'KL':>8s}") for name, r in rows.items(): print_row(name, r) print("\n 每个模型的一条回复：") for name, r in rows.items(): print(f" {name}：{' '.join(r['example'])}") print(""" 预期会看到（编写讲义时只做了冒烟测试，具体数字以你运行的结果为准）： - DPO 之后，金标准奖励和奖励模型的分数都应该比 SFT 模型高：回复更正面，更多 "i loved it" 之类的总结句。 - β 越小，离参考模型越远（KL 越大），奖励通常越高，但多样性也可能下降更多 —— 模型开始反复使用少数几种"讨好"的说法（模式坍缩，讲义第 6.3 节）。 - 注意"长度"一列：偏好数据里，更长的正面评论往往含有更多褒义词，可能被标注者偏爱； 如果 DPO 让回复变长、超过 30 个词的比例上升，金标准奖励就会因为"啰嗦"被扣分 —— 和真实 RLHF 中的长度偏见是同一件事。 - DPO 的训练只用到了偏好数据，没有用到奖励模型；"奖励模型"一列只是用来对照。""") # ---------------- 演示 2 ---------------- print_header("演示 2：best-of-n —— 用奖励模型从 n 条回复里挑最好的") m = 20 if quick else 100 n_max = 16 if quick else 64 ys = P.generate(ref, vocab, P.INSTRUCTIONS["pos"], m * n_max, seed=7) rm = P.rm_score(w_rm, ys, vocab).reshape(m, n_max) gold = np.array([P.gold_reward(y) for y in ys]).reshape(m, n_max) lens = np.array([len(y) for y in ys]).reshape(m, n_max) print(f"\n （从 SFT 模型采样 {m} 组，每组 {n_max} 条，每组用前 n 条做 best-of-n）") print(f" {'n':>4s}{'KL 上界':>9s}{'奖励模型':>10s}{'金标准（用奖励模型选）':>20s}{'长度':>7s}{'金标准（用金标准选）':>20s}") for n in [1, 2, 4, 8, 16, 32, 64]: if n > n_max: break pick = rm[:, :n].argmax(axis=1) rows_idx = np.arange(m) print(f" {n:>4d}{math.log(n) - (n - 1) / n:>9.2f}{rm[rows_idx, pick].mean():>10.2f}" f"{gold[rows_idx, pick].mean():>20.2f}{lens[rows_idx, pick].mean():>9.1f}" f"{gold[:, :n].max(axis=1).mean():>20.2f}") print(""" 预期会看到（具体数字以你运行的结果为准）： - best-of-n 是最简单的"对齐"方法：不训练模型，只是多采样、挑最好的。它的 KL 有一个简单的上界 log n − (n−1)/n。 - 奖励模型的分数随 n 一直上升。用奖励模型挑出的回复往往越来越长（褒义词更多 → 奖励模型更喜欢）， 而金标准对超过 30 个词的回复扣分，所以"用奖励模型选"的金标准奖励涨得比奖励模型慢，n 很大时可能停滞甚至下降。 - 最后一列（用金标准本身来选）一直上升：差距完全来自奖励模型和真实偏好之间的差异 —— 过度优化（03 号文件演示 6）。 - 可以把演示 1 的 DPO 模型按 KL 放在这张表里比较：在差不多的 KL 下，哪种方法的金标准奖励更高？ （DPO 论文图 2 就是这样比较"奖励—KL"前沿的。）""") if __name__ == "__main__": main(quick="--quick" in sys.argv)