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

機器之心報道

編輯：Panda

隨着大型語言模型（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 實現技巧。

有關 Ray、vLLM 和 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 採用了以下技術來獲得進一步的提升：

圖 2 中另外三個模型使用了 ZeRO 的第 3 階段（對模型、梯度和優化器進行分片）。OpenRLHF 使用了英偉達 NCCL 和 vLLM 權重加載器來同步 ZeRO 和 vLLM 引擎的權重，確保實現快速又簡單的集成。

表 2 比較了 OpenRLHF 與該團隊精心微調過的 DSChat 的性能。

訓練穩定性

在訓練大型語言模型（LLM）時，PPO 等強化學習算法容易不穩定。爲了保證穩定，該團隊盡力驗證了 OpenRLHF 的實現細節。圖 2 和圖 3 分別給出了一般的推理和學習流程。

此外，OpenRLHF 還藉助了一些技巧來保證 PPO 實現的訓練穩定，包括：

易用性

爲便於用戶使用，該團隊還爲支持的算法提供了一鍵可用的可訓練腳本（詳見原論文），並且該腳本與 Hugging Face 軟件庫完全兼容。下面給出了 Llama2 70B 模型的 RLHF 訓練的最低配置：

更多技術細節請參閱原論文。