隨著大型語言模型（LLM）規模不斷增大，其效能也在不斷提升。儘管如此，LLM 依然面臨著一個關鍵難題：與人類的價值和意圖對齊。在解決這一難題方面，一種強大的技術是根據人類反饋的強化學習（RLHF）。
但是，隨著模型越來越大，RLHF 通常需要維持多個模型以及越來越複雜的學習流程，這又會導致記憶體和計算資源需求增長。舉個例子，近端策略最佳化（PPO，這是 RLHF 常用的一種演算法）需要在訓練過程中維持四個模型。
由此，當語言模型的引數規模超過 700 億時，為了訓練和協調多個模型，所需的計算資源和排程複雜性會顯著增長 —— 這是當前的架構設計難以滿足的需求。
Transformer 強化學習（TRL）、ColossalChat（CAIChat）和 DeepSpeed-Chat（DSChat）等現有的開源 RLHF 框架是依靠零冗餘最佳化器（Zero Redundancy Optimizer/ZeRO），來將 RLHF 訓練涉及的四個模型配置到同一臺 GPU 上。這個過程被稱為 co-location，即空間並置。
但是，隨著模型引數規模超過 700 億，在記憶體有限的 GPU 上，這種排程方法的效率會越來越低。
為了解決空間並置的限制，TRL 等一些框架選擇在記憶體使用上做出妥協，其做法包括將 actor 和 critic 模型融合起來或採用低秩適應（LoRA）等技術。但是，這些技術會降低模型效能，而且融合 actor-critic 式架構與備受推崇的實踐做法不相容，即使用獎勵模型的權重來初始化 critic 模型的權重。
另一種替代方法是使用來自英偉達 Megatron 的張量並行化和管道並行化技術。但是，Megatron 與人們常用的 Hugging Face 軟體庫不相容，而適應新模型又需要大量修改原始碼，如此就很難使用了。
為了輕鬆實現大規模 RLHF 訓練，OpenLLMAI、位元組跳動、網易伏羲 AI Lab、阿里巴巴的一個聯合團隊提出並開源了 OpenRLHF，其中第一作者為 Jian Hu。該框架使用 Ray、vLLM 和 DeepSpeed 對模型排程進行了重新設計，可支援超 700 億引數的模型的 RLHF 訓練，其優勢包括簡單易用、高效能、實現了分散式 RLHF、整合了 PPO 實現技巧。
論文標題：OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework
論文地址：https://arxiv.org/pdf/2405.11143
開源專案：https://github.com/OpenLLMAI/OpenRLHF
有關 Ray、vLLM 和 DeepSpeed 的具體詳情，請訪問原論文：
Ray: A Distributed Framework for Emerging AI Applications，arXiv:1712.05889
Efficient Memory Management for Large Language Model Serving with PagedAttention，arXiv:2309.06180，也可參看機器之心的報導：《6.7k Star 量的 vLLM 出論文了，讓每個人都能輕鬆快速低成本地部署 LLM 服務》
DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters，https://github.com/microsoft/DeepSpeed
OpenRLHF 可與 Hugging Face Transformer 無縫整合，並且支援混合專家（MoE）、Jamba 和 QLoRA 等常用技術。此外，OpenRLHF 還實現了多個對齊演算法，包括直接偏好最佳化（DPO）和 Kahneman-Tversky 最佳化（KTO）、條件 SFT 和拒絕取樣。
因此，可以說 OpenRLHF 是一個非常全面的 RLHF 訓練框架。
表 1 比較了常用的 RLHF 框架。
OpenRLHF 的設計
排程最佳化
要為更大的模型執行 RLHF 訓練，需要高效地在多臺 GPU 上分配至少四個元件模型（actor、critic、獎勵、參考）。為什麼需要多臺 GPU？因為每臺 GPU 加速器的記憶體有限，比如 NVIDIA A100 的記憶體不到 80GB。OpenRLHF 在模型排程方面創新性地使用了 Ray 來進行模型安放和細粒度的編排。
同時，OpenRLHF 還使用了針對推理最佳化的軟體庫 vLLM 和針對訓練最佳化的軟體庫 DeepSpeed；它們都由基於 Ray 的排程器管理。
OpenRLHF 能將四個模型分配到多臺 GPU 上，而不是將它們並置於同一臺 GPU，如圖 1 所示。
這樣的設計很自然就支援在 RLHF 訓練過程中使用多個獎勵模型，如圖 2 所示，並適用於多種演算法實現。
基於此，演算法工程師無需關心底層資料流的細節，就能快速構建多種對齊策略，比如有用性和有害性分離。這樣的排程器設計還可使用 Ray 和 DeepSpeed 來實現靈活的模型融合或解除安裝策略。比如可以融合 actor - 參考或 critic - 獎勵模型以節省 GPU 資源。
除了能高度定製演算法實現這一優點，該排程器還能以最優方式編排 GPU，從而提升整體訓練效能。
效能最佳化
RLHF 演算法的效能取決於訓練和推理兩方面的效率。從分析結果看，主要瓶頸是在 PPO 樣本生成階段（如圖 2 所示），這個階段佔到了整體訓練時間的 80%。原因是：在生成階段，自迴歸解碼的複雜度為 O (n^2)，並且也受到記憶體限制。
為了進一步加快樣本生成的速度以及支援無法載入到單臺 GPU 的更大型 LLM（比如 700 億引數的模型），OpenRLHF 使用了 vLLM 的張量並行化等先進技術（連續批處理和分頁注意力）來執行生成過程，如圖 1 所示。
在 RLHF 的生成和學習階段，OpenRLHF 採用了以下技術來獲得進一步的提升：
將 Adam 最佳化器狀態解除安裝到 CPU，這能將 GPU 記憶體解放出來用於較大的推理批次大小，這能提升效率以及避免生成的記憶體瓶頸。置頂記憶體和梯度積累，用於降低梯度聚合過程中的 GPU-CPU 通訊負載。
使用 Flash Attention 2 來加速 Transformer 模型訓練。
使用 PyTorch 張量切片移除訓練樣本中的冗餘填充。
圖 2 中另外三個模型使用了 ZeRO 的第 3 階段（對模型、梯度和最佳化器進行分片）。OpenRLHF 使用了英偉達 NCCL 和 vLLM 權重載入器來同步 ZeRO 和 vLLM 引擎的權重，確保實現快速又簡單的整合。
表 2 比較了 OpenRLHF 與該團隊精心微調過的 DSChat 的效能。
訓練穩定性
在訓練大型語言模型（LLM）時，PPO 等強化學習演算法容易不穩定。為了保證穩定，該團隊盡力驗證了 OpenRLHF 的實現細節。圖 2 和圖 3 分別給出了一般的推理和學習流程。
此外，OpenRLHF 還藉助了一些技巧來保證 PPO 實現的訓練穩定，包括：
僅在序列的文字末端 token 上預測獎勵
為語言模型使用 token 層級的強化學習
在 PPO 中使用 KL 散度損失項
在 PPO 中使用已預訓練的損失項，其根據策略損失的相對規模進行調整
為訓練穩定度使用獎勵歸一化
透過全域性統計使用分散式優勢歸一化
使用線性預熱餘弦退火學習率排程器
易用性
為便於使用者使用，該團隊還為支援的演算法提供了一鍵可用的可訓練指令碼（詳見原論文），並且該指令碼與 Hugging Face 軟體庫完全相容。下面給出了 Llama2 70B 模型的 RLHF 訓練的最低配置：
更多技術細節請參閱原論文。