最高提升20倍吞吐量!豆包大模型團(tuán)隊(duì)發(fā)布全新 RLHF 框架,現(xiàn)已開源!
強(qiáng)化學(xué)習(xí)(RL)對(duì)大模型復(fù)雜推理能力提升有關(guān)鍵作用,然而,RL 復(fù)雜的計(jì)算流程以及現(xiàn)有系統(tǒng)局限性,也給訓(xùn)練和部署帶來了挑戰(zhàn)。傳統(tǒng)的 RL/RLHF 系統(tǒng)在靈活性和效率方面存在不足,難以適應(yīng)不斷涌現(xiàn)的新算法需求,無法充分發(fā)揮大模型潛力。
近日,字節(jié)跳動(dòng)豆包大模型團(tuán)隊(duì)與香港大學(xué)聯(lián)合提出 HybridFlow(開源項(xiàng)目名:veRL),一個(gè)靈活且高效的 RL/RLHF 框架。該框架采用混合編程模型,融合單控制器(Single-Controller)的靈活性和多控制器(Multi-Controller)的高效性,可更好實(shí)現(xiàn)和執(zhí)行多種RL算法,顯著提升訓(xùn)練吞吐量,降低開發(fā)和維護(hù)復(fù)雜度。實(shí)驗(yàn)結(jié)果表明,HybridFlow 在運(yùn)行各種 RL(HF) 算法時(shí),吞吐量相較 SOTA 基線提升了 1.5-20 倍。
從 ChatGPT [1] 到 o1 等各種大語言模型,強(qiáng)化學(xué)習(xí)(RL)算法在提升模型性能和適應(yīng)性方面起著至關(guān)重要的作用。在大模型后訓(xùn)練(Post-Training)階段引入 RL 方法,已成為提升模型質(zhì)量和對(duì)齊人類偏好[2, 3]的重要手段。
然而,隨著模型規(guī)模的不斷擴(kuò)大,RL 算法在大模型訓(xùn)練中面臨著靈活性和性能的雙重挑戰(zhàn)。
因此,開發(fā)一個(gè)高效且靈活的大模型 RL 訓(xùn)練框架顯得尤為重要。這不僅需要高效地執(zhí)行復(fù)雜的分布式計(jì)算流程,還要具備適應(yīng)不同 RL 算法的靈活性,以滿足不斷發(fā)展的研究需求。
字節(jié)跳動(dòng)豆包大模型團(tuán)隊(duì)與香港大學(xué)近期公開聯(lián)合研究成果—— HybridFlow ,一個(gè)靈活且高效的大模型 RL 訓(xùn)練框架,兼容多種訓(xùn)練和推理框架,支持靈活的模型部署和多種 RL 算法實(shí)現(xiàn)。
HybridFlow 采用混合編程模型,將單控制器的靈活性與多控制器的高效性相結(jié)合,解耦了控制流和計(jì)算流?;赗ay的分布式編程,動(dòng)態(tài)計(jì)算圖,異構(gòu)調(diào)度能力,通過封裝單模型的分布式計(jì)算、統(tǒng)一模型間的數(shù)據(jù)切分,以及支持異步 RL 控制流,HybridFlow 能夠高效地實(shí)現(xiàn)和執(zhí)行各種 RL 算法,復(fù)用計(jì)算模塊和支持不同的模型部署方式,大大提升了系統(tǒng)的靈活性和開發(fā)效率。
實(shí)驗(yàn)結(jié)果表明,HybridFlow 在各種模型規(guī)模和 RL 算法下,訓(xùn)練吞吐量相比其他框架提升了 1.5 倍至 20 倍。
目前,該論文已被 EuroSys 2025 接收,代碼倉(cāng)庫(kù)也對(duì)外公開。
HybridFlow: A Flexible and Efficient RLHF Framework
論文鏈接:https://team.doubao.com/zh/publication/hybridflow-a-flexible-and-efficient-rlhf-framework?view_from=research
代碼鏈接:https://github.com/volcengine/veRL
1. RL(Post-Training)復(fù)雜計(jì)算流程給 LLM 訓(xùn)練帶來全新的挑戰(zhàn)
在深度學(xué)習(xí)中,數(shù)據(jù)流(DataFlow)是一種重要的計(jì)算模式抽象,用于表示數(shù)據(jù)經(jīng)過一系列復(fù)雜計(jì)算后實(shí)現(xiàn)特定功能。神經(jīng)網(wǎng)絡(luò)的計(jì)算就是典型的 DataFlow ,可以用計(jì)算圖(Computational Graph)來描述,其中節(jié)點(diǎn)代表計(jì)算操作,邊表示數(shù)據(jù)依賴。
大模型 RL 的計(jì)算流程比傳統(tǒng)神經(jīng)網(wǎng)絡(luò)更為復(fù)雜。在 RLHF 中,需要同時(shí)訓(xùn)練多個(gè)模型,如 Actor 、Critic 、參考策略(Reference Policy)和獎(jiǎng)勵(lì)模型(Reward Model),并在它們之間傳遞大量數(shù)據(jù)。這些模型涉及不同的計(jì)算類型(前向反向傳播、優(yōu)化器更新、自回歸生成等),可能采用不同的并行策略。
傳統(tǒng)的分布式 RL 通常假設(shè)模型可在單個(gè) GPU 上訓(xùn)練,或使用數(shù)據(jù)并行方式 [4,5],將控制流和計(jì)算流合并在同一進(jìn)程中。這在處理小規(guī)模模型時(shí)效果良好,但面對(duì)大模型,訓(xùn)練需要復(fù)雜的多維并行,涉及大量分布式計(jì)算,傳統(tǒng)方法難以應(yīng)對(duì)。
2. HybridFlow 解耦控制流和計(jì)算流,兼顧靈活高效
大模型 RL 本質(zhì)上是一個(gè)二維的 DataFlow 問題:high-level 的控制流(描述 RL 算法的流程)+ low-level 的計(jì)算流(描述分布式神經(jīng)網(wǎng)絡(luò)計(jì)算)。
近期開源的 RLHF 框架,如 DeepSpeed-Chat [6]、OpenRLHF [7] 和 NeMo-Aligner [8],采用了統(tǒng)一的多控制器(Multi-Controller)架構(gòu)。各計(jì)算節(jié)點(diǎn)獨(dú)立管理計(jì)算和通信,降低了控制調(diào)度的開銷。然而,控制流和計(jì)算流高度耦合,當(dāng)設(shè)計(jì)新的 RL 算法,組合相同的計(jì)算流和不同的控制流時(shí),需要重寫計(jì)算流代碼,修改所有相關(guān)模型,增加了開發(fā)難度。
與此前框架不同,HybridFlow 采用了混合編程模型,控制流由單控制器(Single-Controller)管理,具有全局視圖,實(shí)現(xiàn)新的控制流簡(jiǎn)單快捷,計(jì)算流由多控制器(Multi-Controller)負(fù)責(zé),保證了計(jì)算的高效執(zhí)行,并且可以在不同的控制流中復(fù)用。
盡管相比純粹的多控制器架構(gòu),這可能帶來一定的控制調(diào)度開銷,但 HybridFlow 通過優(yōu)化數(shù)據(jù)傳輸,降低了控制流與計(jì)算流之間的傳輸量,兼顧了靈活性和高效性。
3. 系統(tǒng)設(shè)計(jì)之一:Hybrid Programming Model (編程模型創(chuàng)新)
- 封裝單模型分布式計(jì)算
在 HybridFlow 中,每個(gè)模型(如 Actor、Critic、參考策略、獎(jiǎng)勵(lì)模型等)的分布式計(jì)算被封裝為獨(dú)立的模塊,稱為模型類。
這些模型類繼承于基礎(chǔ)的并行 Worker 類(如 3DParallelWorker 、FSDPWorker 等),通過抽象的 API 接口,封裝了模型的前向、反向計(jì)算、優(yōu)化器更新和自回歸生成等操作。該封裝方式提高了代碼的復(fù)用性,便于模型的維護(hù)和擴(kuò)展。
對(duì)于不同的 RL 控制流,用戶可以直接復(fù)用封裝好的模型類,同時(shí)自定義部分算法所需的數(shù)值計(jì)算,實(shí)現(xiàn)不同算法。當(dāng)前 HybridFlow 可使用 Megatron-LM [13] 和 PyTorch FSDP [14] 作為訓(xùn)練后端,同時(shí)使用 vLLM [15] 作為自回歸生成后端,支持用戶使用其他框架的訓(xùn)練和推理腳本進(jìn)行自定義擴(kuò)展。
- 靈活的模型部署
HybridFlow 提供了資源池(ResourcePool)概念,可以將一組 GPU 資源虛擬化,并為每個(gè)模型分配計(jì)算資源。不同的資源池實(shí)例可以對(duì)應(yīng)不同設(shè)備集合,支持不同模型在同一組或不同組 GPU 上部署。這種靈活的模型部署方式,滿足了不同算法、模型和硬件環(huán)境下的資源和性能需求。
- 統(tǒng)一模型間的數(shù)據(jù)切分
在大模型 RL 計(jì)算流程中,不同模型之間的數(shù)據(jù)傳輸涉及復(fù)雜的多對(duì)多廣播和數(shù)據(jù)重分片。
為解決該問題,HybridFlow 設(shè)計(jì)了一套通用數(shù)據(jù)傳輸協(xié)議(Transfer Protocol),包括收集(collect)和分發(fā)(distribute)兩個(gè)部分。
通過在模型類的操作上注冊(cè)相應(yīng)的傳輸協(xié)議,比如:@register(transfer_mode=3D_PROTO),HybridFlow 可以在控制器層(Single-Controller)統(tǒng)一管理數(shù)據(jù)的收集和分發(fā),實(shí)現(xiàn)模型間數(shù)據(jù)的自動(dòng)重分片,支持不同并行度下的模型通信。
HybridFlow 框架已經(jīng)支持多種數(shù)據(jù)傳輸協(xié)議,涵蓋大部分?jǐn)?shù)據(jù)重切分場(chǎng)景。同時(shí),用戶可靈活地自定義收集(collect)和分發(fā)(distribute)函數(shù),將其擴(kuò)展到更復(fù)雜的數(shù)據(jù)傳輸場(chǎng)景。
- 支持異步 RL 控制流
在 HybridFlow 中,控制流部分采用單控制器架構(gòu),可靈活實(shí)現(xiàn)異步 RL 控制流。
當(dāng)模型部署在不同設(shè)備集合上時(shí),不同模型計(jì)算可并行執(zhí)行,這提高了系統(tǒng)的并行度和效率。對(duì)于部署在同一組設(shè)備上的模型,HybridFlow 通過調(diào)度機(jī)制實(shí)現(xiàn)了順序執(zhí)行,避免資源爭(zhēng)奪和沖突。
- 少量代碼靈活實(shí)現(xiàn)各種 RL 控制流算法
得益于混合編程模型的設(shè)計(jì),HybridFlow 可以方便地實(shí)現(xiàn)各種 RLHF 算法,如 PPO [9]、ReMax [10]、Safe-RLHF [11]、GRPO [12] 等。用戶只需調(diào)用模型類的 API 接口,按算法邏輯編寫控制流代碼,無需關(guān)心底層的分布式計(jì)算和數(shù)據(jù)傳輸細(xì)節(jié)。
例如,實(shí)現(xiàn) PPO 算法只需少量代碼,通過調(diào)用 actor.generate_sequences 、critic.compute_values 等函數(shù)即可完成。同時(shí),用戶只需要修改少量代碼即可遷移到 Safe-RLHF 、ReMax 以及 GRPO 算法。
4. 系統(tǒng)設(shè)計(jì)之二:3D-HybridEngine (訓(xùn)練推理混合技術(shù))降低通信內(nèi)存開銷
在 Online RL 算法中,Actor 模型需要在訓(xùn)練和生成(Rollout)階段之間頻繁切換,且兩個(gè)階段可能采用不同并行策略。
具體而言,訓(xùn)練階段,需要存儲(chǔ)梯度和優(yōu)化器狀態(tài),模型并行度(Model Parallel Size, MP)可能相應(yīng)增高,而生成階段,模型無需存儲(chǔ)梯度和優(yōu)化器狀態(tài),MP 和數(shù)據(jù)并行度(Data Parallel Size, DP)可能較小。因此,在兩個(gè)階段之間,模型參數(shù)需要重新分片和分配,依賴傳統(tǒng)通信組構(gòu)建方法會(huì)帶來額外通信和內(nèi)存開銷。
此外,為了在新的并行度配置下使用模型參數(shù),通常需要在所有 GPU 之間進(jìn)行全聚合(All-Gather)操作,帶來了巨大的通信開銷,增加了過渡時(shí)間。
為解決這個(gè)問題,HybridFlow 設(shè)計(jì)了 3D-HybridEngine ,提升了訓(xùn)練和生成過程效率。
注:3D-HybridEngine 一次迭代的流程
3D-HybridEngine 通過優(yōu)化并行分組方法,實(shí)現(xiàn)了零冗余的模型參數(shù)重組,具體包括以下步驟:
- 定義不同的并行組
在訓(xùn)練和生成階段,3D-HybridEngine 使用不同的三維并行配置,包括:流水線并行(PP)、張量并行(TP)和數(shù)據(jù)并行(DP)的大小。訓(xùn)練階段的并行配置為 ??-??-?? 。在生成階段,我們新增一個(gè)新的微數(shù)據(jù)并行組(Micro DP Group,????),用于處理 Actor 模型參數(shù)和數(shù)據(jù)的重組。生成階段的并行配置為 ????-????-????-?? 。
- 重組模型參數(shù)過程
通過巧妙地重新定義生成階段的并行分組,可以使每個(gè) GPU 在生成階段復(fù)用訓(xùn)練階段已有的模型參數(shù)分片,避免在 GPU 內(nèi)存中保存額外的模型參數(shù),消除內(nèi)存冗余。
- 減少通信開銷
參數(shù)重組過程中,3D-HybridEngine 僅在每個(gè)微數(shù)據(jù)并行組(Micro DP Group)內(nèi)進(jìn)行 All-Gather 操作,而非所有 GPU 之間進(jìn)行。這大大減少了通信量,降低過渡時(shí)間,提高了整體的訓(xùn)練效率。
5. 實(shí)驗(yàn)結(jié)果:HybridFlow 提供靈活性的同時(shí),加速了訓(xùn)練
團(tuán)隊(duì)在 16 臺(tái) A100 GPU 集群上,對(duì) HybridFlow 和主流 RLHF 框架(DeepSpeed-Chat [6] v0.14.0、OpenRLHF [7] v0.2.5 和 NeMo-Aligner [8] v0.2.0)進(jìn)行對(duì)比實(shí)驗(yàn)。實(shí)驗(yàn)涵蓋了不同模型規(guī)模(7B、13B、34B、70B)的 LLM ,以及不同 RLHF 算法(PPO [9]、ReMax [10]、Safe-RLHF [11])。
所有實(shí)驗(yàn)中,Actor、Critic、參考策略 Reference Policy 和獎(jiǎng)勵(lì)模型 Reward Model 均采用相同規(guī)模模型。更多實(shí)驗(yàn)配置和測(cè)試細(xì)節(jié)請(qǐng)移步完整論文。
- 更高的端到端訓(xùn)練吞吐量
結(jié)果顯示,HybridFlow 在各種模型規(guī)模和 RLHF 算法下,都顯著優(yōu)于其他框架,實(shí)現(xiàn)了更高訓(xùn)練吞吐量。
無論 PPO 、ReMax 還是 Safe-RLHF 算法,HybridFlow 在所有模型規(guī)模下平均訓(xùn)練吞吐量均大幅領(lǐng)先于其他框架,提升幅度在 1.5 倍至 20 倍之間。
隨 GPU 集群規(guī)模擴(kuò)大,HybridFlow 吞吐量也獲得良好擴(kuò)展。這得益于其靈活的模型部署,充分利用硬件資源,實(shí)現(xiàn)高效并行計(jì)算。同時(shí),HybridFlow 能夠支持多種分布式并行框架(Megatron-LM [13]、FSDP [14]、vLLM [15]),滿足不同模型規(guī)模的計(jì)算需求。
- HybridEngine 有效減少開銷
分析 Actor 模型在訓(xùn)練和生成階段的過渡時(shí)間,團(tuán)隊(duì)發(fā)現(xiàn),HybridFlow 的 3D-HybridEngine 的零冗余模型參數(shù)重組技術(shù),有效減少了模型參數(shù)在兩個(gè)階段之間的重分片和通信開銷。
相比其他框架,過渡時(shí)間減少了 55.2% ,在 70B 模型上過渡時(shí)間降低了 89.1% 。
- 不同模型部署方式對(duì)比下的三個(gè)洞察
團(tuán)隊(duì)對(duì)比了不同的模型部署策略,總結(jié)了模型部署和 GPU 分配的三大關(guān)鍵洞察:
- 為 Actor 模型分配更多的 GPU ,可以縮短 critical path ;
- Colocate 模式在相對(duì)小規(guī)模集群中能夠最大化 GPU 利用率;
- 在大規(guī)模集群中將 Actor 和 Critic 模型部署在不同的設(shè)備能夠提升擴(kuò)展率。
值得一提的是,HybridFlow 同樣適用于更廣泛的 RL 訓(xùn)練場(chǎng)景,隨著 o1 模型誕生,業(yè)內(nèi)對(duì) Reasoning 能力、RL 關(guān)注度也在提升,團(tuán)隊(duì)后續(xù)將圍繞相關(guān)場(chǎng)景進(jìn)行探索和實(shí)驗(yàn)。