Hotdry.
ai-engineering

verl 离线 RLHF 工具包:PPO 正则化、KL 散度控制与多 GPU 训练

面向大模型对齐,详解 verl 中 PPO+KL 的离线 RLHF 配置、多 GPU 训练参数与生产监控要点。

verl 是 ByteDance Seed 团队开源的 RL 训练库,专为大语言模型(LLM)后训练设计,支持离线 RLHF(Reinforcement Learning from Human Feedback)管道。通过 HybridFlow 编程模型,它无缝集成 FSDP、Megatron-LM 等训练后端与 vLLM、SGLang 等生成引擎,实现高效 PPO 训练与 KL 散度控制,并在多 GPU 环境下扩展到数百卡规模。

PPO 在 verl 中的核心实现与配置

PPO(Proximal Policy Optimization)是 verl 支持的主要 on-policy RL 算法之一,适用于离线 RLHF 场景。不同于传统 REINFORCE 的高方差问题,PPO 通过 clipped surrogate objective 限制策略更新幅度,确保训练稳定。在 verl 中,PPO 训练流程包括 actor(策略模型)、critic(价值函数)、reference(参考策略)和 rollout(生成)四个模块,使用离线数据集(如 Parquet 格式的 prompt-response 对)驱动。

关键配置参数聚焦 batch 处理与更新策略:

  • data.train_batch_size:全局训练 batch 大小,例如 1024,表示每次迭代采样的 prompt 数量。乘以 actor_rollout_ref.rollout.n(默认 1)得到总轨迹数。该值直接影响样本效率,建议从 512 开始,根据 GPU 内存逐步放大。
  • actor_rollout_ref.actor.ppo_mini_batch_size:PPO 更新时的 mini-batch 大小,如 256。全局视角下,所有 worker/GPU 累加;用于 actor 更新,拆分轨迹集以降低内存压力。
  • actor_rollout_ref.actor.clip_ratio:PPO 裁剪范围,默认 0.2。该阈值防止策略偏离过远,若优势函数(advantage)正向时比率超过此值,则裁剪损失为参考值。实际生产中,可调至 0.1-0.3 以平衡探索与稳定性。
  • actor_rollout_ref.actor.ppo_epochs:每轮轨迹的 PPO 更新 epoch 数,默认 1。增加至 3-4 可提升收敛,但需监控 KL 散度避免过拟合。
  • algorithm.adv_estimator:优势估计器,支持 gae(默认,带 lambda 偏置 - 方差权衡)、reinforce_plus_plus 等。algorithm.lam=0.95 是 GAE 标准值。

这些参数已在 GSM8K 等基准上验证,例如 Qwen2.5-0.5B 通过 PPO 从 36.4% 提升至 56.7% 准确率,仅用 bsz=256。

证据显示,verl 的 PPO 后端如 FSDP 支持 ppo_micro_batch_size_per_gpu=8,模拟梯度累积,避免单次 forward OOM。该设计在离线数据上高效,利用 sequence packing 与 flash attention 2 加速。

KL 散度控制:防止分布偏移的核心机制

RLHF 中,actor 易偏离参考策略(SFT 模型),导致奖励黑客(reward hacking)。verl 通过 KL 散度正则化解决,支持 reward 内罚项与 actor loss 两种方式。

  • Reward 内 KL 罚项algorithm.use_kl_in_reward=Truekl_penalty=kl(或 low_var_kl 以降低方差)。控制器 algorithm.kl_ctrl.type=fixed(固定系数 0.005)或 adaptive(目标 KL 0.1,horizon 10000)。自适应模式动态调整系数,维持 KL 在 [0.01, 0.05] 内。
  • Actor KL lossactor_rollout_ref.actor.use_kl_loss=Truekl_loss_coef=0.001kl_loss_type=low_var_kl。计算 actor 与 ref 的 KL(如 k1/kl 公式),直通估计器(+ 后缀)确保无偏梯度。

生产阈值:监控 kl_mean,若 >0.03 则增大 kl_coef;结合 actor_rollout_ref.actor.entropy_coeff=0.01 鼓励探索。GitHub 示例中,此机制支持 Doubao-1.5-pro 等模型达到 O1 级数学性能。

多 GPU 训练:可扩展对齐管道

verl 的多 GPU 支持通过灵活设备映射实现,HybridEngine 消除 actor resharding 内存冗余,通信开销降至最低。FSDP2 后端推荐,吞吐提升 1.4x。

关键参数:

  • trainer.n_gpus_per_node=8trainer.nnodes=1(单节点起步,多节点用 Ray 调度)。
  • actor_rollout_ref.actor.fsdp_config.param_offload=True:参数 offload 省内存,支持大模型如 DeepSeek-671B。
  • actor_rollout_ref.rollout.gpu_memory_utilization=0.5(vLLM),tensor_model_parallel_size=2(TP)。
  • ppo_max_token_len_per_gpu=16384:每 GPU 最大 token 限,计算公式 n * prompt_len + response_len。

扩展到数百 GPU:用 Megatron-LM 后端,ulysses_sequence_parallel_size=1 开启序列并行;LoRA RL 模式 lora_rank=64 进一步降存。性能调优指南建议 vLLM>=0.8.2,避免 0.7.x OOM bug。

落地清单与监控要点

  1. 安装pip install verl[all],AMD/ROCm 支持 Docker 构建。
  2. 数据准备:Parquet 格式,prompt_key=promptmax_prompt_length=512max_response_length=512。过滤过长 prompt:filter_overlong_prompts=True
  3. 配置模板:从 examples/ppo_trainer/run_qwen2-7b.sh 复制,覆盖 lr=1e-6,warmup_ratio=0.1。
  4. 启动verl-train ppo_trainer.yaml,wandb 日志 logger=['wandb']
  5. 监控:Prometheus+Grafana 追踪 rollout TPS、KL_mean、grad_norm;Ray timeline 分析瓶颈。若 actor/grad_norm 高,检查序列对齐。
  6. 回滚策略:Checkpoint save_freq=100resume_mode=auto;测试前 val_before_train=True
  7. 风险阈值
    指标 正常阈值 异常处理
    KL_mean <0.03 增 kl_coef
    grad_norm <5.0 减 lr 或 clip
    throughput >500 tokens/s/GPU 调 micro_bsz

实际部署中,从 8xA100 起步,逐步 scale 到 64 GPU,确保 rollout_correction.rollout_is_threshold=2.0 修正分布偏移。

资料来源:verl GitHub (https://github.com/volcengine/verl),官方文档 (https://verl.readthedocs.io/en/latest/algo/ppo.html),HybridFlow 论文。

查看归档