近日，摩尔线程发布其面向PyTorch深度学习框架的MUSA扩展库——Torch-MUSA v2.1.1。该版本在v2.1.0的基础上，进一步扩展了对大规模深度学习模型训练与推理的支持能力，并在编译优化、计算性能和算子生态等方面实现显著增强。
Torch-MUSA v2.1.1核心特性
全新的v2.1.1版本集成了v2.1.0版本的多项重要特性，主要包括：
- FSDP2分布式训练支持：引入基于DTensor的每参数分片策略，结合摩尔线程MUSA架构GPU深度优化，大幅提升大模型分布式训练效率。
- AOTInductor集成：支持模型针对摩尔线程计算卡进行提前编译（Ahead-of-Time编译），优化推理部署流程。
- 内存管理优化：新增可插拔的MUSA内存统一系统分配器后端，有效缓解GPU内存碎片化问题，降低训练峰值内存占用。
- Triton-MUSA后端增强：为torch.compile提供完整的Triton-MUSA后端支持，减少设备特定代码，提升编译效率。
- PyTorch3D原生支持：扩展对PyTorch3D计算库的支持，确保3D深度学习工作流在摩尔线程计算卡上的流畅运行。
本次发布的 v2.1.1 版本，聚焦于进一步释放硬件潜力，并在关键计算路径上实现加速：
- 引入MUSA Graphs后端：为torch.compile新增MUSA Graphs后端，进一步降低主机开销，并通过MUSA Graph技术实现端到端的计算加速。
- 集成muSolver计算库：将muSolver深度集成至多个线性代数算子的后端，提升计算效率和数值稳定性。
- 新增融合算子：支持使用MUSA原生的FusedAdamW和FusedAdam，通过算子融合减少内核启动开销，提升训练效率。
- 算子性能优化与扩展：持续优化算子性能，原生支持算子数量超过950个，大幅提升模型的适配性。
- 扩展基准测试模块：加强对更多算子用例的性能评估支持，方便用户进行全面的性能对比与分析。
体验与生态兼容持续优化
从上一版本开始，Torch-MUSA在用户体验和生态兼容性方面实现了显著的改进与优化：
- 简化集成流程：从v2.1.0开始支持torch_musa自动加载，用户无需在Python脚本中显式调用”import torch_musa”导入即可使用，简化了集成流程。
- 生态兼容：v2.1.0和v2.1.1均基于PyTorch 2.5.0构建，确保了与主流PyTorch生态工具链的兼容性。
- 持续性能优化：在FP8混合精度训练（v2.0.0引入的特性）、内核融合、通信优化等方面进行了持续优化，为用户带来更加稳定、高效的开发体验。
Torch-MUSA 开源地址：
https://github.com/MooreThreads/torch_musa
未来计划：
Torch-MUSA将继续跟进PyTorch的版本更新，计划下一版本支持PyTorch 2.7，并持续优化性能与功能，为基于MUSA架构的国产全功能GPU构建更强大的深度学习生态。
版权声明：本文内容转自互联网，本文观点仅代表作者本人。本站仅提供信息存储空间服务，所有权归原作者所有。如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至1393616908@qq.com 举报，一经查实，本站将立刻删除。
