近日，摩尔线程正式推出 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深度学习生态体系。

源码地址：点击下载