9.8 KiB

Raw Blame History

Lesson 23 笔记：PPO 中 Response 采样与梯度问题

核心问题：为什么 y 不能直接求导

response y 是模型采样出来的离散 token 序列，不能像普通 CE loss 那样直接对 token 选择过程求导。

普通 CE Loss 的梯度怎么传

# 普通 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 是"采样"出来的

# PPO 里
y_w = torch.multinomial(probs, num_samples)  # 从模型分布中采样

问题	说明
采样是随机的	`torch.multinomial` 每次结果可能不同
采样不可导	autograd 无法跟踪随机采样过程

梯度流断在这里：

logits → softmax → sample → y (随机!)
                        ↑
                  这个随机采样过程没有梯度

DPO 的解决方案

DPO 绕过了这个问题，不采样，而是直接用 model 输出的 log_probs：

# 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
2	计算 reward r(y)
3	用 ∇ log πθ(y) 把 reward 的信号传回去

关键技巧： 用 log probability 的梯度来"引导"随机采样的方向。

PPO + KL 约束公式：k = r - log r - 1

公式拆解

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 控制

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 中的实现

{%- 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	`<
False	`<

之前阶段有没有训练过 thinking

阶段	有 thinking 训练数据吗
Pretrain	没有
SFT	有字段但实际数据没有（`reasoning_content` 是空的）
RLAIFDataset	有 — 通过 `open_thinking` 参数控制

SFT 中也有 thinking 相关机制

SFTDataset 的数据 features 定义了 reasoning_content 字段，chat template 对 assistant 消息会做：

<|im_start|>assistant
<think>
{reasoning_content}
</think>
{content}<|im_end|>

另外 post_processing_chat() 会随机移除空 thinking 块：

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 详解

公式

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

9.8 KiB Raw Blame History Unescape Escape