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

FP8原生支持，國產GPU的技術突破

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

三大關鍵優(yōu)化，提升AI計算效率

Torch-MUSA v2.0.0在MUSA計算平臺引入多項創(chuàng)新功能，進一步提升深度學習任務的執(zhí)行效率：

▼新增虛擬內存管理支持：

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

▼新增MUSA Graph支持：

MUSA Graph技術將多個MUSA內核整合到一個圖中，通過單次CPU調度大幅減少啟動開銷，提升計算效率，同時與CUDA Graph接口高效兼容。

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

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

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

Torch-MUSA v2.0.0在完整支持PyTorch 2.2.0的基礎上，新增了對PyTorch 2.5.0的支持，使開發(fā)者能夠在基于MUSA Compute Capability 3.1計算架構的全功能GPU上，無縫運行新版本的PyTorch，享受更高效的AI計算體驗。

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

▼Torch-MUSA開源地址：

https://github.com/MooreThreads/torch_musa

▼功能特性：

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

▼版本迭代：

MUSA Graph技術將多個MUSA內核整合到一個圖中，通過單次CPU調度大幅減少啟動開銷，提升計算效率，同時與CUDA Graph接口高效兼容。

v1.1.0：初次發(fā)布，支持PyTorch 2.0，提供基礎張量操作和常見神經網絡層的MUSA加速。

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

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

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

▼未來計劃：

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

▼關于摩爾線程

摩爾線程成立于2020年10月，以全功能GPU為核心，致力于向全球提供加速計算的基礎設施和一站式解決方案，為各行各業(yè)的數智化轉型提供強大的AI計算支持。

我們的目標是成為具備國際競爭力的GPU領軍企業(yè)，為融合人工智能和數字孿生的數智世界打造先進的加速計算平臺。我們的愿景是為美好世界加速。