11 月 28 日消息,近日,摩爾線程正式發佈 PyTorch 深度學習框架的 MUSA 擴展庫 ——Torch-MUSA v2.7.0,新版本在功能集成、性能優化與硬體支持方面實現進一步突破。Torch-MUSA 在短短一個月內,連續完成 v2.5.0 和 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 版本主要更新內容:
新增特性
新增 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,並進一步優化性能與功能。
Torch-MUSA 開源地址:https://github.com/MooreThreads/torch_musa
來源:中國IT之家