RLHF+PPO实战：从奖励模型到策略优化的完整方案

一、引言

大语言模型的训练分为三个阶段：预训练（Pretraining）→ 监督微调（SFT）→ 人类反馈强化学习（RLHF）。RLHF 是让模型对齐人类偏好的关键，也是 ChatGPT 成功的核心秘密。

本文将完整实现 RLHF 全流程：奖励模型训练、PPO 策略优化、训练稳定性技巧，使用 TRL 框架在单卡上完成。

二、RLHF 原理

2.1 三阶段流程

阶段1: SFT → 学习回答格式和内容
阶段2: 奖励模型(RM) → 学习判断回答好坏
阶段3: PPO → 用RM信号优化策略

2.2 PPO 核心公式

L(θ) = E[min(r_t(θ)Â_t, clip(r_t(θ), 1-ε, 1+ε)Â_t)]
r_t(θ) = π_θ(a_t|s_t) / π_old(a_t|s_t)

2.3 奖励函数设计

总奖励 = RM_score + KL_penalty
KL_penalty = β * KL(π_θ || π_ref)

β 控制新策略偏离参考模型的程度，通常在 0.02-0.2 之间。

三、奖励模型训练

3.1 偏好数据集

// 数据格式：每条包含 chosen（优选）和 rejected（劣选）
{
  "prompt": "解释什么是量子计算",
  "chosen": "量子计算利用量子比特的叠加和纠缠特性，实现经典计算机难以完成的计算任务。核心原理包括：量子叠加（qubit同时处于0和1）和量子纠缠（远距离qubit关联）。应用包括：Shor算法破译RSA加密、Grover搜索加速、量子化学模拟。",
  "rejected": "量子计算就是用量子计算机做计算，比普通计算机快很多。"
}

3.2 奖励模型训练代码

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
    Trainer
)
from datasets import load_dataset
import torch
import torch.nn as nn

class RewardModel(nn.Module):
    """基于预训练 LLM 的奖励模型"""
    def __init__(self, base_model="Qwen/Qwen2.5-0.5B"):
        super().__init__()
        self.model = AutoModelForSequenceClassification.from_pretrained(
            base_model,
            num_labels=1,  # 输出标量奖励
            torch_dtype=torch.bfloat16
        )
        self.tokenizer = AutoTokenizer.from_pretrained(base_model)
        self.tokenizer.pad_token = self.tokenizer.eos_token
    
    def forward(self, input_ids, attention_mask):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask
        )
        return outputs.logits  # [B, 1]
    
    def compute_reward(self, texts: list) -> torch.Tensor:
        """计算文本奖励分数"""
        inputs = self.tokenizer(
            texts, padding=True, truncation=True,
            max_length=1024, return_tensors="pt"
        )
        return self.forward(**inputs).squeeze(-1)


def train_reward_model():
    """训练奖励模型"""
    model = RewardModel()
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
    
    # 偏好损失：RM(chosen) > RM(rejected)
    for batch in dataloader:
        # 计算 chosen 和 rejected 的奖励
        rewards_chosen = model.compute_reward(batch["chosen"])
        rewards_rejected = model.compute_reward(batch["rejected"])
        
        # 损失 = -log(σ(r_chosen - r_rejected))
        loss = -torch.nn.functional.logsigmoid(
            rewards_chosen - rewards_rejected
        ).mean()
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # 计算准确率（chosen 奖励 > rejected 奖励的比例）
        accuracy = (rewards_chosen > rewards_rejected).float().mean()
        print(f"Loss: {loss.item():.4f}, Accuracy: {accuracy.item():.2%}")

四、PPO 策略优化

4.1 TRL 框架 PPO 训练

from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from trl.core import LengthSampler
from transformers import AutoTokenizer
import torch

# 1. 配置
config = PPOConfig(
    model_name="Qwen2.5-0.5B-SFT",
    learning_rate=1.41e-5,
    batch_size=8,
    mini_batch_size=4,
    ppo_epochs=4,
    kl_penalty="kl",
    kl_penalty_direction="both",   # 双向KL惩罚
    init_kl_coef=0.2,              # 初始KL系数
    target=6,                      # KL目标值
    horizon=10000,
    gamma=1,
    lam=0.95,
    cliprange=0.2,                 # PPO裁剪范围ε
    cliprange_value=0.2,
    vf_coef=0.1,                   # 价值函数损失权重
)

# 2. 加载模型
model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name,
    torch_dtype=torch.bfloat16
)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(
    config.model_name,
    torch_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)
tokenizer.pad_token = tokenizer.eos_token

# 3. PPO Trainer
ppo_trainer = PPOTrainer(
    config=config,
    model=model,
    ref_model=ref_model,
    tokenizer=tokenizer,
    dataset=dataset,
    data_collator=collator,
)

# 4. 训练循环
generation_kwargs = {
    "min_length": -1,
    "top_k": 0.0,
    "top_p": 1.0,
    "do_sample": True,
    "pad_token_id": tokenizer.eos_token_id,
    "max_new_tokens": 256,
}

reward_model = RewardModel()

for epoch, batch in enumerate(ppo_trainer.dataloader):
    query_tensors = batch["input_ids"]
    
    # === 第一步：采样响应 ===
    response_tensors = ppo_trainer.generate(
        query_tensors,
        return_prompt=False,
        **generation_kwargs
    )
    batch["response"] = [tokenizer.decode(r) for r in response_tensors]
    
    # === 第二步：计算奖励 ===
    # 任务奖励：RM打分
    texts = [
        q + r for q, r in zip(batch["query"], batch["response"])
    ]
    task_rewards = reward_model.compute_reward(texts)
    # 归一化（均值为0，方差为1）
    rewards = (task_rewards - task_rewards.mean()) / (task_rewards.std() + 1e-8)
    
    # === 第三步：PPO更新 ===
    stats = ppo_trainer.step(
        query_tensors,
        response_tensors,
        rewards
    )
    
    # 监控指标
    print(f"Epoch {epoch}: "
          f"reward={rewards.mean():.3f}, "
          f"kl={stats['objective/kl']:.4f}, "
          f"clipfrac={stats['ppo/clipfrac']:.4f}")

4.2 PPO 训练稳定性技巧

class PPOTrainingStabilizer:
    """PPO训练稳定性增强"""
    
    @staticmethod
    def adaptive_kl_penalty(kl_div, target_kl=0.02):
        """自适应KL惩罚系数"""
        if kl_div > target_kl * 2:
            return 2.0  # 增大惩罚
        elif kl_div < target_kl / 2:
            return 0.5  # 减小惩罚
        return 1.0
    
    @staticmethod
    def reward_normalization(rewards, method="zscore"):
        """奖励归一化（防止梯度爆炸）"""
        if method == "zscore":
            return (rewards - rewards.mean()) / (rewards.std() + 1e-8)
        elif method == "minmax":
            r_min, r_max = rewards.min(), rewards.max()
            return (rewards - r_min) / (r_max - r_min + 1e-8)
    
    @staticmethod
    def reward_clipping(rewards, clip_range=5.0):
        """奖励裁剪（防止极端值）"""
        return torch.clamp(rewards, -clip_range, clip_range)
    
    @staticmethod
    def early_stop_check(kl_div, threshold=0.5):
        """KL散度过大时早停"""
        return kl_div > threshold
    
    @staticmethod
    def gradient_clipping(model, max_norm=1.0):
        """梯度裁剪"""
        torch.nn.utils.clip_grad_norm_(
            model.parameters(), max_norm=max_norm
        )

五、DPO vs KTO vs PPO

方法	需要奖励模型	需要偏好数据	训练成本	效果
PPO	✅	❌	高（需在线采样）	最强
DPO	❌	✅	中	强
KTO	❌	部分	低（无需pair）	中
ORPO	❌	❌	低（联合训练）	中上

DPO 实现

def dpo_loss(model, ref_model, batch, beta=0.1):
    """DPO损失函数"""
    # 计算当前模型对 chosen/rejected 的log概率
    logps_chosen = model(**batch["chosen"]).logits
    logps_rejected = model(**batch["rejected"]).logits
    
    # 计算参考模型的log概率
    with torch.no_grad():
        ref_logps_chosen = ref_model(**batch["chosen"]).logits
        ref_logps_rejected = ref_model(**batch["rejected"]).logits
    
    # DPO损失
    log_ratio_chosen = logps_chosen - ref_logps_chosen
    log_ratio_rejected = logps_rejected - ref_logps_rejected
    
    loss = -torch.nn.functional.logsigmoid(
        beta * (log_ratio_chosen - log_ratio_rejected)
    ).mean()
    
    return loss

六、关键调参指南

参数	推荐值	作用
learning_rate	1e-6 ~ 1e-5	PPO学习率应远小于SFT
kl_coef	0.02 ~ 0.2	KL惩罚系数，越大越保守
cliprange (ε)	0.1 ~ 0.2	PPO裁剪范围
mini_batch_size	4 ~ 8	每次更新的样本数
ppo_epochs	4	每批数据重复训练的轮数
target_kl	0.01 ~ 0.06	KL目标值，超过则停止

训练监控仪表盘

Epoch | Reward | KL     | ClipFrac | ValueLoss | PolicyLoss | LR
10    | 0.42   | 0.023  | 0.05     | 0.12      | 0.0015     | 1.4e-5
20    | 0.65   | 0.031  | 0.08     | 0.09      | 0.0012     | 1.3e-5
30    | 0.78   | 0.045  | 0.06     | 0.08      | 0.0008     | 1.2e-5
50    | 0.85   | 0.058  | 0.04     | 0.07      | 0.0006     | 1.0e-5
100   | 0.89   | 0.052  | 0.03     | 0.06      | 0.0005     | 8.0e-6

✅ Reward持续上升
✅ KL在目标范围内(<0.06)
✅ ClipFrac下降(策略收敛)
✅ Loss平稳下降

七、总结

RLHF 的核心要点：

1. 奖励模型质量决定一切 — RM 准确率应 >70%
2. KL惩罚不可忽视 — 防止策略偏离太远导致胡言乱语
3. 奖励归一化 — 防止梯度爆炸
4. 小批量 — PPO 对 batch size 敏感，4-8 即可
5. 监控指标 — reward/kl/clipfrac 三者都要看

DPO 是更简单的替代方案，但 PPO 通过在线采样能更好地探索策略空间。生产环境中优先选择 DPO（简单稳定），追求极致效果时使用 PPO。