近日，摩爾線程正式發(fā)布PyTorch深度學(xué)習(xí)框架的MUSA擴(kuò)展庫——Torch-MUSA v2.7.0，新版本在功能集成、性能優(yōu)化與硬件支持方面實(shí)現(xiàn)進(jìn)一步突破。Torch-MUSA在短短一個(gè)月內(nèi)，連續(xù)完成v2.5.0和v2.7.0兩次版本更新，充分體現(xiàn)了摩爾線程在MUSA生態(tài)建設(shè)上的持續(xù)投入與快速迭代能力。

▼Torch-MUSA開源地址：

https://github.com/MooreThreads/torch_musa

自v2.5.0起，Torch-MUSA版本號(hào)已與PyTorch主版本號(hào)保持同步，便于開發(fā)者進(jìn)行版本識(shí)別與管理。新版本進(jìn)一步集成muSolver與muFFT等計(jì)算加速庫，顯著提升復(fù)雜計(jì)算任務(wù)的執(zhí)行效率；同時(shí)新增對(duì)統(tǒng)一內(nèi)存設(shè)備（Unified Memory）的UMM支持，有效優(yōu)化內(nèi)存使用效率。

此外，新版本繼續(xù)保持與最新MUSA SDK的兼容性，支持使用MUSA SDK 4.2.0至4.3.0及更高版本進(jìn)行編譯。目前Torch-MUSA專屬支持的算子總數(shù)已超過1050個(gè)，系統(tǒng)在性能與穩(wěn)定性方面均實(shí)現(xiàn)進(jìn)一步提升，為大模型訓(xùn)練與推理提供了更高效、更可靠的底層支持。

v2.7.0版本主要更新內(nèi)容

新增特性

▼動(dòng)態(tài)雙精度轉(zhuǎn)換（Dynamic Double Cast）

用戶可通過設(shè)置環(huán)境變量export TORCH_USE_MUSA_DOUBLE_CAST=1，開啟Float64數(shù)據(jù)類型算子的動(dòng)態(tài)轉(zhuǎn)換功能，torch_musa將使用float32作為計(jì)算數(shù)據(jù)類型。

▼分布式檢查點(diǎn)（Distributed Checkpoint）

支持從多個(gè)rank并行加載和保存模型，顯著加速檢查點(diǎn)的保存與加載過程。目前已支持分布式檢查點(diǎn)的異步保存功能。

功能增強(qiáng)

▼新增Poisson、binomial、_standard_gamma、_sample_dirichlet、vdot、upsample（1d、2d、3d、with aa）、flash_attention、transformer_encoder_layer 等多個(gè)實(shí)用算子，MUSA專屬支持的算子總數(shù)突破1050個(gè)。

▼通過升級(jí)PyTorch底層支持，torch.compile與AOTInductor功能進(jìn)一步增強(qiáng)；

▼默認(rèn)啟用TF32計(jì)算模式，提升浮點(diǎn)運(yùn)算效率；

▼優(yōu)化性能分析工具Kineto的穩(wěn)定性，并將其適配版本升級(jí)至2.7.0；

▼繼續(xù)優(yōu)化FSDP2流水線并行策略，進(jìn)一步降低內(nèi)存占用。

v2.5.0版本主要更新內(nèi)容

新增特性

▼新增muFFT與muSolver庫集成，大幅擴(kuò)展計(jì)算能力；

▼在面向邊緣計(jì)算的SoC設(shè)備中支持統(tǒng)一內(nèi)存管理，基于Arm 架構(gòu)的UMA（統(tǒng)一內(nèi)存尋址）設(shè)計(jì)，實(shí)現(xiàn)GPU與CPU共享同一物理內(nèi)存空間，顯著降低模型運(yùn)行過程中的內(nèi)存開銷，具體包括：

消除GPU端重復(fù)內(nèi)存分配；

減少主機(jī)與設(shè)備間的內(nèi)存拷貝；

GPU可直接訪問由CPU分配器申請(qǐng)的內(nèi)存空間。

算子擴(kuò)展與性能優(yōu)化

▼新增支持包括ilshift、irshift、replication_pad1d_bwd、angle、ctcLossTensor、ctcLossTensorBwd、logit、amin/amax/prod.dim_int、glu_bwd等多個(gè)算子；

▼新增基礎(chǔ) Sparse(CSR) 操作支持；

▼擴(kuò)充量化算子支持范圍；

▼修復(fù)torch.norm形狀錯(cuò)誤問題；

▼支持reduce_sum的uint8輸入與int64輸出；

▼C++擴(kuò)展新增支持tensor.is_musa()方法；

▼修復(fù)空輸入下argmax/argmin的異常行為；

▼優(yōu)化var/std、pad、convolution3d、layer_norm等操作的執(zhí)行效率。

系統(tǒng)功能增強(qiáng)

▼開放torch.musa.mccl.version()接口；

▼支持getCurrentMUSABlasHandle與getCurrentMUSABlasLtHandle；

▼優(yōu)化FSDP2流水線并行策略，降低訓(xùn)練內(nèi)存占用。

未來計(jì)劃

Torch-MUSA將繼續(xù)跟進(jìn)PyTorch的版本更新，計(jì)劃下一版本支持PyTorch 2.9.0，并進(jìn)一步優(yōu)化性能與功能，持續(xù)構(gòu)建和完善基于MUSA架構(gòu)國(guó)產(chǎn)全功能GPU的深度學(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)。我們的愿景是為美好世界加速。