Beaver是由北京大学PKU-Alignment团队开发的高度模块化的开源RLHF框架。它旨在为对齐研究提供训练数据和可重复的代码管道,特别是通过安全 RLHF 方法进行约束对齐 LLM 研究。
Beaver 的主要特点是:
2024/06/13
:我们很高兴地宣布 PKU-SafeRLHF 数据集 1.0 版开源。该版本在最初的测试版基础上进行了改进,纳入了人类-人工智能联合注释、扩大了危害类别的范围,并引入了详细的严重性级别标签。有关更多详细信息和访问权限,请访问我们的数据集页面?拥抱脸部:PKU-Alignment/PKU-SafeRLHF。2024/01/16
:我们的方法Safe RLHF已被 ICLR 2024 Spotlight 接受。2023/10/19
:我们在 arXiv 上发布了Safe RLHF 论文,详细介绍了我们新的安全对齐算法及其实现。2023/07/10
:我们很高兴地宣布Beaver-7B v1 / v2 / v3 模型的开源,作为 Safe RLHF 训练系列的第一个里程碑,并辅以相应的奖励模型v1 / v2 / v3 / 统一和成本模型v1/v2/v3/统一检查点?拥抱脸。2023/07/10
:我们扩展了开源安全偏好数据集PKU-Alignment/PKU-SafeRLHF ,它现在包含超过 30 万个示例。 (另请参阅 PKU-SafeRLHF-数据集部分)2023/07/05
:我们增强了对中文预训练模型的支持,并纳入了额外的开源中文数据集。 (另请参阅中文支持 (中文支持) 和自定义数据集 (自定义数据集) 部分)2023/05/15
:首次发布 Safe RLHF 管道、评估结果和训练代码。来自人类反馈的强化学习:通过偏好学习实现奖励最大化
来自人类反馈的安全强化学习:通过偏好学习限制奖励最大化
在哪里
最终目标是找到一个模型
与其他支持 RLHF 的框架相比, safe-rlhf
是第一个支持从 SFT 到 RLHF 以及评估所有阶段的框架。此外, safe-rlhf
是第一个在RLHF阶段考虑安全偏好的框架。它为策略空间中的约束参数搜索提供了更多的理论保证。
快速傅里叶变换 | 偏好模型1培训 | RLHF | 安全RLHF | PTX 损失 | 评估 | 后端 | |
---|---|---|---|---|---|---|---|
海狸 (安全-RLHF) | ✔️ | ✔️ | ✔️ | ✔️ | ✔️ | ✔️ | 深速 |
trlX | ✔️ | 2 | ✔️ | 加速/尼莫 | |||
深度快速聊天 | ✔️ | ✔️ | ✔️ | ✔️ | 深速 | ||
巨大的人工智能 | ✔️ | ✔️ | ✔️ | ✔️ | 巨型人工智能 | ||
羊驼农场 | 3 | ✔️ | ✔️ | ✔️ | 加速 |
PKU-SafeRLHF
数据集是一个人工标记的数据集,包含性能和安全偏好。它包括侮辱、不道德、犯罪、情感伤害、隐私等十多个维度的约束。这些约束是为 RLHF 技术中的细粒度值调整而设计的。
为了方便多轮微调,我们将发布每轮的初始参数权重、所需数据集和训练参数。这确保了科学和学术研究的可重复性。数据集将通过滚动更新逐步发布。
该数据集可在 Hugging Face 上找到:PKU-Alignment/PKU-SafeRLHF。
PKU-SafeRLHF-10K
是PKU-SafeRLHF
的子集,包含具有 10K 实例的第一轮 Safe RLHF 训练数据,包括安全偏好。您可以在 Hugging Face 上找到它:PKU-Alignment/PKU-SafeRLHF-10K。
我们将逐步发布完整的 Safe-RLHF 数据集,其中包括100 万对人类标记的有益和无害的偏好。
Beaver 是一个基于 LLaMA 的大型语言模型,使用safe-rlhf
进行训练。它是在羊驼模型的基础上发展起来的,通过收集人类关于有益和无害的偏好数据,并采用Safe RLHF技术进行训练。在保持羊驼的有益性能的同时,海狸显着提高了其无害性。
海狸被称为“天然大坝工程师”,它们善于利用树枝、灌木、岩石和土壤建造水坝和小木屋,创造适合其他生物栖息的湿地环境,使其成为生态系统中不可或缺的一部分。为了保证大型语言模型(LLM)的安全性和可靠性,同时容纳不同人群的广泛价值观,北京大学团队将他们的开源模型命名为“Beaver”,旨在通过约束价值为LLM筑起一座堤坝。对准(CVA)技术。该技术能够对信息进行细粒度标记,并与安全的强化学习方法相结合,显着减少模型偏差和歧视,从而增强模型的安全性。类似于海狸在生态系统中的作用,Beaver模型将为大型语言模型的发展提供至关重要的支持,并为人工智能技术的可持续发展做出积极贡献。
遵循Vicuna模型的评估方法,我们利用GPT-4来评估Beaver。结果表明,与羊驼毛相比,海狸在与安全性相关的多个维度上都表现出显着的改进。
在 Alpaca-7B 模型上使用 Safe RLHF 管道后,安全偏好的分布发生了显着变化。
从 GitHub 克隆源代码:
git clone https://github.com/PKU-Alignment/safe-rlhf.git
cd safe-rlhf
Native Runner:使用conda
/ mamba
设置 conda 环境:
conda env create --file conda-recipe.yaml # or `mamba env create --file conda-recipe.yaml`
这将自动设置所有依赖项。
Containerized Runner:除了使用带有conda隔离的本机之外,您还可以使用docker镜像来配置环境。
首先,请按照 NVIDIA Container Toolkit:安装指南和 NVIDIA Docker:安装指南来设置nvidia-docker
。然后你可以运行:
make docker-run
此命令将构建并启动一个安装有适当依赖项的 docker 容器。主机路径/
将映射到/host
,当前工作目录将映射到容器内的/workspace
。
safe-rlhf
支持从监督微调 (SFT) 到偏好模型训练再到 RLHF 对齐训练的完整流程。
conda activate safe-rlhf
export WANDB_API_KEY= " ... " # your W&B API key here
或者
make docker-run
export WANDB_API_KEY= " ... " # your W&B API key here
bash scripts/sft.sh
--model_name_or_path < your-model-name-or-checkpoint-path >
--output_dir output/sft
注意:您可能需要根据您的机器设置更新脚本中的一些参数,例如用于训练的 GPU 数量、训练批量大小等。
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.sh
--actor_model_name_or_path output/sft
--reward_model_name_or_path output/rm
--output_dir output/ppo
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 来缓解 GPU 内存使用峰值。
所有训练脚本都可以通过额外的选项--offload
(默认为none
,即禁用 ZeRO-Offload)来将张量(参数和/或优化器状态)卸载到 CPU。例如:
bash scripts/sft.sh
--model_name_or_path ~ /models/llama-7b
--output_dir output/sft
--offload all # or `parameter` or `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
提供了一个抽象来为所有监督微调、偏好模型训练和 RL 训练阶段创建数据集。
class RawSample ( TypedDict , total = False ):
"""Raw sample type.
For SupervisedDataset, should provide (input, answer) or (dialogue).
For PreferenceDataset, should provide (input, answer, other_answer, better).
For SafetyPreferenceDataset, should provide (input, answer, other_answer, safer, is_safe, is_other_safe).
For PromptOnlyDataset, should provide (input).
"""
# Texts
input : NotRequired [ str ] # either `input` or `dialogue` should be provided
"""User input text."""
answer : NotRequired [ str ]
"""Assistant answer text."""
other_answer : NotRequired [ str ]
"""Other assistant answer text via resampling."""
dialogue : NotRequired [ list [ str ]] # either `input` or `dialogue` should be provided
"""Dialogue history."""
# Flags
better : NotRequired [ bool ]
"""Whether ``answer`` is better than ``other_answer``."""
safer : NotRequired [ bool ]
"""Whether ``answer`` is safer than ``other_answer``."""
is_safe : NotRequired [ bool ]
"""Whether ``answer`` is safe."""
is_other_safe : NotRequired [ bool ]
"""Whether ``other_answer`` is safe."""
以下是实现自定义数据集的示例(有关更多示例,请参阅 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 :
# Load a dataset from Hugging Face
self . data = load_dataset ( path or 'my-organization/my-dataset' )[ 'train' ]
def __getitem__ ( self , index : int ) -> RawSample :
data = self . data [ index ]
# Construct a `RawSample` dictionary from your custom dataset item
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 ) # dataset size
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
这将使用随机分割的 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
注意:必须在训练脚本开始解析命令行参数之前导入数据集类。
python3 -m safe_rlhf.serve.cli --model_name_or_path output/sft # or 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 等,为中文提供更好的支持。您只需在训练和推理代码中更新预训练模型的路径即可。
Safe-RLHF 管道康复支持 LLaMA 系列模型,它也支持其他一些对中文支持更好的预训练模型,例如 Baichuan 和 InternLM 等。你只需要在训练和推理的代码中更新预训练模型的路径即可。
# SFT training
bash scripts/sft.sh --model_name_or_path baichuan-inc/Baichuan-7B --output_dir output/baichuan-sft
# Inference
python3 -m safe_rlhf.serve.cli --model_name_or_path output/baichuan-sft
与此同时,我们在原始数据集中添加了对中国数据集(例如 Firefly 和 MOSS 系列)的支持。只需要更改训练代码中的数据集路径,即可使用对应的数据集对中文预训练模型进行微调:
同时,我们也在raw-datasets中增加了支持一些中文数据集,例如Firefly和MOSS系列等。在训练代码中更改数据集路径,你就可以使用相应的数据集来影响中文预训练模型:
# scripts/sft.sh
- --train_datasets alpaca
+ --train_datasets firefly
有关如何添加自定义数据集的说明,请参阅自定义数据集部分。
关于如何添加自定义数据集的方法,请参阅章节自定义数据集(自定义数据集)。
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
# Install 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 evaluation
python3 -m safe_rlhf.evaluate.bigbench
--model_name_or_path output/ppo-lag
--task_name < BIG-bench-task-name >
# Install OpenAI Python API
pip3 install openai
export OPENAI_API_KEY= " ... " # your OpenAI API key here
# GPT-4 evaluation
python3 -m safe_rlhf.evaluate.gpt4
--red_corner_model_name_or_path output/sft
--blue_corner_model_name_or_path output/ppo-lag
如果您发现 Safe-RLHF 有用或在您的研究中使用 Safe-RLHF(模型、代码、数据集等),请考虑在您的出版物中引用以下工作。
@inproceedings { safe-rlhf ,
title = { Safe RLHF: Safe Reinforcement Learning from Human Feedback } ,
author = { Josef Dai and Xuehai Pan and Ruiyang Sun and Jiaming Ji and Xinbo Xu and Mickel Liu and Yizhou Wang and Yaodong Yang } ,
booktitle = { The Twelfth International Conference on Learning Representations } ,
year = { 2024 } ,
url = { https://openreview.net/forum?id=TyFrPOKYXw }
}
@inproceedings { beavertails ,
title = { BeaverTails: Towards Improved Safety Alignment of {LLM} via a Human-Preference Dataset } ,
author = { Jiaming Ji and Mickel Liu and Juntao Dai and Xuehai Pan and Chi Zhang and Ce Bian and Boyuan Chen and Ruiyang Sun and Yizhou Wang and Yaodong Yang } ,
booktitle = { Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track } ,
year = { 2023 } ,
url = { https://openreview.net/forum?id=g0QovXbFw3 }
}
以下所有学生的贡献均等,顺序按字母顺序确定:
全部由王一洲和杨耀东提供建议。致谢:感谢曲毅女士设计的Beaver标志。
该存储库受益于 LLaMA、Stanford Alpaca、DeepSpeed 和 DeepSpeed-Chat。感谢他们的精彩工作以及他们为实现法学硕士研究民主化所做的努力。 Safe-RLHF 及其相关资产是用爱构建和开源的?❤️。
这项工作得到了北京大学的支持和资助。
Safe-RLHF 在 Apache License 2.0 下发布。