9月8日，上海人工智能實驗室(上海AI實驗室)開源書生大模型新一代訓(xùn)練引擎XTuner V1。

XTuner V1是伴隨上海AI實驗室“通專融合”技術(shù)路線的持續(xù)演進(jìn)，以及書生大模型研發(fā)實踐而成長起來的新一代訓(xùn)練引擎。相較于傳統(tǒng)的3D并行訓(xùn)練引擎，XTuner V1不僅能應(yīng)對更加復(fù)雜的訓(xùn)練場景，還具備更快的訓(xùn)練速度，尤其在超大規(guī)模稀疏混合專家(MoE，mixture of experts)模型訓(xùn)練中優(yōu)勢顯著。

同時，為了進(jìn)一步探究XTuner V1訓(xùn)練方案的上限，研究團(tuán)隊與昇騰團(tuán)隊在昇騰384超節(jié)點(Atlas 900 A3 SuperPoD)上進(jìn)行聯(lián)合優(yōu)化，充分利用昇騰384超節(jié)點硬件特性，實現(xiàn)更高的模型算力利用率(MFU，model FLOPS utilization)。相比業(yè)界其他產(chǎn)品，昇騰384超節(jié)點的訓(xùn)練吞吐提升5%以上，MFU提升20%以上，該項研究成果技術(shù)報告也將于近期發(fā)布。

除了訓(xùn)練框架，書生大模型研發(fā)中使用的AIOps工具DeepTrace與ClusterX也將一并開源，為大規(guī)模分布式訓(xùn)練提供全方位保障。

XTuner V1:

https://github.com/InternLM/xtuner

DeepTrace:

https://github.com/DeepLink-org/DeepTrace

ClusterX:

https://github.com/InternLM/clusterx

目前開源社區(qū)主流的訓(xùn)練方案主要分為兩類：

DeepSpeed/PyTorch FSDP(完全分片數(shù)據(jù)并行，F(xiàn)ully Shard Data Parallel)：通信量大但使用簡單，尤其適合稠密型模型訓(xùn)練，開發(fā)者無需具備專業(yè)的AI Infra知識，也能開發(fā)出接近最優(yōu)性能的訓(xùn)練系統(tǒng);

3D并行：通信量小但使用復(fù)雜，開發(fā)者需要具備專業(yè)的AI Infra知識，針對不同硬件和訓(xùn)練場景進(jìn)行針對性調(diào)優(yōu)，尤其適用MoE模型訓(xùn)練。

為了同時兼顧易用性、高性能與超大模型訓(xùn)練，XTuner V1基于PyTorch FSDP進(jìn)行開發(fā)，并針對FSDP通信量大的固有缺陷，進(jìn)行了系列優(yōu)化，可支持1T參數(shù)量級MoE模型訓(xùn)練，并首次在200B以上量級的混合專家模型上，實現(xiàn)訓(xùn)練吞吐超越傳統(tǒng)的3D并行訓(xùn)練方案。

針對當(dāng)前主流的MoE后訓(xùn)練需求，XTuner V1未通過序列并行方式，實現(xiàn)200B量級MoE模型單次forward-backward可處理64k序列長度，更適合當(dāng)下流行的強化學(xué)習(xí)訓(xùn)練場景;對專家并行依賴小，長序列訓(xùn)練時受專家不均衡影響小，200B量級MoE無需專家并行，600B MoE只需節(jié)點內(nèi)專家并行，更適合現(xiàn)代MoE Dropless訓(xùn)練模式;大規(guī)模長短序列混訓(xùn)場景提速2倍以上，數(shù)據(jù)并行負(fù)載均衡，大幅減小因需序列長度不均衡導(dǎo)致的計算空泡。

多維度技術(shù)優(yōu)化

專為“超大模型”而生

XTuner V1之所以能在超大模型訓(xùn)練中展現(xiàn)出卓越的性能，核心在于它在顯存、通信、負(fù)載等多個維度進(jìn)行了系統(tǒng)性優(yōu)化。這些優(yōu)化協(xié)同作用，不僅帶來了性能的跨越式提升，還兼顧了易用性、通用性與擴(kuò)展性。

顯存優(yōu)化：Pytorch FSDP與3D并行最大的差異在于重計算。針對計算損失函數(shù)時的計算圖，XTuner V1基于Liger-Kernel中的Chunk-wise Loss，擴(kuò)展支持了更多種類的損失函數(shù)，能夠支持昇騰NPU;針對重計算保留的激活值，XTuner V1借鑒了昇騰MindSpeed中的Async Checkpointing Swap。最終，無需借助序列并行技術(shù)，實現(xiàn)200B參數(shù)量級MoE模型訓(xùn)練64K長度序列。

通信掩蓋：得益于極致的顯存優(yōu)化，XTuner V1可以讓單次迭代的最大序列長度提升數(shù)倍，從而增加每層計算的耗時，掩蓋參數(shù)聚合的通信耗時。針對因顯存或通信帶寬受限，無法實現(xiàn)通信掩蓋的訓(xùn)練場景，XTuner V1通過Intra-Node Domino-EP來降低每一層聚合參數(shù)的通信量，同時掩蓋因引入專家并行帶來的額外通信開銷。

DP負(fù)載均衡：由于XTuner V1中沒有引入TP、PP等并行策略，相同卡數(shù)下，數(shù)據(jù)并行的維度會遠(yuǎn)大于3D并行。為了緩解變長注意力帶來的計算空泡，并盡可能不影響數(shù)據(jù)的訓(xùn)練順序，會對每n個step內(nèi)的已拼接好的序列進(jìn)行排序，讓每次計算時，不同DP的最長子序列長度是接近的。

基于昇騰384超節(jié)點深度優(yōu)化

訓(xùn)練效率超業(yè)界產(chǎn)品

為了進(jìn)一步探究XTuner V1訓(xùn)練方案的上限，上海人工智能實驗室XTuner團(tuán)隊聯(lián)合昇騰技術(shù)團(tuán)隊在超節(jié)點上進(jìn)行深度優(yōu)化，充分利用昇騰384超節(jié)點硬件特性，實現(xiàn)了更高M(jìn)FU。

昇騰384超節(jié)點通過高速總線連接多顆NPU，突破互聯(lián)瓶頸，讓超節(jié)點像一臺計算機(jī)一樣工作，更加適合FSDP訓(xùn)練：

更高的通信帶寬：最大可實現(xiàn)384顆NPU點到點超大帶寬互聯(lián)，F(xiàn)SDP All Gather耗時僅為業(yè)界產(chǎn)品的1/4~1/3，更容易實現(xiàn)計算-通信掩蓋

計算通信解耦：通過專用硬化調(diào)度和傳輸卸載，實現(xiàn)不占用計算核的高效數(shù)據(jù)通信，F(xiàn)SDP計算通信掩蓋時不會影響計算速度

超節(jié)點高速互連：CPU和NPU通過高速總線互聯(lián)，實現(xiàn)更大帶寬，Checkpointing Swap的開銷更小

除硬件固有優(yōu)勢外，昇騰還從通信、內(nèi)存、計算、框架、工具等維度對基于昇騰384超節(jié)點的MoE訓(xùn)練進(jìn)行了全方位的加持：

Cube調(diào)優(yōu)：對于模型中集中了大量計算任務(wù)的GroupedMatmul算子進(jìn)行分析，發(fā)現(xiàn)內(nèi)部搬運帶寬已經(jīng)擁塞但cube利用率還有提升空間。針對此問題，聯(lián)合研發(fā)團(tuán)隊重點優(yōu)化GroupedMatmul算子分塊邏輯，根據(jù)不同輸入進(jìn)行動態(tài)分塊Tiling策略優(yōu)化搬運效率。同時，根據(jù)場景的不同細(xì)化Cache策略，提高Cache命中率從而提升性能。

QoS調(diào)優(yōu)：QoS(Quality of Service)即服務(wù)質(zhì)量。在有限的帶寬資源下，QoS為各種業(yè)務(wù)分配帶寬，為業(yè)務(wù)提供端到端的服務(wù)質(zhì)量保證。大規(guī)模訓(xùn)練過程中，計算流、通信流、swap流都會存在內(nèi)存訪問，并發(fā)的訪問會導(dǎo)致內(nèi)存帶寬擁塞，從而影響整體性能。通過適當(dāng)調(diào)低通信的內(nèi)存訪存優(yōu)先級，可以減少計算的搬運時間，從而優(yōu)化端到端性能。

跨流內(nèi)存復(fù)用：在FSDP計算流和通信流異步重疊的場景中，Ascend Extension for PyTorch(PTA)中默認(rèn)的跨流內(nèi)存優(yōu)化會導(dǎo)致顯存不能及時釋放，需要開啟PTA中進(jìn)階版的跨流內(nèi)存復(fù)用機(jī)制(MULTI_STREAM_MEMORY_REUSE=2)，可以顯著降低顯存峰值。

集群性能工具高效診斷：借助MindStudio全流程工具鏈中的msprof-analyze性能分析工具與MindStudio Insight可視化工具，開發(fā)者可以充分利用其強大的數(shù)據(jù)分析與可視化能力，在分鐘級時間內(nèi)精準(zhǔn)識別訓(xùn)練過程中的“快慢卡”現(xiàn)象根因，快速定位出性能瓶頸，顯著提升大集群調(diào)優(yōu)效率。

書生大模型工具鏈研發(fā)團(tuán)隊現(xiàn)已將Xtuner V1的工作全部開源，希望為學(xué)術(shù)界與工業(yè)界提供高性能、低門檻、易擴(kuò)展的大模型訓(xùn)練方案，豐富開源社區(qū)的訓(xùn)練工具生態(tài)，為超大模型研發(fā)和應(yīng)用提供堅實易用的基礎(chǔ)設(shè)施。

未來，在研究范式創(chuàng)新及模型能力提升的基礎(chǔ)上，上海AI實驗室將持續(xù)推進(jìn)書生大模型及其全鏈條工具體系的開源，支持免費商用，同時提供線上開放服務(wù)，與各界共同擁抱更廣闊的開源生態(tài)，共促大模型產(chǎn)業(yè)繁榮。