近日，摩尔线程正式推出 PyTorch 深度学习框架的 MUSA 扩展库新版本——Torch-MUSA 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.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（统一内存寻址）设计，实现 CPU 与 GPU 共享物理内存空间，带来以下优势：
- 避免 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深度学习生态体系。
源码地址：点击下载
