作者 | Feng Nie
原文連結：https://zhuanlan.zhihu.com/p/1932578428181279850
本文只做學術分享，如有侵權，請聯繫刪文。
原文連結：
Robin's Home Page：jianzhnie.github.io/llmtech/#/rlhf/infra/RL-Infra_overview
1. 從SFT到強化學習：模型訓練典範的轉變
在2024年OpenAI發布O1系列模型之前，主流的機器學習訓練方式主要依賴於有監督微調（Supervised Fine-Tuning, SFT）。該方法透過讓模型學習“標準答案”，並根據預測與真實標籤之間的損失（loss）來更新模型參數。訓練流程相對簡單，PyTorch 和 TensorFlow 等深度學習框架也圍繞這一典範構建了豐富的訓練加速工具。
然而，隨著O1系列模型的發布，模型訓練的重心逐漸從SFT向強化學習（Reinforcement Learning, RL）轉移。SFT逐漸被視為訓練過程中的“預熱”階段，其作用被弱化為參數初始化或策略引導。取而代之的是，RL在模型能力提升中扮演了越來越關鍵的角色。
1.1. RL演算法的演進與多樣化
RL演算法本身也在不斷迭代與優化。從早期的DPO（Direct Preference Optimization），到經典的PPO（Proximal Policy Optimization），再到近年來湧現出的GRPO、RLOO、Reinforce++、DAPO等新方法，RL演算法在策略更新方式、穩定性、樣本效率等方面持續優化。
儘管DPO因其簡潔性曾一度流行，但隨著任務複雜度和模型規模的提升，其局限性逐漸顯現，目前在實際工程中已較少被採用。儘管如此，主流RL框架的整體結構保持相對一致，核心流程主要包括以下幾個階段：
1.2. RL訓練流程的三大模組
模組一：策略生成（Rollout）
對應“學生自己尋找答案”的過程。這是RL訓練中的推演階段（Rollout），模型基於目前策略生成響應（action），模擬與環境的互動過程。該階段是模型推論過程的擴展，通常需要大量取樣以獲取多樣化的行為軌跡。
模組二：獎勵評估（Reward Evaluation）
對應“給學生答案打分”的過程。傳統上，這一階段依賴於獎勵模型（Reward Model），用於評估生成結果的品質。在目前階段，由於任務複雜度提升，獎勵評估的實現方式也趨於多樣化：
基於規則的評估（Rule-based）：在數學、物理、程式碼等領域，透過結果與規則的匹配度進行打分。
輕量級獎勵模型：訓練一個小型模型（如7B參數）進行打分，成本可控，且效果良好。
在許多研究專案中，這一模組甚至被簡化為Rollout的一部分，未被單獨重視。然而，隨著Agent行為模擬的興起，尤其是在商業應用場景（如電商、客服等）中，獎勵評估的複雜性顯著上升，未來該模組的重要性將不斷提升。
模組三：策略更新（Policy Update）
對應“學生根據打分來學習”的過程。這是RL訓練的核心階段，基於傳統訓練框架（如PyTorch、DeepSpeed等），透過修改損失函數實現策略更新。不同演算法（如PPO、DPO、RLOO等）在此階段的實現邏輯有所不同，但整體結構保持一致。
1.3 總結
從SFT主導的訓練典範到RL驅動的能力提升，大模型的訓練流程正經歷深刻的變革。RL框架的結構雖然保持穩定，但其各模組的功能、實現方式和重要性正在不斷演化。
Rollout模組：面臨長上下文、異構任務帶來的效能挑戰；
獎勵評估模組：從簡單規則向複雜評估演進，未來可能成為RL訓練中的關鍵瓶頸；
策略更新模組：依賴於底層訓練框架的效能優化與演算法迭代。
隨著Agent行為模擬、複雜任務建模、多模態互動等方向的發展，RL框架的設計將更加注重模組間的協同、資源調度的高效性以及演算法與工程實現的統一性。
2. RL訓練框架設計與效能優化挑戰
目前，主流的強化學習（Reinforcement Learning, RL）訓練框架通常被劃分為兩個核心模組：訓練（Training） 和 Rollout（推演）。
在設計一個高效的RL訓練系統時，開發者將面臨一系列關鍵挑戰。以下是我們在技術選型與框架設計過程中總結出的三大核心問題。
2.1 挑戰一：Rollout與訓練模組的協同與資源管理
目前，RL訓練普遍採用On-policy策略，這意味著Rollout與訓練過程必須順序執行。然而，隨著模型規模的持續增長，分散式多卡訓練已成為必然趨勢。
Rollout階段：主要為記憶體密集型任務，尤其在處理長上下文（如Chain-of-Thought）時，需要維護大量的KV Cache（Key-Value Cache）。
訓練階段：則屬於計算密集型任務，涉及大規模的參數更新和梯度計算。
這兩個階段各自已有大量優化手段（如記憶體重複使用、管線平行等），但如何在統一框架中高效管理這兩類異構資源？如何優化兩者之間的參數同步機制？這是構建高效RL系統的關鍵挑戰之一。
2.2 挑戰二：底層訓練與推論框架的多樣性
目前存在多種主流的訓練框架，例如：
Megatron-LM
DeepSpeed（FSDP）
PyTorch FSDP
同時，推論引擎也呈現多樣化趨勢：
vLLM
SGLang
不同訓練框架與推論引擎的架構差異顯著，導致在參數同步、推論排程等環節的實現邏輯差異較大。例如，僅在參數更新部分，不同組合就可能需要完全不同的實現邏輯，這對系統的可維護性與擴展性提出了較高要求。
2.3 挑戰三：異構批次執行帶來的不確定性
Rollout任務通常以批次形式執行，但批次內部任務的複雜度可能存在巨大差異。特別是在引入Agent行為模擬的場景下，這種異構性更加顯著，可能導致整體排程效率下降、資源利用率不均衡等問題。
3. 效能優化分析
3.1 初始實現與效能瓶頸
在RL訓練的早期實現中，整個流程通常分為三個階段：
1. 推論階段（Rollout）：模型根據目前策略生成響應。
2. 評估階段：透過獎勵模型或其他機制對生成結果進行打分。
3. 訓練階段：基於打分結果更新策略模型。
該流程本質上可以基於SFT（Supervised Fine-Tuning）框架實現，區別在於需要初始化多個模型實例（如策略模型、獎勵模型等）。然而，這種實現方式在實際運行中往往存在顯著的效能瓶頸。
3.2 記憶體優化策略
在大規模模型訓練中，顯存佔用主要包括以下幾個部分：
模型參數（Parameters）
梯度（Gradients）
最佳化器狀態（Optimizer States）
激活值（Activations）
以一個7B參數模型為例，在FP32精度下，僅模型參數和梯度就需要約28GB顯存，最佳化器狀態則可能額外佔用28GB×3=84GB，總計高達112GB。顯然，單卡無法承載如此龐大的記憶體需求。
為此，業界提出了多種分散式訓練策略：
資料平行（Data Parallelism, DP）：如 DeepSpeed ZeRO-1/2/3，透過All-Gather操作動態重建完整參數。
張量平行（Tensor Parallelism, TP）與管線平行（Pipeline Parallelism, PP）：如 Megatron-LM，採用參數切分策略，適用於大規模模型。
根據NVIDIA相關論文的研究結論，在千卡以下規模，DP與TP/PP效能相近；但在更大規模下，TP/PP因避免了All-Gather操作的通訊開銷，效能優勢更為明顯。
這個表格比較了資料平行（DP）、張量平行（TP）和管線平行（PP）三種平行策略在不同特性上的表現。
3.3 推論速度優化與引擎選型
目前主流推論引擎（如 vLLM 和 SGLang）在KV Cache重複使用、底層運算子優化等方面已實現顯著效能提升。儘管如此，訓練與推論引擎之間的參數同步仍存在一定挑戰：
推論引擎生成的輸出與訓練引擎在精度上存在差異；
目前主流做法是：在Rollout階段使用推論引擎加速生成，訓練階段再由訓練引擎重新計算logits（僅需prefill階段，計算效率高）。
因此，將高效能推論引擎與訓練框架進行整合，是提升整體RL訓練效率的有效路徑。但如何高效地實現訓練與推論模組的拼接與協同，仍是值得深入研究的問題。
4. 訓練框架與推論引擎的整合
4.1 SPMD和MPMD概念解析
在討論訓練框架與推論引擎如何結合之前，有必要先理解SPMD（Single Program, Multiple Data）和MPMD（Multiple Programs, Multiple Data）的概念。簡而言之，SPMD指的是多個處理單元執行相同的程式但操作不同的資料集，而MPMD則涉及多個處理單元運行不同的程式並處理不同的資料集。前者通常不需要一個中央控制器來協調工作流程，而後者則可能需要以避免混亂。
在討論訓練框架和推論引擎的整合時，首先需要理解兩種平行處理模式：SPMD（Single Program, Multiple Data） 和 MPMD（Multiple Programs, Multiple Data）。這兩種模式也可以被描述為單一控制器與多控制器架構。
單一控制器（SPMD）：所有工作節點執行相同的程式邏輯，適用於資料量大但模型規模較小的場景。
多控制器（MPMD）：每個工作節點可以執行不同的程式，增加了實現複雜度，但無需集中控制，適合特定應用場景。
主流的深度學習訓練框架如DeepSpeed和Megatron都採用了SPMD模式，保證所有進程遵循相同的程式碼邏輯進行運算。然而，對於推論引擎（例如SGLang和vLLM），情況則有所不同。儘管推論引擎（例如SGLang和vLLM）在計算過程中遵循SPMD原則，但在決定下一個token來源或如何處理KV快取等方面，則不完全適用SPMD/MPMD分類。對於這些情況，Google Pathway等系統提供了更靈活的解決方案。
考慮到上述背景，我們更應關注的是訓練框架與推論引擎之間關於訓練資料和模型參數的通訊機制，而非局限於是否採用單一控制器或多控制器架構。
4.2 SLIME的具體實現方法
訓練框架與推論引擎之間的核心挑戰在於訓練資料與模型參數的通訊機制。為了更好地理解這一點，我們可以透過分析slime和roll專案來探討具體實現方案。
SLIME是一個專注於強化學習擴展的後訓練框架，它定義了兩個主要元件：RayTrainGroup用於訓練框架，RolloutGroup用於推論引擎。
4.2.1 資料傳輸機制
SLIME透過定義一個中間件類——Buffer，實現了推論引擎與訓練模組間的資料傳輸。所有的資料都會被儲存在這個Buffer中（甚至可以寫入磁碟），並透過rollout ID進行指定訪問。此外，Buffer類中的資料處理函數以及rollout/eval函數均可以透過命令列參數靈活配置，極大地提高了系統的適應性。
這種設計使得應對業務需求時更加靈活高效，尤其是面對各種特殊需求和資料格式時尤為重要。
Rollout 的generate函數是透過Buffer。
獲取訓練框架所需的資料同樣依賴於這個Buffer：
同步rollout的buffer給actor的過程如下所示：
4.2.2 模型參數同步機制
為了讓rollout引擎能夠在適當的時候正確地同步參數，SLIME將actor的配置資訊傳遞給rollout。這部分涉及到初始化過程組以便在每個訓練階段之後更新權重。
上述過程不僅包括資料緩衝區的同步，還涵蓋了actor間平行配置的協調，保證了參數更新的一致性和準確性。
4.3 ROLL的具體實現方法
ROLL透過叢集（Cluster）的方式定義了多個角色，每個角色負責不同的任務。這種設計方式與演算法層面的認知較為一致，因為從演算法角度來看，訓練框架和推論引擎之間的差異並不明顯，而使用叢集封裝則很好地隱藏了這些複雜性。
4.3.1 資料傳輸機制
類似於Megatron，ROLL允許按照領域（domain）分開取樣，並在pipeline.py文件中進行配置。這使得如果用戶不想編寫資料生成器，ROLL提供了一種更為便捷的解決方案。特別是對於獎勵（reward）模型，理想的狀況是有一個統一的模型，但由於訓練難度大，目前更傾向於針對不同領域使用不同的獎勵模型，並最終進行聚合處理。ROLL支持對不同領域、批次以及查詢進行自定義配置，以適應多樣的應用場景。
4.3.2 模型參數同步機制
ROLL中的模型更新邏輯結合了點對點通訊和集體通訊兩種方式：
點對點通訊：用於同一設備上的參數更新，直接透過worker的node_rank和gpu_rank來判斷是否在同一設備上，從而進行高效的資料交換。
集體通訊：透過廣播參數到目標叢集，只在主進程（rank 0）執行廣播操作，適用於跨設備間的參數同步。
這兩種通訊策略分別對應於colocate和非colocate場景，確保了參數同步的靈活性和效率。
4.3.4 跨機器部署時的考量
當所有元件位於同一台機器上時，硬編碼實現參數同步相對簡單，但當涉及到跨機器部署時，情況變得更加複雜。此時，不僅需要考慮如何有效地管理網路通訊帶來的延遲和頻寬限制，還需要優化分散式環境下的資源分配和負載均衡。此外，單控制器（single controller）模式下，控制器的壓力會隨著叢集規模的擴大而增加，尤其是在處理多媒體資料時，可能需要特別注意效能瓶頸的問題。因此，在跨機器部署的情況下，選擇合適的通訊策略和優化控制器的工作負載變得尤為重要。不過，從SLIME和ROLL的設計來看，參數同步的核心在於通知GPU進行同步操作，中間的通訊過程不依賴於控制器，這為跨機器部署提供了一定的便利性和靈活性。
4.4 Colocation與Ray的應用
將Actor、Ref、Reward、Critic等模型放置在同一張GPU卡上被稱為colocation。然而，正如前文所述，隨著模型規模的增大（例如7B模型已難以在單張卡上訓練），預計下半年會出現多個超過1000B參數量級的模型。這使得平行計算帶來的開銷變得極其顯著。目前，Reward模型普遍較小，7-30B的規模即可滿足需求，因此分開部署往往更具性價比。
為了應對這種複雜性，專案中引入了Ray——一個支持分散式計算的強大框架，它能夠幫助開發者減輕底層邏輯管理的負擔。有關基於Ray的分散式訓練流程和Ray分散式計算框架的詳細介紹，請參閱以下文章： - 圖解OpenRLHF中基於Ray的分散式訓練流程 - Ray分散式計算框架詳解
接下來，我們將比較slime、verl、roll和openrlhf四個框架在colocation與非colocation實現上的差異。
4.4.1 SLIME
SLIME僅定義了兩個主要worker：RayTrainGroup用於訓練，RolloutGroup用於推論。對於colocate，訓練和推論可以分開部署；而在非colocate的情況下，則需要處理分散式通訊以同步參數。這種設計抽象層次高，易於理解，並且能夠很好地適應訓練和推論的不同需求。只需在配置中指定是否colocate，即可自動在所有關鍵環節執行相應操作。
4.4.2 ROLL
對於非colocate場景，ROLL允許細粒度地指定不同worker（例如actor、critic、reward等）部署在不同的顯卡上，甚至可以根據輪次進行配置。若不手動指定，Ray會自動完成部署。鑑於RL任務對資源的高消耗，細粒度的GPU資源配置有助於提高資源利用效率，但這同時也對演算法側的資源排程能力提出了更高要求。顯然，使用Ray來管理這些複雜性更為合適。
4.4.3 Verl
VERL採用了一種獨特的方法來實現colocate和非colocate部署。在非colocation模式下，每個worker（如actor、critic、reward等）作為一個獨立進程運行，依靠Ray來進行排程。而在colocation模式下，多個角色共享同一個Ray actor實例，在同一個進程中實例化多個worker類。透過create_colocated_worker_cls或create_colocated_worker_cls_fused方法動態生成一個多角色類（例如WorkerDict/FusedWorker），該類內部持有多個worker實例。外部可透過統一接口調用不同角色worker的方法，內部則自動分發到對應的worker實例。這種方式使得同進程內的多角色共存成為可能，並在某些場景下能大幅提高效能，比如減少跨進程通訊帶來的延遲和記憶體碎片問題。
4.4.4 OpenRLHF
OpenRLHF提供了靈活的混合部署選項，既支持vLLM引擎、Actor、Reference、Reward和Critic模型節點的共置部署，也支持部分混合部署或完全分離部署，以適應異步訓練的需求。這種靈活性使其能夠應對多樣化的應用場景，但也意味著更複雜的管理和優化需求。
4.4.5 結論
綜上所述，在非colocation情況下，Ray確實可以幫助我們更加輕鬆地管理資源，尤其是在處理複雜的Agent和多輪互動場景時。然而，根據維運團隊的回饋，Ray的設計理念與現有的Kubernetes雲原生生產環境存在一定的衝突，導致在實際生產環境中部署時管理成本較高。不過，Ray團隊也在針對這些問題進行優化，例如使Ray可以直接透過NCCL傳輸tensor資料，從而繞過物件儲存，提高效率。未來，我們可以期待更多來自Ray的更新和改進。
4.5 不同訓練框架與推論引擎的整合
在將不同的訓練框架和推論引擎進行整合時，可能會遇到參數轉換的問題。例如，如果vLLM使用4-維 張量平行（TP），而DeepSpeed分佈在8個GPU上，則需要進行適當的參數轉換以確保資料傳輸的一致性。Megatron-LM也有類似的需求。當存在多個訓練框架和推論引擎時，適配的工作量會成倍增加，這可能導致配置錯誤和效能問題。
4.6 程式碼解耦設計
以Slime為例，其架構分為三層：頂層RolloutGroup負責管理推論引擎的整體流程；中層RolloutRayActor處理具體的推論請求；底層SglangEngine實現具體的推論邏輯。這種分層設計使得替換後端推論引擎變得簡單，只需更改底層實現即可，無需修改上層控制邏輯。同樣，訓練框架也採用了類似的分層結構，保證了系統的靈活性和可維護性。
5. 關於Agentic RL
目前，roll、verl和openrlhf等框架對Agentic RL提供了良好的支持。儘管這樣做可能增加了程式碼複雜度，但隨著技術成熟，預計會有更清晰的設計出現。未來，Agentic RL有望成為主流，現有的RL方法將成為其中的一部分。
6. 框架選擇建議
6.1 框架難點分析
快速發展的技術環境意味著舊框架容易過時，因此保持框架簡潔和高可維護性是關鍵。新框架由於沒有歷史負擔，可以更容易地適應新技術趨勢。
6.2 推薦框架
OpenRLHF：一個高效能的開源RLHF框架，整合了Ray、vLLM、ZeRO-3和HuggingFace Transformers。
slime：新推出的框架，程式碼簡潔，適合想要嘗試大膽框架修改的研究者。
ROLL：強調資料處理和異步操作的支持，特別適用於深入探索Agentic RL的團隊。
verl：穩定且優化良好，適合大規模叢集部署，尤其適合資源豐富的團隊。
根據團隊的具體需求和技術背景，可以選擇最適合的框架來開展工作。對於有特定需求或希望快速擴展的團隊，verl可能是更好的選擇，因為它已經被多個大廠驗證過。而對於追求技術創新和敏捷開發的團隊，slime或ROLL可能更具吸引力。
結尾
在過去半年中，我們深入探討了RL訓練框架、Agent框架以及推論引擎框架。總體而言，程式碼量方面，Agent框架最為龐大，其次是推論引擎和RL訓練框架；而在程式碼難度上，推論引擎居首，隨後是RL訓練框架和Agent框架。值得注意的是，如果排除推論引擎底層運算子的複雜性，RL訓練框架的挑戰主要在於整合各種系統和技術，這要求框架開發者對多種技術和業務邏輯有深刻的理解。
開源框架如verl、slime、roll及openRLHF各具特色，展現了作者們的追求與堅持，並且社群活躍度高。可以說，在開源RL框架領域，中國在技術實力與認知深度方面處於世界領先位置。雖然演算法人才間的差異不大，但在硬體資源（如顯示卡）方面仍存在一定的差距。