項目核心開發(fā)者 Haosheng Zou 本科畢業(yè)于清華大學(xué)電子系，博士畢業(yè)于清華大學(xué)計算機(jī)系朱軍教授組，目前在 360 智腦從事長文本和強(qiáng)化學(xué)習(xí)等后訓(xùn)練工作。開發(fā)者 Xiaowei Lv 目前在人民大學(xué)信息學(xué)院研二在讀。Fenrui Xiao、Junchen Liu、Qi An 和 Xiaodong Sun 等在開發(fā)測試中亦有貢獻(xiàn)。

大模型長序列的處理能力已越來越重要，像復(fù)雜長文本任務(wù)、多幀視頻理解任務(wù)、以及 OpenAI 近期發(fā)布的 o1、o3 系列模型的高計算量模式，需要處理的輸入 + 輸出總 token 數(shù)從幾萬量級上升到了幾百萬量級。面對模型日益增長的長序列需求，在預(yù)訓(xùn)練（Pre-Training）和后訓(xùn)練（Post-Training）階段，所用的平臺框架都需要支持更長序列數(shù)據(jù)的訓(xùn)練。不同于預(yù)訓(xùn)練階段基于 Megatron-LM 定制開發(fā)的常見選擇，后訓(xùn)練階段因后訓(xùn)練算法的多樣性（比如僅 DPO 就有幾十個變種）和訓(xùn)練需求的靈活性，至今沒有一個框架同時在并行策略、后訓(xùn)練算法、GPU 顯存優(yōu)化和簡單易用這 4 個方面上全部做到兼容并包。

在所有開源的后訓(xùn)練框架中，LLaMA-Factory 是用戶最多的框架之一（GitHub star 數(shù)已 37k 多），保持長期迭代更新，支持豐富的模型和后訓(xùn)練算法，有各種 GPU 顯存優(yōu)化技巧和簡單易用的方式。然而，LLaMA-Factory 在長序列后訓(xùn)練上支持仍有所欠缺，尚不支持長序列的關(guān)鍵技術(shù) —— 序列并行。

項目主頁：https://github.com/Qihoo360/360-LLaMA-Factory

最近，360 智腦基于 LLaMA-Factory 開源了 360-LLaMA-Factory，加入了序列并行功能，一行代碼即可支持任意長序列的后訓(xùn)練（Post-Training）—— 僅需額外指定序列并行一個參數(shù)：

sequence_parallel_size: 16

按需增加序列并行的 GPU 卡數(shù)，即可在任意長度的序列上 SFT 或 DPO。

360-LLaMA-Factory 的實現(xiàn)經(jīng)過了嚴(yán)格的正確性驗證，已在主倉 Pull Request 中審核過。正式合并進(jìn) LLaMA-Factory 主倉之前，可先使用 360-LLaMA-Factory。

1、項目背景與項目簡介

360 智腦早在 2023 年就開始了長文本大模型的研發(fā)，到目前為止已經(jīng)成功應(yīng)用于開源并更新了兩個版本的 360Zhinao-7B-Chat-360k 模型，以及近日發(fā)布的長思維鏈推理模型 360gpt2-o1。在 360-LLaMA-Factory 中，我們將 360 智腦內(nèi)部長序列后訓(xùn)練能力系統(tǒng)性地整合進(jìn)了 LLaMA-Factory 中，用戶僅需額外添加一行代碼，即可進(jìn)行理論上任意長度的長序列后訓(xùn)練（增加序列并行的 GPU 卡數(shù)即可）：

sequence_parallel_size: 16

在原先使用 LLaMA-Factory 的基礎(chǔ)上，只需額外增加一個參數(shù)

通過這種方式，360-LLaMA-Factory 將 LLaMA-Factory 的序列并行也做到了簡單易用和兼容并包，和 LLaMA-Factory 的其他功能完全兼容。

粗粒度地測試 8 卡 80G 的全參數(shù)后訓(xùn)練（不考慮除了 zero3-offload 和 gradient checkpointing 外的任何優(yōu)化技巧），360-LLaMA-Factory 至少可以訓(xùn)到 SFT 210k (7B) / 128k (72B) 和 DPO 84k (7B) / 46k (72B)。若加上注掉 logits = logits.float () 和 DPO 預(yù)計算等技巧，2 卡序列并行即可解決諸多常見的訓(xùn)練需求。360-LLaMA-Factory 讓序列并行也真正成為了簡單好用、效果也好的后訓(xùn)練工具。

作為開源社區(qū)的一份子，360-LLaMA-Factory 離不開 LLaMA-Factory、ring-flash-attention 和 EasyContext 等開源項目的開創(chuàng)性工作，我們的底層開發(fā)部分依賴了這些工作，但也有我們自己在具體實現(xiàn)方式上的不同和見解。我們相信我們的代碼實現(xiàn)已做到盡可能好的模塊化和盡可能少的原始代碼修改，且嚴(yán)格檢查過正確性，因此也已向 LLaMA-Factory 主倉提交了 Pull Request，初步審核通過。我們樂于同開源社區(qū)共建完善這項工作。

2、長序列及其后訓(xùn)練

2.1 長序列大模型的訓(xùn)練：預(yù)訓(xùn)練 vs 后訓(xùn)練

隨著大模型訓(xùn)練數(shù)據(jù)長度的增長，預(yù)訓(xùn)練和后訓(xùn)練平臺框架都需要支持長序列數(shù)據(jù)訓(xùn)練。

• 預(yù)訓(xùn)練階段，英偉達(dá)的 Megatron-LM 憑借豐富高效的并行策略與出色的 GPU 顯存優(yōu)化，成為主流框架，基于它的定制開發(fā)往往是最通用的解法， Megatron-LM 本身已實現(xiàn)了序列并行（Megatron-LM 稱之為 context parallelism，其他工作一般稱為 sequence parallelism）。
  
  
• 后訓(xùn)練階段情況相對復(fù)雜。后訓(xùn)練算法多樣，如 DPO 就有諸多變種，且訓(xùn)練需求靈活多變，不同場景對算法、資源、并行性等要求各異。因此，至今沒有一個框架能在并行策略、后訓(xùn)練算法、GPU 顯存優(yōu)化和易用性這四個關(guān)鍵方面做到近乎完美的兼容。雖有框架在部分方面表現(xiàn)尚可，但總體仍存在短板，這也限制了模型在長序列數(shù)據(jù)后訓(xùn)練上的進(jìn)一步發(fā)展。

2.2 長序列的通解 —— 序列并行及其難點

長序列后訓(xùn)練面臨的關(guān)鍵瓶頸是：序列長度增加時，激活顯存會大幅上升。雖然有 unsloth、liger kernel、LoRA 等多種降低顯存占用的技巧，但均未從根本上解決序列長度增加的本質(zhì)問題，其效果存在明確上限。

序列并行（sequence parallelism）被認(rèn)為是解決長序列訓(xùn)練問題的通解，它通過把一條長序列切分到不同的顯卡上進(jìn)行計算，從而避免了每張顯卡處理過長的序列，從根本上解決了 “每張顯卡處理的序列長度增加” 的問題。然而，序列并行的實現(xiàn)難度較大，需要在切分后的序列之間進(jìn)行通信計算 attention，需要侵入修改原始的 attention 函數(shù)。在開源的 Megatron-LM 中，序列并行也是所有并行策略中最后才添加的，LLaMA-Factory 之前還沒有支持序列并行。

2.3 序列并行后訓(xùn)練的相關(guān)工作

我們調(diào)研了其他一些支持序列并行的開源框架，有些實現(xiàn)上有錯或小 bug、導(dǎo)致支持的后訓(xùn)練算法不全；有些更新維護(hù)不及時、訓(xùn)練較新的模型不方便、顯示進(jìn)度條等易用性不足。有的與 LLaMA-Factory 相比繼承依賴更少，支持功能較少但更干凈、更適合定制開發(fā)，有不同的使用場景。此外，各家的序列并行具體實現(xiàn)也不盡相同。詳見下面的表 1 和 GitHub README，有未調(diào)研到的也請包涵并聯(lián)系 360-LLaMA-Factory。

表 1：一些支持序列并行的后訓(xùn)練框架對比

3、360-LLaMA-Factory 框架解析

360-LLaMA-Factory 系統(tǒng)性地為 LLaMA-Factory 增加了序列并行的支持。以下將簡要介紹 360-LLaMA-Factory 框架中的模塊化修改和執(zhí)行流程。

3.1 360-LLaMA-Factory 的框架和模塊化封裝

360-LLaMA-Factory 將序列并行的代碼做到了盡可能好的模塊化和盡可能少的原始代碼修改。

我們認(rèn)為序列并行本質(zhì)上應(yīng)認(rèn)為是對模型的修改，因此在 model_args 中增加了參數(shù)并抽象為 apply_sequence_parallel 修改模型的函數(shù)。

# src/llamafactory/model/loader.py
sequence_parallel_group = apply_sequence_parallel(model_args)  # 序列并行monkey patch，改動attention計算
...
model.sequence_parallel_group = sequence_parallel_group  # 維護(hù)模型的序列并行組，不開則為None

相應(yīng)地，數(shù)據(jù)處理部分也要相應(yīng)地修改，我們將 zigzag ring attention 所需的數(shù)據(jù)處理抽象成了一個 decorator，裝飾原來的數(shù)據(jù)處理函數(shù)。背后，這會將先 shuffle、packing、預(yù)處理好的數(shù)據(jù)進(jìn)一步做好序列并行的準(zhǔn)備：先將每行 pad 或截斷到指定的訓(xùn)練長度，再按 zigzag 切分并按順序?qū)懭霐?shù)據(jù)集，最后在訓(xùn)練時用 SequentialSampler 讀取訓(xùn)練數(shù)據(jù)。

# src/llamafactory/data/loader.py
@sequence_parallel_decorator
def get_dataset(...)

loss 計算則需要在 Trainer 中做序列并行組內(nèi)的 reduce 匯總和計算。

# src/llamafactory/train/sft/trainer.py
dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=sp_group)
dist.all_reduce(label_num, op=dist.ReduceOp.SUM, group=sp_group)
loss /= label_num
# src/llamafactory/train/dpo/trainer.py
dist.all_reduce(policy_chosen_logps, op=dist.ReduceOp.SUM, group=sp_group)
dist.all_reduce(policy_rejected_logps, op=dist.ReduceOp.SUM, group=sp_group)
dist.all_reduce(reference_chosen_logps, op=dist.ReduceOp.SUM, group=sp_group)
dist.all_reduce(reference_rejected_logps, op=dist.ReduceOp.SUM, group=sp_group)

3.2 360-LLaMA-Factory 的 SFT 和 DPOTrainer

除了統(tǒng)一的模塊化抽象，序列并行也需要對 360-LLaMA-Factory 的 Trainer 稍做定制化的修改，以適配各底層庫。針對最普遍的后訓(xùn)練需求 SFT 和 DPO（及其變種），我們對 360-LLaMA-Factory 中的 SFT 和 DPOTrainer 做了盡可能少且清晰的修改。

其中，dummy_forward 是因為我們發(fā)現(xiàn)基于目前的底層序列并行實現(xiàn)，在第一次 forward 時 DPO loss 不等于 log (sigmoid (0))，但學(xué)習(xí)率設(shè)為 0 時之后的 DPO loss 全都等于。因此，訓(xùn)練最開始時先做且僅做一次假前傳，不對正式訓(xùn)練循環(huán)造成任何影響。

從 SFT 和 DPO 的序列并行對比圖中，可以清晰地看出 360-LLaMA-Factory 序列并行帶來的改動。

圖 3：360-LLaMA-Factory SFT 序列并行對比

圖 4：360-LLaMA-Factory DPO 序列并行對比

4、360-LLaMA-Factory 效果驗證

內(nèi)部 360-LLaMA-Factory 的早期版本已訓(xùn)練了開源的 360Zhinao2-7B-Chat-360k。

為驗證本次開源的 360-LLaMA-Factory 的正確性，我們用總量為 30 條的小數(shù)據(jù)集，驗證了序列并行開與不開的對比情況下，訓(xùn)練曲線的差別，以此來確保 360-LLaMA-Factory 所有實現(xiàn)的正確性。從下圖可見，序列并行對訓(xùn)練曲線的影響幾乎可以忽略不計，DPO 稍有一定數(shù)值誤差，但我們也仔細(xì)檢查了該誤差與 DeepSpeed Ulysses 的誤差范圍一致，很可能部分是并行計算本身的隨機(jī)性導(dǎo)致的，亦可參考 ring-flash-attention 的詳細(xì)說明。

圖 5：360-LLaMA-Factory SFT 和 DPO 序列并行開關(guān)對比

為便于對比效果，我們基于第三方全尺寸開源模型粗粒度壓測了最大訓(xùn)練長度，如下表 2、表 3 所示，可見 8 卡 80G 的序列并行上限已可滿足幾十至幾百 k 超長序列的需求：

表 2：第三方開源模型多尺寸 SFT 長度壓測

表 3：第三方開源模型多尺寸 DPO 長度壓測

5、總結(jié)

360 智腦開源了 360-LLaMA-Factory，支持了序列并行，僅需額外 1 個參數(shù)控制?；?LLaMA-Factory 和 ring-flash-attention 開發(fā)，360-LLaMA-Factory 的實現(xiàn)模塊化、效果正確且在長序列上有效。

歡迎開發(fā)者們使用和開發(fā)。在本倉庫（https://github.com/Qihoo360/360-LLaMA-Factory）下提交序列并行相關(guān)的 issue 或 PR 即可。

也歡迎研究者們，尤其是依賴長序列大模型的研究者們，在研究中使用我們的代碼，可以這樣引用我們的工作：

@software{360-llama-factory,
  author = {Haosheng Zou, Xiaowei Lv, Shousheng Jia and Xiangzheng Zhang},
  title = {360-LLaMA-Factory},
  url = {https://github.com/Qihoo360/360-LLaMA-Factory},
  year = {2024}
}