隨著大型語言模型（LLM）規(guī)模不斷增大，其性能也在不斷提升。盡管如此，LLM 依然面臨著一個(gè)關(guān)鍵難題：與人類的價(jià)值和意圖對(duì)齊。在解決這一難題方面，一種強(qiáng)大的技術(shù)是根據(jù)人類反饋的強(qiáng)化學(xué)習(xí)（RLHF）。

但是，隨著模型越來越大，RLHF 通常需要維持多個(gè)模型以及越來越復(fù)雜的學(xué)習(xí)流程，這又會(huì)導(dǎo)致內(nèi)存和計(jì)算資源需求增長。舉個(gè)例子，近端策略優(yōu)化（PPO，這是 RLHF 常用的一種算法）需要在訓(xùn)練過程中維持四個(gè)模型。

由此，當(dāng)語言模型的參數(shù)規(guī)模超過 700 億時(shí)，為了訓(xùn)練和協(xié)調(diào)多個(gè)模型，所需的計(jì)算資源和調(diào)度復(fù)雜性會(huì)顯著增長 —— 這是當(dāng)前的架構(gòu)設(shè)計(jì)難以滿足的需求。

Transformer 強(qiáng)化學(xué)習(xí)（TRL）、ColossalChat（CAIChat）和 DeepSpeed-Chat（DSChat）等現(xiàn)有的開源 RLHF 框架是依靠零冗余優(yōu)化器（Zero Redundancy Optimizer/ZeRO），來將 RLHF 訓(xùn)練涉及的四個(gè)模型配置到同一臺(tái) GPU 上。這個(gè)過程被稱為 co-location，即空間并置。

但是，隨著模型參數(shù)規(guī)模超過 700 億，在內(nèi)存有限的 GPU 上，這種調(diào)度方法的效率會(huì)越來越低。

為了解決空間并置的限制，TRL 等一些框架選擇在內(nèi)存使用上做出妥協(xié)，其做法包括將 actor 和 critic 模型融合起來或采用低秩適應(yīng)（LoRA）等技術(shù)。但是，這些技術(shù)會(huì)降低模型性能，而且融合 actor-critic 式架構(gòu)與備受推崇的實(shí)踐做法不兼容，即使用獎(jiǎng)勵(lì)模型的權(quán)重來初始化 critic 模型的權(quán)重。

另一種替代方法是使用來自英偉達(dá) Megatron 的張量并行化和管道并行化技術(shù)。但是，Megatron 與人們常用的 Hugging Face 軟件庫不兼容，而適應(yīng)新模型又需要大量修改源代碼，如此就很難使用了。

為了輕松實(shí)現(xiàn)大規(guī)模 RLHF 訓(xùn)練，OpenLLMAI、字節(jié)跳動(dòng)、網(wǎng)易伏羲 AI Lab、阿里巴巴的一個(gè)聯(lián)合團(tuán)隊(duì)提出并開源了 OpenRLHF，其中第一作者為 Jian Hu。該框架使用 Ray、vLLM 和 DeepSpeed 對(duì)模型調(diào)度進(jìn)行了重新設(shè)計(jì)，可支持超 700 億參數(shù)的模型的 RLHF 訓(xùn)練，其優(yōu)勢(shì)包括簡(jiǎn)單易用、高性能、實(shí)現(xiàn)了分布式 RLHF、集成了 PPO 實(shí)現(xiàn)技巧。

• 論文標(biāo)題：OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework
• 論文地址：https://arxiv.org/pdf/2405.11143
• 開源項(xiàng)目：https://github.com/OpenLLMAI/OpenRLHF

有關(guān) Ray、vLLM 和 DeepSpeed 的具體詳情，請(qǐng)?jiān)L問原論文：

• Ray: A Distributed Framework for Emerging AI Applications，arXiv:1712.05889
• Efficient Memory Management for Large Language Model Serving with PagedAttention，arXiv:2309.06180，也可參看機(jī)器之心的報(bào)道：《6.7k Star 量的 vLLM 出論文了，讓每個(gè)人都能輕松快速低成本地部署 LLM 服務(wù)》
• 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 等常用技術(shù)。此外，OpenRLHF 還實(shí)現(xiàn)了多個(gè)對(duì)齊算法，包括直接偏好優(yōu)化（DPO）和 Kahneman-Tversky 優(yōu)化（KTO）、條件 SFT 和拒絕采樣。

因此，可以說 OpenRLHF 是一個(gè)非常全面的 RLHF 訓(xùn)練框架。

表 1 比較了常用的 RLHF 框架。

OpenRLHF 的設(shè)計(jì)

調(diào)度優(yōu)化

要為更大的模型執(zhí)行 RLHF 訓(xùn)練，需要高效地在多臺(tái) GPU 上分配至少四個(gè)組件模型（actor、critic、獎(jiǎng)勵(lì)、參考）。為什么需要多臺(tái) GPU？因?yàn)槊颗_(tái) GPU 加速器的內(nèi)存有限，比如 NVIDIA A100 的內(nèi)存不到 80GB。OpenRLHF 在模型調(diào)度方面創(chuàng)新性地使用了 Ray 來進(jìn)行模型安放和細(xì)粒度的編排。

同時(shí)，OpenRLHF 還使用了針對(duì)推理優(yōu)化的軟件庫 vLLM 和針對(duì)訓(xùn)練優(yōu)化的軟件庫 DeepSpeed；它們都由基于 Ray 的調(diào)度器管理。

OpenRLHF 能將四個(gè)模型分配到多臺(tái) GPU 上，而不是將它們并置于同一臺(tái) GPU，如圖 1 所示。

這樣的設(shè)計(jì)很自然就支持在 RLHF 訓(xùn)練過程中使用多個(gè)獎(jiǎng)勵(lì)模型，如圖 2 所示，并適用于多種算法實(shí)現(xiàn)。

基于此，算法工程師無需關(guān)心底層數(shù)據(jù)流的細(xì)節(jié)，就能快速構(gòu)建多種對(duì)齊策略，比如有用性和有害性分離。這樣的調(diào)度器設(shè)計(jì)還可使用 Ray 和 DeepSpeed 來實(shí)現(xiàn)靈活的模型融合或卸載策略。比如可以融合 actor - 參考或 critic - 獎(jiǎng)勵(lì)模型以節(jié)省 GPU 資源。

除了能高度定制算法實(shí)現(xiàn)這一優(yōu)點(diǎn)，該調(diào)度器還能以最優(yōu)方式編排 GPU，從而提升整體訓(xùn)練性能。

性能優(yōu)化

RLHF 算法的性能取決于訓(xùn)練和推理兩方面的效率。從分析結(jié)果看，主要瓶頸是在 PPO 樣本生成階段（如圖 2 所示），這個(gè)階段占到了整體訓(xùn)練時(shí)間的 80%。原因是：在生成階段，自回歸解碼的復(fù)雜度為 O (n^2)，并且也受到內(nèi)存限制。

為了進(jìn)一步加快樣本生成的速度以及支持無法載入到單臺(tái) GPU 的更大型 LLM（比如 700 億參數(shù)的模型），OpenRLHF 使用了 vLLM 的張量并行化等先進(jìn)技術(shù)（連續(xù)批處理和分頁注意力）來執(zhí)行生成過程，如圖 1 所示。

在 RLHF 的生成和學(xué)習(xí)階段，OpenRLHF 采用了以下技術(shù)來獲得進(jìn)一步的提升：

• 將 Adam 優(yōu)化器狀態(tài)卸載到 CPU，這能將 GPU 內(nèi)存解放出來用于較大的推理批量大小，這能提升效率以及避免生成的內(nèi)存瓶頸。置頂內(nèi)存和梯度積累，用于降低梯度聚合過程中的 GPU-CPU 通信負(fù)載。
• 使用 Flash Attention 2 來加速 Transformer 模型訓(xùn)練。
• 使用 PyTorch 張量切片移除訓(xùn)練樣本中的冗余填充。

圖 2 中另外三個(gè)模型使用了 ZeRO 的第 3 階段（對(duì)模型、梯度和優(yōu)化器進(jìn)行分片）。OpenRLHF 使用了英偉達(dá) NCCL 和 vLLM 權(quán)重加載器來同步 ZeRO 和 vLLM 引擎的權(quán)重，確保實(shí)現(xiàn)快速又簡(jiǎn)單的集成。

表 2 比較了 OpenRLHF 與該團(tuán)隊(duì)精心微調(diào)過的 DSChat 的性能。

訓(xùn)練穩(wěn)定性

在訓(xùn)練大型語言模型（LLM）時(shí)，PPO 等強(qiáng)化學(xué)習(xí)算法容易不穩(wěn)定。為了保證穩(wěn)定，該團(tuán)隊(duì)盡力驗(yàn)證了 OpenRLHF 的實(shí)現(xiàn)細(xì)節(jié)。圖 2 和圖 3 分別給出了一般的推理和學(xué)習(xí)流程。

此外，OpenRLHF 還借助了一些技巧來保證 PPO 實(shí)現(xiàn)的訓(xùn)練穩(wěn)定，包括：

• 僅在序列的文本末端 token 上預(yù)測(cè)獎(jiǎng)勵(lì)
• 為語言模型使用 token 層級(jí)的強(qiáng)化學(xué)習(xí)
• 在 PPO 中使用 KL 散度損失項(xiàng)
• 在 PPO 中使用已預(yù)訓(xùn)練的損失項(xiàng)，其根據(jù)策略損失的相對(duì)規(guī)模進(jìn)行調(diào)整
• 為訓(xùn)練穩(wěn)定度使用獎(jiǎng)勵(lì)歸一化
• 通過全局統(tǒng)計(jì)使用分布式優(yōu)勢(shì)歸一化
• 使用線性預(yù)熱余弦退火學(xué)習(xí)率調(diào)度器

易用性

為便于用戶使用，該團(tuán)隊(duì)還為支持的算法提供了一鍵可用的可訓(xùn)練腳本（詳見原論文），并且該腳本與 Hugging Face 軟件庫完全兼容。下面給出了 Llama2 70B 模型的 RLHF 訓(xùn)練的最低配置：

更多技術(shù)細(xì)節(jié)請(qǐng)參閱原論文。