OpenRLHF Trainer — 面向对齐训练的高效 RLHF 工具链

[ English | 中文 | 日本語 ]

OpenRLHF 是首个高性能、生产就绪的开源 RLHF 框架，它将 Ray + vLLM 分布式架构与统一的智能体（Agent）设计范式相结合，为可扩展、可扩展的基于人类反馈的强化学习（RLHF）提供了强大支持。

📚 了解更多：文档 | 幻灯片 | 技术报告 | 视频

🔥 新后端: Molt 为 OpenRLHF 带来了基于 Automodel 的后端，其功能比 DeepSpeed 更强大——可将RL训练扩展到数千亿参数，同时保持熟悉、优雅的 OpenRLHF 工作流程。

📖 目录

• 🗞️ 新闻
• 🏗️ 架构基础 - Ray + vLLM + DeepSpeed 分布式基础设施
• 🎯 设计范式 - 统一的基于智能体的执行流水线
• 🚀 RL 算法 - PPO, REINFORCE++, GRPO, RLOO
• 📋 功能概览 - 完整的 RLHF 流水线能力
• 🎬 快速开始 - 安装与典型工作流程
• 🎓 训练指南 - SFT、Reward Model、RL 训练
• 🎯 单轮智能体 - 自定义奖励函数
• 🤖 多轮智能体 - 复杂环境交互
• 🔧 高级主题 - LoRA、性能调优

新闻

🏗️ 架构基础: Ray + vLLM 分布式架构

OpenRLHF 是首个基于 Ray + vLLM 分布式架构的 RLHF 框架，可跨 GPU 高效编排多个组件：

核心基础设施组件

Ray - 分布式调度器与控制器
OpenRLHF 利用 Ray 实现高效的分布式调度。它将 Actor、Reward、Reference 和 Critic 模型分别部署在不同的 GPU 上，支持高达 700 亿+参数模型的可扩展训练。

混合引擎调度：所有模型和 vLLM 引擎可以共享 GPU 资源—最大限度地减少空闲时间并最大化 GPU 利用率。这使得在有限的硬件资源上运行完整的 RLHF 流水线成为可能。

vLLM - 高性能推理引擎
RLHF 训练80% 的时间用于生成本文。借助 vLLM 及其自动张量并行（AutoTP）和流水线并行（PP）特性，OpenRLHF 能够提供高吞吐、内存高效的生成。

DeepSpeed - 内存高效训练
基于 DeepSpeed ZeRO-3、deepcompile、AutoTP 和 RingAttention。无需重型框架即可实现大模型训练，并可直接使用 HuggingFace 模型。

Transformers - 模型接口
原生集成 HuggingFace Transformers，实现预训练模型的无缝加载、状态管理和微调。

NCCL / CUDA IPC - 高速通信
为分布式训练和推理提供高效的 GPU 间通信。

🎯 设计范式: 基于智能体的执行

在 Ray 分布式架构之上，OpenRLHF 是首个实现统一智能体范式的 RLHF 框架。无论训练是标准 PPO 还是复杂的多轮推理，都遵循一致的智能体执行流水线。

为什么采用智能体范式？

OpenRLHF 通过“令牌进-令牌出”（Token-in-Token-out）的智能体执行方式统一生成与训练，确保完美的一致性、轻松的单/多轮扩展，并消除文本级别的错位。

智能体架构

                 ┌─────────────────────────────┐
                 │    AgentExecutorBase        │
                 │  (Token-in-Token-out 核心) │
                 └─────────────────────────────┘
                              │
                 ┌────────────┴────────────┐
                 ↓                         ↓
         SingleTurnExecutor        MultiTurnExecutor
                 │                         │
      ┌──────────┴──────────┐   ┌─────────┴──────────┐
      ↓                     ↓   ↓                    ↓
  标准RLHF           自定义奖励       多步推理         外部环境
  (单次生成)          函数                           (OpenAI Agent Server)
      ↓                     ↓           ↓                ↓
      └─────────────────────┴───────────┴────────────────┘
                              │
                    一致的 Token 轨迹
                              │
                    ┌─────────┴─────────┐
                    │   RL 算法(解耦)   │
                    │                   │
                    │  PPO, REINFORCE++ │
                    │  GRPO, RLOO, etc. │
                    └───────────────────┘

核心设计原则

两种执行模式 (与 RL 算法正交)

智能体执行模式与您选择的 RL 算法无关。您可以将任何算法（PPO、REINFORCE++、GRPO 等）与任何执行模式结合使用：

🚀 最前沿的 RL 算法

OpenRLHF 实现了 PPO、REINFORCE++、REINFORCE++-baseline、GRPO、RLOO，并参考实践指南和社区最佳实践集成了高级优化技巧。

关键设计：RL 算法与执行模式完全解耦。所有算法均可无缝应用于单轮和多轮智能体执行器，并通过统一的令牌进-令牌出流水线运行，以保证行为一致性。

参考资料: 知乎文章 | Notion 最佳实践

📋 功能概览

OpenRLHF 提供了完整的 RLHF 流水线，并具备基于智能体的灵活性：

🎯 基于智能体的 RL 训练（核心创新）

🎓 监督训练与偏好学习

⚡ 高级能力

🎬 快速开始

安装

推荐方式：使用 Docker 进行无痛安装

# 1. 启动 Docker 容器
docker run --runtime=nvidia -it --rm --shm-size="10g" --cap-add=SYS_ADMIN \
  -v $PWD:/openrlhf nvcr.io/nvidia/pytorch:26.03-py3 bash

# 2. 清理冲突的包
sudo pip uninstall xgboost transformer_engine flash_attn pynvml -y

# 3. 安装 OpenRLHF (选择其一)
pip install openrlhf                    # 基础版
pip install openrlhf[vllm]              # + vLLM 0.27.1 (推荐)
pip install openrlhf[vllm_latest]       # + 最新版 vLLM
pip install openrlhf[vllm,ring,liger]   # + 所有优化项

备选方案：源代码安装

git clone https://github.com/OpenRLHF/OpenRLHF.git
cd OpenRLHF
pip install -e .

[!TIP]
我们推荐使用 vLLM 0.27.1+ 以获得最佳性能。请参阅 Dockerfiles 和 Nvidia-Docker 安装脚本。

准备数据集

OpenRLHF 提供了灵活的数据处理方法：

关键参数：
- --data.input_key: 指定输入数据的 JSON 键名
- --data.apply_chat_template: 使用 HuggingFace 分词器的 聊天模板
- --data.input_template: 自定义模板字符串（聊天模板的替代方案）
- --data.prompt_probs / --data.dataset_probs: 混合多个数据集（例如：0.1,0.4,0.5）
- --eval.dataset: 指定评估数据集路径

聊天模板示例：

dataset = [{"input_key": [
  {"role": "user", "content": "Hello, how are you?"},
  {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
  {"role": "user", "content": "I'd like to show off how chat templating works!"},
]}]

tokenizer.apply_chat_template(dataset[0]["input_key"], tokenize=False)
# 输出: "<s>[INST] Hello, how are you? [/INST]I'm doing great...</s> [INST] I'd like to show off... [/INST]"

[!NOTE]
JSON 键名选项因数据集类型而异。请参阅 奖励数据集， SFT 数据集 和 提示数据集。

监督微调

OpenRLHF 的模型检查点与 HuggingFace 模型完全兼容。您可以使用 --actor.model_name_or_path {name or path}、--reward.model_name_or_path {name or path} 和 --critic.model_name_or_path {name or path} 参数来指定模型名称或路径。我们在 HuggingFace OpenRLHF 上提供了一些预训练检查点和数据集。

然后，您可以使用我们提供的 examples/scripts 目录下的启动脚本，或通过以下命令开始训练。

deepspeed --module openrlhf.cli.train_sft \
   --data.max_len 4096 \
   --data.dataset Open-Orca/OpenOrca \
   --data.input_key question \
   --data.output_key response \
   --data.input_template $'User: {}\nAssistant: ' \
   --train.batch_size 256 \
   --train.micro_batch_size 2 \
   --data.max_samples 500000 \
   --actor.model_name_or_path meta-llama/Meta-Llama-3-8B \
   --ckpt.output_dir ./checkpoint/llama3-8b-sft \
   --ckpt.save_steps -1 \
   --logger.logging_steps 1 \
   --eval.steps -1 \
   --ds.zero_stage 2 \
   --train.max_epochs 1 \
   --ds.packing_samples \
   --ds.param_dtype bf16 \
   --adam.lr 5e-6 \
   --actor.gradient_checkpointing_enable \
   --logger.wandb.key {wandb_token}

# 其他选项：
# --data.apply_chat_template                # 使用 HF tokenizer 对话模板
# --ds.ring_attn_size 2                      # 启用 RingAttention (需先安装 ring_flash_attn)
# --data.multiturn                          # 多轮微调损失
# --actor.pretrain_mode_enable                      # 持续预训练模式

奖励模型训练

deepspeed --module openrlhf.cli.train_rm \
   --ckpt.output_dir ./checkpoint/llama3-8b-rm \
   --ckpt.save_steps -1 \
   --logger.logging_steps 1 \
   --eval.steps -1 \
   --train.batch_size 256 \
   --train.micro_batch_size 1 \
   --actor.model_name_or_path OpenRLHF/Llama-3-8b-sft-mixture \
   --ds.param_dtype bf16 \
   --train.max_epochs 1 \
   --data.max_len 8192 \
   --ds.zero_stage 3 \
   --adam.lr 9e-6 \
   --data.dataset OpenRLHF/preference_dataset_mixture2_and_safe_pku \
   --data.apply_chat_template \
   --chosen_key chosen \
   --rejected_key rejected \
   --ds.packing_samples \
   --actor.gradient_checkpointing_enable \
   --logger.wandb.key {wandb_token}

建议将奖励模型的 --value_prefix_head 选项设置为 score，这样我们就可以使用 AutoModelForSequenceClassification 加载模型：

reward_model = AutoModelForSequenceClassification.from_pretrained(
              reward_model_path,
              num_labels=1,
              torch_dtype=torch.bfloat16,
              attn_implementation="flash_attention_2",
              use_cache=False,
          )
inputs = xxxx (向左填充输入令牌)
reward = reward_model.model(*inputs).last_hidden_state
reward = reward_model.score(reward)[:, -1]

RL 训练: 使用 Ray 和 vLLM 的 PPO/REINFORCE++

OpenRLHF 中的所有 RL 训练均通过智能体执行流水线运行。以下示例展示了单轮智能体执行（默认模式）与混合引擎结合以实现最佳性能：

# 在容器中启动 ray 主节点
ray start --head --node-ip-address 0.0.0.0 --num-gpus 8

# 如果想在更多节点上启动 ray，请使用
ray start --address {MASTER-NODE-ADDRESS}:6379  --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 \
   --train.colocate_all \
   --vllm.gpu_memory_utilization 0.5 \
   --actor.model_name_or_path OpenRLHF/Llama-3-8b-sft-mixture \
   --reward.model_name_or_path OpenRLHF/Llama-3-8b-rm-700k \
   --ckpt.output_dir /openrlhf/examples/test_scripts/final/llama3-8b-rlhf \
   --ckpt.path /openrlhf/examples/test_scripts/ckpt/llama3-8b-rlhf \
   --ckpt.save_hf \
   --train.batch_size 128 \
   --rollout.batch_size 1024 \
   --train.dynamic_batch_enable \
   --rollout.n_samples_per_prompt 1 \
   --train.max_epochs 1 \
   --prompt_max_len 1024 \
   --data.max_samples 100000 \
   --generate_max_len 1024 \
   --ds.zero_stage 3 \
   --ds.param_dtype bf16 \
   --actor.adam.lr 5e-7 \
   --critic.adam.lr 9e-6 \
   --algo.kl.init_coef 0.01 \
   --data.prompt_dataset OpenRLHF/prompt-collection-v0.1 \
   --data.input_key context_messages \
   --data.apply_chat_template \
   --reward.normalize_enable \
   --actor.gradient_checkpointing_enable \
   --ds.packing_samples \
   --vllm.sync_backend nccl \
   --vllm.enforce_eager \
   --vllm.enable_sleep \
   --ds.enable_sleep \
   --logger.wandb.key {wandb_token}

# 算法变体 (均使用单轮智能体执行):
# --algo.advantage.estimator reinforce        # REINFORCE++
# --algo.advantage.estimator rloo             # RLOO
# --algo.advantage.estimator reinforce_baseline  # REINFORCE++-baseline (最适合 RLVR)
# --algo.advantage.estimator group_norm       # GRPO
# --algo.advantage.estimator dr_grpo          # Dr. GRPO

# 高级选项:
# --algo.kl.init_coef 0                                    # 不使用参考模型
# --reward.remote_url http://host:5000/get_reward          # HTTP 奖励模型
# --rollout.n_samples_per_prompt 4                         # 每个 prompt 的多个样本
# --rollout.vllm_generate_batch_size 2048                  # 生成时过采样 (> rollout_batch_size)；需要 --train.async_enable
# --algo.advantage.is_correction_enable                    # 对离策略 (off-policy) rollout 进行 vLLM 重要性采样校正
# --algo.advantage.is_correction_type tis                  # 校正类型: tis (token clamp) | icepop (token filter) | seq-mask-tis (seq-level geom mean)
# --algo.advantage.is_correction_threshold 0.5 5.0         # IS 截断区间: [low, high]
# --ckpt.best_metric_key eval_default_pass1                # 根据 eval 指标保存最佳检查点 (为空则自动检测第一个 pass1， 'none' 为禁用)
# --actor.policy_loss_type gspo                            # 使用 GSPO 策略损失变体 (相对于默认的 'ppo')

[!TIP]
针对推理任务 (RLVR)：使用 --algo.advantage.estimator reinforce_baseline 启用REINFORCE++-baseline——它对不同的奖励尺度具有鲁棒性。

[!NOTE]
Ray 环境设置：使用 --runtime-env-json='{"setup_commands": ["pip install openrlhf[vllm]"]}' 让 Ray 自动部署。

[!NOTE]
GPU 索引错误疑难解答：如果遇到 DeepSpeed GPU 设备设置问题，请设置 export RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1。

📚 更多示例：参见 examples/scripts 与 文档

🎯 单轮智能体: 使用自定义奖励进行强化微调

单轮智能体执行（默认模式）支持自定义奖励函数——非常适合在没有训练好的奖励模型的情况下进行强化微调。您可以提供一个 Python 函数来即时计算奖励，而无需使用预训练的奖励模型。

适用于：
- 基于规则的奖励（长度、格式、代码执行、数学验证）
- 外部 API 奖励（裁判模型、编译器、测试套件）
-