支持序列并行的LLaMA-Factory

原標(biāo)題：一行代碼Post-Train任意長(zhǎng)序列！360智腦開源360-LLaMA-Factory
文章來源：機(jī)器之心
內(nèi)容字?jǐn)?shù)：8718字

360智腦開源360-LLaMA-Factory：輕松實(shí)現(xiàn)大模型長(zhǎng)序列后訓(xùn)練

機(jī)器之心AIxiv專欄報(bào)道了360智腦開源的360-LLaMA-Factory項(xiàng)目。該項(xiàng)目基于LLaMA-Factory，新增序列并行功能，顯著提升了大模型長(zhǎng)序列后訓(xùn)練的能力。此舉解決了現(xiàn)有開源框架在長(zhǎng)序列后訓(xùn)練方面存在的不足，為處理復(fù)雜長(zhǎng)文本任務(wù)、多幀視頻理解任務(wù)以及高計(jì)算量的大模型提供了強(qiáng)有力的支持。

1. 項(xiàng)目背景與簡(jiǎn)介

隨著大模型處理長(zhǎng)序列的需求日益增長(zhǎng)，預(yù)訓(xùn)練和后訓(xùn)練階段都需要支持更長(zhǎng)序列數(shù)據(jù)的訓(xùn)練。雖然預(yù)訓(xùn)練階段已有Megatron-LM等框架支持序列并行，但后訓(xùn)練階段因算法多樣性和需求靈活性，缺乏一個(gè)兼顧并行策略、算法多樣性、GPU顯存優(yōu)化和易用性的框架。LLaMA-Factory作為用戶最多的開源后訓(xùn)練框架，雖然功能豐富，但在長(zhǎng)序列支持上仍有欠缺，尤其缺乏序列并行功能。

360智腦基于LLaMA-Factory開發(fā)了360-LLaMA-Factory，填補(bǔ)了這一空白。該項(xiàng)目通過添加一行代碼sequence_parallel_size: 16即可實(shí)現(xiàn)任意長(zhǎng)度的序列后訓(xùn)練，只需按需增加GPU卡數(shù)即可。該項(xiàng)目已在內(nèi)部經(jīng)過嚴(yán)格驗(yàn)證，并提交了Pull Request給LLaMA-Factory主倉(cāng)。

2. 長(zhǎng)序列及其后訓(xùn)練

長(zhǎng)序列大模型訓(xùn)練面臨的關(guān)鍵挑戰(zhàn)是激活顯存隨序列長(zhǎng)度大幅增加。雖然一些技術(shù)如unsloth、liger kernel和LoRA可以降低顯存占用，但效果有限。序列并行是解決這一問題的有效途徑，它將長(zhǎng)序列切分到多張GPU卡上計(jì)算，避免單卡處理過長(zhǎng)序列。然而，序列并行的實(shí)現(xiàn)難度較大，需要復(fù)雜的通信和attention計(jì)算修改。

360-LLaMA-Factory對(duì)現(xiàn)有的一些支持序列并行的開源框架進(jìn)行了調(diào)研，發(fā)現(xiàn)它們存在一些問題，例如實(shí)現(xiàn)錯(cuò)誤、算法支持不全、更新維護(hù)不及時(shí)等。360-LLaMA-Factory在繼承LLaMA-Factory、ring-flash-attention和EasyContext等開源項(xiàng)目的基礎(chǔ)上，進(jìn)行了模塊化改進(jìn)和優(yōu)化，并對(duì)SFT和DPO訓(xùn)練器進(jìn)行了適配。

3. 360-LLaMA-Factory框架解析

360-LLaMA-Factory通過模塊化封裝，對(duì)原始代碼修改最小化。它在model_args中添加參數(shù)并抽象出apply_sequence_parallel函數(shù)來修改模型，并用decorator抽象數(shù)據(jù)處理函數(shù)，最后在Trainer中進(jìn)行l(wèi)oss計(jì)算的reduce匯總。針對(duì)SFT和DPO，進(jìn)行了最小化的定制化修改，并通過測(cè)試驗(yàn)證了其正確性。

4. 360-LLaMA-Factory效果驗(yàn)證

360-LLaMA-Factory內(nèi)部已用于訓(xùn)練開源模型360Zhinao2-7B-Chat-360k。通過小數(shù)據(jù)集測(cè)試，驗(yàn)證了序列并行對(duì)訓(xùn)練曲線的影響微乎其微。粗粒度壓測(cè)表明，8卡80G的配置下，SFT和DPO訓(xùn)練可支持幾十到幾百k的超長(zhǎng)序列。

5. 總結(jié)

360-LLaMA-Factory通過簡(jiǎn)單易用的方式實(shí)現(xiàn)了長(zhǎng)序列后訓(xùn)練，為大模型應(yīng)用提供了有力支持。其模塊化設(shè)計(jì)、正確性驗(yàn)證以及高效的性能使其成為一個(gè)值得推薦的開源項(xiàng)目。歡迎開發(fā)者使用和貢獻(xiàn)代碼。

聯(lián)系作者

文章來源：機(jī)器之心
作者微信：
作者簡(jiǎn)介：專業(yè)的人工智能媒體和產(chǎn)業(yè)服務(wù)平臺(tái)