摩爾線程近日正式發佈 PyTorch 深度學習框架的 MUSA 擴展庫——Torch-MUSA v2.7.0，新版本在功能集成、性能優化與硬件支持方面實現進一步突破。

自v2.5.0起，Torch-MUSA版本號已與PyTorch主版本號保持同步，便於開發者進行版本識別與管理。新版本進一步集成muSolver與muFFT等計算加速庫，顯著提升複雜計算任務的執行效率；同時新增對統一內存設備（Unified Memory）的UMM支持，有效優化內存使用效率。

此外，新版本繼續保持與最新MUSA SDK的兼容性，支持使用MUSA SDK 4.2.0至4.3.0及更高版本進行編譯。目前Torch-MUSA專屬支持的算子總數已超過1050個，系統在性能與穩定性方面均實現進一步提升。

v2.7.0版本主要更新內容

新增特性

• 動態雙精度轉換（Dynamic Double Cast）

用户可通過設置環境變量export TORCH_USE_MUSA_DOUBLE_CAST=1，開啓Float64數據類型算子的動態轉換功能，torch_musa將使用float32作為計算數據類型。

• 分佈式檢查點（Distributed Checkpoint）

支持從多個rank並行加載和保存模型，顯著加速檢查點的保存與加載過程。目前已支持分佈式檢查點的異步保存功能。

功能增強

• 新增Poisson、binomial、_standard_gamma、_sample_dirichlet、vdot、upsample（1d、2d、3d、with aa）、flash_attention、transformer_encoder_layer 等多個實用算子，MUSA專屬支持的算子總數突破1050個。 
• 通過升級PyTorch底層支持，torch.compile與AOTInductor功能進一步增強； 
• 默認啓用TF32計算模式，提升浮點運算效率； 
• 優化性能分析工具Kineto的穩定性，並將其適配版本升級至2.7.0；
• 繼續優化FSDP2流水線並行策略，進一步降低內存佔用。

v2.5.0版本主要更新內容

新增特性

• 新增muFFT與muSolver庫集成，大幅擴展計算能力；
• 在面向邊緣計算的SoC設備中支持統一內存管理，基於Arm 架構的UMA（統一內存尋址）設計，實現GPU與CPU共享同一物理內存空間，顯著降低模型運行過程中的內存開銷，具體包括：

  • 消除GPU端重複內存分配；

  • 減少主機與設備間的內存拷貝；

  • GPU可直接訪問由CPU分配器申請的內存空間。

算子擴展與性能優化

• 新增支持包括ilshift、irshift、replication_pad1d_bwd、angle、ctcLossTensor、ctcLossTensorBwd、logit、amin/amax/prod.dim_int、glu_bwd等多個算子；
• 新增基礎 Sparse(CSR) 操作支持；
• 擴充量化算子支持範圍；
• 修復torch.norm形狀錯誤問題；
• 支持reduce_sum的uint8輸入與int64輸出；
• C++擴展新增支持tensor.is_musa()方法；
• 修復空輸入下argmax/argmin的異常行為；
• 優化var/std、pad、convolution3d、layer_norm等操作的執行效率。

系統功能增強

• 開放torch.musa.mccl.version()接口；
• 支持getCurrentMUSABlasHandle與getCurrentMUSABlasLtHandle；
• 優化FSDP2流水線並行策略，降低訓練內存佔用。

接下來，Torch-MUSA 計劃繼續跟進PyTorch的版本更新，下一版本支持PyTorch 2.9.0，並進一步優化性能與功能，持續構建和完善基於MUSA架構國產全功能GPU的深度學習生態。