据本网11月28日消息,摩尔线程正式发布了PyTorch深度学习框架的最新版本MUSA扩展库——Torch-MUSA v2.7.0,在功能集成、性能优化、硬件支持等方面取得了进一步突破。
值得一提的是,在短短一个月内,Torch-MUSA就连续完成了v2.5.0、v2.7.0两次版本更新。的另一个变化是,从v2.5.0开始,Torch-MUSA版本号与PyTorch主版本号同步,让开发者更容易识别和管理版本。
v2.7.0版本进一步集成了muSolver、muFFT等计算加速库,显着提升了复杂计算任务的执行效率;
新增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版本主要更新内容新增特性 动态双铸
用户可以通过设置环境变量export TORCH_USE_MUSA_DOUBLE_CAST=1来启用Float64数据类型运算符的动态转换功能。 torch_musa 将使用float32 作为计算数据类型。
分布式检查点
支持多个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、volution3d、layer_norm等操作的执行效率。
系统功能增强 打开torch.musa.mccl.version()接口;
支持getCurrentMUSABlasHandle和getCurrentMUSABlasLtHandle;
优化FSDP2管道并行策略,减少训练内存占用。