OpenRLHF 輕量高效的工業級LLM訓練和對齊框架- 人工智慧

輕量高效率的工業級LLM訓練和對齊框架，OpenRLHF 支援70B模型RLHF全參數全流程訓練！

OpenRLHF 是什麼？自ChatGPT橫空出世以後，大家開始關注到以InstructGPT為代表的RLHF對齊技術，並以此為基礎嘗試復現ChatGPT的訓練流程，逐步出現了ColossalChat、DeepSpeed-Chat等代表性的RLHF復現工作。但彼時大家對對齊技術的理解，基本上都是圍繞著InstructGPT展開的，由於OpenAI最近不太Open，實際上是缺乏第三方的充分驗證的。幸運的是，LLaMA2很快就橫空出世了，不僅充分驗證了RLHF技術的有效性，還有著足夠的創新之處（例如拒絕採樣和多RM等），立刻引爆了整個LLM開源社群。

鑑於InstructGPT和LLaMA2的火爆，我們OpenLLMAI開源社群研究了當前主流的對齊訓練框架，發現大部分框架還缺乏對LLaMA2全流程全參數訓練的支援、缺乏足夠的可擴展性或不夠輕量易用。因此我們決心做一個真正工業級的LLM對齊訓練框架，復現以InstructGPT和LLaMA2為代表的大模型訓練流程，支持主流的RLHF/DPO等對齊技術，幫助大家快速實現自己的對齊想法。

所以，歡迎來到OpenRLHF，快速開啟你的對齊工作吧！

https://github.com/OpenLLMAI/OpenRLHF

OpenRLHF專案介紹

目前OpenLLMAI主要在做的專案有兩個：

LLM框架：https://github.com/OpenLLMAI/OpenRLHF
LLMWiki：https://github.com/OpenLLMAI/OpenLLMWiki大模型煉丹知識庫

其餘項目看後續人力和興趣而定，近期我們可能會啟動一個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 Tricks	34B 全參/4 A100	70B+全參/16 A100	7B 全/4 RTX4090	QLoRA	Mixtral 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 RLHF	13B 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應該會快很多，後續會在正式的技術報告中補充更多的模型並更新性能數據）：

model	SFT	RM	PPO	Notes
Baichuan2-7B	1h	4h	71h
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 # weTP5T launch nvidia container ./docker_run.sh # we need con. 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 ninaging ninja pip3 install torch # check ninjia nin> 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_openhf.shp

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 con讓./build_openrlhf.sh # huggingface login ~/.local/bin/huggingface-cli login # continue pretrain ./train_continue_pretrain_llama.sh # train SFT modelin_continue_pretrain_llama.sh # train SFT model ./train_shftsft ./train_rm_llama.sh # train PPO model ./train_ppo_llama.sh # train DPO model ./train_dpo_llama.sh # train KTO model ./train_dpo_llama.sh # train KTO model ./train_kto_llama.shin ./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 (ie, pip install) ./build_openhIAf.shue P. image pip uninstall xgboost transformer_engine -y # huggingface login ~/.local/bin/huggingface-cli login # launch the master node of ray in container ray。 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 baseB53553555556p (pixpem. 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、Dylan cer1998等同學也對專案的發展做出了重要貢獻（這裡無法一一列出，後續所有的開發者都會在正式的技術報告/論文中說明；還有很多的同學和老師雖然沒有直接參與貢獻，但是提出了很多寶貴的意見，真的非常感謝大家）。也歡迎越來越多志同道合的朋友們加入我們，願OpenLLMAI與大家一起成長！

有意參與貢獻的同學可直接在git上參與開發、聯絡相關負責人或官方信箱。

RL：hijkzzz
Ray：wuxibin
NLP：Xianyu
官方信箱：[email protected]

贊助我們

目前OpenLLMAI是一個純開源組織，無論是OpenRLHF/OpenLLMWiki等項目，還是OpenLLM Talk和技術交流群，都是完全開源開放的。但長遠來看，沒有資金支持注定難以為繼，用愛發電走到今天並不容易，感謝大家一路的支持。最後，求贊助呀，歡迎大家有錢的捧個錢場（算力！！！），有人的捧個人場（參與開發或其他貢獻）！贊助或合作請聯絡[email protected]。