支持序列并行的LLaMA-Factory
原標題:一行代碼Post-Train任意長序列!360智腦開源360-LLaMA-Factory
文章來源:機器之心
內容字數:8718字
360智腦開源360-LLaMA-Factory:輕松實現大模型長序列后訓練
機器之心AIxiv專欄報道了360智腦開源的360-LLaMA-Factory項目。該項目基于LLaMA-Factory,新增序列并行功能,顯著提升了大模型長序列后訓練的能力。此舉解決了現有開源框架在長序列后訓練方面存在的不足,為處理復雜長文本任務、多幀視頻理解任務以及高計算量的大模型提供了強有力的支持。
1. 項目背景與簡介
隨著大模型處理長序列的需求日益增長,預訓練和后訓練階段都需要支持更長序列數據的訓練。雖然預訓練階段已有Megatron-LM等框架支持序列并行,但后訓練階段因算法多樣性和需求靈活性,缺乏一個兼顧并行策略、算法多樣性、GPU顯存優化和易用性的框架。LLaMA-Factory作為用戶最多的開源后訓練框架,雖然功能豐富,但在長序列支持上仍有欠缺,尤其缺乏序列并行功能。
360智腦基于LLaMA-Factory開發了360-LLaMA-Factory,填補了這一空白。該項目通過添加一行代碼sequence_parallel_size: 16
即可實現任意長度的序列后訓練,只需按需增加GPU卡數即可。該項目已在內部經過嚴格驗證,并提交了Pull Request給LLaMA-Factory主倉。
2. 長序列及其后訓練
長序列大模型訓練面臨的關鍵挑戰是激活顯存隨序列長度大幅增加。雖然一些技術如unsloth、liger kernel和LoRA可以降低顯存占用,但效果有限。序列并行是解決這一問題的有效途徑,它將長序列切分到多張GPU卡上計算,避免單卡處理過長序列。然而,序列并行的實現難度較大,需要復雜的通信和attention計算修改。
360-LLaMA-Factory對現有的一些支持序列并行的開源框架進行了調研,發現它們存在一些問題,例如實現錯誤、算法支持不全、更新維護不及時等。360-LLaMA-Factory在繼承LLaMA-Factory、ring-flash-attention和EasyContext等開源項目的基礎上,進行了模塊化改進和優化,并對SFT和DPO訓練器進行了適配。
3. 360-LLaMA-Factory框架解析
360-LLaMA-Factory通過模塊化封裝,對原始代碼修改最小化。它在model_args
中添加參數并抽象出apply_sequence_parallel
函數來修改模型,并用decorator抽象數據處理函數,最后在Trainer中進行loss計算的reduce匯總。針對SFT和DPO,進行了最小化的定制化修改,并通過測試驗證了其正確性。
4. 360-LLaMA-Factory效果驗證
360-LLaMA-Factory內部已用于訓練開源模型360Zhinao2-7B-Chat-360k。通過小數據集測試,驗證了序列并行對訓練曲線的影響微乎其微。粗粒度壓測表明,8卡80G的配置下,SFT和DPO訓練可支持幾十到幾百k的超長序列。
5. 總結
360-LLaMA-Factory通過簡單易用的方式實現了長序列后訓練,為大模型應用提供了有力支持。其模塊化設計、正確性驗證以及高效的性能使其成為一個值得推薦的開源項目。歡迎開發者使用和貢獻代碼。
聯系作者
文章來源:機器之心
作者微信:
作者簡介:專業的人工智能媒體和產業服務平臺