20 KiB

Raw Blame History

第 22 课：DPO loss 源码

这一课只解决一个问题：MiniMind 的 train_dpo.py 如何把 policy/reference 对 chosen/rejected 的 token logprob 变成一个可反向传播的 DPO loss。

0. 本节主线

DPO loss 的计算链路是：

policy/ref 分别 forward 同一批 x
-> logits_to_log_probs 取出每个 y token 的 logprob
-> mask 去掉 prompt/padding，只保留 assistant 回复 token
-> 每条序列沿 seq 维求和，得到 sequence logprob
-> batch 前半是 chosen，后半是 rejected
-> policy chosen - policy rejected 得到 policy logratio
-> ref chosen - ref rejected 得到 reference logratio
-> 两个 logratio 相减，得到 policy 相对 reference 的偏好提升
-> loss = -logsigmoid(beta * improvement)

一句话：

DPO 奖励的是：policy 相比 reference，更明显地偏向 chosen 而不是 rejected。

这节课不是重新讲 DPODataset。第 21 课已经解决：

chosen/rejected -> x/y/mask
policy/ref -> token logprob

第 22 课只处理：

token logprob -> sequence logprob -> logratio -> DPO loss

1. 本节要懂的 6 个原理

原理	要理解什么	源码位置
token logprob 来自 gather	从 `[batch, seq, vocab]` 里取出目标 token 的 log probability	`trainer/train_dpo.py:25-31`
sequence logprob 是 masked sum	只把 assistant 区域 token logprob 加起来	`trainer/train_dpo.py:34-37`
batch 前半/后半有语义	前半是 chosen，后半是 rejected	`trainer/train_dpo.py:39-44`
policy logratio 表示 policy 偏好	`chosen_policy - rejected_policy` 越大，policy 越偏向 chosen	`trainer/train_dpo.py:46`
reference logratio 是基准偏好	DPO 看 policy 相对 ref 的偏好变化，不只看绝对概率	`trainer/train_dpo.py:47-48`
loss 用 `-logsigmoid`	improvement 越大，loss 越小；beta 控制强度	`trainer/train_dpo.py:49-50`, `README.md:951-956`

学完本节，你应该能手算一条偏好样本的：

chosen_policy_logprob
rejected_policy_logprob
chosen_ref_logprob
rejected_ref_logprob
pi_logratio
ref_logratio
logits = pi_logratio - ref_logratio
loss = -logsigmoid(beta * logits)

2. DPO loss 完整原理

2.1 为什么不是只最大化 chosen

如果只最大化 chosen，就会退回到类似 SFT：

让 chosen 的概率变高

但 DPO 的数据里还有 rejected。偏好信息不是：

chosen 是唯一正确答案

而是：

chosen 比 rejected 更好

所以 DPO 关注的是两条回答的相对概率：

log p(chosen) - log p(rejected)

这就是 logratio。

2.2 为什么还要 reference model

只看 policy 的 logratio：

policy_logratio = log pi(chosen) - log pi(rejected)

会鼓励 policy 强烈偏向 chosen。但这没有约束 policy 偏离原始 SFT 模型多远。

reference model 提供一个基准：

ref_logratio = log ref(chosen) - log ref(rejected)

DPO 看的是：

policy_logratio - ref_logratio

含义是：

policy 是否比 reference 更偏向 chosen。

如果 reference 本来就更喜欢 chosen，而 policy 只是保持原样，那么提升不大；如果 policy 在训练后更明显地喜欢 chosen，提升变大，loss 变小。

2.3 为什么 sequence logprob 要用 sum

一条回复有多个 token。

假设 chosen 回复 token 是：

y = [y1, y2, y3]

模型给这条回复的条件概率可以写成：

p(y|x) = p(y1|x) * p(y2|x,y1) * p(y3|x,y1,y2)

取 log 后乘法变加法：

log p(y|x) = log p(y1|...) + log p(y2|...) + log p(y3|...)

所以源码里：

(log_probs * mask).sum(dim=1)

得到的是每条回复的 sequence logprob。

注意：这里用 sum，不是 mean。sum 对应整段序列的 log probability。第 22 课先按源码理解，不展开不同长度归一化的变体。

2.4 `-logsigmoid` 在优化什么

DPO 源码最终计算：

loss = -logsigmoid(beta * logits)

其中：

logits = policy_logratio - ref_logratio

如果 logits 很大：

policy 相对 ref 更偏向 chosen
sigmoid(beta * logits) 接近 1
logsigmoid 接近 0
loss 接近 0

如果 logits 很小甚至为负：

policy 没有比 ref 更偏向 chosen，甚至更偏向 rejected
sigmoid(beta * logits) 小
logsigmoid 是更大的负数
loss 变大

所以 DPO loss 的梯度会推动 policy：

提高 chosen sequence logprob
降低 rejected sequence logprob
同时相对 reference 控制变化

2.5 beta 控制偏好优化强度

beta 乘在：

policy_logratio - ref_logratio

前面。

直观理解：

beta 越大，模型越强烈响应偏好差异；
beta 越小，更新更温和。

MiniMind 默认：

beta = 0.15

第 22 课只要求理解 beta 是 DPO loss 的强度系数，不要求从 KL 约束推导它。

3. 源码阅读顺序图

建议按这个顺序读：

1. README.md:948-956
   先看 DPO 公式长什么样。

2. trainer/train_dpo.py:25-31
   看 logits_to_log_probs 如何得到 token logprob。

3. trainer/train_dpo.py:34-37
   看 token logprob 如何通过 mask 变成 sequence logprob。

4. trainer/train_dpo.py:39-44
   看 chosen/rejected 如何按 batch 前后半拆开。

5. trainer/train_dpo.py:46-50
   看 policy/ref logratio 和最终 loss。

6. trainer/train_dpo.py:73-85
   看 train_epoch 如何把 ref/policy logprob 接到 dpo_loss。

和前几课的关系：

第 10 课：logits、labels、next-token 对齐
第 21 课：DPO x/y/mask 和 reference model
第 22 课：logprob 和 DPO loss
第 23 课：PPO 训练链路

4. MiniMind 源码走读

4.1 README 里的 DPO 公式

源码证据：DPO 数学形式

文件：README.md:948-956

看它是为了理解：代码里的 pi_logratios - ref_logratios 对应公式里的哪一部分。

公式摘录：


\mathcal{L}_{\mathrm{DPO}}
= -\mathbb{E}\left[
\log \sigma\left(
\beta \left[
\log \frac{\pi_{\theta}(y_w \mid x)}{\pi_{\mathrm{ref}}(y_w \mid x)}
-
\log \frac{\pi_{\theta}(y_l \mid x)}{\pi_{\mathrm{ref}}(y_l \mid x)}
\right]
\right)
\right]

把分式里的 log 展开：


\log \frac{\pi_{\theta}(y_w \mid x)}{\pi_{\mathrm{ref}}(y_w \mid x)}
=
\log \pi_{\theta}(y_w \mid x)
-
\log \pi_{\mathrm{ref}}(y_w \mid x)

所以中括号里的量可以改写成：


\left[
\log \pi_{\theta}(y_w \mid x)
-
\log \pi_{\theta}(y_l \mid x)
\right]
-
\left[
\log \pi_{\mathrm{ref}}(y_w \mid x)
-
\log \pi_{\mathrm{ref}}(y_l \mid x)
\right]

这段公式说明：

y_w 是 winner，也就是 chosen。
y_l 是 loser，也就是 rejected。
policy 是 pi_theta。
reference 是 pi_ref。
DPO 比的是 chosen/rejected 在 policy 和 reference 下的相对概率。

把公式映射到源码变量：

公式符号	源码变量	含义
`\log \pi_{\theta}(y_w \mid x)`	`chosen_policy_log_probs`	policy 对 chosen 回复的 sequence logprob
`\log \pi_{\theta}(y_l \mid x)`	`reject_policy_log_probs`	policy 对 rejected 回复的 sequence logprob
`\log \pi_{\mathrm{ref}}(y_w \mid x)`	`chosen_ref_log_probs`	reference 对 chosen 回复的 sequence logprob
`\log \pi_{\mathrm{ref}}(y_l \mid x)`	`reject_ref_log_probs`	reference 对 rejected 回复的 sequence logprob
`\log \pi_{\theta}(y_w \mid x) - \log \pi_{\theta}(y_l \mid x)`	`pi_logratios`	policy 对 chosen 相比 rejected 的偏好
`\log \pi_{\mathrm{ref}}(y_w \mid x) - \log \pi_{\mathrm{ref}}(y_l \mid x)`	`ref_logratios`	reference 对 chosen 相比 rejected 的基准偏好
`pi_logratios - ref_logratios`	`logits`	policy 相对 reference 的偏好提升
`-\log\sigma(\beta \cdot \text{logits})`	`loss`	DPO loss

对应到代码就是：

pi_logratios = chosen_policy_log_probs - reject_policy_log_probs
ref_logratios = chosen_ref_log_probs - reject_ref_log_probs
logits = pi_logratios - ref_logratios
loss = -F.logsigmoid(beta * logits)

4.2 logits_to_log_probs 取每个 token 的 logprob

源码证据：token logprob

文件：trainer/train_dpo.py:25-31

看它是为了理解：模型输出 logits 后如何取出目标 token 的 logprob。

代码摘录：

def logits_to_log_probs(logits, labels):
    log_probs = F.log_softmax(logits, dim=2)
    log_probs_per_token = torch.gather(log_probs, dim=2, index=labels.unsqueeze(2)).squeeze(-1)
    return log_probs_per_token

这段代码说明：

logits shape 是 [batch, seq, vocab]。
labels shape 是 [batch, seq]。
F.log_softmax(logits, dim=2) 在 vocab 维得到 log probability。
labels.unsqueeze(2) 变成 [batch, seq, 1]，作为 gather 索引。
返回值 shape 是 [batch, seq]，每个位置是目标 token 的 logprob。

注意：

DPO dataset 已经把 x/y 对齐好；
这里不再做 logits[..., :-1, :] 和 labels[..., 1:]。

如果你在这里再 shift 一次，就会错位。

4.3 mask 后求和得到 sequence logprob

源码证据：masked sum

文件：trainer/train_dpo.py:34-37

看它是为了理解：token-level logprob 如何变成一条回复的 logprob。

代码摘录：

def dpo_loss(ref_log_probs, policy_log_probs, mask, beta):
    ref_log_probs = (ref_log_probs * mask).sum(dim=1)
    policy_log_probs = (policy_log_probs * mask).sum(dim=1)

这段代码说明：

输入的 ref_log_probs 和 policy_log_probs 都是 [batch, seq]。
mask 也是 [batch, seq]。
prompt/padding 区域 mask 为 0，不参与求和。
assistant 回复区域 mask 为 1，参与求和。
sum(dim=1) 后变成 [batch]，每个元素是一条序列的 logprob。

这里覆盖第 21 课的一个关键点：

DPO 的忽略区域靠 mask，不靠 label=-100。

4.4 按 batch 前后半拆 chosen/rejected

源码证据：batch split

文件：trainer/train_dpo.py:39-44

看它是为了理解：为什么 train_epoch 拼接顺序不能乱。

代码摘录：

batch_size = ref_log_probs.shape[0]
chosen_ref_log_probs = ref_log_probs[:batch_size // 2]
reject_ref_log_probs = ref_log_probs[batch_size // 2:]
chosen_policy_log_probs = policy_log_probs[:batch_size // 2]
reject_policy_log_probs = policy_log_probs[batch_size // 2:]

这段代码说明：

dpo_loss 假设 batch 前半是 chosen。
batch 后半是 rejected。
chosen/ref 和 policy/ref 都按同样位置切分。
如果 train loop 拼接顺序错，loss 语义就错。

对应第 21 课的拼接：

x = torch.cat([x_chosen, x_rejected], dim=0)
y = torch.cat([y_chosen, y_rejected], dim=0)
mask = torch.cat([mask_chosen, mask_rejected], dim=0)

4.5 logratio 表示 chosen 相对 rejected 的偏好

源码证据：policy/ref logratio

文件：trainer/train_dpo.py:46-48

看它是为了理解：DPO 为什么不是单独看 chosen 概率。

代码摘录：

pi_logratios = chosen_policy_log_probs - reject_policy_log_probs
ref_logratios = chosen_ref_log_probs - reject_ref_log_probs
logits = pi_logratios - ref_logratios

这段代码说明：

pi_logratios 是 policy 对 chosen 相比 rejected 的偏好。
ref_logratios 是 reference 对 chosen 相比 rejected 的基准偏好。
logits 是 policy 相对 reference 的偏好提升。
如果 logits > 0，policy 比 reference 更偏向 chosen。
如果 logits < 0，policy 比 reference 更偏向 rejected，或偏好提升不够。

变量名里的 logits 这里不是 vocab logits，而是 DPO 二分类目标里的标量分数。它 shape 是：

[batch_size / 2]

不要和模型输出的：

[batch, seq, vocab]

混淆。

4.6 `-logsigmoid` 把偏好提升变成 loss

源码证据：最终 loss

文件：trainer/train_dpo.py:49-50

看它是为了理解：为什么 logits 越大 loss 越小。

代码摘录：

loss = -F.logsigmoid(beta * logits)
return loss.mean()

这段代码说明：

beta * logits 是带强度系数的偏好提升。
logsigmoid 越接近 0，说明 sigmoid 越接近 1。
前面加负号后，偏好提升越大，loss 越小。
mean() 对 batch 里的偏好对取平均。

手算时可以记住：

logits 很大 -> loss 接近 0
logits = 0 -> loss = -log(0.5) = 0.693...
logits 为负 -> loss 大于 0.693...

4.7 train_epoch 如何接入 DPO loss

源码证据：训练步骤

文件：trainer/train_dpo.py:73-85

看它是为了理解：DPO loss 如何进入普通训练循环。

代码摘录：

with torch.no_grad():
    ref_outputs = ref_model(x)
    ref_logits = ref_outputs.logits
ref_log_probs = logits_to_log_probs(ref_logits, y)

outputs = model(x)
logits = outputs.logits
policy_log_probs = logits_to_log_probs(logits, y)

dpo_loss_val = dpo_loss(ref_log_probs, policy_log_probs, mask, beta=beta)
loss = dpo_loss_val + outputs.aux_loss
loss = loss / args.accumulation_steps

这段代码说明：

ref model forward 不构建梯度。
policy model forward 正常构建梯度。
两边 logprob 的 shape 都是 [2 * batch, seq]。
dpo_loss 返回标量。
如果 policy 是 MoE，还会加 outputs.aux_loss。
后续 backward/clip/step 和普通训练一样。

5. 本节必须会写 / 暂时不要求

必须会写：

1. 从 logits 和 labels 得到 token logprob：
   log_probs = log_softmax(logits, dim=-1)
   token_logps = gather(log_probs, labels)

2. 从 token logprob 得到 sequence logprob：
   seq_logps = (token_logps * mask).sum(dim=1)

3. 拆 chosen/rejected：
   chosen = seq_logps[:batch_size//2]
   rejected = seq_logps[batch_size//2:]

4. 计算 logratio：
   pi_logratio = chosen_policy - rejected_policy
   ref_logratio = chosen_ref - rejected_ref

5. 计算 DPO loss：
   logits = pi_logratio - ref_logratio
   loss = -logsigmoid(beta * logits).mean()

暂时不要求：

1. 从 PPO KL 约束推导 DPO。
2. beta 的系统调参。
3. reference logprob 预缓存。
4. IPO / KTO / ORPO 等变体。
5. 长度归一化、reward margin、label smoothing 等 DPO 变体。

6. 手写模块

本节对应教学版文件：

course/impl/core/losses.py

6.1 补 `sequence_log_probs`

接口：

def sequence_log_probs(logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    ...

对齐源码：

trainer/train_dpo.py::logits_to_log_probs
trainer/train_dpo.py::dpo_loss 的 masked sum

行为要求：

输入:
  logits: [batch, seq, vocab]
  labels: [batch, seq]
  mask:   [batch, seq]

输出:
  seq_logps: [batch]

内部步骤：

log_softmax
gather labels
乘 mask
sum(dim=1)

6.2 补 `dpo_loss`

接口：

def dpo_loss(
    policy_chosen_logps: torch.Tensor,
    policy_rejected_logps: torch.Tensor,
    ref_chosen_logps: torch.Tensor,
    ref_rejected_logps: torch.Tensor,
    beta: float,
) -> torch.Tensor:
    ...

对齐源码：

trainer/train_dpo.py::dpo_loss

行为要求：

pi_logratios = policy_chosen_logps - policy_rejected_logps
ref_logratios = ref_chosen_logps - ref_rejected_logps
logits = pi_logratios - ref_logratios
loss = -logsigmoid(beta * logits).mean()

教学版接口把 chosen/rejected 拆开传入，比原源码更显式：

原源码：输入拼接后的 ref_log_probs / policy_log_probs / mask，再内部拆前后半。
教学版：先在调用处得到四个 sequence logprob，再传给 dpo_loss。

这样更容易测试和手算。

7. 实验验证

本节新增实验：

course/labs/trace_dpo_loss.py

它复用：

course/labs/tiny_dpo.jsonl

实验验证：

1. 从 policy/ref logits 得到 token logprob。
2. mask 后得到 chosen/rejected sequence logprob。
3. 手算 pi_logratio 和 ref_logratio。
4. 手算 logits = pi_logratio - ref_logratio。
5. 手算 -logsigmoid(beta * logits)。
6. 和 train_dpo.py::dpo_loss 输出对齐。

运行命令：

cd /home/sun/minimind
python course/labs/trace_dpo_loss.py \
  --data_path course/labs/tiny_dpo.jsonl \
  --max_length 96 \
  --beta 0.15

重点记录：

chosen_policy_logp =
rejected_policy_logp =
chosen_ref_logp =
rejected_ref_logp =
pi_logratio =
ref_logratio =
dpo_logits =
manual_dpo_loss =
source_dpo_loss =
abs_diff =

初始 policy/ref 权重相同时，常见现象是：

pi_logratio == ref_logratio
dpo_logits 接近 0
loss 接近 0.693147

因为训练还没开始，policy 相对 reference 没有偏好提升。

8. 阶段组装

DPO 阶段现在有两课：

第 21 课：
  chosen/rejected 数据
  DPODataset
  policy/ref model
  token logprob 输入

第 22 课：
  sequence logprob
  policy/ref logratio
  DPO loss

后续教学版 DPO 最小闭环可以这样组装：

CourseDPODataset
-> policy/ref model
-> sequence_log_probs
-> dpo_loss
-> optimizer 更新 policy
-> 保存 dpo 权重

8.1 阶段验收顺序

建议先跑两个 lab：

cd /home/sun/minimind
python course/labs/trace_dpo_dataset_reference.py --data_path course/labs/tiny_dpo.jsonl --max_length 96
python course/labs/trace_dpo_loss.py --data_path course/labs/tiny_dpo.jsonl --max_length 96 --beta 0.15

如果补了 course/impl/core/losses.py，后续可新增：

course/impl/tests/test_dpo.py

测试内容：

sequence_log_probs 对齐源码 gather + mask sum
dpo_loss 对齐 train_dpo.py::dpo_loss

8.2 Portfolio 记录

完成本节后，可以在 course/portfolio/experiments.md 记录：

DPO loss trace:
- 用 tiny chosen/rejected 样本构造 x/y/mask。
- 计算 policy/ref token logprob。
- mask 后求 sequence logprob。
- 拆 chosen/rejected，计算 policy/ref logratio。
- 手算 DPO loss 并和源码函数输出对齐。

在 course/notes/mistakes.md 记录容易错的点：

1. 把 DPO logits 和模型 vocab logits 混淆。
2. 对 DPO 的 x/y 又额外 shift 一次。
3. 忘记 mask，导致 prompt/padding 进入 sequence logprob。
4. chosen/rejected 前后半拆反。
5. 把 alpha/beta 和蒸馏课的 alpha 混淆。

9. 本节检查

logits_to_log_probs 的输出 shape 是什么？
为什么 sequence logprob 用 (token_logps * mask).sum(dim=1)？
DPO 源码为什么假设 batch 前半是 chosen，后半是 rejected？
pi_logratios 和 ref_logratios 分别代表什么？
logits = pi_logratios - ref_logratios 里的 logits 为什么不是 vocab logits？
当 policy 和 reference 完全一样时，DPO loss 大约是多少？为什么？
beta 变大时，-logsigmoid(beta * logits) 对同一个 logits 会发生什么变化？

10. 下一课

第 23 课进入 PPO 训练链路。

下一课要解决：

PPO 为什么需要 rollout；
actor、critic、reference、reward model 分别是什么；
response logprob、reward、advantage 如何进入 PPO loss；
clip ratio 和 KL penalty 在源码里如何实现；
PPO 和 DPO 的训练数据、模型数量、更新方式有什么不同。

20 KiB Raw Blame History Unescape Escape

第 22 课：DPO loss 源码

目录

0. 本节主线

1. 本节要懂的 6 个原理

2. DPO loss 完整原理

2.1 为什么不是只最大化 chosen

2.2 为什么还要 reference model

2.3 为什么 sequence logprob 要用 sum

2.4 -logsigmoid 在优化什么

2.5 beta 控制偏好优化强度

3. 源码阅读顺序图

4. MiniMind 源码走读

4.1 README 里的 DPO 公式

源码证据：DPO 数学形式

4.2 logits_to_log_probs 取每个 token 的 logprob

源码证据：token logprob

4.3 mask 后求和得到 sequence logprob

源码证据：masked sum

4.4 按 batch 前后半拆 chosen/rejected

源码证据：batch split

4.5 logratio 表示 chosen 相对 rejected 的偏好

源码证据：policy/ref logratio

4.6 -logsigmoid 把偏好提升变成 loss

源码证据：最终 loss

4.7 train_epoch 如何接入 DPO loss

源码证据：训练步骤

5. 本节必须会写 / 暂时不要求

6. 手写模块

6.1 补 sequence_log_probs

6.2 补 dpo_loss

7. 实验验证

8. 阶段组装

8.1 阶段验收顺序

8.2 Portfolio 记录

9. 本节检查

10. 下一课

20 KiB

Raw Blame History

2.4 `-logsigmoid` 在优化什么

4.6 `-logsigmoid` 把偏好提升变成 loss

6.1 补 `sequence_log_probs`

6.2 补 `dpo_loss`