一、RLHF 是什么

1. LLM 訓練流程

借用 Andrej Karpathy 在 Microsoft build 2023 上的大語言模型訓練的圖。

模型訓練總體上分為四個階段：預訓練（Pre-training）、微調（Fine-tuning）、獎勵建模（Reward Modeling）和強化學習訓練（Reinforcement Learning Training）。

• 預訓練階段Pre-training）。從互聯網收集大量數據，如維基百科、雜志、報紙和小說等，以擬合下一個 token。預訓練的模型一般不可以直接用，如果和他對話，大概率會重復輸入內容。
• 微調階段（Fine-tuning）。即 SFT，收集多人對話和示范，以提升預訓練模型的能力。微調后的模型即可部署使用。

今天重點介紹的是強化學習階段，包括獎勵建模和強化學習訓練。

• 獎勵建模階段（Reward Modeling）。比如同一個 prompt，可以收集多個輸出，并由人類標注出好的、壞的，組成一個個 pair 對。收集到多個這樣的 pair 后，就可以對獎勵模型進行訓練。
• 強化學習訓練（Reinforcement Learning Training）。有了獎勵模型后，可以使用 PPO、DPO 等算法進行強化學習訓練。這些算法的輸入是 prompt。

2. 為什么需要 RLHF？

使用 RLHF 可以通過降低模型輸出不正確或有害內容的概率，來提升模型輸出符合人類偏好的概率。互聯網上有海量的數據，通過預訓練或者微調可以擬合下一個 token，也就是模型可以判斷輸出什么，但是沒有負反饋來表達不應輸出什么。例如，當詢問 GPT 如何制造炸彈時，如果沒有 RLHF，模型就會輸出制造炸彈的步驟，而有了 RLHF 之后，模型就可以拒絕回答這類有害內容。

3. 如何將 LLM 訓練表述為強化學習問題？

強化學習問題有五個要素：策略、狀態、動作、環境和獎勵，可以將大語言模型映射到這些概念上，即：

• 策略：LLM 模型，輸入是狀態，輸出是所有可選的動作(token)的概率。
• 狀態：用戶輸入的 prompt 和 LLM 已經解碼出來的那部分內容的組合。
• 動作：模型輸出的 token(動作空間即詞表大小)。
• 環境：reward model，輸入是狀態，輸出是獎勵分數 r。
• 獎勵：reward model 給出的評價分數，也可以看做是人類對當前模型輸出的評價，是一個客觀的值。

因此，LLM 接受一個 prompt，輸出一個長度為 256 response 的過程，可視為智能體在環境中行動了 256 步，最終由 reward model 給出相應的獎勵。基于以上表述，我們就可以利用 PPO 算法對策略(LLM)進行優化。

4.  LLM PPO 算法流程

如上圖，PPO 算法的工作流程分為推理階段和訓練階段，推理階段負責生成經驗，訓練階段負責經驗回放，這兩個階段交替進行。

左邊推理階段，給定一個 prompt，首先經過 rollout 模型生成 response。然后，將 prompt 和 response 拼成一個 sequence，這個 sequence 同時過 actor、initial、critic、 reward 4 個模型進行 inference，生成對應 logits、value、score。其中，actor、initial、critic 這三個模型是直接從微調模型SFT初始化的，reword 模型是從 SFT 訓練出來的。這幾個模型的作用是：

• actor 模型是我們最終要訓練的目標模型；
• initial 模型是一個參考模型，它是 freeze 的，不會用作訓練；
• critic 和 reward 模型，是用來評判 actor 模型的；
• rollout 模型，它和 actor 模型是同參數的，也就是他們參數量完全一樣，可以共享一份參數，也可以進行分離部署。rollout 模型是負責生成的，可以使用現在很多的推理框架來單獨部署。

右邊訓練階段，可以看到推理階段生成的這些橙色框的內容，就是我們所說的經驗。根據這經驗五元組，來對 critic 和 actor 模型進行更新。

RLHF 算法整體比較復雜：多階段，有推理有訓練；多模型，有 5 個模型，甚至可能更多；框架中需要重點關注 actor 模型，它與 rollout 模型同參數，在訓練完成后，要做參數同步，更新給推理階段的 rollout 模型。

二、RLHF 的挑戰與動機

RLHF 面臨的挑戰包括模型規模的增大、模型參數同步的耗時增加以及研究對靈活性的要求。大型模型的特點和新挑戰使得我們難以直接將現有的強化學習框架應用于語言模型的 RLHF。

1. 現有的 RL 框架是否能用于 LLM RLHF?

Ray 里有一個傳統的 RL 框架 Rlib，它可以直接用于大語言模型的 RLHF 嗎？答案是可以，但是也不一定可以。

說可以是因為將大語言模型的強化學習訓練，對齊到傳統的強化學習的五要素中，都能映射上。如上圖右邊所示，gym.vector，有人或車走動，映射到 RL 就是一個自回歸生成過程（Autoregressive Generation）；獎勵 reward，可以通過一個模型以及當前模型和初始模型的 KL 散度來計算。

說不可以是因為大語言模型的一些特點或者一些新的挑戰，使得難以將 Rlib 應用于 RLHF。具體的挑戰總結起來有三點。

2. 挑戰 1：模型規模的增大

模型規模變大是最直觀的一個挑戰。傳統的 RL 算法中，模型大小從 100K~100M。但是，LLM 模型大小遠超傳統模型：7B~405B。在這么大的模型規模下，對每一個計算點，都需要單獨針對性地做優化。比如：

• rollout：用來做生成 generate，需要一些專門針對推理做優化的框架，如 vLLM、TensorRT-LLM、SGLang；
• actor/cirtic：用來進行訓練的，需要針對訓練的框架，如 Megatron-LM、Deepspeed、FSDP 等；

這樣在這個算法中，需要用到若干種框架的組合：比如，推理時用 vLLM；訓練時，大一點的模型用 TP、PP、DP 并行技術，也可以用 Deepspeed、FSDP 等更 native 的方式。

3. 挑戰 2：模型參數同步的耗時增加

隨著模型規模增大，模型參數的同步變得更加耗時。比如，英偉達 A100 顯卡的跨機器的 RDMA 帶寬可達 50GB/s，它同步一個 Llama 405B 耗時仍需 60 秒左右。這在 PPO 算法中，占比較高。在傳統模型中，這個同步耗時可以忽略不計，但在大語言模型中，需要針對性優化。

4. 挑戰 3：對靈活性的要求

由于 LLM RLHF 領域還處于高速發展中，過去一年涌現了非常多的新算法，包括 PPO 算法的一些變種，比如 ReMax、Safe-RLHF、GRPO；還有針對 PPO 的改進，比如 DPO、KTO、ReST 等。其中，DPO 將 RLHF 簡化成兩個模型；KTO 更進一步，不需要有 pair 對；Google 提出的 Rest，進行模型自學習，而不需要外部獎勵模型。

在這樣的背景下，對于 RL 算法研究者來說，一個僅僅可以調整模型參數，喂數據進行訓練的算法框架并不能滿足需求。研究者需要一個編程框架，以支持實現各種不同的算法，也就是需要可以復用的算法模塊，并通過算法模塊靈活組合來實現不同的 RL 算法。

三、HybridFlow 設計與實現

HybridFlow 針對前文所述挑戰進行了設計和實現。

1. HybridFlow 核心設計思想

在分布式系統中采用 Single-Controller 還是 Multi-Controller 是個重要命題。Google Pathways 在論文中對這兩點做了非常詳細的闡述。

Single-Controller 由一個中心控制器驅動所有 worker 工作，這個中心控制器向所有 worker 發送數據和計算指令，worker 可以運行不同的程序(MPMD)。它的優點是靈活、異構，即每個 worker 可以執行不同的算法；缺點是因為是單點控制，所以會引入額外的通信開銷。

Multi-Controller 無中心控制器，各個 worker 自驅，各自計算，通過一些通信原語進行同步，worker 運行相同的程序(SPMD)。它的優點是高效；缺點是因為同構所以不夠靈活。

典型的 Single-Controller，比如 Spark、TensorFlow1.0 的 non-SPMD 的模式。

Multi-Controller 則更普遍，基本上主流的開源 RLHF 分布式訓練框架，比如 Megatron-LM、Deepspeed、FSDP，以及 vLLM 都是 Multi-Controller 模式的。這是因為主流 RLHF 框架基本上是從預訓練框架演進而來，而預訓練基本上都是 Multi-Controller 模式。

那么，Multi-Controller 模式在處理 RLHF 的多模型的異構場景時，有哪些缺點呢？

首先，Multi-Controller 是 SPMD 模式的，所有 worker 運行相同的程序，這時推理和訓練的框架以及工作流都不一樣，要實現這樣的 SPMD 算法難度就很高。

其次，靈活性不足。比如算法研究者想改一個算法，需要改每個模塊中 3D 并行的策略，設計同步點。

通過我們的觀察分析，可以將 RL 的數據流拆成兩部分：

第一部分是 RL 的數據流，就是在不同的組件間交換數據。如前文提到的，先在 rollout 模型做一些推理，拿到結果后，再把數據發送給其它 4 個模型做推理，然后再把經驗給到訓練側的 actor 和 critic。

第二部分是組件內部的數據流。每個組件都是分布式的，組件內部是 SPMD 的工作模式，通過 NCCL 進行數據交換。

因為這樣兩層的數據流架構，所以我們將其命名為 HybridFlow，即通過 Single-Controller 實現 RL 的數據流，通過 Multi-Controller 實現各個組件的計算和通信流。

HybridFlow 對于各個組件，可復用已有框架，不用重復造輪子。比如，訓練使用 Megatron-LM/Deepspeed/FSDP，推理使用 vLLM/TensorRT-LLM/SGLang。同時，通過組件的抽象和封裝，使得 Single-Controller 可以像單進程一樣去調用組件，而不需要關注組件內部的并行策略。

2. HybridFlow 實現

HybridFlow 的核心是 worker group 的概念，它的本質是將一組 SPMD 進程進行抽象，使其可以像單進程一樣被調用。

舉例來說，如上圖，用 Tensor 并行加數據并行來實現一個簡單的 Transformer 的 FFN layer。假設 TP、DP 各 2 個：首先，在 DP 維度進行數據拆分，拆成 batch 0 和 batch 1。然后，在 DP group 內部，把相同數據發給所有的 rank。所有 worker 收到數據后，就進行推理。有結果后，再把數據聚合回來，也就是從每一個 DP 的 rank 0 去把數據取回。

Worker group 負責了數據分發（data dispatch）和聚合（data aggregation）兩個重要功能。

如前文所介紹，第一層RL數據流是 Single-Controller 模式的，由 driver 進程去驅動所有的 worker group 工作。即數據的分發和聚合都在 driver 側，通過 ray object store 來完成，并提供了簡單的 dispatch function 和 collect function。

根據整體測試的結果，在千卡集群下，ray object store 基本可以滿足要求。

HybridFlow 封裝的一些組件接口如上圖所示，包括 actor/rollout 的混合訓練推理引擎，reward、reference 以及 critic 模型。

在 HybridFlow 框架基礎上，實現 RLHF 算法就比較簡單了。對每個模型，每個組件創建對應的 work group，共創建 4 個 group；然后，對于每個算法，調用對應 group 的接口，進行訓練。

3. HybridFlow 優點

有了 HybridFlow 的框架和流程后，算法工程師再對算法做微調就比較簡單了。上圖右邊展示了一個完整的 PPO 算法流程。當進行 PPO 算法變種時，只需要改其中一些步驟：比如 ReMax、GRPO 算法，不需要 critic 計算 advantages，把這行刪掉就可以；Safe-RLHF 算法需要多個 reward 模型，增加對應的 group 就可以。

WorkerGroup 可以進行靈活組合，使用 Ray PlacementGroup 調度，通過設置 actor 的 placement_group、GPU 數量等，靈活控制 WorkerGroup 是否共享同一個 placement_group， 從而實現了不同 RLHF 組件邏輯隔離、物理共享，提高了資源利用率。

實驗結果表明，相比業界其他 RLHF 框架（Deepspeed-Chat、OpenRLHF、Nemo-Aligner)，HybridFlow 在 7B 到 70B，模型吞吐量提高了 1.6 倍到 20 倍不等。

我們的論文已被 USENIX ATC 2024 接收，代碼已在 GitHub 上開源，歡迎大家關注。

• Paper: https://arxiv.org/abs/2409.19256, EuroSys’ 25
• Github: https://github.com/volcengine/verl  HybridFlow 在 GitHub 上開源項目名 veRL。