基于安全RLHF的约束型价值对齐大模型

Beaver是由北京大学PKU-Alignment团队开发的一个高度模块化的开源RLHF框架。其目标是为对齐研究，尤其是通过安全RLHF方法进行的约束型对齐大模型研究，提供训练数据和可复现的代码流水线。

Beaver的主要特点包括：

• 支持主流预训练模型的SFT、RLHF和安全RLHF训练：LLaMA、OPT、百川等。
• 提供大规模人工标注数据集**（多达100万条）**，包含有益和无害偏好，以支持可复现的RLHF研究。
• 支持奖励模型和成本模型的训练，并提供预训练检查点。
• 支持针对SFT和RLHF的自定义参数和数据集。
• 提供多尺度的安全约束验证指标，例如BIG-bench、GPT-4评估等。

🦫 最新动态？

• 🎉 2024/06/13: 我们很高兴地宣布开源我们的PKU-SafeRLHF数据集1.0版本。该版本在初始测试版的基础上进行了改进，加入了人机联合标注，扩展了危害类别范围，并引入了详细的严重程度标签。更多详情及访问方式，请前往Hugging Face上的数据集页面：PKU-Alignment/PKU-SafeRLHF。
• 🎉 2024/01/16: 我们的安全RLHF方法已被ICLR 2024 Spotlight接收。
• 📄 2023/10/19: 我们已在arXiv上发布了安全RLHF论文，详细介绍了我们的新型安全对齐算法及其实现。
• 🚀 2023/07/10: 我们非常高兴地宣布开源Beaver-7B v1 / v2 / v3模型，作为安全RLHF训练系列的第一个里程碑，并配套提供了相应的奖励模型 v1 / v2 / v3 / 统一版以及成本模型 v1 / v2 / v3 / 统一版检查点，均已发布在Hugging Face平台上。
• 🔥 2023/07/10: 我们扩展了开源的安全偏好数据集，即PKU-Alignment/PKU-SafeRLHF，目前包含超过30万条样本。（详见PKEU-SafeRLHF数据集部分）
• ⚙ 2023/07/05: 我们增强了对中国预训练模型的支持，并整合了更多开源中文数据集。（详见中文支持和自定义数据集部分）
• ⭐️ 2023/05/15: 首次发布安全RLHF流水线、评估结果及训练代码。

目录

• 基于安全RLHF的约束型价值对齐
• 与其他RLHF库的比较
• PKU-SafeRLHF数据集
  • PKU-SafeRLHF-10K
  • PKU-SafeRLHF-1M
• 为什么叫“Beaver”？
• Beaver与Alpaca对比
• 安装
• 训练
• 自定义数据集
• 推理
  • 交互式CLI演示
  • 交互式竞技场
• 中文支持
• 基准测试与评估
  • 基于奖励和成本模型的竞技场
  • BIG-bench
  • GPT-4评估
• 未来计划
• 引用
• PKU-Alignment团队
• 许可证

基于安全RLHF的约束型价值对齐

从人类反馈中学习强化学习：通过偏好学习实现奖励最大化

从人类反馈中学习安全强化学习：通过偏好学习实现受约束的奖励最大化

其中$R (\cdot)$和$C (\cdot)$分别为奖励函数和成本函数。它们是基于人类偏好训练的人类代理神经网络。

最终目标是找到一个既有益（高奖励）又无害（低成本）的模型$\pi_{\theta}$。

与其他 RLHF 库的比较

与支持 RLHF 的其他框架相比，safe-rlhf 是首个支持从 SFT 到 RLHF 再到评估全流程的框架。此外，safe-rlhf 还是第一个在 RLHF 阶段就将安全偏好纳入考量的框架，它在策略空间中的约束参数搜索方面具有更强的理论保证。

	SFT	偏好模型1训练	RLHF	安全 RLHF	PTX 损失	评估	后端
Beaver (Safe-RLHF)	✔️	✔️	✔️	✔️	✔️	✔️	DeepSpeed
trlX	✔️	❌2	✔️	❌	❌	❌	Accelerate / NeMo
DeepSpeed-Chat	✔️	✔️	✔️	❌	✔️	❌	DeepSpeed
Colossal-AI	✔️	✔️	✔️	❌	✔️	❌	ColossalAI
AlpacaFarm	❌3	✔️	✔️	❌	❌	✔️	Accelerate

<sup>1. 在 RLHF 的语境中，“偏好模型”即为“奖励模型”。而在 Safe RLHF 中，“偏好模型”则同时指代“奖励模型”和“成本模型”。
2. trlX 仓库的 examples 目录中确实有一个用于奖励模型训练的示例，但该示例并未得到官方支持，也未集成到 trlX 库中。
3. Alpaca 的监督微调支持是在 tatsu-lab/stanford_alpaca 仓库中提供的。</sup>

PKU-SafeRLHF 数据集

PKU-SafeRLHF 数据集是一个由人类标注的、同时包含性能偏好和安全偏好的数据集。它涵盖了十多个维度的约束条件，例如侮辱性内容、不道德行为、犯罪、情感伤害以及隐私等，旨在实现 RLHF 技术中的细粒度价值观对齐。

为了便于多轮微调，我们将逐步发布每一轮所需的初始参数权重、所需数据集以及训练参数，以确保科学和学术研究的可重复性。该数据集将通过滚动更新的方式逐步发布。

数据集已在 Hugging Face 上开放：PKU-Alignment/PKU-SafeRLHF。

PKU-SafeRLHF-10K

PKU-SafeRLHF-10K 是 PKU-SafeRLHF 的一个子集，包含了第一轮安全 RLHF 训练的 1 万条数据，其中包括安全偏好信息。你可以在 Hugging Face 上找到它：PKU-Alignment/PKU-SafeRLHF-10K。

PKU-SafeRLHF-1M

我们将逐步发布完整的安全 RLHF 数据集，其中包含 100 万对由人类标注的 有益性和无害性偏好数据。

为什么叫“Beaver”

Beaver 是一款基于 LLaMA 的大型语言模型，使用 safe-rlhf 技术进行训练。它建立在 Alpaca 模型的基础上，通过收集与有益性和无害性相关的用户偏好数据，并采用安全 RLHF 技术进行训练。在保持 Alpaca 有益性能的同时，Beaver 显著提升了其无害性。

海狸素有“天然筑坝工程师”的美誉，它们善于利用树枝、灌木、岩石和泥土建造水坝及小型木屋，从而创造出适合其他生物栖息的湿地环境，成为生态系统中不可或缺的一部分。为了确保大型语言模型（LLMs）的安全可靠，同时满足不同人群的多样化价值取向，北京大学团队将其开源模型命名为“Beaver”，并计划通过约束性价值观对齐（CVA）技术为 LLMs 筑起一道“大坝”。这项技术能够对信息进行细粒度标注，并结合安全强化学习方法，显著降低模型的偏见和歧视，从而提升模型的安全性。正如海狸在生态系统中所扮演的角色一样，Beaver 模型将为大型语言模型的发展提供重要支持，并为人工智能技术的可持续发展作出积极贡献。

Beaver 与 Alpaca 对比

参照 Vicuna 模型的评估方法，我们使用 GPT-4 对 Beaver 进行了评估。结果显示，与 Alpaca 相比，Beaver 在多项安全相关指标上均有显著提升。

在对 Alpaca-7B 模型应用安全 RLHF 流程后，安全偏好分布发生了显著变化。

安装

从 GitHub 克隆源代码：

git clone https://github.com/PKU-Alignment/safe-rlhf.git
cd safe-rlhf

原生运行环境： 使用 conda 或 mamba 设置 Conda 环境：

conda env create --file conda-recipe.yaml  # 或 `mamba env create --file conda-recipe.yaml`

这将自动安装所有依赖项。

容器化运行环境： 除了使用带有 Conda 隔离的本地机器外，您也可以选择使用 Docker 镜像来配置环境。

首先，请按照 NVIDIA Container Toolkit：安装指南 和 NVIDIA Docker：安装指南 设置 nvidia-docker。然后您可以运行：

make docker-run

该命令将构建并启动一个安装了所需依赖项的 Docker 容器。宿主机的根目录 / 将被映射到容器内的 /host，而当前工作目录则会被映射到容器内的 /workspace。

训练

safe-rlhf 支持从监督微调（SFT）到偏好模型训练，再到 RLHF 对齐训练的完整流程。

0. 按照安装部分的说明正确设置训练环境。

conda activate safe-rlhf
export WANDB_API_KEY="..."  # 在此处填写您的 W&B API 密钥

或者

make docker-run
export WANDB_API_KEY="..."  # 在此处填写您的 W&B API 密钥

1. 监督微调（SFT）

bash scripts/sft.sh \
    --model_name_or_path <您的模型名称或检查点路径> \
    --output_dir output/sft

注意：您可能需要根据自己的机器配置调整脚本中的一些参数，例如用于训练的 GPU 数量、训练批次大小等。

2. 价值模型（奖励模型和成本模型）

bash scripts/reward-model.sh \
    --model_name_or_path output/sft \
    --output_dir output/rm

bash scripts/cost-model.sh \
    --model_name_or_path output/sft \
    --output_dir output/cm

3. RLHF（可选）

bash scripts/ppo.sh \
    --actor_model_name_or_path output/sft \
    --reward_model_name_or_path output/rm \
    --output_dir output/ppo

4. 安全 RLHF

bash scripts/ppo-lag.sh \
    --actor_model_name_or_path output/sft \
    --reward_model_name_or_path output/rm \
    --cost_model_name_or_path output/cm \
    --output_dir output/ppo-lag

以下是一个使用 LLaMA-7B 运行整个流程的示例命令：

conda activate safe-rlhf
bash scripts/sft.sh --model_name_or_path ~/models/llama-7b --output_dir output/sft
bash scripts/reward-model.sh --model_name_or_path output/sft --output_dir output/rm
bash scripts/cost-model.sh --model_name_or_path output/sft --output_dir output/cm
bash scripts/ppo-lag.sh \
    --actor_model_name_or_path output/sft \
    --reward_model_name_or_path output/rm \
    --cost_model_name_or_path output/cm \
    --output_dir output/ppo-lag

计算资源需求

上述所有训练过程均已在配备 8 块 NVIDIA A800-80GB GPU 的云服务器上，使用 LLaMA-7B 进行测试。

对于 GPU 显存不足的用户，可以启用 DeepSpeed ZeRO-Offload 来缓解峰值显存占用。

所有训练脚本都支持通过额外的 --offload 参数（默认为 none，即禁用 ZeRO-Offload）将张量（参数和/或优化器状态）卸载到 CPU 上。例如：

bash scripts/sft.sh \
    --model_name_or_path ~/models/llama-7b \
    --output_dir output/sft \
    --offload all  # 或 `parameter` 或 `optimizer`

在多节点环境下，用户可以参考 DeepSpeed：资源配置（多节点） 文档以获取更多详细信息。以下是一个在 4 个节点（每个节点有 8 个 GPU）上启动训练过程的示例：

# myhostfile
worker-1 slots=8
worker-2 slots=8
worker-3 slots=8
worker-4 slots=8

然后使用以下命令启动训练脚本：

bash scripts/sft.sh \
    --hostfile myhostfile \
    --model_name_or_path ~/models/llama-7b \
    --output_dir output/sft

自定义数据集

safe-rlhf 提供了一个抽象层，用于为监督微调、偏好模型训练和强化学习训练的所有阶段创建数据集。

class RawSample(TypedDict, total=False):
    """原始样本类型。

    对于监督数据集，应提供 (input, answer) 或 (dialogue)。
    对于偏好数据集，应提供 (input, answer, other_answer, better)。
    对于安全偏好数据集，应提供 (input, answer, other_answer, safer, is_safe, is_other_safe)。
    对于仅提示数据集，应提供 (input)。
    """

    # 文本
    input: NotRequired[str]  # either `input` or `dialogue` should be provided
    """用户输入文本。"""
    answer: NotRequired[str]
    """助手回答文本。"""
    other_answer: NotRequired[str]
    """通过重采样得到的另一条助手回答文本。"""
    dialogue: NotRequired[list[str]]  # either `input` or `dialogue` should be provided
    """对话历史。"""

    # 标记
    better: NotRequired[bool]
    """是否 ``answer`` 比 ``other_answer`` 更好。"""
    safer: NotRequired[bool]
    """是否 ``answer`` 比 ``other_answer`` 更安全。"""
    is_safe: NotRequired[bool]
    """是否 ``answer`` 是安全的。"""
    is_other_safe: NotRequired[bool]
    """是否 ``other_answer`` 是安全的。"""

以下是一个实现自定义数据集的示例（更多示例请参见 safe_rlhf/datasets/raw）：

import argparse
from datasets import load_dataset
from safe_rlhf.datasets import RawDataset, RawSample, parse_dataset


class MyRawDataset(RawDataset):
    NAME = 'my-dataset-name'

    def __init__(self, path=None) -> None:
        # 从 Hugging Face 加载数据集
        self.data = load_dataset(path or 'my-organization/my-dataset')['train']

    def __getitem__(self, index: int) -> RawSample:
        data = self.data[index]
        # 从自定义数据集条目中构造一个 `RawSample` 字典
        return RawSample(
            input=data['col1'],
            answer=data['col2'],
            other_answer=data['col3'],
            better=float(data['col4']) > float(data['col5']),
            ...
        )

    def __len__(self) -> int:
        return len(self.data)  # 数据集大小


def parse_arguments():
    parser = argparse.ArgumentParser(...)
    parser.add_argument(
        '--datasets',
        type=parse_dataset,
        nargs='+',
        metavar='DATASET[:PROPORTION[:PATH]]',
    )
    ...
    return parser.parse_args()


def main():
    args = parse_arguments()
    ...


if __name__ == '__main__':
    main()

然后你可以将这个数据集传递给训练脚本，如下所示：

python3 train.py --datasets my-dataset-name

你也可以传递多个数据集，并可选地指定每个数据集的比例（用冒号 : 分隔）。例如：

python3 train.py --datasets alpaca:0.75 my-dataset-name:0.5

这将会随机选取斯坦福 Alpaca 数据集的 75% 和你自定义数据集的 50% 进行使用。

此外，如果已经从 Hugging Face 克隆了数据集仓库，数据集参数后面还可以跟本地路径（同样用冒号 : 分隔）。

git lfs install
git clone https://huggingface.co/datasets/my-organization/my-dataset ~/path/to/my-dataset/repository
python3 train.py --datasets alpaca:0.75 my-dataset-name:0.5:~/path/to/my-dataset/repository

注意：在训练脚本开始解析命令行参数之前，必须先导入数据集类。

推理

交互式 CLI 演示

python3 -m safe_rlhf.serve.cli --model_name_or_path output/sft  # 或 output/ppo-lag

交互式竞技场

python3 -m safe_rlhf.serve.arena --red_corner_model_name_or_path output/sft --blue_corner_model_name_or_path output/ppo-lag

中文支持

Safe-RLHF 流程不仅支持 LLaMA 系列模型，还支持其他一些对中文支持更好的预训练模型，例如 Baichuan 和 InternLM 等。你只需要在训练和推理代码中更新预训练模型的路径即可。

# SFT 训练
bash scripts/sft.sh --model_name_or_path baichuan-inc/Baichuan-7B --output_dir output/baichuan-sft

# 推理
python3 -m safe_rlhf.serve.cli --model_name_or_path output/baichuan-sft

同时，我们也在 raw-datasets 中增加了支持一些中文数据集，例如 Firefly 和 MOSS 系列等。在训练代码中更改数据集路径，你就可以使用相应的数据集来微调中文预训练模型：

# scripts/sft.sh
-	--train_datasets alpaca \
+	--train_datasets firefly \

关于如何添加自定义数据集的方法，请参阅章节 Custom Datasets (自定义数据集)。

基准测试与评估

通过奖励和成本模型进行竞技场评估

scripts/arena-evaluation.sh \
    --red_corner_model_name_or_path output/sft \
    --blue_corner_model_name_or_path output/ppo-lag \
    --reward_model_name_or_path output/rm \
    --cost_model_name_or_path output/cm \
    --output_dir output/arena-evaluation

BIG-bench

# 安装 BIG-bench
git clone https://github.com/google/BIG-bench.git
(
    cd BIG-bench
    python3 setup.py sdist
    python3 -m pip install -e .
)

# BIG-bench 评估
python3 -m safe rlhf.evaluate.bigbench \
    --model_name_or_path output/ppo-lag \
    --task_name <BIG-bench-task-name>

GPT-4 评估

# 安装 OpenAI Python API
pip3 install openai
export OPENAI_API_KEY="..."  # 在此处填写你的 OpenAI API 密钥

# GPT-4 评估
python3 -m safe rlhf.evaluate.gpt4 \
    --red_corner_model_name_or_path output/sft \
    --blue_corner_model_name_or_path output/ppo-lag

未来计划

• Beaver-7B 检查点已在 Hugging Face 上发布。
• 发布 Safe RLHF 论文预印本。
• 我们将逐步发布完整的 Safe-RLHF 数据集。
• 使用 Safe-RLHF 训练更大的 LLM。
• 支持内存高效的训练方法，如 LoRA、PEFT 等。

引用

如果您认为 Safe-RLHF 有用，或在您的研究中使用了 Safe-RLHF（模型、代码、数据集等），请考虑在您的出版物中引用以下文献。

@inproceedings{safe-rlhf,
  title={Safe RLHF: 基于人类反馈的安全强化学习},
  author={Josef Dai 和 Xuehai Pan 和 Ruiyang Sun 和 Jiaming Ji 和 Xinbo Xu 和 Mickel Liu 和 Yizhou Wang 和 Yaodong Yang},
  booktitle={第十二届国际学习表示会议},
  year={2024},
  url={https://openreview.net/forum?id=TyFrPOKYXw}
}
@inproceedings{beavertails,
  title={BeaverTails：通过人类偏好数据集实现 {LLM} 更佳的安全对齐},
  author={Jiaming Ji 和 Mickel Liu 和 Juntao Dai 和 Xuehai Pan 和 Chi Zhang 和 Ce Bian 和 Boyuan Chen 和 Ruiyang Sun 和 Yizhou Wang 和 Yaodong Yang},
  booktitle={第三十七届神经信息处理系统大会数据集与基准赛道},
  year={2023},
  url={https://openreview.net/forum?id=g0QovXbFw3}
}

北大对齐团队

以下所有同学贡献均等，顺序按姓名字母排序：

• Juntao Dai
• Jiaming Ji
• Xuehai Pan
• Ruiyang Sun

全体由 Yizhou Wang 和 Yaodong Yang 指导。
特别感谢：我们感谢 Ms. Yi Qu 设计的海狸标志。

致谢

本仓库受益于 LLaMA、斯坦福 Alpaca、DeepSpeed 和 DeepSpeed-Chat 的优秀工作。感谢他们为推动 LLM 研究的民主化所做出的努力。Safe-RLHF 及其相关资源以爱构建并开源 🤗❤️。

本项目得到北京大学的支持与资助。

许可证

Safe-RLHF 采用 Apache License 2.0 开源。

safe-rlhf 快速上手指南

safe-rlhf (Beaver) 是由北京大学 PKU-Alignment 团队开发的高度模块化开源框架，旨在通过 Safe RLHF（安全强化学习人类反馈）技术，训练既“有帮助”又“无害”的大语言模型。它支持从 SFT、奖励/成本模型训练到 Safe RLHF 全流程，并提供大规模人工标注的安全偏好数据集。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU (建议显存 24GB+ 以运行 7B 模型训练)
• 前置依赖:
  • PyTorch (与您的 CUDA 版本匹配)
  • DeepSpeed (框架后端依赖)
  • Hugging Face Transformers, Datasets, Accelerate

国内加速建议： 推荐使用国内镜像源安装 Python 依赖，并配置 Hugging Face 镜像以加速模型和数据集下载。

# 设置 Hugging Face 镜像 (可选)
export HF_ENDPOINT=https://hf-mirror.com

安装步骤

1. 克隆项目代码

git clone https://github.com/PKU-Alignment/safe-rlhf.git
cd safe-rlhf

2. 创建虚拟环境并安装依赖

建议使用 conda 或 venv 创建独立环境。

# 创建 conda 环境
conda create -n safe-rlhf python=3.9 -y
conda activate safe-rlhf

# 安装 PyTorch (请根据实际 CUDA 版本选择，此处以 cu118 为例)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# 安装项目依赖
# 国内用户可使用清华源加速
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 safe-rlhf 包本身
pip install -e .

基本使用

safe-rlhf 的核心流程通常包含三个阶段：SFT (有监督微调) -> Reward/Cost Model Training (奖励/成本模型训练) -> Safe RLHF (安全强化学习)。

以下是最简单的命令行使用示例，展示如何启动训练任务。

1. 数据准备

框架会自动从 Hugging Face 下载 PKU-SafeRLHF 数据集。如果您网络受限，可手动下载后本地加载，或直接使用默认配置（需配置好镜像）。

2. 运行 SFT (有监督微调)

使用提供的配置文件对预训练模型（如 LLaMA-7B）进行微调。

deepspeed --num_gpus=8 \
    safe_rlhf/train/sft.py \
    --model_name_or_path meta-llama/Llama-2-7b-hf \
    --dataset_names PKU-Alignment/PKU-SafeRLHF \
    --split train \
    --output_dir ./output/sft_model \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --learning_rate 2e-5 \
    --num_train_epochs 3 \
    --deepspeed ds_configs/stage3.json

3. 训练奖励模型 (Reward Model) 与 成本模型 (Cost Model)

Safe RLHF 需要两个关键模型：一个评估“有用性”(Reward)，一个评估“有害性”(Cost)。

训练 Reward Model:

deepspeed --num_gpus=8 \
    safe_rlhf/train/reward.py \
    --model_name_or_path ./output/sft_model \
    --dataset_names PKU-Alignment/PKU-SafeRLHF \
    --split train \
    --output_dir ./output/reward_model \
    --per_device_train_batch_size 4 \
    --learning_rate 1e-5 \
    --num_train_epochs 1 \
    --deepspeed ds_configs/stage3.json

训练 Cost Model:

deepspeed --num_gpus=8 \
    safe_rlhf/train/cost.py \
    --model_name_or_path ./output/sft_model \
    --dataset_names PKU-Alignment/PKU-SafeRLHF \
    --split train \
    --output_dir ./output/cost_model \
    --per_device_train_batch_size 4 \
    --learning_rate 1e-5 \
    --num_train_epochs 1 \
    --deepspeed ds_configs/stage3.json

4. 运行 Safe RLHF

结合上述模型进行最终的安全对齐训练。

deepspeed --num_gpus=8 \
    safe_rlhf/train/rl.py \
    --model_name_or_path ./output/sft_model \
    --reward_model_name_or_path ./output/reward_model \
    --cost_model_name_or_path ./output/cost_model \
    --dataset_names PKU-Alignment/PKU-SafeRLHF \
    --split train \
    --output_dir ./output/safe_rlhf_model \
    --per_device_prompt_batch_size 4 \
    --per_device_train_batch_size 4 \
    --learning_rate 1e-6 \
    --num_train_epochs 1 \
    --deepspeed ds_configs/stage3.json

5. 交互式推理

训练完成后，可以使用 CLI 进行对话测试：

python safe_rlhf/inference/generate.py \
    --model_name_or_path ./output/safe_rlhf_model \
    --max_new_tokens 256 \
    --temperature 0.7 \
    --top_p 0.9