OpenRLHF 轻量高效的工业级LLM训练和对齐框架

轻量高效的工业级LLM训练和对齐框架,OpenRLHF 支持70B模型RLHF全参数全流程训练!

?自横空出世以后,大家开始关注到以InstructGPT为代表的RLHF对齐技术,并以此为基础尝试复现ChatGPT的训练流程,逐步出现了 ColossalChat、DeepSpeed-Chat等代表性的RLHF复现工作。但彼时大家对对齐技术的理解,基本都是围绕着InstructGPT展开的,由于OpenAI最近不太Open,实际上是缺乏第三方的充分验证的。幸运的是,很快就横空出世了,不仅充分验证了RLHF技术的有效性,还有着足够的创新之处(比如拒绝采样和多RM等),立马引爆了整个开源社区。

鉴于InstructGPT和LLaMA2的火爆,我们OpenLLMAI开源社区调研了当前主流的对齐训练框架,发现大部分框架还缺乏对LLaMA2全流程全参数训练的支持、缺乏足够的可扩展性或者不够轻量易用。因此我们决心做一个真正工业级的LLM对齐训练框架,复现以InstructGPT和LLaMA2为代表的大模型训练流程,支持主流的RLHF/DPO等对齐技术,帮助大家快速实现自己的对齐想法。

所以,欢迎来到,快速开启你的对齐工作吧!

https://github.com/OpenLLMAI/OpenRLHF

OpenRLHF项目介绍

目前OpenLLMAI主要在做的项目有两个:

其余项目看后续人力和兴趣而定,近期我们可能会启动一个KD或者SE的框架,暂时可能对训练一个通用的小模型兴趣不大,主要原因是时间、资金和精力都很有限,用爱发电有点难以为继,所以很多时候更多会考虑兴趣导向。但是,兴趣毕竟不能当饭吃,所以我们近期花了比较大的精力来准备这份宣传材料(以前羡鱼同学太过佛系/强迫症/忙,毛病很多hh)。诚然,OpenLLMAI还很懵懂,OpenRLHF也还不够完善,但是我们尽了最大的诚意,也希望可以得到社区更广泛的认可和支持,一群人可以走的更远!

 

OpenRLHF设计思路

1.设计目标:轻量高效工业级LLM训练和对齐框架

由于目前业界缺乏真正工业级别的LLM对齐框架,大部分厂商可能会选择自己实现(感谢OpenAI开了个好头),短期来看这无可厚非,但长远来看终归难免重复造轮子的问题。

因此,我们的目标是做出一个轻量高效的工业级LLM训练和对齐框架。为了实现这一目标,我们一方面在第一个版本中做了比较审慎的开发和测试,力求第一个版本的可用性;另一方面在此正式开源,以吸引更多志同道合者来参与共建。对于框架,我们坚信开源才有生命力!

2.设计思想:简单易用、高性能、可扩展、探索性

  • 简单易用:易用性是我们设计OpenRLHF框架的第一个指导思想,因为高性能是一个合格框架的应有之义,所以我们并不会过多的强调这个事情,在保证高性能的前提下,提升易用性是我们的第一目标。

  • 可扩展:以7B为基础,向下兼容1-2B小模型的训练,向上逐步支持不断增长的模型规模,比如34B/70B/170B的训练。

  • 探索性:在保证基础的框架功能之外,我们会保持对齐技术的前沿性,跟踪最新进展并快速实现,同时也会提供我们团队开发的最新的对齐算法。后续我们还会开发LLMPipeline模块,以提供主流对齐算法或者主流模型训练技术的快速实践和公平比较。

3.实现思路

  • 易用性:在基础大模型框架方面,我们调研了DeepSpeed/Megatron-LM等LLM训练框架,在第一个版本中选择了更简洁易用的DeepSpeed;在模型库方面,我们毫不犹豫的选择了拥抱抱抱脸;在分布式扩展方面,我们选择了ray,别问,问就是祥ray!(主要用于资源调度)

  • 可扩展和高性能:使用ray进行合理的GPU资源调度,将Actor、Reward、Reference 和 Critic 模型分配到单独的 GPU 上,将训练和推理隔离以充分利用推理社区的优秀工具,同时配合offload、PEFT等显存节省技术,实现大模型的规模扩展和高效训练。

  • 探索性:第一个版本我们完整的复现了InstructGPT和LLaMA2的训练流程,并支持了DPO等更新的对齐技术,未来也将继续保持探索性,并开发pipeline模块,以支持InstructGPTPipeline和LLaMA2Pipeline等主流模型的pipeline,帮助社区进行更科学的比较和研究。

OpenRLHF主要亮点

OpenRLHF主要特性

  • 首个开源的全面复现LLaMA2和InstructGPT的RLHF对齐框架;

    • 支持SFT/RM/PPO全流程训练;

    • 支持拒绝采样、多个RM;

  • 简单易用:OpenRLHF 是目前最简单的高性能 RLHF 库之一,只需单个8卡 DGXA100 节点即可实现 34B 模型 RLHF 训练,可通过脚本一键启动训练

  • 训推分离,分布式可扩展的RLHF;

    • 使用多卡24GB RTX 4090 GPU 进行7B 模型的全流程训练

    • 使用多卡 A100 80G GPU 和 vLLM 进行70B+ 模型的全流程训练

    • 训推分离:训练和推理分离,以复用社区良好的推理工具(我们最终使用了vLLM)降低推理延时;

    • 分布式可扩展:通过ray/deepspeed/vLLM的加持,在合理的资源调度下,我们实现了高效可扩展的训练,以下为两个示例:

  • 高性能:得益于ray/deepspeed和其他的显存节省技术、推理加速框架,我们在13B LLaMA2 模型上的训练性能是 DeepSpeedChat 的 4 倍以上;

    • zero系列

    • FlashAttention2

    • LoRA、QLoRA

    • offload

    • gradient checkpointing

    • 推理加速:vLLM

    • 显存节省技巧:

  • 前沿性:紧跟前沿进展,目前支持主流的对齐技术、主流的大模型;

    • LLaMA

    • baichuan

    • qwen

    • Mixtral 8*7b

    • 最前沿的模型:

    • 标准的RLHF:SFT/RM/PPO;

    • Rejection Sampling;

    • DPO (direct-preference-optimization)/IPO/cDPO;

    • Kahneman-Tversky optimization (KTO);

    • Conditional SFT (https://arxiv.org/abs/2308.12050);

    • 最先进的对齐技术:

  • 强化学习技巧:We integrated the implementation tricks for PPO to improve the training stability, referencing Implementation Matters in Deep Policy Gradients and ppo-implementation-details.

 

OpenRLHF性能展示

支持矩阵:

下面的支持矩阵展示了OpenRLHF与主流LLM对齐训练框架的比较(调研可能存在延迟,错漏之处请联系我们修正):

 PPO Tricks34B 全参/4 A10070B+全参/16 A1007B 全/4 RTX4090QLoRAMixtral MOE 8*7b
OpenRLHF
DeepSpeedChat✖️✖️✖️✖️✖️✖️
ColossalAIChat✖️✖️✖️✖️✖️✖️
TRL✖️✖️✖️✖️
LLaMA-Factory✖️✖️✖️✖️✔(QLoRA)

OpenRLHF的主要优势在于良好的可扩展性高效的性能,可以支持70B模型的全流程全参数高效训练,也可以应对未来更大规模的扩展。而LLaMA-Factory/trl/trlx 等框架都存在类似的问题, 不支持 70B 全参数RLHF训练,有的框架主打Lora 微调 13b 级别的模型,一般采样合并 actor critic 的方案(节省显存,这是小规模上进行RLHF的权宜之计,但并不符合标准RLHF的实现,而且可扩展性很差,总有放不下的时候)。当然了,OpenRLHF也存在一些劣势,比如文档和benchmark不够完善,易用性还有待提高。具体而言,就OpenRLHF与各流行RLHF框架的对比我们做如下说明(错漏之处,欢迎大家指正),更详细和全面的对比后续可以在我们正式的技术报告中找到。

  • LLaMA-Factory:优势是高效微调和易用性(这一点非常值得我们学习,甚至有web-ui),使用merged actor-critic,无法支持70B 全参数PPO训练,也不便于扩展模型规模;

  • Colossal-Chat:使用single-step RL,而我们的框架使用的是step-wise RL。详见OpenRLHF vs Colossal-Chat;

  • trl/trlx:优势是与Hugging Face的生态兼容的非常好,但可能存在封装过深不易修改的问题,同样的,目前暂不支持70B 全参数PPO训练;而且使用的是merged actor-critic以节省显存,但这与标准实现不符;

  • NeMo-Aligner:基于Megatron的生成目前效率不高,影响了整体训练效率,与Hugging Face的生态兼容性不太好,模型可能需要做专门的修改;

性能数据:

根据现有测试,我们的OpenRLHF框架在13B模型上的训练效率大约是DeepSpeedChat的4倍左右(人力所限,测试可能存在延迟,大家可以向我们报告其他框架的性能数据以做修正)。

 7B llama2 RLHF13B llama2 RLHF (50k samples)
OpenRLHF-17 hours with 8 A100
DeepSpeedChat-48 hours with 16 A100

训练吞吐:

  • 默认配置:

  • 4 A100 80G for Actor, 2 A100 80G for Critic, 1 A100 80G for RM, and 1 A100 80G for InitPolicy

  • ZeRO2 with Adam Offload

  • Max Sequence Length: 2048

  • 性能吞吐(默认配置下samples/s,后续会换成tokens/s):

  • 7B llama2: 0.105 samples/gpu/secsmicro_batch_size = 16/8 (rollout/train), generation_length = 100~300

  • 13B llama2: 0.04 samples/gpu/secsmicro_batch_size = 8/4 (rollout/train), generation_length = 200~400

  • 34B codellama: 0.007 samples/gpu/secsmicro_batch_size = 2/1 (rollout/train), generation_length = 300~800

主流模型性能数据(人力原因,暂时来不及重新测试,这里报告的是当时支持该模型时的测试数据,当前版本PPO应该会快很多,后续会在正式的技术报告中补充更多的模型并更新性能数据):

modelSFTRMPPONotes
Baichuan2-7B1h4h71h 
Qwen-7B--- 

OpenRLHF使用方法

官方文档:

包括本文在内的官方正式文档都将在Github上维护,提升文档质量以改善易用性也是我们后续工作的重点方向之一(人力原因,文档目前比较粗糙,欢迎大家参与贡献):

  • 项目主页

  • 官方文档

OpenRLHF安装

我们支持nvidia-docker(推荐,以避免潜在的环境问题)或者conda环境安装(后续可以提供配置好的conda环境或者镜像)

首先,clone仓库:

Clone the repository:
git clone https://github.com/openllmai/OpenRLHF.git

 

然后,安装nv-docker或者conda环境:

#安装nv-docker
cd examples/scripts

# install nvidia-docker (Optional)
./nvidia_docker_install.sh

# launch nvidia container
./docker_run.sh
# we need conda
conda create -n openrlhf python=3.10
# so, we need install some package manually: when installing torch, you may need to match the corresponding cuda version.
pip install packaging ninja
pip3 install torch
# check ninjia
ninja --version
echo $? # output: 0
# install flash-attn: may take some time.
# For network error: you can download specified version from https://github.com/Dao-AILab/flash-attention/releases.
pip install flash-attn==2.4.2
./build_openrlhf.sh
# enjoy it!

conda activate openrlhf

 

 

OpenRLHF训练

训练脚本:

配置好环境之后,进入/openrlhf/examples/scripts目录,根据自己的需求修改训练脚本,一键即可启动训练,支持单机和多机训练,支持7B-70B+的模型全量全流程训练。以下为部分重要参数,用户可以根据情况进行修改,以支持自己的模型训练:

  • -pretrain:预训练模型地址, 抱抱脸格式

  • -dataset :数据集地址, 抱抱脸格式

  • -dataset_probs:多个数据集混合的采样概率比如: 0.5,0.4,0.1

  • -save_path:模型保存地址, 抱抱脸格式

  • -max_epochs:训练 epoch 次数

  • -micro_train_batch_size:单 GPU batch_size

  • -train_batch_size:全局 batch_size

  • -learning_rate:学习速率

单机训练脚本:

cd examples/scripts

# install nvidia-docker (Optional)
./nvidia_docker_install.sh

# launch nvidia container
./docker_run.sh

# cd in container
cd /openrlhf/examples/scripts

# build OpenRLHF (i.e, pip install)
./build_openrlhf.sh

# huggingface login
~/.local/bin/huggingface-cli login

# continue pretrain
./train_continue_pretrain_llama.sh

# train SFT model
./train_sft_llama.sh

# train RM model
./train_rm_llama.sh

# train PPO model
./train_ppo_llama.sh

# train DPO model
./train_dpo_llama.sh

# train KTO model
./train_kto_llama.sh

# train Rejection Sampling model
./train_rejection_sampling_llama.sh

# train Conditional SFT model
./train_conditional_llama.sh

 

多机训练脚本,16卡A100 70B模型全参数RLHF训练

cd examples/scripts

# launch nvidia container
./docker_run.sh

# cd in container
cd /openrlhf/examples/scripts

# build OpenRLHF (i.e, pip install)
./build_openrlhf.sh
# due to the compatibility of nVIDIA PyTorch image
pip uninstall xgboost transformer_engine -y

# huggingface login
~/.local/bin/huggingface-cli login

# launch the master node of ray in container
ray start --head --node-ip-address 0.0.0.0 --num-gpus 8
# if you want to launch ray on more nodes, use
ray start --address {MASTER-NODE-ADDRESS}:6379 --num-gpus 8


# train ray PPO model, requires 8 gpus in default config
./train_ppo_llama_ray.sh

# for 70B models and vLLM-based RLHF (important!)
pip install vllm==0.3.2
# due to the compatibility of vLLM
pip uninstall flash_attn -y

./train_ppo_llama_ray_70b.sh

 

推理

推理和评估我们推荐复用业界开源的工具或者代码,可以参考以下脚本:

未来工作

OpenRLHF未来的开发工作将以易用性和实用性(文档、教程、实践经验等)、前沿性(新算法、模型pipeline等)和稳定性为主,具体而言,有以下潜在的工作,希望大家可以一起参与:

  • 文档:中英文版本

    • 教程:提供良好的教程

    • 环境:提供配置好的镜像或conda环境

  • 性能测试、benchmark;

    • 基础功能的测试

    • 与其他框架的对比

    • 支持模型的测试

    • 对齐算法的测试

  • 进一步的性能优化;

  • 稳定性提升:定期代码review;

  • 新功能、新算法;

  • 新模型支持:Google的新模型等;

  • evaluation模块提供更全面的评估能力;

组织介绍

组织简介

OpenLLMAI: Open AI for everyone.

我们或许离OpenAI很远,但我们离Open很近。所以。我们对组织成员的要求也只有两个,那就是希望大家足够Open也足够自信。我们的态度:”有一分热发一分光“。愿与大家同走AI之毅行路,士不可以不弘毅,任重而道远!

大家因热爱相聚,主要想做的事情就两个:1.交流LLM的技术(技术分享、知识传播);2.开发LLM的工具(训练框架、模型、数据工程等),欢迎感兴趣的同学加入我们!组织的详细介绍见知乎旧文OpenLLMAI组织介绍。

开发者们

一路走来,OpenRLHF项目吸引了20+的贡献者,贡献了130+的提交,收获了800+ stars。在此感谢所有的贡献者,尤其是hijkzzz、wuxibin和Xianyu等同学为项目的发展做出了突出贡献,其中hijkzzz和Xianyu同学是本项目的发起人,hijkzzz作为项目的管理员提交了本项目的第一版代码,长期投入了大量精力进行维护,为项目的发展做出了不可替代的贡献;wuxibin作为项目的核心开发者,主要负责基于Ray对框架进行大规模扩展,并且长期投入了大量精力进行日常维护;Xianyu作为项目的管理员负责了NLP部分的开发工作和一些项目规划工作;此外还有pikaqqqqqq、li-plus、wwxFromTju、jovany-wang、xffxff、dabney777、suc16、Dylancer1998等同学也对项目的发展做出了重要贡献(这里无法一一列出,后续所有的开发者都会在正式的技术报告/论文中说明;还有很多的同学和老师虽然没有直接参与贡献,但是提出了很多宝贵的意见,真的非常感谢大家)。也欢迎越来越多志同道合的朋友们加入我们,愿OpenLLMAI与大家一起成长!

有意参与贡献的同学可直接在git上参与开发、联系相关负责人或者官方邮箱。

  • RL:hijkzzz

  • Ray:wuxibin

  • NLP:Xianyu

  • 官方邮箱:xianyuai@openllmai.top

赞助我们

目前OpenLLMAI是一个纯开源组织,无论是OpenRLHF/OpenLLMWiki等项目,还是OpenLLM Talk和技术交流群,都是完全开源开放的。但长远来看,没有资金支持注定难以为继,用爱发电走到今天并不容易,感谢大家一路的支持。最后,求赞助呀,欢迎大家有钱的捧个钱场(算力!!!),有人的捧个人场(参与开发或者其他贡献)!赞助或合作请联系xianyuai@openllmai.top。

参考资料

https://github.com/OpenLLMAI/OpenRLHF

https://github.com/NVIDIA/Megatron-LM

https://chat.openai.com/

InstructGPT

LLaMA2

https://github.com/facebookresearch/llama

Hugging Face Transformers

DeepSpeed

https://github.com/microsoft/DeepSpeed/tree/master/blogs/deepspeed-chat

Ray

https://github.com/hpcaitech/ColossalAI/tree/main/applications/Chat

https://github.com/CarperAI/trlx

https://github.com/NVIDIA/NeMo-Aligner

https://github.com/hiyouga/LLaMA-Factory

https://github.com/OpenLLMAI/OpenLLMWiki

【OpenLLMAI】相信开源的力量:我们有自己的组织了!任重道远,行则将至! - OpenLLMAI的文章 - 知乎https://zhuanlan.zhihu.com/p/647882819

如何正确复现 Instruct GPT / RLHF? - 蜗牛在花园跑酷的文章 - 知乎 https://zhuanlan.zhihu.com/p/622134699

开启训练之旅: 基于Ray和vLLM构建70B+模型的开源RLHF全量训练框架 - 蜗牛在花园跑酷的文章 - 知乎https://zhuanlan.zhihu.com/p/678828949

【OpenLLM 006】LoRA:大模型的低秩适配-最近大火的lora到底是什么东西?为啥stable diffusion和开源ChatGPT复现都在用? - OpenLLMAI的文章 - 知乎

https://zhuanlan.zhihu.com/p/620327907

https://arxiv.org/abs/2005.12729

https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/

https://github.com/baichuan-inc/Baichuan2

https://github.com/QwenLM/Qwen

https://mistral.ai/news/mixtral-of-experts/

https://github.com/OpenLLMAI/OpenRLHF/issues/221

 

评分

留下评论

您的邮箱地址不会被公开。 必填项已用 * 标注