如何基于 OpenRLHF 定制內部 RFT 訓練框架

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

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

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

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

強化學習在AI訓練框架中的挑戰

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

2. **參數調優的復雜性:** RL算法參數眾多，對精度要求高，缺乏RL經驗的工程師難以判斷訓練失敗是框架問題還是超參數問題。

OpenRLHF框架的優勢及推薦

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

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

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

3. **高可靠性和優秀的默認超參數:** 經過大量驗證，默認超參數表現出色，減少了調參的工作量。

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

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

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

OpenRLHF框架代碼解讀

文章對OpenRLHF框架的關鍵文件進行了簡要分析：

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

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

3. openrlhf/trainer/ppo_utils/experience_maker.py: 生成rollout數據，包括調用VLLM和計算logprob、KL、reward等。

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

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

實現OpenAI RFT可能遇到的問題及解決方法

文章還指出了在使用OpenRLHF實現OpenAI RFT時可能遇到的問題以及相應的解決方法：

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

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

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

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

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

聯系作者

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