如何基于 OpenRLHF 定制內(nèi)部 RFT 訓練框架

在這里做一個簡單的code walkthroug，以release的0.5.3版本為基準。

原標題：如何基于 OpenRLHF 定制內(nèi)部 RFT 訓練框架
文章來源：智猩猩GenAI
內(nèi)容字數(shù)：3218字

第四屆全球自動駕駛峰會及OpenRLHF框架解讀

本文首先簡要介紹了將于1月14日在北京舉辦的第四屆全球自動駕駛峰會，并列舉了部分已確認參會的嘉賓及峰會日程安排。隨后，文章重點關(guān)注如何將強化學習融入AI訓練框架，特別是針對LLM+RL的訓練。

強化學習在AI訓練框架中的挑戰(zhàn)

1. **多模型協(xié)調(diào)與通信:** RL算法（如PPO、Reinforce）通常涉及多個模型間的協(xié)作，其通信方式與傳統(tǒng)的預訓練/微調(diào)(pretrain/SFT)中的MPI集合通信有所不同，增加了實現(xiàn)難度。

2. **參數(shù)調(diào)優(yōu)的復雜性:** RL算法參數(shù)眾多，對精度要求高，缺乏RL經(jīng)驗的工程師難以判斷訓練失敗是框架問題還是超參數(shù)問題。

OpenRLHF框架的優(yōu)勢及推薦

文章推薦使用OpenRLHF框架進行LLM+RL訓練，并闡述了其優(yōu)勢：

1. **基于Ray的多模型調(diào)度:** 支持任意擴展每個模型的GPU數(shù)量，提高可擴展性。

2. **模塊化設(shè)計:** 訓練和推理模塊分離，方便替換和定制。

3. **高可靠性和優(yōu)秀的默認超參數(shù):** 經(jīng)過大量驗證，默認超參數(shù)表現(xiàn)出色，減少了調(diào)參的工作量。

4. **輕量級和易讀性:** 代碼簡潔易懂，方便二次開發(fā)。

5. **功能全面:** 支持主流的LLM+RL算法（PPO、Reinforce）以及關(guān)鍵優(yōu)化功能，例如打包樣本(packing samples)。

文章建議將現(xiàn)有RL訓練框架中的LLM部分替換成自己的框架，并復用OpenRLHF的多模型調(diào)度機制，從而實現(xiàn)事半功倍的效果。

OpenRLHF框架代碼解讀

文章對OpenRLHF框架的關(guān)鍵文件進行了簡要分析：

1. openrlhf/cli/train_ppo_ray.py: 訓練入口，包含配置項和模型初始化。

2. openrlhf/trainer/ppo_trainer.py: 包含PPO訓練的基本流程，包括rollout生成和模型訓練。

3. openrlhf/trainer/ppo_utils/experience_maker.py: 生成rollout數(shù)據(jù)，包括調(diào)用VLLM和計算logprob、KL、reward等。

4. openrlhf/trainer/ray/ppo_actor.py: Ray版本的PPOTrainer，包含通信同步邏輯。

5. openrlhf/utils/deepspeed/deepspeed.py: 進行deepspeed初始化，需要替換deepspeed時需關(guān)注此文件。

實現(xiàn)OpenAI RFT可能遇到的問題及解決方法

文章還指出了在使用OpenRLHF實現(xiàn)OpenAI RFT時可能遇到的問題以及相應(yīng)的解決方法：

1. **Verifier接入:** 使用remote_rm_url接入Verifier，參考openrlhf/cli/serve_rm.py。

2. **數(shù)據(jù)多樣性:** 通過--n_samples_per_prompt參數(shù)設(shè)置每個問題采樣多個回答。

3. **Critic模型缺失:** 使用--freezing_actor_steps預訓練Critic，或使用無需Critic的算法（Reinforce、RLOO）。

4. **Reward后處理:** 在process_experiences方法中注冊自定義的處理方法。

總而言之，文章推薦OpenRLHF框架作為LLM+RL訓練的理想選擇，并提供了詳細的代碼解讀和問題解決方法，方便工程師進行二次開發(fā)和應(yīng)用。

聯(lián)系作者

文章來源：智猩猩GenAI
作者微信：
作者簡介：智猩猩旗下公眾號之一，深入關(guān)注大模型與AI智能體，及時搜羅生成式AI技術(shù)產(chǎn)品。