支持序列并行的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通過簡單易用的方式實現了長序列后訓練，為大模型應用提供了有力支持。其模塊化設計、正確性驗證以及高效的性能使其成為一個值得推薦的開源項目。歡迎開發者使用和貢獻代碼。

聯系作者

文章來源：機器之心
作者微信：
作者簡介：專業的人工智能媒體和產業服務平臺