openrlhf-training

OpenRLHF - 高性能 RLHF 训练

快速开始

OpenRLHF 是一个基于 Ray 的 RLHF 框架，针对使用 vLLM 推理加速的分布式训练进行了优化。

安装：

# Launch Docker container
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN \
  -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:25.02-py3 bash

# Uninstall conflicts
sudo pip uninstall xgboost transformer_engine flash_attn pynvml -y

# Install OpenRLHF with vLLM
pip install openrlhf[vllm]

PPO 训练（混合引擎）：

ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

ray job submit --address="http://127.0.0.1:8265" \
  --runtime-env-json='{"working_dir": "/openrlhf"}' \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  --ref_num_nodes 1 --ref_num_gpus_per_node 8 \
  --reward_num_nodes 1 --reward_num_gpus_per_node 8 \
  --critic_num_nodes 1 --critic_num_gpus_per_node 8 \
  --actor_num_nodes 1 --actor_num_gpus_per_node 8 \
  --vllm_num_engines 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --vllm_gpu_memory_utilization 0.5 \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
  --save_path ./output/llama3-8b-rlhf \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --rollout_batch_size 1024 \
  --max_epochs 1 --prompt_max_len 1024 --generate_max_len 1024 \
  --zero_stage 3 --bf16 \
  --actor_learning_rate 5e-7 --critic_learning_rate 9e-6 \
  --init_kl_coef 0.01 --normalize_reward \
  --gradient_checkpointing --packing_samples \
  --vllm_enable_sleep --deepspeed_enable_sleep

GRPO 训练（组归一化策略优化）：

# Same command as PPO, but add:
--advantage_estimator group_norm

常见工作流

工作流 1：完整的 RLHF 流程（SFT → 奖励模型 → PPO）

步骤 1：训练奖励模型（DPO）：

deepspeed --module openrlhf.cli.train_rm \
  --save_path ./output/llama3-8b-rm \
  --save_steps -1 --logging_steps 1 \
  --eval_steps -1 --train_batch_size 256 \
  --micro_train_batch_size 1 --pretrain meta-llama/Meta-Llama-3-8B \
  --bf16 --max_epochs 1 --max_len 8192 \
  --zero_stage 3 --learning_rate 9e-6 \
  --dataset OpenRLHF/preference_dataset_mixture2_and_safe_pku \
  --apply_chat_template --chosen_key chosen \
  --rejected_key rejected --flash_attn --gradient_checkpointing

步骤 2：PPO 训练：

ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

ray job submit --address="http://127.0.0.1:8265" \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  --ref_num_nodes 1 --ref_num_gpus_per_node 8 \
  --reward_num_nodes 1 --reward_num_gpus_per_node 8 \
  --critic_num_nodes 1 --critic_num_gpus_per_node 8 \
  --actor_num_nodes 1 --actor_num_gpus_per_node 8 \
  --vllm_num_engines 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain ./output/llama3-8b-rm \
  --save_path ./output/llama3-8b-ppo \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --rollout_batch_size 1024 \
  --max_epochs 1 --prompt_max_len 1024 --generate_max_len 1024 \
  --zero_stage 3 --bf16 \
  --actor_learning_rate 5e-7 --critic_learning_rate 9e-6 \
  --init_kl_coef 0.01 --normalize_reward \
  --vllm_enable_sleep --deepspeed_enable_sleep

工作流 2：GRPO 训练（无需 Critic 模型）

PPO 的内存高效替代方案：

ray job submit --address="http://127.0.0.1:8265" \
  -- python3 -m openrlhf.cli.train_ppo_ray \
  --advantage_estimator group_norm \
  --ref_num_nodes 1 --ref_num_gpus_per_node 8 \
  --reward_num_nodes 1 --reward_num_gpus_per_node 8 \
  --actor_num_nodes 1 --actor_num_gpus_per_node 8 \
  --vllm_num_engines 4 --vllm_tensor_parallel_size 2 \
  --colocate_all_models \
  --pretrain OpenRLHF/Llama-3-8b-sft-mixture \
  --reward_pretrain OpenRLHF/Llama-3-8b-rm-700k \
  --save_path ./output/llama3-8b-grpo \
  --micro_train_batch_size 8 --train_batch_size 128 \
  --micro_rollout_batch_size 16 --rollout_batch_size 1024 \
  --max_epochs 1 --bf16 \
  --actor_learning_rate 5e-7 \
  --init_kl_coef 0.01 --use_kl_loss --kl_estimator k3 \
  --normalize_reward --no_advantage_std_norm

GRPO 关键参数：

• --advantage_estimator group_norm - 启用 GRPO
• --use_kl_loss - 使用 GRPO 论文中的 KL 损失
• --kl_estimator k3 - 损失函数（k2 ≈ k1）
• --no_advantage_std_norm - 禁用标准差归一化

工作流 3：DPO 训练（偏好优化）

无需奖励模型的更简单替代方案：

deepspeed --module openrlhf.cli.train_dpo \
  --save_path ./output/llama3-8b-dpo \
  --save_steps -1 --logging_steps 1 \
  --eval_steps -1 --train_batch_size 256 \
  --micro_train_batch_size 2 --pretrain meta-llama/Meta-Llama-3-8B \
  --bf16 --max_epochs 1 --max_len 8192 \
  --zero_stage 3 --learning_rate 5e-7 --beta 0.1 \
  --dataset OpenRLHF/preference_dataset_mixture2_and_safe_pku \
  --apply_chat_template --chosen_key chosen \
  --rejected_key rejected --flash_attn --gradient_checkpointing

何时使用 OpenRLHF 与替代方案

在以下情况下使用 OpenRLHF：

• 使用 RL 训练大型模型（7B-70B+）
• 需要 vLLM 推理加速
• 需要基于 Ray 的分布式架构
• 拥有多节点 GPU 集群
• 需要在一个框架中使用 PPO/GRPO/RLOO/DPO

算法选择：

• PPO：控制力最强，适用于复杂奖励
• GRPO：内存高效，无需 Critic
• RLOO：采用逐 Token KL 的改进版 PPO
• REINFORCE++：比 GRPO 更稳定，比 PPO 更快
• DPO：最简单，无需奖励模型

改用其他方案：

• TRL：单节点训练，API 更简单
• veRL：ByteDance 针对 671B 模型的框架
• DeepSpeedChat：集成于 DeepSpeed 生态系统

常见问题

问题：大型模型导致 GPU 内存溢出

禁用模型共置：

# Remove --colocate_all_models flag
# Allocate separate GPUs for each model
--actor_num_gpus_per_node 8 \
--critic_num_gpus_per_node 8 \
--reward_num_gpus_per_node 8 \
--ref_num_gpus_per_node 8

问题：DeepSpeed GPU 索引超出范围

设置环境变量：

export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1

问题：训练不稳定

使用混合引擎而非异步模式：

--colocate_all_models \
--vllm_enable_sleep \
--deepspeed_enable_sleep

调整 KL 系数：

--init_kl_coef 0.05  # Increase from 0.01

问题：PPO 生成速度慢

启用 vLLM 加速：

--vllm_num_engines 4 \
--vllm_tensor_parallel_size 2 \
--vllm_gpu_memory_utilization 0.5

高级主题

混合引擎 GPU 共享：有关 vLLM 休眠模式、DeepSpeed 休眠模式和最佳节点分配，请参阅 references/hybrid-engine.md。

算法比较：有关 PPO、GRPO、RLOO 和 REINFORCE++ 的基准测试与超参数，请参阅 references/algorithm-comparison.md。

多节点设置：有关 Ray 集群配置和容错，请参阅 references/multi-node-training.md。

自定义奖励函数：有关强化微调和 Agent RLHF，请参阅 references/custom-rewards.md。

硬件要求

• GPU：推荐使用 NVIDIA A100/H100
• 显存：
  • 7B 模型：8× A100 40GB（混合引擎）
  • 70B 模型：48× A100 80GB（vLLM:Actor:Critic = 1:1:1）
• 多节点：推荐使用配备 InfiniBand 的 Ray 集群
• Docker：NVIDIA PyTorch 容器 25.02+

性能：

• 比 DeepSpeedChat 快 2×
• vLLM 推理加速
• 混合引擎最大限度减少 GPU 空闲时间

资源

• 文档：https://github.com/OpenRLHF/OpenRLHF
• 论文：https://arxiv.org/abs/2405.11143
• 示例：https://github.com/OpenRLHF/OpenRLHF/tree/main/examples
• Discord：社区支持

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确认输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集、少量步骤，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请执行此操作
CUDA 显存不足	减小批次大小，启用梯度累积/检查点或混合精度，或对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否存在无效值，并启用梯度裁剪。
已安装版本中缺少这些说明里的函数、标志或端点	查阅已安装版本自身的文档（help()、--help、官方文档），进行适配并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果无法运行或验证某项内容，请如实说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前先估算计算成本。
• 此处的版本特定信息可能已过时：请根据已安装版本的官方文档进行确认。
• 对于需要付费、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• fine-tuning-with-trl：使用 TRL 进行大模型微调和强化学习——用于指令微调的 SFT、用于偏好对齐的 DPO、用于奖励优化的 PPO/GRPO…
• verl-rl-training：提供使用 verl（Volcano Engine RL）进行大模型强化学习训练的指导。
• simpo-training：用于大模型对齐的简单偏好优化。