11月28日,摩尔线程正式推出PyTorch深度学习框架的最新版本——Torch-MUSA v2.7.0扩展库,这次更新在功能集成、运行效率与硬件适配方面均有显著提升。
值得一提的是,短短一个月内,Torch-MUSA已经接连完成了v2.5.0和v2.7.0两次版本迭代。
另一个重要变化是,从v2.5.0版本开始,Torch-MUSA的版本号与PyTorch主版本保持同步,方便开发者进行版本识别与管理。
v2.7.0版本进一步整合了muSolver、muFFT等计算加速库,使复杂计算任务的执行效率得到明显提升;
新增的统一内存设备(Unified Memory)支持UMM功能,有效优化了内存使用效率。
继续保持与最新MUSA SDK的兼容性,支持使用MUSA SDK 4.2.0至4.3.0及更高版本进行编译。
目前,Torch-MUSA专属支持的算子总数已突破1050个,系统在性能与稳定性方面均有进一步提升,为大模型训练与推理提供了更高效、更可靠的底层支持。
下一次版本升级将迭代至v2.9.0,届时将进一步完善性能表现与功能特性,持续推进基于MUSA架构的国产全功能GPU深度学习生态建设。
▼ Torch-MUSA开源地址:
https://github.com/MooreThreads/torch_musa

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流水线并行策略,降低训练内存占用。
