Hotdry.
归档
ai-systems

verl中KL正则化PPO的多GPU数据并行离线RLHF偏好对齐蒸馏实现

详解verl框架下KL-regularized PPO的多GPU数据并行配置，用于离线RLHF偏好对齐蒸馏，提供工程参数清单、监控要点与稳定性优化策略。

在 verl 框架中，KL 正则化 PPO（KL-regularized PPO）是实现离线 RLHF（Reinforcement Learning from Human Feedback）偏好对齐蒸馏的核心算法，通过多 GPU 数据并行显著提升 LLM 训练的扩展性和稳定性。该方法利用预处理偏好数据集，避免在线采样开销，同时 KL 散度约束防止策略过度偏离参考模型，确保训练收敛可靠。

verl 的 HybridFlow 编程模型将 RL 数据流解耦为控制流和计算流，支持 FSDP2 和 Megatron-LM 后端，实现高效多 GPU 并行。核心优势在于 3D-HybridEngine：在 actor rollout（生成）和训练阶段动态重分片模型，消除内存冗余，减少 AllReduce 通信开销达 40%。例如，在 8 卡 A100 集群上训练 Qwen2.5-7B，verl PPO 吞吐量较传统框架提升 1.5 倍。

离线 RLHF 流程聚焦偏好对齐蒸馏：首先准备 Parquet 格式数据集（prompt/response pairs），包含人类偏好排名或规则奖励（如 GSM8K 数学准确率）。Reward 模型从偏好数据训练，PPO 使用 actor（策略）、critic（价值）和 reference（KL 基准）模型迭代优化。偏好蒸馏通过组采样（group sampling）实现，如 GRPO 变体，但标准 KL-PPO 适用于排名数据：优势估计 A_t = r_t + γV (s_{t+1}) - V (s_t)，损失 L = E [min (ratio A, clip (ratio,1-ε,1+ε) A)] - β KL (π||π_ref)。“verl 是 HybridFlow 论文的开源实现，支持 PPO 等算法。” 配置中设置 algorithm.adv_estimator=gae，kl_ctrl.type=moving_average，target_kl=0.1。

多 GPU 数据并行关键参数清单：

  • 并行策略：actor_rollout_ref.actor.strategy=fsdp2（推荐，内存降 7%，吞吐升 1.5%）；tensor_model_parallel_size=2（张量并行，平衡 DP/TP）；n_gpus_per_node=8，nnodes=4（扩展至数百 GPU）。
  • 批处理优化：data.train_batch_size=2048（全局）；ppo_mini_batch_size=512；ppo_micro_batch_size_per_gpu=16（动态调整）；use_dynamic_bsz=True，ppo_max_token_len_per_gpu=16384（actor 3x (prompt+response)，critic 2x）。
  • KL 正则：kl_ctrl.kl_coef=0.002（初始 0.0010.005，监控 target_kl=0.050.15）；kl_loss_type=low_var_kl（低方差估计）；use_kl_loss=True，kl_loss_coef=0.01。
  • PPO 核心：ppo_epochs=4~8；clip_ratio=0.2；entropy_coeff=0.01（探索）；grad_clip=1.0；lr=1e-6（actor），1e-5（critic）。
  • 序列优化：use_remove_padding=True（打包减填充 20%）；ulysses_sequence_parallel_size=2（长上下文 > 32k）。

稳定性优化与监控要点：

  • KL 监控：实时追踪 KL divergence，若 > 0.2 则减 lr 或增 kl_coef；使用 wandb/mlflow 记录 policy_loss、value_loss、kl_div、explained_var。
  • 内存管理：actor_rollout_ref.rollout.gpu_memory_utilization=0.6（vLLM/SGLang 推理）；enable_gradient_checkpointing=True；fsdp_config.offload_policy=True（CPU 卸载）。
  • 回滚策略：每 epoch 验证集评估，若 perf 降 > 5% 回滚 checkpoint；多阶段训练：先 SFT，再 RM 训练，最后 PPO（total_epochs=10~20）。
  • 风险阈值：熵崩（entropy<0.1）增 entropy_coeff；reward hacking 监控 reward 均值波动 < 10%。

实战脚本示例（run_qwen2-7b_rlhf.sh）：

python3 -m verl.trainer.main_ppo \
  data.train_files=~/data/rlhf/train.parquet \
  data.train_batch_size=2048 \
  actor_rollout_ref.actor.strategy=fsdp2 \
  actor_rollout_ref.rollout.name=vllm tensor_model_parallel_size=2 \
  algorithm.kl_ctrl.kl_coef=0.002 target_kl=0.1 \
  trainer.n_gpus_per_node=8 total_epochs=15

在 DeepSeek-671B MoE 模型上，verl 支持 expert parallelism，结合 LoRA RL 进一步降内存，支持 multi-turn 工具调用偏好蒸馏。

该实现优化了 LLM 扩展性：动态批处理提升 GPU 利用 85%，序列并行处理长 CoT 推理。实际部署中，优先 FSDP2+vLLM 组合，Nsight Systems 分析瓶颈（steps=[1,5,10]）。

资料来源：verl GitHub（https://github.com/volcengine/verl），HybridFlow 论文（https://arxiv.org/abs/2409.19256），官方文档（https://verl.readthedocs.io）。

查看归档