近日，摩爾線程正式發(fā)布Torch-MUSA v2.0.0版本，這是其面向PyTorch深度學(xué)習(xí)框架的MUSA擴(kuò)展庫(kù)的重要升級(jí)。新版本基于MUSA Compute Capability 3.1計(jì)算架構(gòu)，支持原生FP8數(shù)據(jù)類型，支持PyTorch 2.5.0，并通過(guò)多項(xiàng)針對(duì)MUSA計(jì)算平臺(tái)的性能優(yōu)化，進(jìn)一步提升了對(duì)AI模型和大規(guī)模數(shù)據(jù)處理的支持能力。

FP8原生支持，國(guó)產(chǎn)GPU的技術(shù)突破

作為本次升級(jí)的核心亮點(diǎn)，Torch-MUSA v2.0.0率先在國(guó)產(chǎn)GPU上實(shí)現(xiàn)了對(duì)FP8數(shù)據(jù)類型的完整支持。FP8是當(dāng)前AI計(jì)算的一種前沿低精度格式，在支持原生FP8的GPU上，大語(yǔ)言模型（LLM）訓(xùn)練采用FP8混合精度可大幅提高GPU算力，顯著降低顯存占用。摩爾線程基于新一代 MUSA Compute Capability 3.1 計(jì)算架構(gòu)的全功能 GPU 原生支持 FP8 計(jì)算，這為 Torch-MUSA v2.0.0 實(shí)現(xiàn) FP8 矩陣乘法和分布式通信優(yōu)化提供了堅(jiān)實(shí)的基礎(chǔ)。依托這一底層架構(gòu)優(yōu)勢(shì)，Torch-MUSA v2.0.0 能夠充分發(fā)揮 FP8 的計(jì)算效能，顯著提升大語(yǔ)言模型訓(xùn)練和推理的效率。

三大關(guān)鍵優(yōu)化，提升AI計(jì)算效率

Torch-MUSA v2.0.0在MUSA計(jì)算平臺(tái)引入多項(xiàng)創(chuàng)新功能，進(jìn)一步提升深度學(xué)習(xí)任務(wù)的執(zhí)行效率：

▼新增虛擬內(nèi)存管理支持：

MUSA虛擬內(nèi)存管理技術(shù)能夠有效緩解GPU內(nèi)存碎片化問(wèn)題，降低模型訓(xùn)練過(guò)程中的峰值內(nèi)存占用，特別適用于FSDP、DeepSpeed和Megatron-LM等主流大模型訓(xùn)練框架。

▼新增MUSA Graph支持：

MUSA Graph技術(shù)將多個(gè)MUSA內(nèi)核整合到一個(gè)圖中，通過(guò)單次CPU調(diào)度大幅減少啟動(dòng)開(kāi)銷，提升計(jì)算效率，同時(shí)與CUDA Graph接口高效兼容。

▼torch.compile增加Triton后端支持：

為torch.compile提供了Triton-MUSA后端支持，開(kāi)發(fā)者可以直接使用PyTorch原生接口，獲得更高效的性能表現(xiàn)。

支持PyTorch 2.5.0，生態(tài)兼容性更強(qiáng)

Torch-MUSA v2.0.0在完整支持PyTorch 2.2.0的基礎(chǔ)上，新增了對(duì)PyTorch 2.5.0的支持，使開(kāi)發(fā)者能夠在基于MUSA Compute Capability 3.1計(jì)算架構(gòu)的全功能GPU上，無(wú)縫運(yùn)行新版本的PyTorch，享受更高效的AI計(jì)算體驗(yàn)。

Torch-MUSA已完全開(kāi)源，開(kāi)發(fā)者可通過(guò)訪問(wèn)GitHub獲取源代碼。摩爾線程鼓勵(lì)開(kāi)發(fā)者積極參與該項(xiàng)目的開(kāi)發(fā)與改進(jìn)，通過(guò)提交問(wèn)題報(bào)告（issue）或代碼修改申請(qǐng)（pull request）等方式，共同推動(dòng)Torch-MUSA以及MUSA軟件生態(tài)的持續(xù)進(jìn)步與創(chuàng)新。

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

https://github.com/MooreThreads/torch_musa

▼功能特性：

在Torch-MUSA中，用戶只需指定torch.device("musa")，即可輕松將現(xiàn)有的PyTorch模型遷移到MUSA架構(gòu)的GPU上運(yùn)行，無(wú)需大幅修改代碼。Torch-MUSA完全兼容PyTorch的自動(dòng)微分和動(dòng)態(tài)圖機(jī)制，支持多種常用的神經(jīng)網(wǎng)絡(luò)模塊及優(yōu)化算法，并加速了關(guān)鍵深度學(xué)習(xí)算子的計(jì)算。此外，Torch-MUSA還支持多種PyTorch特性，包括DDP、JIT、FSDP、Profiler、Extension等。

▼版本迭代：

MUSA Graph技術(shù)將多個(gè)MUSA內(nèi)核整合到一個(gè)圖中，通過(guò)單次CPU調(diào)度大幅減少啟動(dòng)開(kāi)銷，提升計(jì)算效率，同時(shí)與CUDA Graph接口高效兼容。

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專有特性如compare_tool、musa_converter，幫助用戶更快的定位模型精度問(wèn)題。

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

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

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

Torch-MUSA將繼續(xù)跟進(jìn)PyTorch的版本更新，計(jì)劃支持更高版本的PyTorch。摩爾線程期待與廣大開(kāi)發(fā)者和研究人員共同完善Torch-MUSA的功能，持續(xù)優(yōu)化性能，為基于MUSA架構(gòu)的國(guó)產(chǎn)全功能GPU構(gòu)建更強(qiáng)大的深度學(xué)習(xí)生態(tài)。

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

摩爾線程成立于2020年10月，以全功能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)。我們的愿景是為美好世界加速。