騰訊近日發佈了基於Megatron-Core和SGLang/vLLM研發的大模型訓練庫WeChat-YATT(Yet Another Transformer Trainer),內部項目代號爲gCore。這一訓練庫專注於強化學習和多模態模型訓練,致力於爲開發者提供易擴展、簡潔、高效、可靠的大模型訓練解決方案。

WeChat-YATT通過定製化的並行計算策略,能夠有效處理大尺寸模型、長序列輸入和大數據集等複雜場景,成功解決了微信內部多個實際業務場景的關鍵痛點問題,顯著提升了大模型訓練效率。該工具爲研究人員和開發者提供了靈活且可擴展的技術方案,有望推動多模態和強化學習領域的創新發展。

image.png

WeChat-YATT重點解決了大模型分佈式訓練過程中面臨的兩大核心技術瓶頸。

首先是多模態場景下的可擴展性瓶頸問題。隨着圖像、視頻等多模態數據規模不斷增長,傳統架構中由單一控制器進行數據管理的方式容易成爲通信和內存瓶頸,導致系統吞吐量受限,甚至引發訓練流程異常中斷。WeChat-YATT通過引入並行控制器(Parallel Controller)的並行管理機制,有效分散了系統壓力,大幅提升了系統的可擴展性和穩定性,更好地應對多模態、大數據量的複雜應用場景。

其次是動態採樣與生成式獎勵計算下的效率短板。在需要頻繁動態採樣或生成式獎勵計算的訓練流程中,模型頻繁切換和"長尾"任務容易產生大量額外開銷,導致GPU算力無法充分利用,嚴重影響整體訓練效率。WeChat-YATT通過部分共存策略和異步交互機制,大幅減輕了模型切換損耗和長尾任務影響,實現了訓練過程中的高吞吐量和高資源利用率,更好地支撐大規模RLHF任務的高效迭代。

針對不同業務場景的需求,WeChat-YATT支持全員共存和部分共存兩種資源放置模式,以最大化提升集羣資源利用率。

全員共存模式採用串行調度機制,Actor Rollouts、GenRM(生成式獎勵模型)與Train依次串行執行。每個角色完成任務後主動釋放計算資源,系統隨即加載下一個任務所需模型。該策略適用於絕大多數常規訓練場景。值得注意的是,在每個階段,相關組件均可獨佔全部GPU資源,這極大縮短了資源空閒"氣泡"時間,顯著提升總體訓練吞吐量和效率。

部分共存模式下,Actor Rollouts與GenRM獨立部署,並通過異步方式進行高效交互。Actor訓練階段佔用全部GPU資源,在Rollouts生成階段釋放GPU資源並喚醒Actor Rollouts及GenRM兩大組件協同工作。系統通過動態負載評估進行資源分配與均衡,當Rollouts生成完畢後,這兩者釋放資源,Actor隨之重新加載到GPU上進入下一輪訓練流程。部分共存模式特別適合Rollouts與GenRM需要高頻交互、動態採樣的複雜任務場景。

WeChat-YATT還具備多項技術特色。在內存利用方面,項目採用並行控制器架構,有效降低了單節點內存消耗,更適合多模態場景下的大模型訓練,提升了系統擴展性和穩定性。在GenRM支持方面,針對生成式獎勵模型場景實現了不同的資源放置策略,用戶可根據具體場景選擇最優的訓練方案。

智能檢查點策略是另一大亮點。WeChat-YATT支持異步Checkpoint保存,並針對微信業務場景特點,根據調度流程實現斷點自動保存,進一步保障了訓練安全性和高可用性。此外,系統在訓練過程中實現了各個數據並行組間的負載均衡,有效減少資源空閒時間,顯著提升整體訓練吞吐量。

這一訓練庫的發佈標誌着騰訊在大模型技術基礎設施建設方面取得重要進展,也爲業界提供了處理複雜多模態訓練場景的有效解決方案。