大模型(LLM)如GPT、Llama等在人工智能領域掀起了一場革命,但如何高效地訓練這些龐大的模型並使其符合人類價值觀仍然是一個難題。
強化學習與人類反饋(RLHF)作爲一種重要的LLM訓練方法,近年來得到廣泛應用,但傳統的RLHF框架在靈活性、效率和可擴展性方面存在侷限性。

爲解決這些問題,字節跳動豆包大模型團隊開源了名爲HybridFlow的RLHF框架,爲LLM訓練帶來了新的可能性。
RLHF通常包含三個階段:
首先,actor模型根據輸入的提示生成文本;然後,critic模型、reference模型和reward模型對生成的文本進行評估,並計算出相應的價值、參考概率和獎勵值;

最後,利用這些評估結果對actor模型進行訓練,使其生成更符合人類偏好的文本。傳統的RLHF框架通常採用單一控制器來管理整個數據流,但這對於需要分佈式計算的LLM來說效率低下。
HybridFlow框架創新性地結合了單控制器和多控制器模式,並通過分層的API設計將複雜的計算和數據依賴關係解耦,從而實現RLHF數據流的靈活表示和高效執行。

HybridFlow的優勢主要體現在以下三個方面:
靈活支持多種RLHF算法和模型: HybridFlow提供了模塊化的API,用戶可以輕鬆地實現和擴展各種RLHF算法,例如PPO、ReMax和Safe-RLHF等。
高效的模型權重重組:3D-HybridEngine組件支持actor模型在訓練和生成階段高效地進行模型權重重組,最大限度地減少內存冗餘和通信開銷。
自動化的模型部署和並行策略選擇: Auto Mapping組件可以根據模型負載和數據依賴關係自動將模型映射到不同的設備,並選擇最佳的並行策略,從而簡化模型部署流程並提升訓練效率。
實驗結果表明,HybridFlow在運行各種RLHF算法時,吞吐量提升顯著,最高可達20.57倍。HybridFlow的開源將爲RLHF研究和開發提供強大的工具,推動未來LLM技術的發展。
論文地址:https://arxiv.org/pdf/2409.19256
