国产GPU突破!摩尔线程发布Torch-MUSA v2.0.0:FP8原生支持

pepsi 科技 5

5月9日消息,近日,国产GPU厂商摩尔线程正式发布Torch-MUSA v2.0.0版本,这是其面向PyTorch深度学习框架的MUSA扩展库的重要升级。

在Torch-MUSA中,用户只需指定torch.device("musa"),即可将现有的PyTorch模型迁移到MUSA架构的GPU上运行,无需大幅修改代码,目前Torch-MUSA已完全开源,可通过GitHub获取源代码。

作为本次升级的核心亮点,Torch-MUSA v2.0.0率先在国产GPU上实现了对FP8数据类型的完整支持。

FP8是当前AI计算的一种低精度格式,在支持原生FP8的GPU上,大语言模型训练采用FP8混合精度可大幅提高GPU算力,降低显存占用。

摩尔线程基于新一代MUSA Compute Capability 3.1计算架构的全功能GPU原生支持FP8计算,为Torch-MUSA v2.0.0实现FP8矩阵乘法和分布式通信优化提供了基础。

依托这一底层架构优势,Torch-MUSA v2.0.0能够充分发挥FP8的计算效能,显著提升大语言模型训练和推理的效率。

Torch-MUSA v2.0.0在MUSA计算平台引入多项创新功能,进一步提升深度学习任务的执行效率,主要包括:

1、新增虚拟内存管理支持:

MUSA虚拟内存管理技术能够有效缓解GPU内存碎片化问题,降低模型训练过程中的峰值内存占用,特别适用于FSDP、DeepSpeed和Megatron-LM等主流大模型训练框架。

2、新增MUSA Graph支持:

MUSA Graph技术将多个MUSA内核整合到一个图中,通过单次CPU调度大幅减少启动开销,提升计算效率,同时与CUDA Graph接口高效兼容。

3、torch.compile增加Triton后端支持:

为torch.compile提供了Triton-MUSA后端支持,开发者可以直接使用PyTorch原生接口,获得更高效的性能表现。

不仅如此,Torch-MUSA v2.0.0在完整支持PyTorch 2.2.0的基础上,还新增了对PyTorch 2.5.0的支持,使开发者能够在基于MUSA Compute Capability 3.1计算架构的全功能GPU上,无缝运行新版本的PyTorch。

未来Torch-MUSA还将继续跟进PyTorch的版本更新,计划支持更高版本的PyTorch。

国产GPU突破!摩尔线程发布Torch-MUSA v2.0.0:FP8原生支持-第1张图片-芙蓉之城

标签: 摩尔线程

抱歉,评论功能暂时关闭!