大規模語言模型(LLM)的訓練和最佳化是人工智慧領域的關鍵挑戰。高效率的訓練方法不僅需要確保模型的效能,更要確保其符合人類價值。強化學習與人類回饋(RLHF)作為一種有效的LLM訓練方法,近年來已廣泛應用,但其效率和可擴展性仍有待提高。位元組跳動豆包大模型團隊為此開源了名為HybridFlow的RLHF框架,旨在解決傳統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
Downcodes小編總結:HybridFlow框架的開源,為大規模語言模型的訓練提供了新的想法和工具,其高效性和靈活性有望推動LLM技術進一步發展,值得關注和深入研究。 期待未來能看到更多基於HybridFlow的創新應用。