Hotdry.
归档
ai-engineering

VERL PPO 中微调 KL 散度阈值：稳定 RLHF 训练的工程实践

VERL PPO RLHF 中 KL 阈值调控，平衡策略偏离与奖励黑客，提供参数配置、自适应机制与监控清单。

在 VERL 框架的 PPO 算法中，KL 散度（Kullback-Leibler divergence）作为核心正则化机制，用于约束当前策略（actor）与参考策略（reference model，通常为 SFT 模型）之间的分布差异。这有助于在 RLHF 训练中防止 “奖励黑客”（reward hacking），即模型为了最大化奖励模型分数而生成无意义或有害输出，同时避免策略过度保守导致学习停滞。工程实践中，KL 阈值的微调直接影响训练稳定性：阈值过低易引发策略崩溃，过高则抑制探索。VERL 通过固定或自适应 KL 控制器，提供灵活调控，实现高效平衡。

KL 散度在 PPO 中的数学作用源于其作为代理目标函数的惩罚项。PPO 优化目标为 (L^{PPO} = \mathbb {E} \left [ \min \left ( r_t (\theta) \hat {A}t, \clip(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}t \right) - \beta D{KL}(\pi\theta || \pi_{ref}) \right] )，其中 ( r_t (\theta) = \frac {\pi_\theta (a_t|s_t)}{\pi_{old}(a_t|s_t)} ) 为重要性采样比率，( \hat {A}t ) 为 GAE 优势估计，(\beta) 为 KL 系数。KL 项确保策略更新幅度受控，典型值为 ( D{KL} \approx 0.01 \sim 0.1 )。VERL 支持两种集成方式：一是嵌入奖励函数（algorithm.use_kl_in_reward: true，KL_penalty 类型如 'kl' 或 'low_var_kl'）；二是独立损失（actor_rollout_ref.actor.use_kl_loss: true，kl_loss_coef: 0.001）。GitHub 仓库中 PPO 示例默认禁用 KL 损失，但推荐在长序列任务中启用以防分布漂移。

VERL 的配置参数高度模块化，便于工程落地。核心参数包括：

  • algorithm.kl_ctrl.kl_coef: 0.001：初始 KL 系数，控制惩罚强度，小模型用 0.0005，大模型调至 0.002。
  • algorithm.kl_ctrl.target_kl: 0.05：目标 KL 值，自适应模式下以此为基准，数学任务推荐 0.03~0.08，代码生成用 0.1。
  • algorithm.kl_ctrl.horizon: 10000：调整周期，步数累积达此值时重置比例误差。
  • actor_rollout_ref.actor.kl_loss_type: 'low_var_kl'：KL 类型，低方差估计（K3 形式）适用于 GRPO/PPO 混合，避免梯度爆炸。 示例配置：
algorithm:
  kl_ctrl:
    type: adaptive  # 或 fixed
    kl_coef: 0.001
    target_kl: 0.05
    horizon: 10000
actor_rollout_ref:
  actor:
    use_kl_loss: true
    kl_loss_coef: 0.001

这些参数在 VERL 文档中被验证为 SOTA 基线，能在 Qwen2.5-7B 上将 GSM8K 准确率从 36% 提升至 56%，KL 均值稳定在 0.04 附近。

自适应 KL 控制器是 VERL 的工程亮点，实现公式：kl_coef *= (1 + clip((current_kl / target_kl - 1), -0.2, 0.2) * n_steps / horizon)。每步计算当前 KL（参考 core_algos.py 中的 low_var_kl：(D_{KL}(\pi_\theta || \pi_{ref}) \approx \mathbb {E} [\frac {\pi_{ref}}{\pi_\theta} - \log (\frac {\pi_{ref}}{\pi_\theta}) - 1] )，无偏低方差），若当前 KL > target，则增大系数抑制更新；反之减小促进探索。实践阈值：启动阶段 target_kl=0.1 鼓励学习，中期降至 0.03 锁定收敛。监控日志中 actor/kl_divergence 应波动 <0.2，避免>0.3 触发 OOM 或 NaN。

调优实践分三阶段：

  1. 初始化：小数据集验证，kl_coef=0.001，target_kl=0.08，观察 10 epochs 内 KL 均值趋近目标。
  2. 迭代优化：W&B 日志追踪 train/kl_divergence、reward/mean 与 policy_loss。若奖励飙升但 KL<0.01，降 target_kl；KL>0.15 且熵崩塌（entropy<0.1），升 kl_coef 20%。
  3. 生产部署：A/B 测试，多机 8xA100，启用 FSDP2（strategy: fsdp2），KL 阈值与基线对齐。

监控清单：

指标 健康阈值 异常处理
KL 均值 0.02~0.08 >0.15: 增 kl_coef 50%；<0.01: 减 target_kl
KL 方差 <0.05 高：启用 low_var_kl+；检查 ref 模型一致性
策略熵 >0.2 低：增 entropy_coeff=0.01
奖励方差 <2.0 高：审视 RM 噪声，回滚至上 checkpoint

风险与回滚：KL 失效常见于长上下文（>2048 tokens），表现为 reward hacking（奖励 > 5 但生成乱码）。回滚策略：固定 KL 模式 + 早停（KL>0.2 连续 5 epochs）；备用 GRPO（无 critic，adv_estimator: grpo）。VERL recipe/dapo 示例证实，KL 阈值优化后 AIME 分数超 DeepSeek GRPO 3 分。

落地参数清单：

  • 保守模式（稳定优先）：target_kl=0.03, kl_coef=0.002, horizon=5000。
  • 激进模式（快速收敛）：target_kl=0.1, kl_coef=0.0005, horizon=20000。
  • 监控脚本：trainer.logger: ["wandb"]，追踪 kl_divergence、explained_variance（>0.8）。
  • 验证：每 10 epochs eval GSM8K，KL 与准确率正相关。

通过上述微调，RLHF 训练失败率降 70%，适用于 MoE/671B 规模。来源：VERL GitHub（PPO 示例配置），文档（algo/ppo.html）。

查看归档