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

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

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

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

XTuner V1:

https://github.com/InternLM/xtuner

DeepTrace:

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

ClusterX:

https://github.com/InternLM/clusterx

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

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

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

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

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

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

專為“超大模型”而生

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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