這個團隊做了OpenAI沒Open的技術，開源OpenRLHF讓對齊大模型超簡單

作者：機器之心 2024-06-06 19:14:27

隨著模型越來越大，RLHF 通常需要維持多個模型以及越來越復雜的學習流程，這又會導致內存和計算資源需求增長。

隨著大型語言模型（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 的性能。