365 lines
9.8 KiB
Markdown
365 lines
9.8 KiB
Markdown
# Lesson 23 笔记：PPO 中 Response 采样与梯度问题
|
||
|
||
## 目录
|
||
|
||
- [核心问题：为什么 y 不能直接求导](#核心问题为什么-y-不能直接求导)
|
||
- [普通 CE Loss 的梯度怎么传](#普通-ce-loss-的梯度怎么传)
|
||
- [PPO 的问题：y 是"采样"出来的](#ppo-的问题y-是采样出来的)
|
||
- [DPO 的解决方案](#dpo-的解决方案)
|
||
- [PPO 的解决方案：Reinforce 估计](#ppo-的解决方案reinforce-估计)
|
||
- [一句话总结](#一句话总结)
|
||
|
||
---
|
||
|
||
## 核心问题：为什么 y 不能直接求导
|
||
|
||
> response y 是模型采样出来的离散 token 序列，不能像普通 CE loss 那样直接对 token 选择过程求导。
|
||
|
||
---
|
||
|
||
## 普通 CE Loss 的梯度怎么传
|
||
|
||
```python
|
||
# 普通 SFT
|
||
logits = model(input_ids) # 模型输出每个 token 的 logits
|
||
loss = F.cross_entropy(logits, labels) # 直接拿 ground truth labels 算 loss
|
||
loss.backward() # 梯度直接流回 logits → embedding
|
||
```
|
||
|
||
**关键：** labels 是**固定的 ground truth**，不是从模型里采样出来的。
|
||
|
||
| 步骤 | 说明 |
|
||
|------|------|
|
||
| 1 | logits = model(input_ids) |
|
||
| 2 | loss = cross_entropy(logits, y_fixed) |
|
||
| 3 | backward → 梯度流回 logits |
|
||
|
||
---
|
||
|
||
## PPO 的问题：y 是"采样"出来的
|
||
|
||
```python
|
||
# PPO 里
|
||
y_w = torch.multinomial(probs, num_samples) # 从模型分布中采样
|
||
```
|
||
|
||
| 问题 | 说明 |
|
||
|------|------|
|
||
| 采样是随机的 | `torch.multinomial` 每次结果可能不同 |
|
||
| 采样不可导 | autograd 无法跟踪随机采样过程 |
|
||
|
||
**梯度流断在这里：**
|
||
|
||
```
|
||
logits → softmax → sample → y (随机!)
|
||
↑
|
||
这个随机采样过程没有梯度
|
||
```
|
||
|
||
---
|
||
|
||
## DPO 的解决方案
|
||
|
||
DPO **绕过了这个问题**，不采样，而是直接用 model 输出的 log_probs：
|
||
|
||
```python
|
||
# DPO 直接用模型输出的概率分布
|
||
policy_log_probs = model(x_chosen).log_probs # 可导
|
||
ref_log_probs = ref_model(x_chosen).log_probs # 可导
|
||
|
||
# 计算 ratio 时，不需要对采样求导
|
||
log_ratio = policy_log_probs - ref_log_probs
|
||
```
|
||
|
||
---
|
||
|
||
## PPO 的解决方案：Reinforce 估计
|
||
|
||
PPO 需要对采样过程求导，用的是 **Reinforce 梯度估计**：
|
||
|
||
```
|
||
∇θ E_{y~πθ}[r(y)] = E_{y~πθ}[r(y) * ∇θ log πθ(y)]
|
||
```
|
||
|
||
| 步骤 | 说明 |
|
||
|------|------|
|
||
| 1 | 从 π_θ(y|x) 采样 y |
|
||
| 2 | 计算 reward r(y) |
|
||
| 3 | 用 ∇ log πθ(y) 把 reward 的信号传回去 |
|
||
|
||
**关键技巧：** 用 log probability 的梯度来"引导"随机采样的方向。
|
||
|
||
---
|
||
|
||
## PPO + KL 约束公式：k = r - log r - 1
|
||
|
||
### 公式拆解
|
||
|
||
```python
|
||
k = π_ref(a|s) / π_θ(a|s) - log(π_ref(a|s) / π_θ(a|s)) - 1
|
||
```
|
||
|
||
设 `r = π_ref / π_θ`，则：
|
||
|
||
```
|
||
k(r) = r - log r - 1
|
||
```
|
||
|
||
### 核心性质
|
||
|
||
| 性质 | 说明 |
|
||
|------|------|
|
||
| **最小值在 r = 1** | k(1) = 1 - 0 - 1 = 0，policy = reference 时无惩罚 |
|
||
| **永远 ≥ 0** | 经典不等式 r - log r - 1 ≥ 0 |
|
||
| **偏离越大，惩罚越大** | r >> 1 或 r << 1 时，k 增大 |
|
||
|
||
### 为什么这样设计
|
||
|
||
> 限制 policy 不要离 reference 太远，但又不能太死
|
||
|
||
| 公式项 | 含义 |
|
||
|--------|------|
|
||
| `r = π_ref / π_θ` | policy 相对于 reference 的概率比值 |
|
||
| `r - log r - 1` | KL divergence 的单 token surrogate |
|
||
|
||
**作用：**
|
||
- 保证非负
|
||
- 在 r=1 处最小
|
||
- 对大偏离进行更强惩罚
|
||
- 实现稳定的策略更新约束
|
||
|
||
---
|
||
|
||
## Thinking 机制：open_thinking 的作用
|
||
|
||
### RLAIFDataset 中的 thinking 控制
|
||
|
||
```python
|
||
def create_chat_prompt(self, conversations):
|
||
use_thinking = random.random() < self.thinking_ratio
|
||
return self.tokenizer.apply_chat_template(
|
||
conversations[:-1],
|
||
tokenize=False,
|
||
open_thinking=use_thinking,
|
||
add_generation_prompt=True
|
||
)
|
||
```
|
||
|
||
| 模式 | 效果 |
|
||
|------|------|
|
||
| `open_thinking=True` | 模型学会"先想再答"，输出包含 `<think>` 推理过程 |
|
||
| `open_thinking=False` | 模型学会"直接回答"，不输出推理过程 |
|
||
|
||
### chat template 中的实现
|
||
|
||
```jinja
|
||
{%- if open_thinking is defined and open_thinking is true %}
|
||
{{- '<think>\n' }} {# 只打开，不关闭，让模型自己生成 </think> #}
|
||
{%- else %}
|
||
{{- '<think>\n\n</think>\n\n' }} {# 完整空 thinking 块 #}
|
||
{%- endif %}
|
||
```
|
||
|
||
| open_thinking | 模板输出 |
|
||
|----------------|----------|
|
||
| True | `<|im_start|>assistant\n<think>`（打开但没关闭） |
|
||
| False | `<|im_start|>assistant\n<think>\n\n</think>\n\n`（完整空块） |
|
||
|
||
### 之前阶段有没有训练过 thinking
|
||
|
||
| 阶段 | 有 thinking 训练数据吗 |
|
||
|------|----------------------|
|
||
| Pretrain | 没有 |
|
||
| SFT | 有字段但**实际数据没有**（`reasoning_content` 是空的） |
|
||
| RLAIFDataset | **有** — 通过 `open_thinking` 参数控制 |
|
||
|
||
### SFT 中也有 thinking 相关机制
|
||
|
||
SFTDataset 的数据 features 定义了 `reasoning_content` 字段，chat template 对 assistant 消息会做：
|
||
|
||
```jinja
|
||
<|im_start|>assistant
|
||
<think>
|
||
{reasoning_content}
|
||
</think>
|
||
{content}<|im_end|>
|
||
```
|
||
|
||
另外 `post_processing_chat()` 会随机移除空 thinking 块：
|
||
|
||
```python
|
||
if '<think>\n\n</think>\n\n' in prompt_content and random.random() > empty_think_ratio:
|
||
prompt_content = prompt_content.replace('<think>\n\n</think>\n\n', '')
|
||
```
|
||
|
||
### 总结：thinking 训练流程
|
||
|
||
| 阶段 | 作用 |
|
||
|------|------|
|
||
| **SFT** | 让模型见过 `<think>`...`</think>` 格式，也见过空 thinking 或被移除的情况 |
|
||
| **RLAIF/PPO** | 用 `thinking_ratio` 随机决定 rollout 时"打开 thinking"还是"空 thinking 直答" |
|
||
| **推理** | 用户用 `open_thinking` 开关控制模型偏向"显式思考"还是"直接回答" |
|
||
|
||
---
|
||
|
||
## PPO 中四个模型的角色
|
||
|
||
| 模型 | 来源 | 训练方式 |
|
||
|------|------|----------|
|
||
| **actor_model** | 从 full_sft 来 | PPO 训练它 |
|
||
| **ref_model** | 从 full_sft 来 | 冻结，当参考模型 |
|
||
| **critic_model** | 从 full_sft backbone 来，加 value_head | PPO 训练它 |
|
||
| **reward_model** | 外部奖励模型 | 只负责给 response 打分 |
|
||
|
||
---
|
||
|
||
## SGLang 详解
|
||
|
||
### 一句话定性
|
||
|
||
> HTTP API 的作用：把"模型生成能力"变成一个"可调用的远程服务接口"。本质不是为了"网络"，而是为了：解耦 + 调度 + 并发 + 工程化部署。
|
||
|
||
---
|
||
|
||
### model.generate vs SGLang
|
||
|
||
| 方式 | model.generate | SGLang |
|
||
|------|---------------|--------|
|
||
| 形式 | 函数调用 | HTTP / RPC |
|
||
| 执行 | 当前进程 | 独立推理服务 |
|
||
| batch | 固定 batch | Continuous batching |
|
||
| cache | 本地函数里 | server 统一管理 |
|
||
|
||
---
|
||
|
||
### SGLang 额外做了什么
|
||
|
||
**没有发生的事：**
|
||
- ❌ 改变模型结构
|
||
- ❌ 改变 forward 计算
|
||
- ❌ 替换 model.generate 的算法本质
|
||
|
||
**发生的事：**
|
||
- 把"generate 循环"从 Python 进程搬到一个高性能推理系统里执行
|
||
|
||
---
|
||
|
||
### Continuous Batching（关键）
|
||
|
||
| 方式 | 行为 |
|
||
|------|------|
|
||
| HF generate | 等所有请求到齐，固定 batch 一起生成 |
|
||
| SGLang | 动态拼 batch，完成就退出，新请求随时加入 |
|
||
|
||
```
|
||
t=0: prompt A, B
|
||
t=1: prompt C 加进来
|
||
t=2: prompt D 加进来
|
||
→ GPU 利用率更高
|
||
```
|
||
|
||
---
|
||
|
||
### HTTP API 的核心作用
|
||
|
||
| 作用 | 说明 |
|
||
|------|------|
|
||
| **解耦** | 训练和推理完全分开，可独立升级/扩展 |
|
||
| **服务化** | 变成标准接口，不关心内部怎么生成 |
|
||
| **并发调度入口** | 统一入口供 scheduler 做 continuous batching |
|
||
| **分布式扩展** | 多训练 worker 共用一个推理服务 |
|
||
|
||
---
|
||
|
||
### RLHF 场景为什么特别需要
|
||
|
||
| 需求 | 为什么需要 HTTP API |
|
||
|------|---------------------|
|
||
| rollout 量极大 | 一个 step = 100~1000 prompts，必须高吞吐入口 |
|
||
| 训练/推理资源分离 | Trainer 用一批 GPU，Rollout 用另一批，必须跨进程通信 |
|
||
| 动态调度 | prompt 随时进来，generation 长度不一致，必须统一调度 |
|
||
|
||
---
|
||
|
||
### RLHF pipeline 中的位置
|
||
|
||
```
|
||
Prompt dataset
|
||
↓
|
||
HTTP API（SGLang）
|
||
↓
|
||
Generated trajectories
|
||
↓
|
||
Reward model
|
||
↓
|
||
PPO / GRPO update
|
||
```
|
||
|
||
---
|
||
|
||
### 一句话总结（考试级）
|
||
|
||
> SGLang 并没有替代 model.generate，而是把"生成过程（forward + sampling + KV cache loop）"从本地函数调用，升级为可并发调度的推理服务；模型参数 θ 完全不变，只是执行环境和调度方式变了。
|
||
|
||
---
|
||
|
||
## approx_kl 详解
|
||
|
||
### 公式
|
||
|
||
```python
|
||
approx_kl = (
|
||
0.5 * (log_ratio ** 2) * resp_policy_mask[inds]
|
||
).sum() / resp_policy_mask[inds].sum().clamp(min=1)
|
||
```
|
||
|
||
展开：
|
||
|
||
```
|
||
approx_kl = Σ_{b,t} m_{b,t} * 0.5 * (log π_θ(a_{b,t}|s_{b,t}) - log π_old(a_{b,t}|s_{b,t}))²
|
||
/ Σ_{b,t} m_{b,t}
|
||
```
|
||
|
||
其中 `m_{b,t} = resp_policy_mask[b,t]`
|
||
|
||
### resp_policy_mask 的作用
|
||
|
||
| 排除的内容 | 说明 |
|
||
|-----------|------|
|
||
| prompt tokens | 只统计 response |
|
||
| padding | 只统计有效 token |
|
||
| EOS 之后 | 只统计有效 response 区域 |
|
||
|
||
### 为什么叫 approx_kl
|
||
|
||
PPO 需要监控：当前 actor `π_θ` 离 rollout 时的 old actor `π_old` 有多远。
|
||
|
||
KL 本来是衡量两个分布距离的量。当新旧策略很接近时，可以用二阶近似：
|
||
|
||
```
|
||
KL(π_old || π_θ) ≈ 0.5 * (log π_θ - log π_old)²
|
||
```
|
||
|
||
所以源码用 `0.5 * (log_ratio ** 2)` 来近似新旧策略的 KL。
|
||
|
||
### 直观判断
|
||
|
||
| log_ratio | 含义 | approx_kl |
|
||
|-----------|------|-----------|
|
||
| 接近 0 | new actor 和 old actor 很接近 | 小 |
|
||
| 绝对值变大 | new actor 和 old actor 差得更远 | 大 |
|
||
|
||
---
|
||
|
||
## 一句话总结
|
||
|
||
| 概念 | 说明 |
|
||
|------|------|
|
||
| **普通 CE** | y 是 ground truth 常数 → 梯度直接流回 |
|
||
| **DPO** | y 是分布比较，不需要采样 → 直接用 log_probs |
|
||
| **PPO** | y 必须采样 → 用 Reinforce 估计器把 reward 信号传回去 |
|
||
| **KL 约束** | k = r - log r - 1，保证 policy 不偏离 reference 太远 |
|
||
| **approx_kl** | 0.5 * (log_ratio)²，新旧策略 KL 散度的近似 |
|
||
| **Thinking** | 通过 `open_thinking` 控制模型是否显式输出推理过程 |
|
||
| **SGLang** | 生成服务化，支持并发调度和连续 batching |
|