近日，摩爾線程發(fā)布其面向PyTorch深度學(xué)習(xí)框架的MUSA擴(kuò)展庫(kù)——Torch-MUSA v2.1.1。該版本在v2.1.0的基礎(chǔ)上，進(jìn)一步擴(kuò)展了對(duì)大規(guī)模深度學(xué)習(xí)模型訓(xùn)練與推理的支持能力，并在編譯優(yōu)化、計(jì)算性能和算子生態(tài)等方面實(shí)現(xiàn)顯著增強(qiáng)。

Torch-MUSA v2.1.1核心特性

全新的v2.1.1版本集成了v2.1.0版本的多項(xiàng)重要特性，主要包括：

▼FSDP2分布式訓(xùn)練支持：引入基于DTensor的每參數(shù)分片策略，結(jié)合摩爾線程MUSA架構(gòu)GPU深度優(yōu)化，大幅提升大模型分布式訓(xùn)練效率。

▼AOTInductor集成：支持模型針對(duì)摩爾線程計(jì)算卡進(jìn)行提前編譯（Ahead-of-Time編譯），優(yōu)化推理部署流程。

▼內(nèi)存管理優(yōu)化：新增可插拔的MUSA內(nèi)存統(tǒng)一系統(tǒng)分配器后端，有效緩解GPU內(nèi)存碎片化問(wèn)題，降低訓(xùn)練峰值內(nèi)存占用。

▼Triton-MUSA后端增強(qiáng)：為torch.compile提供完整的Triton-MUSA后端支持，減少設(shè)備特定代碼，提升編譯效率。

▼PyTorch3D原生支持：擴(kuò)展對(duì)PyTorch3D計(jì)算庫(kù)的支持，確保3D深度學(xué)習(xí)工作流在摩爾線程計(jì)算卡上的流暢運(yùn)行。

本次發(fā)布的 v2.1.1 版本，聚焦于進(jìn)一步釋放硬件潛力，并在關(guān)鍵計(jì)算路徑上實(shí)現(xiàn)加速：

▼引入MUSA Graphs后端：為torch.compile新增MUSA Graphs后端，進(jìn)一步降低主機(jī)開(kāi)銷(xiāo)，并通過(guò)MUSA Graph技術(shù)實(shí)現(xiàn)端到端的計(jì)算加速。

▼集成muSolver計(jì)算庫(kù)：將muSolver深度集成至多個(gè)線性代數(shù)算子的后端，提升計(jì)算效率和數(shù)值穩(wěn)定性。

▼新增融合算子：支持使用MUSA原生的FusedAdamW和FusedAdam，通過(guò)算子融合減少內(nèi)核啟動(dòng)開(kāi)銷(xiāo)，提升訓(xùn)練效率。

▼算子性能優(yōu)化與擴(kuò)展：持續(xù)優(yōu)化算子性能，原生支持算子數(shù)量增加至950+。

▼擴(kuò)展基準(zhǔn)測(cè)試模塊：加強(qiáng)對(duì)更多算子用例的性能評(píng)估支持，方便用戶進(jìn)行全面的性能對(duì)比與分析。

體驗(yàn)與生態(tài)兼容持續(xù)優(yōu)化

從上一版本開(kāi)始，Torch-MUSA在用戶體驗(yàn)和生態(tài)兼容性方面實(shí)現(xiàn)了顯著的改進(jìn)與優(yōu)化：

▼簡(jiǎn)化集成流程：從v2.1.0開(kāi)始支持torch_musa自動(dòng)加載，用戶無(wú)需在Python腳本中顯式調(diào)用"import torch_musa"導(dǎo)入即可使用，簡(jiǎn)化了集成流程。

▼生態(tài)兼容：v2.1.0和v2.2.0均基于PyTorch 2.5.0構(gòu)建，確保了與主流PyTorch生態(tài)工具鏈的兼容性。

▼持續(xù)性能優(yōu)化：在FP8混合精度訓(xùn)練（v2.0.0引入的特性）、內(nèi)核融合、通信優(yōu)化等方面進(jìn)行了持續(xù)優(yōu)化，為用戶帶來(lái)更加穩(wěn)定、高效的開(kāi)發(fā)體驗(yàn)。

▼Torch-MUSA 開(kāi)源地址：

https://github.com/MooreThreads/torch_musa

▼版本迭代：

v1.1.0：初次發(fā)布，支持PyTorch 2.0，提供基礎(chǔ)張量操作和常見(jiàn)神經(jīng)網(wǎng)絡(luò)層的MUSA加速。

v1.2.0：進(jìn)一步擴(kuò)展算子支持，支持了完整功能的Profiler、MUSA Extension，并增加了Torch-MUSA專(zhuān)有特性如compare_tool、musa_converter，幫助用戶更快的定位模型精度問(wèn)題。

v1.3.0：支持PyTorch 2.2.0，性能進(jìn)一步提升，支持FSDP，支持更復(fù)雜的模型和更大規(guī)模的數(shù)據(jù)處理。

v2.0.0：首次原生支持FP8數(shù)據(jù)類(lèi)型，支持PyTorch 2.5.0，新增MUSA虛擬內(nèi)存管理技術(shù)優(yōu)化大模型訓(xùn)練內(nèi)存效率，通過(guò)集成Triton-MUSA后端顯著提升torch.compile編譯效率，支持MUSA Graph技術(shù)。

v2.1.0：支持FSDP2，集成AOTInductor，支持PyTorch3D，為torch.compile提供完整的Triton-MUSA后端支持，支持torch_musa自動(dòng)加載。

v2.1.1：torch.compile支持MUSA Graphs后端，原生支持算子超950個(gè)，并優(yōu)化了大量算子性能。

▼未來(lái)計(jì)劃：

Torch-MUSA將繼續(xù)跟進(jìn)PyTorch的版本更新，計(jì)劃下一版本支持PyTorch 2.7，并持續(xù)優(yōu)化性能與功能，為基于MUSA架構(gòu)的國(guó)產(chǎn)全功能GPU構(gòu)建更強(qiáng)大的深度學(xué)習(xí)生態(tài)。

▼ 關(guān)于摩爾線程

摩爾線程以全功能GPU為核心，致力于向全球提供加速計(jì)算的基礎(chǔ)設(shè)施和一站式解決方案，為各行各業(yè)的數(shù)智化轉(zhuǎn)型提供強(qiáng)大的AI計(jì)算支持。

我們的目標(biāo)是成為具備國(guó)際競(jìng)爭(zhēng)力的GPU領(lǐng)軍企業(yè)，為融合人工智能和數(shù)字孿生的數(shù)智世界打造先進(jìn)的加速計(jì)算平臺(tái)。我們的愿景是為美好世界加速。