摩尔线程发布Torch-MUSA v2.7.0,持续增强AI模型训练与推理支持
近日,摩尔线程正式发布PyTorch深度学习框架的MUSA扩展库——Torch-MUSA v2.7.0,新版本在功能集成、性能优化与硬件支持方面实现进一步突破。Torch-MUSA在短短一个月内,连续完成v2.5.0和v2.7.0两次版本更新,充分体现了摩尔线程在MUSA生态建设上的持续投入与快速迭代能力。 ▼ Torch-MUSA开源地址: https://github.com/MooreThreads/torch_musa 自v2.5.0起,Torch-MUSA版本号已与PyTorch主版本号保持同步,便于开发者进行版本识别与管理。新版本进一步集成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(统一内存寻址)设计,实现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流水线并行策略,降低训练内存占用。 未来计划 Torch-MUSA将继续跟进PyTorch的版本更新,计划下一版本支持PyTorch 2.9.0,并进一步优化性能与功能,持续构建和完善基于MUSA架构国产全功能GPU的深度学习生态。 ▼ 关于摩尔线程 摩尔线程以全功能GPU为核心,致力于向全球提供加速计算的基础设施和一站式解决方案,为各行各业的数智化转型提供强大的AI计算支持。 我们的目标是成为具备国际竞争力的GPU领军企业,为融合人工智能和数字孪生的数智世界打造先进的加速计算平台。我们的愿景是为美好世界加速。 更多信息,请长按二维码关注了解