IT之家11月28日消息,近日,摩尔线程正式发布PyTorch深度学习框架的MUSA扩展库——Torch-MUSAv2.7.0,新版本在功能集成、性能优化与硬件支持方面实现进一步突破。Torch-MUSA在短短一个月内,连续完成v2.5.0和v2.7.0两次版本更新。
据介绍,自v2.5.0起,Torch-MUSA版本号已与PyTorch主版本号保持同步,便于开发者进行版本识别与管理。新版本进一步集成muSolver与muFFT等计算加速库,显著提升复杂计算任务的执行效率;同时新增对统一内存设备(UnifiedMemory)的UMM支持,有效优化内存使用效率。
此外,新版本继续保持与最新MUSASDK的兼容性,支持使用MUSASDK4.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流水线并行策略,降低训练内存占用。
IT之家从官方获悉,Torch-MUSA将继续跟进PyTorch的版本更新,计划下一版本支持PyTorch2.9.0,并进一步优化性能与功能。
Torch-MUSA开源地址:https://github.com/MooreThreads/torch_musa