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

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

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

項(xiàng)目主頁(yè)：https://github.com/Qihoo360/360-LLaMA-Factory

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

sequence_parallel_size: 16

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

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

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

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

sequence_parallel_size: 16

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

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

粗粒度地測(cè)試 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ù)計(jì)算等技巧，2 卡序列并行即可解決諸多常見(jiàn)的訓(xùn)練需求。360-LLaMA-Factory 讓序列并行也真正成為了簡(jiǎn)單好用、效果也好的后訓(xùn)練工具。

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

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

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

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

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

2.2 長(zhǎng)序列的通解 —— 序列并行及其難點(diǎn)

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

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

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

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

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

3、360-LLaMA-Factory 框架解析

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

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

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

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

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

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

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

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

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

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

從 SFT 和 DPO 的序列并行對(duì)比圖中，可以清晰地看出 360-LLaMA-Factory 序列并行帶來(lái)的改動(dòng)。

圖 3：360-LLaMA-Factory SFT 序列并行對(duì)比

圖 4：360-LLaMA-Factory DPO 序列并行對(duì)比

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

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

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

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

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

表 2：第三方開(kāi)源模型多尺寸 SFT 長(zhǎng)度壓測(cè)

表 3：第三方開(kāi)源模型多尺寸 DPO 長(zhǎng)度壓測(cè)

5、總結(jié)

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

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

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

@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}
}